Repository: pola-rs/polars (language: python)

Issue:
top_k is slow and uses a lot of memory
### Checks

- [X] I have checked that this issue has not already been reported.
- [X] I have confirmed this bug exists on the [latest version](https://pypi.org/project/polars/) of Polars.

### Reproducible example

The issue is best shown with quite large data, so I've made a dummy 500 million record dataset 
```python
import numpy as np
import polars as pl

np.random.seed(43)

TOTAL_ROWS = 500_000_000
INITIAL_ROWS = int(TOTAL_ROWS ** 0.5)

# Create the cross product of two smaller datasets
lf_a = pl.LazyFrame({'a': np.random.randint(0, 1_000_000, size=INITIAL_ROWS)})
lf_b = pl.LazyFrame({'b': np.random.randint(0, 1_000_000, size=INITIAL_ROWS)})
cross_product = lf_a.join(lf_b, how="cross")

# Sink into a parquet file
cross_product.sink_parquet("data.parquet", compression="snappy")
```

Select the top 5 `a` values

```python
import polars as pl
res = pl.scan_parquet("data.parquet").select("a").top_k(5, by='a').collect()
print(res)
```

Find mean of `a`

```python
import polars as pl
res = pl.scan_parquet("data.parquet").select("a").mean().collect()
print(res)
```

Sort `a`

```python
import polars as pl
res = pl.scan_parquet("data.parquet").select("a").sort(by="a").collect()
print(res)
```

Read `a`

```python
import polars as pl
df = pl.scan_parquet("data.parquet").select("a").collect()
print(df)
```

### Log output


> [!NOTE]
> I have not set `POLARS_VERBOSE=1` in these outputs as it does not write anything to stderr in this case and potentially interferes with performance.

The logs are edited for clarity (removed extra lines and converted memory to GB).

```shell
/usr/bin/time -lh python top_5.py
shape: (5, 1)
┌────────┐
│ a      │
│ ---    │
│ i64    │
╞════════╡
│ 999938 │
│ 999938 │
│ 999938 │
│ 999938 │
│ 999938 │
└────────┘
	20.87s real		5.22s user		11.11s sys
        28.7 GB  peak memory footprint
```

```shell
/usr/bin/time -lh python mean.py
shape: (1, 1)
┌───────────────┐
│ a             │
│ ---           │
│ f64           │
╞═══════════════╡
│ 502017.772764 │
└───────────────┘
	0.56s real		0.40s user		1.15s sys
        4.2 GB peak memory footprint
```

```shell
/usr/bin/time -lh python sort.py
shape: (499_969_600, 1)
┌────────┐
│ a      │
│ ---    │
│ i64    │
╞════════╡
│ 19     │
│ 19     │
│ 19     │
│ 19     │
│ 19     │
│ …      │
│ 999938 │
│ 999938 │
│ 999938 │
│ 999938 │
│ 999938 │
└────────┘
	3.38s real		8.32s user		1.94s sys
        12.2 GB peak memory footprint
```

```shell
/usr/bin/time -lh python read_a.py
shape: (499_969_600, 1)
┌────────┐
│ a      │
│ ---    │
│ i64    │
╞════════╡
│ 276292 │
│ 276292 │
│ 276292 │
│ 276292 │
│ 276292 │
│ …      │
│ 624089 │
│ 624089 │
│ 624089 │
│ 624089 │
│ 624089 │
└────────┘
	0.39s real		0.29s user		1.06s sys
        4.2 GB  peak memory footprint
```

Equivalent to top_5 in duckdb

```shell
/usr/bin/time -lh duckdb -c "select a from 'data.parquet' order by a desc limit 5;"
┌────────┐
│   a    │
│ int64  │
├────────┤
│ 999938 │
│ 999938 │
│ 999938 │
│ 999938 │
│ 999938 │
└────────┘
	0.58s real		2.04s user		1.54s sys
        0.026 GB peak memory footprint
```

### Issue description

Top k uses a lot more memory and takes much longer than expected. The large memory use means I get OOM errors and my actual queries fail to execute with polars.

### Expected behavior

According to the the [top_k docs](https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.top_k.html) the time complexity is $O(n)$, so I would expect it to perform like `.mean`. To be honest, I don't see how top_k is able to perform in $O(n)$. In algorithms where you keep the top k contenders in a heap, I believe you have $O(n \cdot log(k))$ time. But as shown above, the runtime is much worse than sorting the whole `a` column, so I'm not sure what is going on.

For memory use, I was really hoping that top_k would scale with the size of k. I have a large dataset where it is not feasible to load the whole thing into memory. It would be a great addition if polars could make optimizations like that. I have tried to enable the `collect(streaming=True)` flag, but that does not seem to affect the memory use of this query.

I included the max and count stats here, as that is probably related to improving the top_k performance. I would not expect a simple max and count to load the whole dataset into memory. For those two statistics, all the information is actually available in the rowgroup metadata so there are opportunities for optimizing them much further.

### Installed versions

<details>

```
--------Version info---------
Polars:              1.9.0
Index type:          UInt32
Platform:            macOS-14.7-arm64-arm-64bit
Python:              3.12.3 (main, Aug  1 2024, 12:12:15) [Clang 15.0.0 (clang-1500.3.9.4)]

----Optional dependencies----
adbc_driver_manager  <not installed>
altair               <not installed>
cloudpickle          <not installed>
connectorx           <not installed>
deltalake            <not installed>
fastexcel            <not installed>
fsspec               <not installed>
gevent               <not installed>
great_tables         <not installed>
matplotlib           <not installed>
nest_asyncio         <not installed>
numpy                2.1.2
openpyxl             <not installed>
pandas               <not installed>
pyarrow              17.0.0
pydantic             <not installed>
pyiceberg            <not installed>
sqlalchemy           <not installed>
torch                <not installed>
xlsx2csv             <not installed>
xlsxwriter           <not installed>
```

</details>



Discussion / hints:
> To be honest, I don't see how top_k is able to perform in O(n)

Easy, here's what immediately comes to mind (using k = 5):

1. Read the first five values. These are your current top 5.
2. Iterate through remaining items. If the item is larger than the minimum of the 5, swap them.

There are probably tricks to speed this up. For example, if you sort your initial 5 items, then you can perform five "mindless" swaps before you need to re-sort your top 5 again, since after swapping the first, you know your new top-5-minimum is either the newly-swapped item or the first or second value. I bet someone has come up with a faster algo, but this one is O(n).
> 2. Iterate through remaining items. If the item is larger than the minimum of the 5, swap them.

I think you have to account for
- keeping track of what the minimum is (no impact on O)
- Finding/keeping track of which value is the minimum so that you can swap it. That should take $log(k)$ with an efficient algorithm.

For k=5 that is not a lot of book keeping... but you get the point. Also, I guess you would need to swap less frequently as you go through and already have high values in your top list. Not sure how to account for that in O. Worst case is you have to swap for every record.
quickselect is average case $O(n)$: given some fixed $k$, [`select_nth_unstable`](https://doc.rust-lang.org/beta/std/primitive.slice.html#method.select_nth_unstable) (which is what polars uses) I think uses an introspective version of quickselect (to defeat bad cases that would result in $O(n^2)$ behaviour).

Idea, you partition your data into `[0..k, k, k+1..N]` such that all entries in `A[i < k] < A[k]`. You can that do that with quicksort by recursing only onto one side of your pivot. This is $O(n)$. Once you've done that, you then in general must sort the `A[:k]` entries, which is $O(k \log k)$, but $k$ is constant. So in the complexity analysis, the whole algorithm is still $O(n)$: you don't do more work for a fixed $k$ as you increase $n$.
You're right, it's O(N log k). I see now you included that in your initial post. If k is considered a constant then you can say the algorithm is O(n).

It looks like polars uses Rust's [`select_nth_unstable_by`](https://doc.rust-lang.org/beta/std/primitive.slice.html#method.select_nth_unstable) and @orlp helped identify/speed up their implementation [here](https://github.com/rust-lang/rust/issues/91644), so he probably has a lot more knowledge of what's going on here.
Regarding running-time, you don't need all this code.
Look at:

```python
import numpy as np
import polars as pl
import datetime

np.random.seed(43)

TOTAL_ROWS = 500_000_000

df = pl.DataFrame({'a': np.random.randint(0, 1_000_000, size=TOTAL_ROWS)})

t1 = datetime.datetime.now()
res = df.select("a").top_k(5, by='a')
t2 = datetime.datetime.now()
print((t2 - t1).total_seconds())

t1 = datetime.datetime.now()
res = df.select("a").sort(by="a")
t2 = datetime.datetime.now()
print((t2 - t1).total_seconds())

```

Sorting is quicker, and I wonder why.

By the way, this code:

```python
res = df.select(pl.col("a").top_k(k=5))
```
Works as expected.

I think they have different implementations. but I need to verify it.

**Edit:**
It seems indeed `df.select("a").top_k(5, by='a')` uses different implementation - sort, followed by slice.
**Edit:** I was dense and just realized now that @barak1412 is saying that `top_k` works as expected in expression context i.e,  does not do sort followed by slice.

-----------

As @wence- mentioned, you get $O(n)$ for fixed $k$ with quickselect. However, looking through the code, It looks like a sort followed by a slice for both `DataFrame` and `LazyFrame` I'm not savvy enough to gauge if that is somehow optimized away.

https://github.com/pola-rs/polars/blob/8a76daddccfe06d789312ec6e077dd230a6722c3/crates/polars-lazy/src/frame/mod.rs#L347-L359

My other point was about space complexity. Maybe it is $O(n \cdot log(n))$ at the moment. I'm not sure why we get a lot of re-allocation during the sort, but I think we do. If the full sort is avoided, it should become $O(n)$. What I would really like to see is $O(1)$ space complexity, i.e, that the dataset size does not matter. (Considering $k$ a variable rather than a constant, the space complexity is $O(k)$ ).

However, any work to achieve that, will require the https://github.com/pola-rs/polars/labels/new-streaming engine so it is possible to iterate over larger than memory lazyframes https://github.com/pola-rs/polars/issues/10683 .

In that future, I think quickselect is not the right algorithm, because it can backtrack basically across the whole dataset, requiring the dataset to be in memory (or loaded again). I think an algorithm that keeps candidate rows around as it iterates through the dataset will perform much better – at least for $k << n$.

### Work around
For my actual use case, I have more columns and they are all loaded into memory in the current top_k implementation, which is what tips the memory use over the edge for my system. I found that you can leverage that `quantile` has a more efficient implementation, which can be used to find the minimum value of interest and use that in a filter. Because of predicate pushdown optimization, only rowgroups of interest are then loaded into memory.

Given these columns

```python
a = pl.col("a")
b = pl.col("b")
a_plus_b = (a + b).alias('a_plus_b')
a_squared = (a**2).alias('a_squared')
```

Straight top_k

```python
res = (pl
    .scan_parquet("data.parquet")
    .select(a, b, a_plus_b, a_squared)
    .top_k(5, by='a')
    .collect()
)
```
takes 34.2 seconds, 52.7 GB memory.


Finding the quantile

```python
# Percentile for the 5th largest value given a 500 million dataset
percentile_fraction = 1 - 1e-8

res = (pl
    .scan_parquet("data.parquet")
    .select("a")
    .quantile(percentile_fraction)
    .collect()
)
```
takes 2.7 seconds, 8.2 GB memory.


Combining finding quantile and filtering before top_k

```python
percentile_fraction = 1 - 1e-8
lf = pl.scan_parquet("data.parquet")

percentile = lf.select(a).quantile(percentile_fraction).collect()
res = (lf
    .select(a, b, a_plus_b, a_squared)
    .filter( a >= percentile)
    .top_k(5, by='a')
    .collect()
)
```
takes 2.8 seconds, 8.2 GB memory.






@foldager You can see in the comment in the code you pasted it is written "this will optimize to top-k".

So this will be changed in the future.

Try measure the memory consumption using this code instead:
```python
res = df.select(pl.col("a").top_k(k=5))
```
Thanks @barak1412. Expression `top_k` has a fast running time in both dataframe and lazy frame context. It improves my work-around as I can replace `quantile()` with the expression `top_k`.

```python
res = pl.scan_parquet("data.parquet").select(pl.col("a").top_k(k=5)).collect()
```

takes 1.6 seconds, 8.2 GB memory.

```python
a = pl.col("a")
b = pl.col("b")
a_plus_b = (a + b).alias('a_plus_b')
a_squared = (a**2).alias('a_squared')

# Base lazy frame
lf = pl.scan_parquet("data.parquet")

min_a = lf.select(a.top_k(k=5).min()).collect().item()
res = (lf
    .filter(a >= min_a)
    .select(a, b, a_plus_b, a_squared)
    .top_k(5, by=a)
    .collect()
)
```
takes 1.6 seconds, 8.2 GB memory. The same as above because the second step is fast.

A python program with the second step alone
```
lf = pl.scan_parquet("data.parquet")

res = (lf
    .filter(a >= 999938)
    .select(a, b, a_plus_b, a_squared)
    .top_k(5, by=a)
    .collect()
)
```
takes 0.12 seconds 71.9 MB memory.

Looking into it.

Fix patch:
diff --git a/crates/polars-core/src/chunked_array/ops/sort/arg_bottom_k.rs b/crates/polars-core/src/chunked_array/ops/sort/arg_bottom_k.rs
index 7f257f23f59e..7787ef28076f 100644
--- a/crates/polars-core/src/chunked_array/ops/sort/arg_bottom_k.rs
+++ b/crates/polars-core/src/chunked_array/ops/sort/arg_bottom_k.rs
@@ -39,6 +39,11 @@ pub fn _arg_bottom_k(
     _broadcast_bools(by_column.len(), &mut sort_options.descending);
     _broadcast_bools(by_column.len(), &mut sort_options.nulls_last);
 
+    // Don't go into row encoding.
+    if by_column.len() == 1 && sort_options.limit.is_some() && !sort_options.maintain_order {
+        return Ok(NoNull::new(by_column[0].arg_sort((&*sort_options).into())));
+    }
+
     let encoded = _get_rows_encoded(
         by_column,
         &sort_options.descending,
diff --git a/crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs b/crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs
index ca34d37318a7..4f9a1ff9e9b3 100644
--- a/crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs
+++ b/crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs
@@ -18,7 +18,7 @@ pub(super) fn arg_sort<I, J, T>(
     iters: I,
     options: SortOptions,
     null_count: usize,
-    len: usize,
+    mut len: usize,
 ) -> IdxCa
 where
     I: IntoIterator<Item = J>,
@@ -49,14 +49,46 @@ where
         vals.extend(iter);
     }
 
-    sort_impl(vals.as_mut_slice(), options);
+    let vals = if let Some((limit, desc)) = 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()
+                .select_nth_unstable_by(limit, |a, b| a.1.tot_cmp(&b.1));
+            lower
+        };
+
+        sort_impl(out, options);
+        out
+    } else {
+        sort_impl(vals.as_mut_slice(), options);
+        vals.as_slice()
+    };
 
-    let iter = vals.into_iter().map(|(idx, _v)| idx);
+    let iter = vals.iter().map(|(idx, _v)| idx).copied();
     let idx = if nulls_last {
         let mut idx = Vec::with_capacity(len);
         idx.extend(iter);
-        idx.extend(nulls_idx);
+
+        let nulls_idx = if options.limit.is_some() {
+            &nulls_idx[..len - idx.len()]
+        } else {
+            &nulls_idx
+        };
+        idx.extend_from_slice(nulls_idx);
         idx
+    } else if options.limit.is_some() {
+        nulls_idx.extend(iter.take(len - nulls_idx.len()));
+        nulls_idx
     } else {
         let ptr = nulls_idx.as_ptr() as usize;
         nulls_idx.extend(iter);
@@ -90,9 +122,29 @@ where
         }));
     }
 
