fix: Fix input replacement logic for slice (#11631)

pola-rs · Oct 10, 2023 · 543cea5 · 543cea5
1 parent 9a78bfa
commit 543cea5
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/crates/polars-plan/src/logical_plan/aexpr/mod.rs b/crates/polars-plan/src/logical_plan/aexpr/mod.rs
@@ -340,7 +340,7 @@ impl AExpr {
             Column(_) | Literal(_) | Wildcard | Count | Nth(_) => return self,
             Alias(input, _) => input,
             Cast { expr, .. } => expr,
-            Explode(input) | Slice { input, .. } => input,
+            Explode(input) => input,
             BinaryExpr { left, right, .. } => {
                 *right = inputs[0];
                 *left = inputs[1];
@@ -390,6 +390,16 @@ impl AExpr {
                 input.extend(inputs.iter().rev().copied());
                 return self;
             },
+            Slice {
+                input,
+                offset,
+                length,
+            } => {
+                *length = inputs[0];
+                *offset = inputs[1];
+                *input = inputs[2];
+                return self;
+            },
             Window {
                 function,
                 partition_by,

diff --git a/py-polars/tests/unit/test_cse.py b/py-polars/tests/unit/test_cse.py
@@ -498,6 +498,18 @@ def test_cse_slice_11594() -> None:
         "2": [2, 1, 2, 1, 2],
     }
 
+    q = df.select(
+        pl.col("a").slice(offset=1, length=pl.count() - 1).alias("1"),
+        pl.col("a").slice(offset=0, length=pl.count() - 1).alias("2"),
+    )
+
+    assert "__POLARS_CSE" in q.explain(comm_subexpr_elim=True)
+
+    assert q.collect(comm_subexpr_elim=True).to_dict(False) == {
+        "1": [2, 1, 2, 1, 2],
+        "2": [1, 2, 1, 2, 1],
+    }
+
 
 def test_cse_is_in_11489() -> None:
     df = pl.DataFrame(