Support slice pushdown with horizontal concatenation

pola-rs · Dec 19, 2023 · 66b1978 · 66b1978
1 parent 9c85968
commit 66b1978
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/crates/polars-plan/src/logical_plan/optimizer/slice_pushdown_lp.rs b/crates/polars-plan/src/logical_plan/optimizer/slice_pushdown_lp.rs
@@ -350,6 +350,11 @@ impl SlicePushDown {
                     self.no_pushdown_restart_opt(lp, state, lp_arena, expr_arena)
                 }
             }
+            (HConcat {inputs, schema, options}, _) => {
+                // Slice can always be pushed down for horizontal concatenation
+                let lp = HConcat {inputs, schema, options};
+                self.pushdown_and_continue(lp, state, lp_arena, expr_arena)
+            }
             (catch_all, state) => {
                 self.no_pushdown_finish_opt(catch_all, state, lp_arena)
             }

diff --git a/py-polars/tests/unit/operations/test_slice.py b/py-polars/tests/unit/operations/test_slice.py
@@ -121,3 +121,22 @@ def test_hstack_slice_pushdown() -> None:
     plan = out.explain()
 
     assert not plan.startswith("SLICE")
+
+
+def test_hconcat_slice_pushdown() -> None:
+    num_dfs = 3
+    lfs = [
+        pl.LazyFrame({f"column_{i}": list(range(i, i + 10))}) for i in range(num_dfs)
+    ]
+
+    out = pl.concat(lfs, how="horizontal").slice(2, 3)
+    plan = out.explain()
+
+    assert not plan.startswith("SLICE")
+
+    expected = pl.DataFrame(
+        {f"column_{i}": list(range(i + 2, i + 5)) for i in range(num_dfs)}
+    )
+
+    df_out = out.collect()
+    assert_frame_equal(df_out, expected)