-    sort_impl(vals.as_mut_slice(), options);
+    let vals = if let Some((limit, desc)) = 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()
+                .select_nth_unstable_by(limit, |a, b| a.1.tot_cmp(&b.1));
+            lower
+        };
+        sort_impl(out, options);
+        out
+    } else {
+        sort_impl(vals.as_mut_slice(), options);
+        vals.as_slice()
+    };
 
-    let iter = vals.into_iter().map(|(idx, _v)| idx);
+    let iter = vals.iter().map(|(idx, _v)| idx).copied();
     let idx: Vec<_> = iter.collect_trusted();
 
     ChunkedArray::with_chunk(name, IdxArr::from_data_default(Buffer::from(idx), None))
diff --git a/crates/polars-core/src/chunked_array/ops/sort/categorical.rs b/crates/polars-core/src/chunked_array/ops/sort/categorical.rs
index 5dd71a7b1eb8..c89e0790f251 100644
--- a/crates/polars-core/src/chunked_array/ops/sort/categorical.rs
+++ b/crates/polars-core/src/chunked_array/ops/sort/categorical.rs
@@ -53,6 +53,7 @@ impl CategoricalChunked {
             descending,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         })
     }
 
diff --git a/crates/polars-core/src/chunked_array/ops/sort/mod.rs b/crates/polars-core/src/chunked_array/ops/sort/mod.rs
index 727f2ace15a8..add7e8b696a4 100644
--- a/crates/polars-core/src/chunked_array/ops/sort/mod.rs
+++ b/crates/polars-core/src/chunked_array/ops/sort/mod.rs
@@ -335,6 +335,7 @@ impl ChunkSort<StringType> for StringChunked {
             nulls_last: false,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         })
     }
 
