Top-K without a sort: how DataFusion keeps a bounded heap of the k smallest rows

TL;DR: SortExec with fetch = Some(k) does not sort. It converts sort keys to Arrow’s row format, keeps the k best rows in a BinaryHeap, and evicts the current worst row whenever a better one arrives. Memory is O(k) rows instead of O(N), time is O(N log k) instead of O(N log N) — and the heap’s boundary is published as a dynamic filter so the scan underneath stops reading rows that cannot possibly make the cut.

The problem: `ORDER BY x LIMIT 3` on 100 million rows

This is the single most common analytical query shape after GROUP BY, and the naive plan for it is absurd:

sql
-- datafusion/physical-plan/src/topk/mod.rs:71-79
/// ```sql
/// > explain SELECT customer_id, revenue FROM sales ORDER BY revenue DESC limit 3;
/// +--------------+----------------------------------------+
/// | plan_type    | plan                                   |
/// +--------------+----------------------------------------+
/// | logical_plan | Limit: 3                               |
/// |              |   Sort: revenue DESC NULLS FIRST       |
/// |              |     Projection: customer_id, revenue   |
/// |              |       TableScan: sales                 |
/// +--------------+----------------------------------------+
/// ```

The sort materializes 100 million rows, sorts them (O(N log N) comparisons, O(N) memory), and then throws away 99,999,997 of them. The sort is the entire cost of the query and it is almost entirely wasted work.

The information-theoretic observation is simple: to know the top k rows you only need to compare each new row against the current k-th best. That’s a selection problem, not a sorting problem, and it admits an O(N log k) / O(k) solution.

DataFusion implements exactly that, and then goes further: it turns the running threshold into a filter that the scan can use.

Where it lives

There is no separate TopKExec node in this tree. Top-K is a mode of SortExec, selected at execute time by the presence of fetch:

rust
// datafusion/physical-plan/src/sorts/sort.rs:1609-1645
match (sort_satisfied, self.fetch.as_ref()) {
    (true, Some(fetch)) => Ok(Box::pin(LimitStream::new(
        input, 0, Some(*fetch),
        BaselineMetrics::new(&self.metrics_set, partition),
    ))),
    (true, None) => Ok(input),
    (false, Some(fetch)) => {
        let filter = self.filter.clone();
        let mut topk = TopK::try_new(
            partition,
            input.schema(),
            self.common_sort_prefix.clone(),
            self.expr.clone(),
            *fetch,
            context.session_config().batch_size(),
            context.runtime_env(),
            &self.metrics_set,
            Arc::clone(&unwrap_or_internal_err!(filter)),
        )?;
        Ok(Box::pin(RecordBatchStreamAdapter::new(
            self.schema(),
            futures::stream::once(async move {
                while let Some(batch) = input.next().await {
                    let batch = batch?;
                    topk.insert_batch(batch)?;
                    if topk.finished {
                        break;
                    }
                }
                drop(input);
                topk.emit()
            })
            .try_flatten(),
        )))
    }
    (false, None) => {
        // ... full ExternalSorter path
    }
}

Three cases, and note the third: with no fetch, SortExec builds an ExternalSorter (a spilling, multi-run merge sort). With a fetch, it builds a TopK. With a fetch and an already-sorted input, it just takes the first k rows.

The heap

rust
// datafusion/physical-plan/src/topk/mod.rs:925-953
/// This structure keeps at most the *smallest* k items, using the
/// [arrow::row] format for sort keys. While it is called "topK" for
/// values like `1, 2, 3, 4, 5` the "top 3" really means the
/// *smallest* 3 , `1, 2, 3`, not the *largest* 3 `3, 4, 5`.
///
/// Using the `Row` format handles things such as ascending vs
/// descending and nulls first vs nulls last.
struct TopKHeap {
    /// The maximum number of elements to store in this heap.
    k: usize,
    /// Storage for up at most `k` items using a BinaryHeap. Reversed
    /// so that the smallest k so far is on the top
    inner: BinaryHeap<TopKRow>,
    /// Storage the original row values (TopKRow only has the sort key)
    store: RecordBatchStore,
    /// The size of all owned data held by this heap
    owned_bytes: usize,
}

