Skip to content

Commit

Permalink
fix: Only slice after sort when slice is smaller than frame length (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
mcrumiller authored Dec 2, 2024
1 parent ac4822b commit 57de43f
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
16 changes: 9 additions & 7 deletions crates/polars-core/src/frame/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2029,13 +2029,15 @@ impl DataFrame {
return Ok(out);
}
if let Some((0, k)) = slice {
let desc = if sort_options.descending.len() == 1 {
sort_options.descending[0]
} else {
false
};
sort_options.limit = Some((k as IdxSize, desc));
return self.bottom_k_impl(k, by_column, sort_options);
if k < self.len() {
let desc = if sort_options.descending.len() == 1 {
sort_options.descending[0]
} else {
false
};
sort_options.limit = Some((k as IdxSize, desc));
return self.bottom_k_impl(k, by_column, sort_options);
}
}

#[cfg(feature = "dtype-struct")]
Expand Down
11 changes: 11 additions & 0 deletions py-polars/tests/unit/operations/test_slice.py
Original file line number Diff line number Diff line change
Expand Up @@ -288,3 +288,14 @@ def test_slice_first_in_agg_18551() -> None:
"x": ["B", "C"],
"y": ["A", None],
}


def test_slice_after_sort_with_nulls_20079() -> None:
df = pl.LazyFrame({"a": [None, 1.2, None]})
out = df.sort("a", nulls_last=True).slice(0, 10).collect()
expected = pl.DataFrame({"a": [1.2, None, None]})
assert_frame_equal(out, expected)

out = df.sort("a", nulls_last=False).slice(0, 10).collect()
expected = pl.DataFrame({"a": [None, None, 1.2]})
assert_frame_equal(out, expected)

0 comments on commit 57de43f

Please sign in to comment.