@@ -406,6 +407,7 @@ impl ChunkSort<BinaryType> for BinaryChunked {
             nulls_last: false,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         })
     }
 
@@ -536,6 +538,7 @@ impl ChunkSort<BinaryOffsetType> for BinaryOffsetChunked {
             nulls_last: false,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         })
     }
 
@@ -672,6 +675,7 @@ impl ChunkSort<BooleanType> for BooleanChunked {
             nulls_last: false,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         })
     }
 
@@ -797,6 +801,7 @@ mod test {
             nulls_last: false,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         });
         assert_eq!(
             Vec::from(&out),
@@ -816,6 +821,7 @@ mod test {
             nulls_last: true,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         });
         assert_eq!(
             Vec::from(&out),
@@ -925,6 +931,7 @@ mod test {
             nulls_last: false,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         });
         let expected = &[None, None, Some("a"), Some("b"), Some("c")];
         assert_eq!(Vec::from(&out), expected);
@@ -934,6 +941,7 @@ mod test {
             nulls_last: false,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         });
 
         let expected = &[None, None, Some("c"), Some("b"), Some("a")];
@@ -944,6 +952,7 @@ mod test {
             nulls_last: true,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         });
         let expected = &[Some("a"), Some("b"), Some("c"), None, None];
         assert_eq!(Vec::from(&out), expected);
