Polars Finds the Top K With a Bounded Heap

TL;DR: There are two top-k implementations and they are not the same algorithm. The in-memory engine materialises the column and calls select_nth_unstable_by — a quickselect partition, O(n) expected, no ordering among the winners. The streaming engine cannot materialise anything, so it keeps a BinaryHeap holding exactly k keys per pipeline, admits a new key only when it is strictly below the current threshold, and merges the per-pipeline heaps pairwise when the input ends. Ties are broken by arrival order, not by row index, and the threshold is published as a filter so that rows which cannot win are never even read.

top_k looks like a trivial operation. Sort the column, take the first k, done. That is fine when the column already fits in memory and you are willing to pay O(n log n) for a full sort to look at k rows — which is almost never a good trade when k is 10 and n is a billion. The interesting question is what a dataframe engine does when it cannot afford the sort.

In Polars there are two answers, and they live in different crates.

The in-memory answer: partition, don’t sort

The eager engine has the whole column in RAM, so it can afford to copy it and then throw most of it away. This is the numeric path:

rust
fn top_k_num_impl<T>(ca: &ChunkedArray<T>, k: usize, descending: bool) -> ChunkedArray<T>
where
    T: PolarsNumericType,
{
    if k >= ca.len() && ca.null_count() == 0 {
        return ca.clone();
    }

    // Get rid of all the nulls and transform into Vec<T::Native>.
    let mut nnca = ca.drop_nulls();
    nnca.rechunk_mut();
    let chunk = nnca.downcast_into_iter().next().unwrap();
    let (_, buffer, _) = chunk.into_inner();
    let mut vec = buffer.to_vec();

    // Partition.
    if k < vec.len() {
        if descending {
            vec.select_nth_unstable_by(k, TotalOrd::tot_cmp);
        } else {
            vec.select_nth_unstable_by(k, |a, b| TotalOrd::tot_cmp(b, a));
        }
    }

    // Reconstruct output (with nulls at the end).
    let out_len = k.min(ca.len());
    let non_null_count = ca.len() - ca.null_count();
    vec.resize(out_len, T::Native::default());
    let validity = first_n_valid_mask(non_null_count, out_len);

    let arr = PrimitiveArray::from_vec(vec).with_validity_typed(validity);
    ChunkedArray::with_chunk_like(ca, arr)
}

crates/polars-ops/src/chunked_array/top_k.rs:66

select_nth_unstable_by is quickselect: it partitions the slice so that the element at index k is the one that would be there in sorted order, and it does that in O(n) expected time instead of O(n log n). The first k elements are then the winners. The result is not sorted — the winners come out in whatever order the partition left them — and ties are broken arbitrarily, because the algorithm is unstable.

The multi-column case cannot use quickselect on a single column, so it row-encodes the keys into bytes and partitions those:

rust
    let sorted = if k >= from_n_rows {
        match (sort_options.multithreaded, sort_options.maintain_order) {
            (true, true) => RAYON.install(|| {
                rows.par_sort();
            }),
            (true, false) => RAYON.install(|| {
                rows.par_sort_unstable();
            }),
            (false, true) => rows.sort(),
            (false, false) => rows.sort_unstable(),
        }
        &rows
    } else if sort_options.maintain_order {
        // todo: maybe there is some more efficient method, comparable to select_nth_unstable
        if sort_options.multithreaded {
            RAYON.install(|| {
                rows.par_sort();
            })
        } else {
            rows.sort();
        }
        &rows[..k]
    } else {
        // todo: possible multi threaded `select_nth_unstable`?
        let (lower, _el, _upper) = rows.select_nth_unstable(k);
        if sort_options.multithreaded {
            RAYON.install(|| {
                lower.par_sort_unstable();
            })
        } else {
            lower.sort_unstable();
        }
        &*lower
    };

crates/polars-core/src/chunked_array/ops/sort/arg_bottom_k.rs:63

Note the escape hatch: maintain_order forces a full sort, because quickselect cannot promise anything about the order of the winners. That flag is the difference between “the k largest values” and “the k largest values, in the order they appeared”.

There is a third path, used by arg_sort with a limit, which is the one place in Polars where a top-k keeps a stable order:

