Partial Sort: How DataFusion Sorts a Stream That Is Already Sorted by a Prefix

TL;DR: A full sort has to buffer everything before it can emit anything, because the last row might be the smallest. But if the input is already sorted by a prefix of the required ordering — say the data is ordered by (a, b) and you need (a, b, c) — then only the rows that share a prefix value need to be sorted together. DataFusion’s PartialSortExec buffers one prefix group at a time, sorts just that group, and emits it as soon as it can prove the group is complete. Peak memory drops from O(dataset) to O(largest group), and the first rows come out immediately.

A SortExec is a barrier. It consumes the entire input, sorts it, and then produces output. For a 100 GB table that means 100 GB of memory (or a spill to disk), and it means the first output row arrives only after the last input row has been read.

That is unavoidable for a full sort. If the data is completely unsorted, the smallest row could be anywhere, including the very end.

But real data is rarely completely unsorted. Parquet files written in order, time-series data, data that came out of a previous sort — all of it has some structure. And when the input is already ordered by a prefix of the required ordering, the problem decomposes.

The idea

The module’s header gives the example:

Consider a sort plan having an input with ordering a ASC, b ASC

text Copy

+---+---+---+
| a | b | d |
+---+---+---+
| 0 | 0 | 3 |
| 0 | 0 | 2 |
| 0 | 1 | 1 |
| 0 | 2 | 0 |
+---+---+---+

and required ordering for the plan is a ASC, b ASC, d ASC. The first 3 rows(segment) can be sorted as the segment already has the required information for the sort, but the last row requires further information as the input can continue with a batch with a starting row where a and b does not change

So the input is sorted by (a, b), and we need (a, b, d). The rows with the same (a, b) prefix form a group, and within a group only the d column is out of order. Sort each group by d and concatenate — the result is fully sorted by (a, b, d).

The groups are:

• (0, 0): rows with d = 3, 2, 1 → sorted to 1, 2, 3
• (0, 1): one row
• (0, 2): one row

And the groups appear in the input in prefix order, so the output is globally sorted.

The buffering rule

The subtle part is knowing when a group is complete. You cannot emit the group for prefix (0, 0) until you have seen a row with a different prefix — otherwise more rows with prefix (0, 0) might arrive in the next batch.

The rule, from the source:

[PartialSortExec] buffers rows only until it can prove a prefix group will never be seen again, then sorts and emits buffered rows. A group is guaranteed to never be seen again once a row with a different prefix value arrives. This relies on the input’s existing ordering guarantees.

The doc then walks through the buffer state:

           Buffer
       +---+---+---+
       | a | b | c |
       +---+---+---+
       | 0 | 0 | 3 |
       | 0 | 0 | 2 |
       | 0 | 0 | 1 |
       +---+---+---+

“As long as the (a, b) prefix keeps repeating, more rows are buffered.”

Then, when a batch arrives with a new prefix:

           Buffer
       +---+---+---+
       | a | b | c |
       +---+---+---+
       | 0 | 0 | 3 |
       | 0 | 0 | 2 |
       | 0 | 0 | 1 |
       | 0 | 1 | 1 |  <-- first row of new batch, new prefix
       | 0 | 2 | 4 |  <-- new prefix
       | 0 | 2 | 0 |
       | 1 | 0 | 5 |  <-- last row of new batch, new prefix
       +---+---+---+

“Once known complete, each prefix group is sorted by the remaining suffix ordering and emitted as a RecordBatch. Any rows from the most recently seen prefix remain buffered, as more rows with the same prefix may arrive in future batches.”

So the algorithm is:

1. Accumulate rows into a buffer.
2. When a new prefix value appears, the previous prefix group is closed: sort it by the suffix and emit it.
3. The last (still open) group stays buffered.
4. At end of input, close and emit the final group.

The peak memory is therefore the size of the largest prefix group, not the size of the dataset. For data clustered by a low-cardinality key, that is a factor of thousands.

Why it is a “partial” sort

The word “partial” refers to the ordering guarantee: the input only partially satisfies the required ordering. The operator requires:

• the input is ordered by a prefix of the required ordering, and
• the operator sorts the suffix within each group.

So PartialSortExec is the composition of a “group by the prefix” and a “sort by the suffix within the group”. The prefix is a prefix of the required ordering, not an arbitrary grouping — that is what makes the concatenation of the sorted groups globally sorted.

The comparison is the expensive part

Within a group, the sort is by the remaining columns only. Arrow’s lexsort_to_indices does the work, but the comparator only needs the suffix columns, which is both less work and a smaller comparator.

The implementation uses Arrow’s make_comparator to build a comparator over the suffix columns and finds the group boundaries by comparing consecutive rows on the prefix columns:

use arrow_ord::ord::make_comparator;

That is the whole trick, expressed in two comparators over disjoint column sets: one to find where the groups are, one to sort within a group.

The catch

• The input’s ordering guarantee is a promise. “This relies on the input’s existing ordering guarantees.” If the input is not actually sorted by the prefix, the output is silently wrong. This is the same class of trust as the range-partitioning contract we saw earlier.
• A single huge group means no memory saving. If all the rows share the same prefix value, the operator degenerates to a full sort.
• It only helps if the prefix is actually a prefix. Sorting by (a, b) when the input is sorted by (a) and then by (c) does not decompose.
• The operator has to be placed correctly. It is a physical operator, and the planner has to know that the input’s ordering is a prefix of the required one. That is exactly the kind of fact EquivalenceProperties tracks.

Why it matters

The payoff is not just memory. It is latency: the first output batch comes out as soon as the first group closes, rather than after the entire input has been read and sorted. For a dashboard query that streams results, or for a LIMIT 10 on top of a sort, the difference is the difference between a second and a minute.

And it is a nice illustration of a general idea: sorting is only irreducible when the data is entirely unordered. Every bit of known structure in the input is a chance to break the problem into smaller problems — which is exactly what a partial sort does, one prefix group at a time.