@@ -953,6 +962,7 @@ mod test {
             nulls_last: true,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         });
         let expected = &[Some("c"), Some("b"), Some("a"), None, None];
         assert_eq!(Vec::from(&out), expected);
diff --git a/crates/polars-core/src/chunked_array/ops/sort/options.rs b/crates/polars-core/src/chunked_array/ops/sort/options.rs
index 046d0b251b04..95bff0b1b47a 100644
--- a/crates/polars-core/src/chunked_array/ops/sort/options.rs
+++ b/crates/polars-core/src/chunked_array/ops/sort/options.rs
@@ -41,6 +41,10 @@ pub struct SortOptions {
     /// If true 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)>,
 }
 
 /// Sort options for multi-series sorting.
@@ -96,6 +100,10 @@ pub struct SortMultipleOptions {
     pub multithreaded: bool,
     /// 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)>,
 }
 
 impl Default for SortOptions {
@@ -105,6 +113,7 @@ impl Default for SortOptions {
             nulls_last: false,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         }
     }
 }
@@ -116,6 +125,7 @@ impl Default for SortMultipleOptions {
             nulls_last: vec![false],
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         }
     }
 }
@@ -224,6 +234,7 @@ impl From<&SortOptions> for SortMultipleOptions {
             nulls_last: vec![value.nulls_last],
             multithreaded: value.multithreaded,
             maintain_order: value.maintain_order,
+            limit: value.limit,
         }
     }
 }