rust
fn sort_limit_impl<T>(
    vals: &mut [(IdxSize, T)],
    limit: usize,
    options: SortOptions,
) -> &[(IdxSize, T)]
where
    T: TotalOrd + Send + Sync,
{
    if limit >= vals.len() {
        sort_impl(vals, options);
        return vals;
    }
    let cmp = |a: &(IdxSize, T), b: &(IdxSize, T)| {
        let ord = if options.descending {
            b.1.tot_cmp(&a.1)
        } else {
            a.1.tot_cmp(&b.1)
        };
        ord.then(a.0.cmp(&b.0))
    };
    let (lower, _el, _upper) = vals.select_nth_unstable_by(limit, cmp);
    if options.multithreaded {
        RAYON.install(|| lower.par_sort_unstable_by(cmp));
    } else {
        lower.sort_unstable_by(cmp);
    }
    lower
}

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

The .then(a.0.cmp(&b.0)) is the tie-break: when the values compare equal, the original row index decides. That is what makes the result stable, and it is why the comparison carries the index along.

The streaming answer: a heap of exactly k keys

The streaming engine cannot afford to materialise a sorted copy of the column. It receives morsels, and it must decide, row by row, whether a row is worth keeping. That is a job for a bounded heap: keep the k best keys seen so far, throw everything else away.

The state is a struct in crates/polars-stream/src/nodes/top_k.rs:

rust
struct BottomKWithPayload<P> {
    k: usize,
    heap: BinaryHeap<Priority<P, (DfsKey, RowIdxKey)>>,
    df_subsets: SlotMap<DfsKey, DfSubset>,
    row_idxs: SlotMap<RowIdxKey, IdxSize>,
    marked_for_pruning: SecondaryMap<DfsKey, ()>,
    to_prune: Vec<DfsKey>,
    gather_idxs: Vec<IdxSize>,
    shared_optimum: Arc<RwLock<Option<P>>>,
    /// The best `k` keys of all pipelines, when they share them.
    shared_keys: Option<Arc<Mutex<BinaryHeap<P>>>>,
}

crates/polars-stream/src/nodes/top_k.rs:78

The heap holds Priority<P, (DfsKey, RowIdxKey)>. P is the key type — a ReorderWithNulls<TotalOrdWrap<f64>> for numeric keys, a ReorderWithNulls<Vec<u8>> for strings, a row encoding for multi-column keys — and the payload is a pair of handles: which data frame the row came from, and which row inside it.

The Priority wrapper is the trick that makes this work. It is a newtype that compares only on the key:

rust
use std::cmp::Ordering;

/// A pair which is ordered exclusively by the first element.
#[derive(Copy, Clone, Debug)]
pub struct Priority<P, T>(pub P, pub T);

impl<P: Ord + Eq, T> Ord for Priority<P, T> {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.cmp(&other.0)
    }
}

crates/polars-utils/src/priority.rs:7

So the heap’s ordering is the ordering of the keys, and the payload is carried along without influencing it. That is what allows the heap to be a plain BinaryHeap<Priority<...>> instead of a hand-rolled structure.

The bounded part

The whole algorithm is add_one:

rust
    fn add_one<Q>(
        &mut self,
        dfs_key: DfsKey,
        row_idx: IdxSize,
        key: Q,
        is_less: impl Fn(&Q, &P) -> bool,
        to_owned: impl Fn(Q) -> P,
    ) -> bool {
        // We use a max-heap for our bottom k. This means the top element in our heap (peek())
        // is the first to be replaced.
        if !(self.heap.len() < self.k || is_less(&key, self.peek_optimum().unwrap())) {
            return false;
        }
        let opt = self.heap.peek().map(|p| p.1);
        let row_idx_key = self.row_idxs.insert(row_idx);
        let df_subset = &mut self.df_subsets[dfs_key];
        df_subset.subset_len += 1;
        df_subset.rows.push(row_idx_key);
        self.heap
            .push(Priority(to_owned(key), (dfs_key, row_idx_key)));

        if self.heap.len() > self.k {
            let (dfs_key, row_idx_key) = self.heap.pop().unwrap().1;
            self.row_idxs[row_idx_key] = IdxSize::MAX;
            self.df_subsets[dfs_key].subset_len -= 1;
            self.mark_if_sparse(dfs_key);
        }

        opt != self.heap.peek().map(|p| p.1)
    }

crates/polars-stream/src/nodes/top_k.rs:215

Read it as a max-heap holding the k smallest keys seen so far (Polars calls this “bottom k”, and reverses the comparison when the user asked for a top-k). peek() is the largest key in the heap, i.e. the current threshold: the k-th best key so far. A new key is admitted only if the heap is not full yet, or if the new key is strictly below that threshold. If it is admitted, it is pushed and the largest element is popped, so the heap size never exceeds k + 1 for an instant and never exceeds k afterwards.

