don't cache in-expression window functions (#3840)

polars/polars-lazy/src/physical_plan/expressions/binary.rs

@@ -92,10 +92,13 @@ impl PhysicalExpr for BinaryExpr {
     }
 
     fn evaluate(&self, df: &DataFrame, state: &ExecutionState) -> Result<Series> {
+        let mut state = state.clone();
+        // don't cache window functions as they run in parallel
+        state.cache_window = false;
         let (lhs, rhs) = POOL.install(|| {
             rayon::join(
-                || self.left.evaluate(df, state),
-                || self.right.evaluate(df, state),
+                || self.left.evaluate(df, &state),
+                || self.right.evaluate(df, &state),
             )
         });
         apply_operator_owned(lhs?, rhs?, self.op)

polars/polars-lazy/src/physical_plan/expressions/ternary.rs

@@ -87,11 +87,14 @@ impl PhysicalExpr for TernaryExpr {
         &self.expr
     }
     fn evaluate(&self, df: &DataFrame, state: &ExecutionState) -> Result<Series> {
-        let mask_series = self.predicate.evaluate(df, state)?;
+        let mut state = state.clone();
+        // don't cache window functions as they run in parallel
+        state.cache_window = false;
+        let mask_series = self.predicate.evaluate(df, &state)?;
         let mut mask = mask_series.bool()?.clone();
 
-        let op_truthy = || self.truthy.evaluate(df, state);
-        let op_falsy = || self.falsy.evaluate(df, state);
+        let op_truthy = || self.truthy.evaluate(df, &state);
+        let op_falsy = || self.falsy.evaluate(df, &state);
 
         let (truthy, falsy) = POOL.install(|| rayon::join(op_truthy, op_falsy));
         let mut truthy = truthy?;

polars/polars-lazy/src/physical_plan/expressions/window.rs

@@ -256,6 +256,9 @@ impl PhysicalExpr for WindowExpr {
             // and every partition run the cache should be empty so we expect a max of 1.
             debug_assert!(gt_map.len() <= 1);
             if let Some(gt) = gt_map.get_mut(&cache_key) {
+                // We take now, but it is important that we set this before we return!
+                // a next windows function may get this cached key and get an empty if this
+                // does not happen
                 (std::mem::take(gt), true, cache_key)
             } else {
                 (create_groups()?, false, cache_key)

py-polars/tests/db-benchmark/various.py

@@ -52,3 +52,27 @@
     )
     assert computed[0, "min"] == minimum
     assert computed[0, "max"] == maximum
+
+
+def test_windows_not_cached() -> None:
+    ldf = (
+        pl.DataFrame(
+            [
+                pl.Series("key", ["a", "a", "b", "b"]),
+                pl.Series("val", [2, 2, 1, 3]),
+            ]
+        )
+        .lazy()
+        .filter(
+            (pl.col("key").cumcount().over("key") == 0)
+            | (pl.col("val").shift(1).over("key").is_not_null())
+            | (pl.col("val") != pl.col("val").shift(1).over("key"))
+        )
+    )
+    # this might fail if they are cached
+    for i in range(1000):
+        ldf.collect()
+
+
+if __name__ == "__main__":
+    test_windows_not_cached()