@@ -235,6 +246,7 @@ impl From<&SortMultipleOptions> for SortOptions {
             nulls_last: value.nulls_last.first().copied().unwrap_or(false),
             multithreaded: value.multithreaded,
             maintain_order: value.maintain_order,
+            limit: value.limit,
         }
     }
 }
diff --git a/crates/polars-core/src/frame/mod.rs b/crates/polars-core/src/frame/mod.rs
index e3b969a81756..7e2d7b050dcf 100644
--- a/crates/polars-core/src/frame/mod.rs
+++ b/crates/polars-core/src/frame/mod.rs
@@ -1989,6 +1989,12 @@ 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);
         }
 
@@ -2012,6 +2018,7 @@ impl DataFrame {
                     nulls_last: sort_options.nulls_last[0],
                     multithreaded: sort_options.multithreaded,
                     maintain_order: sort_options.maintain_order,
+                    limit: sort_options.limit,
                 };
                 // fast path for a frame with a single series
                 // no need to compute the sort indices and then take by these indices
diff --git a/crates/polars-expr/src/expressions/sortby.rs b/crates/polars-expr/src/expressions/sortby.rs
index 1624d7c9bcd6..fad081cb49ed 100644
--- a/crates/polars-expr/src/expressions/sortby.rs
+++ b/crates/polars-expr/src/expressions/sortby.rs
@@ -160,6 +160,7 @@ fn sort_by_groups_multiple_by(
                 nulls_last: nulls_last.to_owned(),
                 multithreaded,
                 maintain_order,
+                limit: None,
             };
 
             let sorted_idx = groups[0]