The “smallest k” wording is the key to the whole design. Arrow’s row format encodes a value so that byte-wise memcmp order equals SQL sort order — including descending order and nulls-first/last, because the SortOptions are baked into the encoder. So for ORDER BY revenue DESC, the smallest row bytes correspond to the highest revenue. “Keep the smallest k row-encodings” is exactly “keep the top k rows”, with no per-comparison special-casing for direction or nulls.

That is a genuinely nice trick. A naive implementation would have to write a custom comparator per sort direction, handle nulls, and probably get NULLS FIRST vs NULLS LAST wrong. Arrow’s RowConverter gives you memcmp semantics for free.

Each heap entry is deliberately not the row itself:

rust
// datafusion/physical-plan/src/topk/mod.rs:1111-1126
/// Represents one of the top K rows held in this heap. Orders
/// according to memcmp of row (e.g. the arrow Row format, but could
/// also be primitive values)
///
/// Reuses allocations to minimize runtime overhead of creating new Vecs
#[derive(Debug, PartialEq)]
struct TopKRow {
    /// the value of the sort key for this row. This contains the
    /// bytes that could be stored in `OwnedRow` but uses `Vec<u8>` to
    /// reuse allocations.
    row: Vec<u8>,
    /// the RecordBatch this row came from: an id into a [`RecordBatchStore`]
    batch_id: u32,
    /// the index in this record batch the row came from
    index: usize,
}

So the heap stores sort keys plus a pointer (batch id, row index) to the payload. The payload stays in the original RecordBatch, held alive by RecordBatchStore.

Admission: one comparison per row

rust
// datafusion/physical-plan/src/topk/mod.rs:515-535
fn find_new_topk_items(
    &mut self,
    items: impl Iterator<Item = usize>,
    batch_entry: &mut RecordBatchEntry,
) -> usize {
    let mut replacements = 0;
    let rows = &mut self.scratch_rows;
    for (index, row) in items.zip(rows.iter()) {
        match self.heap.max() {
            // heap has k items, and the new row is greater than the
            // current max in the heap ==> it is not a new topk
            Some(max_row) if row.as_ref() >= max_row.row() => {}
            // don't yet have k items or new item is lower than the currently k low values
            None | Some(_) => {
                self.heap.add(batch_entry, row, index);
                replacements += 1;
            }
        }
    }
    replacements
}

heap.max() returns self.inner.peek() — the largest of the k retained rows in row-byte order, which is the k-th best row and therefore the current admission threshold. A new row is admitted only if it is strictly smaller than that threshold. One memcmp per input row, and O(log k) work only for the rows that actually make the cut.

The eviction path is where the “reuse allocations” comment pays off:

rust
// datafusion/physical-plan/src/topk/mod.rs:978-1016
fn add(
    &mut self,
    batch_entry: &mut RecordBatchEntry,
    row: impl AsRef<[u8]>,
    index: usize,
) {
    let batch_id = batch_entry.id;
    batch_entry.uses += 1;

    assert!(self.inner.len() <= self.k);
    let row = row.as_ref();

    // Reuse storage for evicted item if possible
    if self.inner.len() == self.k {
        let mut prev_min = self.inner.peek_mut().unwrap();

        // Update batch use
        if prev_min.batch_id == batch_entry.id {
            batch_entry.uses -= 1;
        } else {
            self.store.unuse(prev_min.batch_id);
        }

        // update memory accounting
        self.owned_bytes -= prev_min.owned_size();

        prev_min.replace_with(row, batch_id, index);

        self.owned_bytes += prev_min.owned_size();
    } else {
        let new_row = TopKRow::new(row, batch_id, index);
        self.owned_bytes += new_row.owned_size();
        // put the new row into the heap
        self.inner.push(new_row);
    }
}

