From df036a81bcdf7023c6a4e76b7ec3fc1700dbb231 Mon Sep 17 00:00:00 2001 From: ritchie Date: Mon, 22 Jan 2024 08:24:41 +0100 Subject: [PATCH] feat: register 'set_sorted' as batch/elementwise --- crates/polars-plan/src/dsl/mod.rs | 3 ++- py-polars/tests/unit/operations/test_slice.py | 8 ++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/crates/polars-plan/src/dsl/mod.rs b/crates/polars-plan/src/dsl/mod.rs index 633017ccbb9a..b17a64452f70 100644 --- a/crates/polars-plan/src/dsl/mod.rs +++ b/crates/polars-plan/src/dsl/mod.rs @@ -1608,7 +1608,8 @@ impl Expr { /// This can lead to incorrect results if this `Series` is not sorted!! /// Use with care! pub fn set_sorted_flag(self, sorted: IsSorted) -> Expr { - self.apply_private(FunctionExpr::SetSortedFlag(sorted)) + // This is `map`. If a column is sorted. Chunks of that column are also sorted. + self.map_private(FunctionExpr::SetSortedFlag(sorted)) } #[cfg(feature = "row_hash")] diff --git a/py-polars/tests/unit/operations/test_slice.py b/py-polars/tests/unit/operations/test_slice.py index c080f991acdd..0aa4fea66209 100644 --- a/py-polars/tests/unit/operations/test_slice.py +++ b/py-polars/tests/unit/operations/test_slice.py @@ -160,3 +160,11 @@ def test_slice_nullcount(ref: list[int | None]) -> None: assert s.null_count() == sum(x is None for x in ref) assert s.slice(64).null_count() == sum(x is None for x in ref[64:]) assert s.slice(50, 60).slice(25).null_count() == sum(x is None for x in ref[75:110]) + + +def test_slice_pushdown_set_sorted() -> None: + ldf = pl.LazyFrame({"foo": [1, 2, 3]}) + ldf = ldf.set_sorted("foo").head(5) + plan = ldf.explain() + # check the set sorted is above slice + assert plan.index("set_sorted") < plan.index("SLICE")