@@ -180,6 +181,7 @@ fn sort_by_groups_multiple_by(
                 nulls_last: nulls_last.to_owned(),
                 multithreaded,
                 maintain_order,
+                limit: None,
             };
             let sorted_idx = groups[0]
                 .as_materialized_series()
diff --git a/crates/polars-ops/src/chunked_array/top_k.rs b/crates/polars-ops/src/chunked_array/top_k.rs
index 9caf861b6cd9..ef37e267c10f 100644
--- a/crates/polars-ops/src/chunked_array/top_k.rs
+++ b/crates/polars-ops/src/chunked_array/top_k.rs
@@ -285,6 +285,7 @@ fn top_k_by_impl(
         nulls_last: vec![true; by.len()],
         multithreaded,
         maintain_order: false,
+        limit: None,
     };
 
     let idx = _arg_bottom_k(k, by, &mut sort_options)?;
diff --git a/crates/polars-ops/src/frame/join/hash_join/sort_merge.rs b/crates/polars-ops/src/frame/join/hash_join/sort_merge.rs
index fce2f2bf6cf0..95cde8387733 100644
--- a/crates/polars-ops/src/frame/join/hash_join/sort_merge.rs
+++ b/crates/polars-ops/src/frame/join/hash_join/sort_merge.rs
@@ -225,6 +225,7 @@ pub(crate) fn _sort_or_hash_inner(
                 nulls_last: false,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             });
             let s_right = unsafe { s_right.take_unchecked(&sort_idx) };
             let ids = par_sorted_merge_inner_no_nulls(s_left, &s_right);
