Top K per Partition: How DataFusion Replaces a Full Sort with N Small Heaps

TL;DR: SELECT * FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY store ORDER BY revenue DESC) rn FROM sales) WHERE rn <= 3 compiles, naively, to a full sort of the entire table, a window aggregation over every row, and then a filter. But the answer only needs three rows per store. DataFusion replaces the whole thing with one bounded heap per partition key — and the interesting engineering is in making N heaps share one row converter, one memory reservation, one row store, and one set of metrics, instead of paying that setup N times.

The “top K per group” query is one of the most common analytical patterns there is:

SELECT * FROM (
    SELECT *, ROW_NUMBER() OVER (PARTITION BY store ORDER BY revenue DESC) as rn
    FROM sales
) WHERE rn <= 3;

“Top 3 products by revenue for each store.” “Most recent 5 events per user.” “Latest record per device.”

The naive physical plan is brutal:

FilterExec: rn <= 3
  BoundedWindowAggExec: ROW_NUMBER() PARTITION BY [store] ORDER BY [revenue DESC]
    SortExec: expr=[store ASC, revenue DESC]
      DataSourceExec

That sorts the entire dataset — O(N log N) — computes a row number for every row, and then throws away all but 3 per store. The sort is the dominant cost, and almost all of its work is discarded.

The replacement

PartitionedTopKExec (datafusion/physical-plan/src/sorts/partitioned_topk.rs) replaces the sort-then-window-then-filter with a single pass that keeps, per partition key, only the K best rows seen so far.

The state is a bounded heap per partition. For each incoming row:

1. Compute the partition key and the ordering key.
2. Find the heap for that partition.
3. If the heap is not full, or the row beats the heap’s worst entry, insert it and evict the worst if necessary.

At the end, each heap has at most K rows, and the answer is the concatenation of all the heaps. The cost is O(N log K), and the memory is O(number_of_partitions × K) instead of O(N).

The three ranking functions

Not all three window functions have the same retention rule, and getting this right is subtle:

pub enum WindowFnKind {
    /// `ROW_NUMBER()` — keep exactly K rows per partition.
    RowNumber,
    /// `RANK()` — keep K rows plus any rows tied at the boundary.
    Rank,
    /// `DENSE_RANK()` — keep every row whose ob value is among the K
    /// distinct-smallest ob values seen in the partition.
    DenseRank,
}

ROW_NUMBER is exactly K rows, no ties, no ambiguity.

RANK is K rows plus any rows tied with the K-th row at the boundary. So WHERE rk <= 3 might keep 5 rows if rows 3, 4 and 5 are tied. The source says so: “RANK semantics — WHERE rk <= K may keep more than K rows when ties straddle the boundary”.

DENSE_RANK is different again: it keeps every row whose ordering value is among the K distinct smallest values. So if the top value appears a million times, all a million rows are kept. “total kept rows is unbounded in rows-per-distinct-value”.

The admission test reflects this:

fn qualifies(&self, k: usize, key: &[u8]) -> bool {
    matches!(self.classify(k, key), None | Some(Ordering::Less))
}

/// How `key` compares with the worst retained row, or `None` while the
/// heap holds fewer than `k` rows and so admits anything.
///
/// The worst retained row is the root, so this is one comparison. `RANK`
/// needs all three outcomes, since `Equal` is a tie at the boundary;
/// `ROW_NUMBER` only needs [`Self::qualifies`].
fn classify(&self, k: usize, key: &[u8]) -> Option<Ordering> {
    if self.inner.len() < k {
        return None;
    }
    self.inner.peek().map(|worst| key.cmp(worst.key.as_slice()))
}

classify is one comparison against the heap’s root — the worst retained row. ROW_NUMBER only needs to know “better or not” (Less), so it uses qualifies and never materialises the Equal case. RANK needs all three outcomes, because Equal is a tie at the boundary that must be kept.

Sharing the setup

Now the part that makes the implementation interesting.

The obvious implementation is a HashMap<PartitionKey, TopK> — one full TopK per partition. But each TopK owns a RowConverter, a MemoryConsumer registration, a RecordBatchStore, and a set of metric counters. With 100,000 distinct partitions, you pay 100,000 times the fixed setup cost — and most of it is identical.

PartitionedTopK is built specifically to avoid that:

/// Top-K-per-partition operator state.
///
/// Sibling to [`TopK`]. Where `TopK` maintains a single global heap,
/// `PartitionedTopK` maintains one [`PartitionHeap`] per distinct partition
/// key while sharing everything else in [`PartitionedTopKCore`] across all
/// partitions.
///
/// This sharing is the point of the type: with N distinct partition
/// keys, a naive `HashMap<_, TopK>` pays N × constant overhead for
/// `RowConverter::new`, `MemoryConsumer::register`, `RecordBatchStore::new`,
/// and metric counter setup. `PartitionedTopK` pays it once.
pub(crate) struct PartitionedTopK {
    core: PartitionedTopKCore<PartitionHeap>,
}

The core holds the shared machinery; the heap holds only the per-partition data. The PartitionedTopKCore<PartitionHeap> is generic over the heap type, which is how the three ranking variants share the same core with different heaps.

Row format keys

The ordering keys are compared as Arrow row format bytes:

use arrow::row::SortField;

The row format is Arrow’s order-preserving byte encoding: two values encode to bytes that compare the same way the values do. That means the heap’s comparison is a plain memcmp on byte slices, with no Arrow comparator, no null handling in the comparison loop, and no per-column branching.

The comparison is literally:

self.inner.peek().map(|worst| key.cmp(worst.key.as_slice()))

[u8]::cmp — a byte comparison. For a BinaryHeap of K entries, that is O(log K) memcmps per row.

One pass, not grouped

A detail that looks like a micro-optimisation but is actually a complexity issue:

// Rows are visited in batch order rather than grouped by partition.
// Interleaving groups cannot change a decision — a heap is only
// read and written by rows of its own partition, which still reach
// it in ascending row order — and it costs no pass over the
// partitions seen so far, which would be quadratic in the input
// whenever partition count grows with it.

If you grouped the rows of a batch by partition before feeding the heaps, you would have to scan the batch’s distinct partition keys — and if the number of distinct partitions grows with the batch size, that grouping pass is quadratic in the input. Visiting in batch order is linear, and the interleaving is harmless because each partition’s rows still reach its heap in ascending row order.

Interning the partition keys

The other detail is the partition key’s identity. The map owns the key, not the row:

let mut interned = false;
let heap = heaps.entry_ref(pk.as_ref()).or_insert_with(|| {
    interned = true;
    PartitionHeap::default()
});
if interned {
    // The map has just taken the only copy of this key.
    interned_bytes += pk.as_ref().len();
}

entry_ref is the hashbrown-style API that lets you look up a key without constructing an owned one — the map keeps the key it already has. So the row’s partition key bytes are only allocated the first time a partition is seen:

One probe per row, and entry_ref owns the key only on Vacant, so a key is allocated once per partition for the operator’s lifetime — not once per row.

Without this, every row would allocate a copy of its partition key bytes, and with 100 million rows that is 100 million allocations of the same handful of distinct values. The interned_bytes accounting is then folded into the memory reservation — note the comment: “Accumulated locally and folded in once, so the running totals are not touched per row.”

The rows are gathered once

The heap stores keys, not full rows. The actual row data is kept in a shared RecordBatchStore, and the admission records only which row it kept:

// 2. One pass over the rows: find the row's partition and admit it to
//    that partition's heap if it qualifies. Only the key is copied
//    here; an admission records which row it kept and phase 3 gathers
//    them all in one `take`, so a batch is materialized once no matter
//    how many partitions it touched.

“a batch is materialized once no matter how many partitions it touched” — so a batch that touches 10,000 partitions is still gathered with a single take, not 10,000.

The catch

• The operator only fires for a specific plan shape. It is a physical optimizer rule that recognises FilterExec: rn <= K on top of a BoundedWindowAggExec with ROW_NUMBER/RANK/DENSE_RANK and replaces the sort. If the query is written any other way, you get the sort.
• RANK and DENSE_RANK can keep unboundedly many rows. The “top K” is really “top K values”, and a value can have a million rows. The memory bound is not K rows per partition.
• The partition key is materialised. Every row’s partition key is encoded into row-format bytes. That is a per-row cost that the sort also pays, but it is not free.
• The Equal case is genuinely ambiguous. For RANK, the rows tied at the boundary are all kept, and which of them “wins” the boundary slot is not defined — so the set is right and the count may exceed K.

Why it matters

This is a textbook example of an optimisation that is worth far more than its implementation complexity suggests. The sort-then-window-then-filter plan does O(N log N) work and O(N) memory; the partitioned-top-k plan does O(N log K) work and O(partitions × K) memory. On a table with a billion rows and a thousand partitions, that is the difference between a query that needs a terabyte of memory and one that needs a megabyte.

And the way it is implemented — a shared core, interned partition keys, row-format comparisons, one gather per batch — is a good illustration of what it takes to make a per-group operator actually scale.