Note peek_mut(): the heap’s worst element is overwritten in place with the new row’s bytes, reusing the Vec<u8> allocation. No allocation, no deallocation, no sift_down from scratch — BinaryHeap::peek_mut handles the reordering. In a steady-state top-K over millions of rows, this is the difference between a garbage-generating loop and an allocation-free one.

Also note the reference counting: batch_entry.uses is incremented when a row from that batch enters the heap and decremented when it leaves. A batch whose rows have all been evicted is dropped from the store entirely:

rust
// datafusion/physical-plan/src/topk/mod.rs:1244-1253
/// Insert a record batch entry into this store, tracking its
/// memory use, if it has any uses
pub fn insert(&mut self, entry: RecordBatchEntry) {
    // uses of 0 means that none of the rows in the batch were stored in the topk
    if entry.uses > 0 {
        self.batches_size += get_record_batch_memory_size(&entry.batch);
        self.total_rows += entry.batch.num_rows();
        self.batches.insert(entry.id, entry);
    }
}

Emitting: gather the survivors

rust
// datafusion/physical-plan/src/topk/mod.rs:1024-1051
fn emit_with_state(&mut self) -> Result<(Option<RecordBatch>, Vec<TopKRow>)> {
    // generate sorted rows
    let topk_rows = std::mem::take(&mut self.inner).into_sorted_vec();

    if self.store.is_empty() {
        return Ok((None, topk_rows));
    }

    // The batch ids are not continuous, so `indices` below is built through
    // the store's "batch_id -> array_pos" mapping.
    let (record_batches, batch_id_array_pos) = self.store.positional();

    let indices: Vec<_> = topk_rows
        .iter()
        .map(|k| (batch_id_array_pos[&k.batch_id], k.index))
        .collect();

    // At this point `indices` contains indexes within the
    // rows and `input_arrays` contains a reference to the
    // relevant RecordBatch for that index. `interleave_record_batch` pulls
    // them together into a single new batch
    let new_batch = interleave_record_batch(&record_batches, &indices)?;

    Ok((Some(new_batch), topk_rows))
}

BinaryHeap::into_sorted_vec() gives the k survivors already in sorted order — one O(k log k) operation at the very end, not per batch. Then interleave_record_batch gathers the payloads from the (possibly many) retained batches into one output batch, in sorted order.

Compaction

The store holds whole RecordBatches, but only k of their rows are actually wanted. Over a long scan the store can accumulate many batches each contributing one or two surviving rows, wasting memory. maybe_compact fixes this:

rust
// datafusion/physical-plan/src/topk/mod.rs:1053-1069
/// Compact this heap, rewriting all stored batches into a single
/// input batch
pub fn maybe_compact(&mut self) -> Result<()> {
    // Don't compact if there's only one batch (compacting into itself is pointless)
    if self.store.len() <= 1 {
        return Ok(());
    }

    let total_rows = self.store.total_rows;
    let num_rows = self.inner.len();

    // Compact when the store holds more than STORE_COMPACTION_RATIO x what
    // the compacted result would need. The multiplier avoids compacting when
    // the savings would be marginal.
    if total_rows <= num_rows * STORE_COMPACTION_RATIO {
        return Ok(());
    }

with

rust
// datafusion/physical-plan/src/topk/mod.rs:918-923
/// length in [`TopKHeap::maybe_compact`], every partition's retained rows
/// (heap slots, plus ties for `RANK`) in [`RecordBatchStore::compact`]. The
/// multiplier avoids compacting when the savings would be marginal.
const STORE_COMPACTION_RATIO: usize = 2;

So compaction fires only when the store holds more than twice as many rows as the heap actually needs — amortizing the rewrite cost against the savings, and never doing a pointless rewrite.

The dynamic filter: pushing the threshold down

The heap knows a threshold: “no row at or above this row-encoding can enter the top k”. DataFusion publishes that as a dynamic filter expression that the scan below can evaluate, so the scan can skip rows early:

rust
// datafusion/physical-plan/src/topk/mod.rs:559-568
/// Update the filter representation of our TopK heap.
/// For example, given the sort expression `ORDER BY a DESC, b ASC LIMIT 3`,
/// and the current heap values `[(1, 5), (1, 4), (2, 3)]`,
/// the filter will be updated to:
///
/// ```sql
/// (a > 1 OR (a = 1 AND b < 5)) AND
/// (a > 1 OR (a = 1 AND b < 4)) AND
/// (a > 2 OR (a = 2 AND b < 3))
/// ```
fn update_filter(&mut self) -> Result<()> {

The predicate is built outside the lock, then installed with a double-check, because multiple partitions share one filter and any of them may tighten it:

rust
// datafusion/physical-plan/src/topk/mod.rs:587-610
// Extract scalar values BEFORE acquiring lock to reduce critical section
let thresholds = boundary.threshold_values(&self.expr)?;

// Build the filter expression OUTSIDE any synchronization
let predicate = Self::build_filter_expression(&self.expr, &thresholds)?;
let new_threshold =
    boundary.threshold(self.encode_topk_common_prefix_row(boundary)?);

// update the threshold. Since there was a lock gap, we must check if it is still the best
// may have changed while we were building the expression without the lock
let mut filter = self.filter.write();
let still_needs_update = filter
    .shared_threshold
    .as_ref()
    .map(|current| new_threshold.is_more_selective_than(current))
    .unwrap_or(true);
if !still_needs_update {
    // some other thread updated the threshold to a better one while we
    // were building so there is no need to update the filter
    return Ok(());
}

Note the multi-column predicate construction: for ORDER BY a DESC, b ASC, the predicate on b is only applied when a equals the threshold’s a — otherwise a row with a better a would be wrongly rejected. The builder threads a running equality conjunction (prev_sort_expr) through the columns to get this right (topk/mod.rs:679-699).

Early termination

There’s an even stronger trick. If the input is already sorted on a common prefix of the top-K’s ordering, then once the last row of a batch is strictly past the heap boundary, no subsequent row can ever enter the top k — so the query can stop reading:

rust
// datafusion/physical-plan/src/topk/mod.rs:708-717
/// If input ordering shares a common sort prefix with the TopK,
/// check if the computation can be finished early.
///
/// This is the case if the last row of the current batch is strictly
/// greater than either the shared dynamic-filter threshold prefix or the max
/// row in the local heap, comparing only on the shared prefix columns.
fn attempt_early_completion(&mut self, batch: &RecordBatch) -> Result<()> {
    let last_row_prefix = self.last_row_common_prefix(batch)?;
    self.finish_if_prefix_past_boundary(last_row_prefix.as_ref())
}

For SELECT day, ts FROM sensor ORDER BY day DESC, ts DESC LIMIT 10 where sensor is already sorted by day, the scan reads only the first few “days” of data and stops. The finished flag then breaks the input loop (sorts/sort.rs:1636-1638).

Parallel top-K

Top-K has an additional trick that only works because of how the plan is built: k rows per partition is enough, because the merge takes the global top k afterwards.

The optimizer’s parallelize_sorts phase rewrites CoalescePartitions → Sort into Sort(preserve_partitioning) → SortPreservingMerge:

rust
// datafusion/physical-optimizer/src/utils.rs:54-58
let mut new_sort = SortExec::new(ordering, Arc::clone(&node.plan)).with_fetch(fetch);
if node.plan.output_partitioning().partition_count() > 1 {
    new_sort = new_sort.with_preserve_partitioning(true);
}

and

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs:351-361
requirements = add_sort_above_with_check(requirements, sort_reqs, fetch)?;

let spm =
    SortPreservingMergeExec::new(sort_exprs, Arc::clone(&requirements.plan));
Ok(Transformed::yes(
    PlanWithCorrespondingCoalescePartitions::new(
        Arc::new(spm.with_fetch(fetch)),
        false,
        vec![requirements],
    ),
))

So the plan becomes:

text

SortPreservingMergeExec(fetch=k)      ← global top k
  SortExec(fetch=k, preserve_partitioning=true)   ← per-partition top k
    ├── partition 0  ─┐
    ├── partition 1   ├─ N parallel scans
    └── partition 2  ─┘

Each of the N partitions keeps only k rows, so total memory is N × k rows instead of the whole dataset — and the merge only has to merge N × k rows, not N full streams. This is exactly the win described in the EnsureRequirements module docs:

text
/// Each input partition is sorted in parallel, then a `SortPreservingMergeExec`
/// at the top performs a cheap merge of pre-sorted streams. For TopK queries
/// (`fetch=K`), each parallel sort only keeps K rows per partition, so total
/// memory is `N × K` rather than coalescing the entire stream first.

The fetch is propagated to the merge as well, so the merge itself stops after emitting k rows and the upstream partitions get cancelled.

Sort pushdown: where the `fetch` travels

The fetch is a piece of metadata that the sort-pushdown rule threads through the plan. The rule lives in datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs and carries three things per node:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:60-80
#[derive(Clone, Debug)]
pub struct ParentRequirements {
    ordering_requirement: Option<OrderingRequirements>,
    fetch: Option<usize>,
    /// The distribution required by the consumer above any SortExec we insert.
    /// When this is `SinglePartition` and the input has multiple partitions,
    /// `add_sort_above_with_distribution` wraps the sort in `SortPreservingMergeExec`.
    distribution_requirement: Distribution,
}

pub type SortPushDown = PlanContext<ParentRequirements>;

The fetch arithmetic is where the subtlety lives. A LIMIT with an OFFSET needs skip more input rows than it returns:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:103-112
fn input_fetch(
    plan: &Arc<dyn ExecutionPlan>,
    parent_fetch: Option<usize>,
) -> Option<usize> {
    let fetch = min_fetch(plan.fetch(), parent_fetch)?;
    let skip = plan
        .downcast_ref::<GlobalLimitExec>()
        .map_or(0, |limit| limit.skip());
    Some(fetch.saturating_add(skip))
}

The comment explains the trap: min(fetch + skip, parent_fetch) would be too small when the parent fetch is the tighter bound, and using the bare fetch would turn LIMIT 10 OFFSET 5 into TopK(10) below the limit — yielding 5 result rows instead of 10.

Pushing a sort through a filter: refused, on purpose

The pushdown rule decides, per node, whether the parent’s ordering requirement can be pushed into the children. For most operators the answer is “yes, if the operator preserves input order”. For a FilterExec, the answer is a hard no:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:637-646
} else if maintains_input_order.is_empty()
    || !maintains_input_order.iter().any(|o| *o)
    || plan.is::<RepartitionExec>()
    || plan.is::<FilterExec>()
    || pushdown_would_violate_requirements(&parent_required, plan.as_ref())
{
    // If the current plan is a leaf node or can not maintain any of the input ordering, can not pushed down requirements.
    // For RepartitionExec, we always choose to not push down the sort requirements even the RepartitionExec(input_partition=1) could maintain input ordering.
    // Pushing down is not beneficial
    Ok(None)
}

Returning None means “cannot push down” — and the caller then inserts a SortExec above the filter:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:435-444
} else {
    // Can not push down requirements, add new `SortExec` (distribution-aware):
    sort_push_down = add_sort_above_with_distribution(
        sort_push_down,
        parent_requirement.into_single(),
        parent_fetch,
        &parent_distribution,
    );
    assign_initial_requirements(&mut sort_push_down);
}