@@ -252,6 +253,7 @@ pub(crate) fn _sort_or_hash_inner(
                 nulls_last: false,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             });
             let s_left = unsafe { s_left.take_unchecked(&sort_idx) };
             let ids = par_sorted_merge_inner_no_nulls(&s_left, s_right);
@@ -323,6 +325,7 @@ pub(crate) fn sort_or_hash_left(
                 nulls_last: false,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             });
             let s_right = unsafe { s_right.take_unchecked(&sort_idx) };
 
diff --git a/crates/polars-pipe/src/executors/sinks/sort/source.rs b/crates/polars-pipe/src/executors/sinks/sort/source.rs
index 1c1fa2984a0e..6f544e8e6ef1 100644
--- a/crates/polars-pipe/src/executors/sinks/sort/source.rs
+++ b/crates/polars-pipe/src/executors/sinks/sort/source.rs
@@ -101,6 +101,7 @@ impl SortSource {
                     nulls_last: self.nulls_last,
                     multithreaded: true,
                     maintain_order: false,
+                    limit: None,
                 },
             ),
             Some((offset, len)) => {
@@ -119,6 +120,7 @@ impl SortSource {
                             nulls_last: self.nulls_last,
                             multithreaded: true,
                             maintain_order: false,
+                            limit: None,
                         },
                     );
                     *len = len.saturating_sub(df_len);
diff --git a/crates/polars-python/src/expr/general.rs b/crates/polars-python/src/expr/general.rs
index 7125388e88cd..fe5fdafdbbb8 100644
--- a/crates/polars-python/src/expr/general.rs
+++ b/crates/polars-python/src/expr/general.rs
@@ -260,6 +260,7 @@ impl PyExpr {
                 nulls_last,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             })
             .into()
     }
@@ -272,6 +273,7 @@ impl PyExpr {
                 nulls_last,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             })
             .into()
     }
@@ -349,6 +351,7 @@ impl PyExpr {
                     nulls_last,
                     multithreaded,
                     maintain_order,
+                    limit: None,
                 },
             )
             .into()
diff --git a/crates/polars-python/src/functions/lazy.rs b/crates/polars-python/src/functions/lazy.rs
index d3ebb376d10f..1c4e738ea69c 100644
--- a/crates/polars-python/src/functions/lazy.rs
+++ b/crates/polars-python/src/functions/lazy.rs
@@ -75,6 +75,7 @@ pub fn arg_sort_by(
             nulls_last,
             multithreaded,
             maintain_order,
+            limit: None,
         },
     )
     .into()
diff --git a/crates/polars-python/src/lazyframe/general.rs b/crates/polars-python/src/lazyframe/general.rs
index fd89884ece82..f9fb740d4cae 100644
--- a/crates/polars-python/src/lazyframe/general.rs
+++ b/crates/polars-python/src/lazyframe/general.rs
@@ -539,6 +539,7 @@ impl PyLazyFrame {
                 nulls_last: vec![nulls_last],
                 multithreaded,
                 maintain_order,
+                limit: None,
             },
         )
         .into()
@@ -561,6 +562,7 @@ impl PyLazyFrame {
                 nulls_last,
                 maintain_order,
                 multithreaded,
+                limit: None,
             },
         )
         .into()
diff --git a/crates/polars-python/src/series/general.rs b/crates/polars-python/src/series/general.rs
index b14285e77aa0..7312995d7606 100644
--- a/crates/polars-python/src/series/general.rs
+++ b/crates/polars-python/src/series/general.rs
@@ -457,6 +457,7 @@ impl PySeries {
             nulls_last,
             multithreaded: true,
             maintain_order: false,
+            limit: None,
         };
         Ok(self.series.is_sorted(options).map_err(PyPolarsErr::from)?)
     }