The comparison is a strict <, which is where the tie behaviour comes from: a key equal to the threshold is not admitted. Once the heap holds k keys, the first k arrivals of a group of equal keys are the ones that stay. (And within the heap, BinaryHeap::pop removes the last element of its vector, so among equal keys the most recently pushed is the one discarded.) The result is that top_k on ties keeps the earliest rows, deterministically, even though the heap itself is not a stable structure.

The heap is also bounded in memory, not just in size:

rust
impl<P: Ord + Clone> BottomKWithPayload<P> {
    pub fn new(k: usize, num_pipelines: usize) -> Self {
        let share_keys = num_pipelines > 1 && k <= MAX_SHARED_K;
        Self::with_shared(
            k,
            Arc::default(),
            share_keys.then(|| Arc::new(Mutex::new(BinaryHeap::with_capacity(k + 1)))),
        )
    }

    fn with_shared(
        k: usize,
        shared_optimum: Arc<RwLock<Option<P>>>,
        shared_keys: Option<Arc<Mutex<BinaryHeap<P>>>>,
    ) -> Self {
        let initial_capacity = get_ideal_morsel_size().min(k + 1);
        Self {
            k,
            heap: BinaryHeap::with_capacity(initial_capacity),
            df_subsets: SlotMap::with_key(),
            row_idxs: SlotMap::with_key(),
            marked_for_pruning: SecondaryMap::new(),
            to_prune: Vec::new(),
            gather_idxs: Vec::new(),
            shared_optimum,
            shared_keys,
        }
    }

crates/polars-stream/src/nodes/top_k.rs:97

The heap is pre-allocated for at most a morsel’s worth of keys or k + 1, whichever is smaller, so a pipeline that never sees a competitive key never grows a large allocation.

The rows are not copied

The heap stores handles, not data. A row that is still in the heap lives in the data frame it came from, and the RowIdxKey slotmap remembers where. Rows that fall out of the heap are marked, and the data frames are pruned periodically:

rust
    /// Schedules a subset for pruning once none of its rows are in the heap, or once at most half
    /// of its rows are in the heap and a gather would remove enough rows to be worthwhile.
    fn mark_if_sparse(&mut self, dfs_key: DfsKey) {
        let df_subset = &self.df_subsets[dfs_key];
        let in_heap = df_subset.subset_len;
        let removed = df_subset.df.height() - in_heap;
        let should_prune =
            in_heap == 0 || (removed >= in_heap && removed >= MIN_GATHER_REMOVED_ROWS);
        if should_prune && self.marked_for_pruning.insert(dfs_key, ()).is_none() {
            self.to_prune.push(dfs_key);
        }
    }