This is the right call, and the reason is a nice illustration of “the optimizer must think about data volume, not just ordering”. Pushing a sort below a filter would mean sorting the unfiltered rows and then discarding most of them — strictly more work. Keeping the sort above the filter means filtering first and sorting only the survivors.

The same logic blocks pushing a limit below a filter:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:455-460
// A parent fetch can be pushed through a plan only if the plan explicitly
// supports limit pushdown and does not change cardinality. For example, if
// LIMIT 2 is applied below a FilterExec that removes half the rows, only one
// row may remain. Applied after the filter, it correctly returns two rows
// when at least two are available.

FilterExec doesn’t advertise supports_limit_pushdown(), so a LIMIT cannot be pushed through it either. And there’s a second, subtler guard for operators that do support limit pushdown but change cardinality:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:462-477
// UnionExec advertises `supports_limit_pushdown() == true` because it can
// forward a LIMIT k to each of its children—i.e. apply "LIMIT k" separately
// on each branch before merging them together.
//
// However, UnionExec's `cardinality_effect() == GreaterEqual` (it sums up
// all child row counts), so pushing a global TopK/LIMIT through it would
// break the semantics of "take the first k rows of the combined result."

So can_push_fetch_through requires both:

rust
// datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:151-155
fn can_push_fetch_through(plan: &Arc<dyn ExecutionPlan>) -> bool {
    plan.supports_limit_pushdown()
        && matches!(plan.cardinality_effect(), CardinalityEffect::Equal)
}

