Skip to content

Commit

Permalink
feat: register 'set_sorted' as batch/elementwise (#13896)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Jan 22, 2024
1 parent 3dc9018 commit a75ebdf
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 1 deletion.
3 changes: 2 additions & 1 deletion crates/polars-plan/src/dsl/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down
8 changes: 8 additions & 0 deletions py-polars/tests/unit/operations/test_slice.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

0 comments on commit a75ebdf

Please sign in to comment.