fix: Don't deduplicate sort that has slice pushdown (#15784)

pola-rs · Apr 19, 2024 · 4a870cc · 4a870cc
1 parent e2845ae
commit 4a870cc
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/crates/polars-plan/src/logical_plan/optimizer/collapse_and_project.rs b/crates/polars-plan/src/logical_plan/optimizer/collapse_and_project.rs
@@ -131,7 +131,11 @@ impl OptimizationRule for SimpleProjectionAndCollapse {
                 slice,
                 sort_options,
             } => match lp_arena.get(*input) {
-                Sort { input: inner, .. } => Some(Sort {
+                Sort {
+                    input: inner,
+                    slice: None,
+                    ..
+                } => Some(Sort {
                     input: *inner,
                     by_column: by_column.clone(),
                     slice: *slice,

diff --git a/py-polars/tests/unit/operations/test_slice.py b/py-polars/tests/unit/operations/test_slice.py
@@ -240,3 +240,9 @@ def test_slice_lazy_frame_raises_proper(input_slice: tuple[int | None]) -> None:
     s = slice(*input_slice)
     with pytest.raises(ValueError, match="not supported"):
         ldf[s].collect()
+
+
+def test_double_sort_slice_pushdown_15779() -> None:
+    assert (
+        pl.LazyFrame({"foo": [1, 2]}).sort("foo").head(0).sort("foo").collect()
+    ).shape == (0, 1)