From a413c4a868e29aacb6f2b21e51533fc98fe380e9 Mon Sep 17 00:00:00 2001 From: Marshall Date: Mon, 2 Dec 2024 03:07:33 -0500 Subject: [PATCH] fix: Only slice after sort when slice is smaller than frame length (#20084) --- crates/polars-core/src/frame/mod.rs | 16 +++++++++------- py-polars/tests/unit/operations/test_slice.py | 11 +++++++++++ 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/crates/polars-core/src/frame/mod.rs b/crates/polars-core/src/frame/mod.rs index 34692949adf3..b524bbd06d0f 100644 --- a/crates/polars-core/src/frame/mod.rs +++ b/crates/polars-core/src/frame/mod.rs @@ -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")] diff --git a/py-polars/tests/unit/operations/test_slice.py b/py-polars/tests/unit/operations/test_slice.py index 94dc1e3283ff..b148b2443421 100644 --- a/py-polars/tests/unit/operations/test_slice.py +++ b/py-polars/tests/unit/operations/test_slice.py @@ -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)