    pub fn prune(&mut self) {
        for dfs_key in self.to_prune.drain(..) {
            self.marked_for_pruning.remove(dfs_key);
            if self.df_subsets[dfs_key].subset_len == 0 {
                let df_subset = self.df_subsets.remove(dfs_key).unwrap();
                for row_idx in df_subset.rows {
                    self.row_idxs.remove(row_idx);
                }
            } else {
                self.df_subsets[dfs_key].gather(&mut self.row_idxs, &mut self.gather_idxs);
            }
        }
    }

crates/polars-stream/src/nodes/top_k.rs:243

Pruning is not free — it is a gather over the surviving rows — so it is only done when at least half the rows are gone and at least sixteen rows would be removed. The constant MIN_GATHER_REMOVED_ROWS = 16 is there because a gather of a handful of rows costs more than the memory it saves.

Merging the per-pipeline heaps

Each pipeline of the streaming engine gets its own reducer with its own heap, because the pipelines run concurrently and cannot share a mutex on the hot path. When the input ends, the heaps are merged pairwise:

rust
    pub fn combine(&mut self, other: &BottomKWithPayload<P>) {
        let mut new_df_keys =
            SecondaryMap::<DfsKey, DfsKey>::with_capacity(other.df_subsets.capacity());
        for (dfs_key, dfs) in &other.df_subsets {
            if dfs.subset_len > 0 {
                let subset = DfSubset {
                    df: dfs.df.clone(),
                    rows: Vec::new(),
                    subset_len: 0,
                };
                new_df_keys.insert(dfs_key, self.df_subsets.insert(subset));
            }
        }
        for prio in &other.heap {
            let (dfs_key, row_idx_key) = prio.1;
            self.add_one(
                new_df_keys[dfs_key],
                other.row_idxs[row_idx_key],
                prio.0.clone(),
                |l, r| l < r,
                |x| x,
            );
        }
        for &dfs_key in new_df_keys.values() {
            self.mark_if_sparse(dfs_key);
        }
        self.prune();
    }

crates/polars-stream/src/nodes/top_k.rs:265

Merging two heaps of k keys each is O(k log k): every key of the second heap is offered to the first through the same add_one admission test, so most of them are rejected on the first comparison. The merge is done in parallel with a rayon reduction over the per-pipeline reducers.

So the complexity of the whole streaming top-k is:

Time: O(n log k) for n rows, versus O(n log n) for a full sort. With k = 10 that is roughly log2(10) = 3.3 comparisons per row instead of log2(n), which for a billion rows is the difference between a few seconds and a few minutes.
Memory: O(k) for the heap, plus the rows that are still in it. The heap never holds more than k keys, but the payload of those keys has to survive until the end of the query, which is why the pruning machinery exists.

Publishing the threshold

The heap’s threshold is worth more than the answer. Once the engine knows the k-th best key, every row that is worse than that key can be discarded immediately, wherever it comes from. Polars publishes it as a dynamic predicate that is pushed back into the scan:

rust
            let keys = columns[0].as_materialized_series();
            let key_ca: &ChunkedArray<T> = keys.as_phys_any().downcast_ref().unwrap();
            let pred: BooleanChunked = key_ca
                .iter()
                .map(|opt_x| ReorderWithNulls(opt_x.map(TotalOrdWrap)) < v)
                .collect_ca(PlSmallStr::EMPTY);
            Ok(Some(Column::from(pred.into_series())))

crates/polars-stream/src/nodes/top_k.rs:359

For integer and temporal keys the predicate also exposes a value range, so a parquet reader can skip entire row groups using statistics:

rust
    fn runtime_range(&self) -> RuntimeRange {
        // Floats order NaN apart from statistics, enums by their categories.
        if !(self.dtype.is_integer() || self.dtype.is_temporal() || self.dtype.is_decimal()) {
            return RuntimeRange::Disabled;
        }
        let Some(ReorderWithNulls(Some(TotalOrdWrap(threshold)))) = *self.shared_optimum.read()
        else {
            return RuntimeRange::Pending;
        };

crates/polars-stream/src/nodes/top_k.rs:438

That is why the streaming implementation is worth the extra machinery: the threshold does not just reduce what is kept, it reduces what is read.

For small k the pipelines go further and share their best keys, so that the published threshold is the k-th key over all pipelines rather than one pipeline’s k-th key:

rust
/// Up to this `k` several pipelines share their best `k` keys, so the published optimum is
/// the k-th key over all pipelines instead of the best k-th key of one pipeline.
const MAX_SHARED_K: usize = 128;

crates/polars-stream/src/nodes/top_k.rs:39

The bottom line

top_k in Polars is two algorithms wearing the same name:

In memory, it is a quickselect partition over a materialised copy of the values. O(n) expected time, O(n) extra memory, winners in arbitrary order. maintain_order turns it back into a full sort.
Streaming, it is a BinaryHeap of exactly k keys per pipeline, with the payload carried along out of band so the heap never holds data. O(n log k) time, O(k) heap memory, and — the real prize — a threshold that is published back into the scan so that losing rows are never read.

The tie-breaking rules differ, and that is worth knowing: the in-memory quickselect breaks ties arbitrarily, the streaming heap keeps the earliest arrivals, and the arg_sort-with-limit path keeps the earliest rows by index. If your query depends on which of several equal values comes back, ask for maintain_order and pay for the sort.

And this is the query that makes the streaming engine reach for the heap:

python
import polars as pl

lf = pl.scan_parquet("big.parquet")
result = lf.sort("x", descending=True).head(10).collect()

The streaming engine does not lower that into “sort everything, then slice”. It lowers it into a top-k node — the bounded heap — followed by a sort of the ten survivors:

rust
/// Keeps the first `k` rows in sort order, in no particular order.
pub(crate) fn build_top_k_stream(
    input: PhysStream,
    by_column: Vec<ExprIR>,
    sort_options: &SortMultipleOptions,
    k: u64,
    dyn_pred: Option<DynamicPred>,
    expr_arena: &mut Arena<AExpr>,
    phys_sm: &mut DenseSlotMap<PhysNodeKey, PhysNode>,
) -> PhysStream {

crates/polars-stream/src/physical_plan/lower_ir.rs:67

The k rows that come out of the heap are then sorted, so the answer is correctly ordered; what the heap avoids is sorting the other billion rows. And because the heap’s threshold is published as a dynamic predicate, the parquet reader can skip row groups whose statistics cannot contain a winner — so the losing rows are not merely discarded, they are never read.