fix: slice expr can be taken in cse

pola-rs · Oct 10, 2023 · 1febff3 · 1febff3
1 parent 3b3c4a0
commit 1febff3
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/crates/polars-plan/src/logical_plan/visitor/expr.rs b/crates/polars-plan/src/logical_plan/visitor/expr.rs
@@ -154,6 +154,7 @@ impl AexprNode {
                 | (Filter { .. }, Filter { .. })
                 | (Ternary { .. }, Ternary { .. })
                 | (Count, Count)
+                | (Slice { .. }, Slice { .. })
                 | (Explode(_), Explode(_)) => true,
                 (SortBy { descending: l, .. }, SortBy { descending: r, .. }) => l == r,
                 (Agg(l), Agg(r)) => l.equal_nodes(r),

diff --git a/py-polars/tests/unit/test_cse.py b/py-polars/tests/unit/test_cse.py
@@ -483,6 +483,22 @@ def test_no_cse_in_with_context() -> None:
     }
 
 
+def test_cse_slice_11594() -> None:
+    df = pl.LazyFrame({"a": [1, 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=1, 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": [2, 1, 2, 1, 2],
+    }
+
+
 def test_cse_is_in_11489() -> None:
     df = pl.DataFrame(
         {"cond": [1, 2, 3, 2, 1], "x": [1.0, 0.20, 3.0, 4.0, 0.50]}