Trade-offs and failure modes

The row-format conversion is O(N) no matter what. Even though the heap only does O(log k) work per admitted row, every row must be encoded into the row format first (self.row_converter.append(rows, &sort_keys)?, topk/mod.rs:471-473). That’s a full pass over the sort key columns, comparable in cost to a comparison sort’s first phase. The win comes from avoiding the O(N log N) comparison phase and the O(N) memory, not from avoiding all work.

k in the millions degenerates. The heap is O(log k) per admission and holds k rows. If k is a large fraction of N, a plain sort (which is O(N log N) but with excellent locality and no per-row encoding) may actually win. DataFusion doesn’t switch strategies based on k; the fetch is either present or not.

Ties are resolved arbitrarily, and differently across partitions. The admission test is row < max_row — strictly smaller. A row equal to the threshold is rejected. That means the k retained rows are “the k strictly-smallest rows”, and among equal rows the choice is arbitrary (it depends on arrival order). SQL leaves tie-breaking unspecified for LIMIT without ORDER BY on a unique key, so this is legal — but it means LIMIT 10 over rows with duplicate sort keys may return a different subset on different runs or with different partition counts. For deterministic results you need a total order in the ORDER BY clause.

The dynamic filter is a double-edged sword. It lets the scan skip rows, but building and updating the predicate costs a lock acquisition per tightening and a BinaryExpr tree rebuild. The code mitigates this by building the expression outside the lock and re-checking under it, and by only updating when the new threshold is strictly more selective.

Memory is O(k) rows but not O(k) bytes. The heap holds k keys plus a reference to the whole RecordBatch each key came from. If the k survivors are spread across k different 8192-row batches, the store holds k full batches — k × 8192 rows of payload — until compaction fires. That’s why maybe_compact exists, and why it fires at a ratio of 2 rather than 1: compaction is itself a rewrite, and doing it eagerly would cost more than it saves.

Early termination only works on a common prefix. attempt_early_completion compares only the shared prefix columns of the sort key. If the input ordering shares nothing with the top-K’s ordering (the common case), the check is a no-op and the query reads everything. The optimization pays off for time-series-ish data that arrives pre-sorted.

Bottom line

The top-K heap is a small, self-contained piece of systems engineering that pays for itself many times over. Three ideas are doing the work:

Encode once, compare with memcmp. Arrow’s row format turns “compare two sort keys according to these SortOptions” into “compare two byte strings”, with null ordering and direction already baked in.
Keep only the boundary. The heap’s worst element is the admission threshold, so the data structure that implements the selection algorithm is also the thing that implements the filter.
Publish the threshold. Turning the boundary into a dynamic filter expression means the rows that can’t win are never read from storage at all — the top-K operator stops being a consumer-only optimization and starts pruning the scan.

The result is a query plan where ORDER BY revenue DESC LIMIT 3 over a hundred million rows reads a few thousand rows and holds three in memory.

The problem: ORDER BY x LIMIT 3 on 100 million rows