Test patch:
diff --git a/crates/polars-lazy/src/tests/aggregations.rs b/crates/polars-lazy/src/tests/aggregations.rs
index 6b2d8cb05da0..2ab337ef51e9 100644
--- a/crates/polars-lazy/src/tests/aggregations.rs
+++ b/crates/polars-lazy/src/tests/aggregations.rs
@@ -450,6 +450,7 @@ fn take_aggregations() -> PolarsResult<()> {
                             nulls_last: false,
                             multithreaded: true,
                             maintain_order: false,
+                            limit: None,
                         })
                         .head(Some(2)),
                 )
@@ -489,6 +490,7 @@ fn test_take_consistency() -> PolarsResult<()> {
                 nulls_last: false,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             })
             .get(lit(0))])
         .collect()?;
@@ -507,6 +509,7 @@ fn test_take_consistency() -> PolarsResult<()> {
                 nulls_last: false,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             })
             .get(lit(0))])
         .collect()?;
@@ -526,6 +529,7 @@ fn test_take_consistency() -> PolarsResult<()> {
                     nulls_last: false,
                     multithreaded: true,
                     maintain_order: false,
+                    limit: None,
                 })
                 .get(lit(0))
                 .alias("1"),
@@ -537,6 +541,7 @@ fn test_take_consistency() -> PolarsResult<()> {
                             nulls_last: false,
                             multithreaded: true,
                             maintain_order: false,
+                            limit: None,
                         })
                         .get(lit(0)),
                 )
diff --git a/crates/polars-lazy/src/tests/queries.rs b/crates/polars-lazy/src/tests/queries.rs
index 95cbf586be67..0c4e518b5042 100644
--- a/crates/polars-lazy/src/tests/queries.rs
+++ b/crates/polars-lazy/src/tests/queries.rs
@@ -1666,6 +1666,7 @@ fn test_single_group_result() -> PolarsResult<()> {
                 nulls_last: false,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             })
             .over([col("a")])])
         .collect()?;
diff --git a/py-polars/tests/unit/operations/test_top_k.py b/py-polars/tests/unit/operations/test_top_k.py
index debb3e729274..866ef88e6e10 100644
--- a/py-polars/tests/unit/operations/test_top_k.py
+++ b/py-polars/tests/unit/operations/test_top_k.py
@@ -397,3 +397,37 @@ def test_bottom_k_nulls(s: pl.Series, should_sort: bool) -> None:
 def test_top_k_descending_deprecated() -> None:
     with pytest.deprecated_call():
         pl.col("a").top_k_by("b", descending=True)  # type: ignore[call-arg]
+
+
+def test_top_k_df() -> None:
+    df = pl.LazyFrame({"a": [3, 4, 1, 2, 5]})
+    expected = [5, 4, 3]
+    assert df.sort("a", descending=True).limit(3).collect()["a"].to_list() == expected
+    assert df.top_k(3, by="a").collect()["a"].to_list() == expected
+    expected = [1, 2, 3]
+    assert df.sort("a", descending=False).limit(3).collect()["a"].to_list() == expected
+    assert df.bottom_k(3, by="a").collect()["a"].to_list() == expected
+
+    df = pl.LazyFrame({"a": [1, None, None, 4, 5]})
+    expected2 = [5, 4, 1, None]
+    assert (
+        df.sort("a", descending=True, nulls_last=True).limit(4).collect()["a"].to_list()
+        == expected2
+    )
+    assert df.top_k(4, by="a").collect()["a"].to_list() == expected2
+    expected2 = [1, 4, 5, None]
+    assert (
+        df.sort("a", descending=False, nulls_last=True)
+        .limit(4)
+        .collect()["a"]
+        .to_list()
+        == expected2
+    )
+    assert df.bottom_k(4, by="a").collect()["a"].to_list() == expected2
+
+    assert df.sort("a", descending=False, nulls_last=False).limit(4).collect()[
+        "a"
+    ].to_list() == [None, None, 1, 4]
+    assert df.sort("a", descending=True, nulls_last=False).limit(4).collect()[
+        "a"
+    ].to_list() == [None, None, 5, 4]