fix: Fix incorrect top-k by sorted column, fix head() returning extra rows

crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs

@@ -105,8 +105,7 @@ where
     {
         len = options
             .limit
-            .map(|(limit, _)| std::cmp::min(limit as usize, len))
-            .unwrap_or(len);
+            .map_or(len, |limit| std::cmp::min(limit.try_into().unwrap(), len));
         return ChunkedArray::with_chunk(
             name,
             IdxArr::from_data_default(
@@ -136,17 +135,12 @@ where
         vals.extend(iter);
     }
 
-    let vals = if let Some((limit, desc)) = options.limit {
+    let vals = if let Some(limit) = options.limit {
         let limit = limit as usize;
         // Overwrite output len.
         len = limit;
         let out = if limit >= vals.len() {
             vals.as_mut_slice()
-        } else if desc {
-            let (lower, _el, _upper) = vals
-                .as_mut_slice()
-                .select_nth_unstable_by(limit, |a, b| b.1.tot_cmp(&a.1));
-            lower
         } else {
             let (lower, _el, _upper) = vals
                 .as_mut_slice()
@@ -205,8 +199,7 @@ where
     if is_sorted_flag != IsSorted::Not {
         let len_final = options
             .limit
-            .map(|(limit, _)| std::cmp::min(limit as usize, len))
-            .unwrap_or(len);
+            .map_or(len, |limit| std::cmp::min(limit.try_into().unwrap(), len));
         if (options.descending && is_sorted_flag == IsSorted::Descending)
             || (!options.descending && is_sorted_flag == IsSorted::Ascending)
         {
@@ -237,15 +230,10 @@ where
         }));
     }
 
-    let vals = if let Some((limit, desc)) = options.limit {
+    let vals = if let Some(limit) = options.limit {
         let limit = limit as usize;
         let out = if limit >= vals.len() {
             vals.as_mut_slice()
-        } else if desc {
-            let (lower, _el, _upper) = vals
-                .as_mut_slice()
-                .select_nth_unstable_by(limit, |a, b| b.1.tot_cmp(&a.1));
-            lower
         } else {
             let (lower, _el, _upper) = vals
                 .as_mut_slice()

crates/polars-core/src/chunked_array/ops/sort/options.rs

@@ -42,9 +42,7 @@ pub struct SortOptions {
     /// Default `false`.
     pub maintain_order: bool,
     /// Limit a sort output, this is for optimization purposes and might be ignored.
-    /// - Len
-    /// - Descending
-    pub limit: Option<(IdxSize, bool)>,
+    pub limit: Option<IdxSize>,
 }
 
 /// Sort options for multi-series sorting.
@@ -101,9 +99,7 @@ pub struct SortMultipleOptions {
     /// Whether maintain the order of equal elements. Default `false`.
     pub maintain_order: bool,
     /// Limit a sort output, this is for optimization purposes and might be ignored.
-    /// - Len
-    /// - Descending
-    pub limit: Option<(IdxSize, bool)>,
+    pub limit: Option<IdxSize>,
 }
 
 impl Default for SortOptions {

crates/polars-core/src/frame/column/mod.rs

@@ -1041,14 +1041,9 @@ impl Column {
 
         // @NOTE: This can theoretically be pushed into the previous operation but it is really
         // worth it... probably not...
-        if let Some((limit, limit_dsc)) = options.limit {
+        if let Some(limit) = options.limit {
             let limit = limit.min(length);
-
-            if limit_dsc {
-                values = values.drain((length - limit) as usize..).collect();
-            } else {
-                values.truncate(limit as usize);
-            }
+            values.truncate(limit as usize);
         }
 
         IdxCa::from_vec(self.name().clone(), values)

crates/polars-core/src/frame/mod.rs

@@ -2078,14 +2078,9 @@ impl DataFrame {
             set_sorted(&mut out);
             return Ok(out);
         }
+
         if let Some((0, k)) = slice {
             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);
             }
         }

py-polars/tests/unit/operations/test_sort.py

@@ -1087,3 +1087,47 @@ def test_sort_literals() -> None:
 
     with pytest.raises(pl.exceptions.ShapeError):
         df.sort(pl.Series(values=[1, 2]))
+
+
+def test_sorted_slice_after_function_20712() -> None:
+    assert_frame_equal(
+        pl.LazyFrame({"a": 10 * ["A"]})
+        .with_columns(b=pl.col("a").str.extract("(.*)"))
+        .sort("b")
+        .head(2)
+        .collect(),
+        pl.DataFrame({"a": ["A", "A"], "b": ["A", "A"]}),
+    )
+
+
+@pytest.mark.slow
+def test_sort_into_function_into_dynamic_groupby_20715() -> None:
+    assert (
+        pl.select(
+            time=pl.datetime_range(
+                pl.lit("2025-01-13 00:01:00.000000").str.to_datetime(
+                    "%Y-%m-%d %H:%M:%S%.f"
+                ),
+                pl.lit("2025-01-17 00:00:00.000000").str.to_datetime(
+                    "%Y-%m-%d %H:%M:%S%.f"
+                ),
+                interval="64m",
+            )
+            .cast(pl.String)
+            .reverse(),
+            val=pl.Series(range(90)),
+            cat=pl.Series(list(range(2)) * 45),
+        )
+        .lazy()
+        .with_columns(
+            pl.col("time")
+            .str.to_datetime("%Y-%m-%d %H:%M:%S%.f", strict=False)
+            .alias("time2")
+        )
+        .sort("time2")
+        .group_by_dynamic("time2", every="1m", group_by=["cat"])
+        .agg(pl.sum("val"))
+        .sort("time2")
+        .collect()
+        .shape
+    ) == (90, 3)

py-polars/tests/unit/operations/test_top_k.py

@@ -458,3 +458,70 @@ def test_top_k_df(
     assert df.sort("a", descending=True, nulls_last=False).limit(4).collect()[
         "a"
     ].to_list() == [None, None, 5, 4]
+
+
+@pytest.mark.parametrize("descending", [True, False])
+def test_sorted_top_k_20719(descending: bool) -> None:
+    df = pl.DataFrame(
+        [
+            {"a": 1, "b": 1},
+            {"a": 5, "b": 5},
+            {"a": 9, "b": 9},
+            {"a": 10, "b": 20},
+        ]
+    ).sort(by="a", descending=descending)
+
+    # Note: Output stability is guaranteed by the input sortedness as an
+    # implementation detail.
+
+    for func, reverse in [
+        [pl.DataFrame.top_k, False],
+        [pl.DataFrame.bottom_k, True],
+    ]:
+        assert_frame_equal(
+            df.pipe(func, 2, by="a", reverse=reverse),  # type: ignore[arg-type]
+            pl.DataFrame(
+                [
+                    {"a": 10, "b": 20},
+                    {"a": 9, "b": 9},
+                ]
+            ),
+        )
+
+    for func, reverse in [
+        [pl.DataFrame.top_k, True],
+        [pl.DataFrame.bottom_k, False],
+    ]:
+        assert_frame_equal(
+            df.pipe(func, 2, by="a", reverse=reverse),  # type: ignore[arg-type]
+            pl.DataFrame(
+                [
+                    {"a": 1, "b": 1},
+                    {"a": 5, "b": 5},
+                ]
+            ),
+        )
+
+
+@pytest.mark.parametrize(
+    ("func", "reverse", "expect"),
+    [
+        (pl.DataFrame.top_k, False, pl.DataFrame({"a": [2, 2]})),
+        (pl.DataFrame.bottom_k, True, pl.DataFrame({"a": [2, 2]})),
+        (pl.DataFrame.top_k, True, pl.DataFrame({"a": [1, 2]})),
+        (pl.DataFrame.bottom_k, False, pl.DataFrame({"a": [1, 2]})),
+    ],
+)
+@pytest.mark.parametrize("descending", [True, False])
+def test_sorted_top_k_duplicates(
+    func: Callable[[pl.DataFrame], pl.DataFrame],
+    reverse: bool,
+    expect: pl.DataFrame,
+    descending: bool,
+) -> None:
+    assert_frame_equal(
+        pl.DataFrame({"a": [1, 2, 2]})  # type: ignore[call-arg]
+        .sort("a", descending=descending)
+        .pipe(func, 2, by="a", reverse=reverse),
+        expect,
+    )