Opensource internals

Polars Joins Two Tables with a Piecewise Merge Join

12 min read Jay Hubbert

When both sides of an equi-join are sorted, Polars uses a merge join: a linear sweep that advances two cursors and emits matches, cut into independently mergeable pieces.

TL;DR: An equi-join on two sorted tables does not need a hash table. It needs a merge: one cursor on the left key column, one on the right, and a comparison that decides which side advances. Polars implements exactly that in match_keys_impl, and the streaming engine goes one step further: it cuts both inputs into pieces that are mergeable independently, so the sweep can be suspended at a morsel boundary and resumed later.

The choice is made from the optimizer’s sortedness metadata, not by inspecting the data at run time.

The default equi-join in Polars is a hash join, and a hash join has a build/probe split: one side is turned into a hash table, the other side is streamed through it. That is a good algorithm for arbitrary data and a poor one for sorted data, because it throws away the ordering that the input already has and replaces it with random access into a table.

If both sides are sorted, the join is a merge, and the merge is a linear scan.

The hash join, for contrast

Here is the hash join’s structure. The build side becomes a partitioned hash table:

rust 
// crates/polars-ops/src/frame/join/hash_join/single_keys.rs:22
pub(crate) fn build_tables<T, I>(
    keys: Vec<I>,
    nulls_equal: bool,
) -> Vec<PlHashMap<<T as ToTotalOrd>::TotalOrdItem, IdxVec>>
where
    T: TotalHash + TotalEq + ToTotalOrd,
    <T as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + Hash + Eq + DirtyHash + IsNull,
    I: IntoIterator<Item = T> + Send + Sync + Clone,
{

and the probe side hashes each key, picks the partition that owns it, and looks it up:

rust 
// crates/polars-ops/src/frame/join/hash_join/single_keys_inner.rs:31
    probe.into_iter().enumerate_idx().for_each(|(idx_a, k)| {
        let k = k.to_total_ord();
        let idx_a = idx_a + local_offset;
        // probe table that contains the hashed value
        let current_probe_table =
            unsafe { hash_tbls.get_unchecked(hash_to_partition(k.dirty_hash(), n_tables)) };

        let value = current_probe_table.get(&k);

        if let Some(indexes_b) = value {
            let tuples = indexes_b.iter().map(|&idx_b| swap_fn(idx_a, idx_b));
            results.extend(tuples);
        }
    });

Two hash computations per row (one when building, one when probing), a hash table allocation, and one random memory access per probe. A merge join does none of that.

The merge sweep

The merge join kernel is short, and it is the whole algorithm:

rust 
// crates/polars-ops/src/frame/join/merge_join.rs:139
    while *build_row_offset < build_key.len() {
        if gather_build.len() >= limit_results {
            return;
        }

        let build_keyval = unsafe { build_key.get_unchecked(*build_row_offset) };
        let build_keyval = build_keyval.as_ref();
        let mut build_keyval_matched = false;

        if nulls_equal || build_keyval.is_some() {
            for probe_idx in *probe_row_offset..probe_key.len() {
                let probe_keyval = unsafe { probe_key.get_unchecked(probe_idx) };
                let probe_keyval = probe_keyval.as_ref();

                let mut ord: Ordering = match (&build_keyval, &probe_keyval) {
                    (None, None) if nulls_equal => Ordering::Equal,
                    (Some(l), Some(r)) => TotalOrd::tot_cmp(*l, *r),
                    _ => continue,
                };
                if descending {
                    ord = ord.reverse();
                }

                match ord {
                    Ordering::Equal => {
                        if let Some(probe_unmatched) = gather_probe_unmatched.as_mut() {
                            // All probe keys up to and *excluding* this matched key are unmatched
                            probe_unmatched
                                .extend(*probe_first_unmatched as IdxSize..probe_idx as IdxSize);
                            *probe_first_unmatched = (*probe_first_unmatched).max(probe_idx + 1);
                        }
                        gather_build.push(*build_row_offset as IdxSize);
                        gather_probe.push(probe_idx as IdxSize);
                        build_keyval_matched = true;
                    },
                    Ordering::Greater => {
                        if let Some(probe_unmatched) = gather_probe_unmatched.as_mut() {
                            // All probe keys up to and *including* this matched key are unmatched
                            probe_unmatched
                                .extend(*probe_first_unmatched as IdxSize..=probe_idx as IdxSize);
                            *probe_first_unmatched = (*probe_first_unmatched).max(probe_idx + 1);
                        }
                        *probe_row_offset = probe_idx + 1;
                    },
                    Ordering::Less => {
                        break;
                    },
                }
            }
        }
        if build_emit_unmatched && !build_keyval_matched {
            gather_build.push(*build_row_offset as IdxSize);
            gather_probe.push(IdxSize::MAX);
        }
        *build_row_offset += 1;
    }

The shape of it:

text 
   build (sorted)                 probe (sorted)
   ------------------------       ------------------------
   [ 3 3 3 5 5 9 ]                [ 3 8 ]

   build cursor -> 0 (key 3)      probe cursor -> 0 (key 3)
     3 == 3  -> emit (0, 0)
     3 == 3  -> emit (1, 0)   # the probe cursor did NOT advance
     3 == 3  -> emit (2, 0)
     3 <  8  -> break, advance the build cursor (key 5)
     5 >  3  -> advance the probe cursor (key 8)
     5 <  8  -> break, advance the build cursor (key 5)
     5 <  8  -> break, advance the build cursor (key 9)
     9 >  8  -> advance the probe cursor, past the end
     key 9 is unmatched

Three rules, and they are the whole of a merge join:

  • Equal: emit the pair. Note that the probe cursor does not advance. That is what handles duplicate keys on the probe side: the next probe row with the same key is compared again, and matches too. Duplicates on the build side are handled by the outer loop: the build cursor advances one row at a time, so each duplicate build row gets its own pass over the same run of probe rows.
  • Greater (build key > probe key): the probe row is behind and can never match this or any later build row, so the probe cursor advances.
  • Less (build key < probe key): the probe side has already moved past this build key, so no later probe row can match it either; break out of the inner loop and advance the build cursor.

Both cursors only ever move forward, so the whole join is O(n + m) in the number of key comparisons, plus the number of emitted pairs.

Nulls are handled explicitly rather than by accident. A null build key with nulls_equal == false skips the inner loop entirely (the if nulls_equal || build_keyval.is_some() guard), so it matches nothing and is emitted as unmatched if the join type requires that. A null probe key falls into the _ => continue arm, which skips it without advancing the cursor.

Descending keys are handled by reversing the comparison rather than by flipping the comparison sites:

rust 
// crates/polars-ops/src/frame/join/merge_join.rs:131
                if descending {
                    ord = ord.reverse();
                }

That works because both sides are sorted in the same direction, so the merge structure is symmetric.

Making it piecewise

The kernel above is written as a loop with two cursors that are passed in and returned out (build_row_offset, probe_row_offset). That signature is the interesting part: the sweep is resumable.

The streaming engine uses it that way. Because the output of a join can be much larger than either input (a many-to-many join), the engine emits results in morsels, and the sweep has to stop at a morsel boundary and pick up where it left off:

rust 
// crates/polars-stream/src/nodes/joins/merge_join.rs:498
    let mut build_row_offset = 0;
    let mut probe_row_offset = 0;
    let mut probe_last_matched = 0;
    arenas.gather_probe_unmatched.clear();
    while build_row_offset < build.height() || probe_row_offset < probe.height() {
        arenas.gather_build.clear();
        arenas.gather_probe.clear();
        let gather_probe_unmatched = params
            .probe_params()
            .emit_unmatched
            .then_some(&mut arenas.gather_probe_unmatched);
        match_keys(
            &build_key,
            &probe_key,
            &mut arenas.gather_build,
            &mut arenas.gather_probe,
            gather_probe_unmatched,
            params.build_params().emit_unmatched,
            params.key_descending,
            params.args.nulls_equal,
            morsel_size,
            &mut build_row_offset,
            &mut probe_row_offset,
            &mut probe_last_matched,
        );

The limit_results argument is the morsel size, and the kernel returns as soon as it has produced that many pairs. The cursors keep their values, so the next call resumes exactly where it stopped. That is the “piecewise” in the name: the join is computed as a sequence of independently emitted pieces, but the sweep is one continuous traversal of the two sorted inputs.

There is a second, coarser kind of piece: whole chunks of the input. The streaming node cannot assume it has both sides fully in memory, so it has to decide, given what it has buffered so far, which rows of the build side can be merged with what it has of the probe side. That decision is a binary search, because both sides are sorted:

rust 
// crates/polars-stream/src/nodes/joins/merge_join.rs:697
    let build_last_idx = usize::min(build.height(), search_limit) - 1;
    let build_last = get_key(build, build_params, params, build_last_idx).await;
    let build_first_incomplete = match build_done {
        false => {
            binary_search_lower(build, &build_last, ..=build_last_idx, params, build_params).await
        },
        true => build.height(),
    };

    let probe_last_idx = usize::min(probe.height(), search_limit) - 1;
    let probe_last = get_key(probe, probe_params, params, probe_last_idx).await;
    let probe_first_incomplete = match probe_done {
        false => {
            binary_search_lower(probe, &probe_last, ..=probe_last_idx, params, probe_params).await
        },
        true => probe.height(),
    };

build_first_incomplete is the index of the first build row whose group might continue into the next morsel — the last group in the buffered window is incomplete if the input has not ended yet, because the next morsel might contain more rows with the same key. The same is computed for the probe side, and then the two are reconciled:

rust 
// crates/polars-stream/src/nodes/joins/merge_join.rs:729
    let build_mergeable_until; // bound is *exclusive*
    let probe_mergeable_until;
    match keys_cmp(&build_last_completed_val, &probe_last_completed_val, params) {
        Ordering::Equal => {
            build_mergeable_until = build_first_incomplete;
            probe_mergeable_until = probe_first_incomplete;
        },
        Ordering::Less => {
            build_mergeable_until = build_first_incomplete;
            let build_last_mergeable =
                get_key(build, build_params, params, build_mergeable_until - 1).await;
            probe_mergeable_until = binary_search_upper(
                probe,
                &build_last_mergeable,
                ..probe_first_incomplete,
                params,
                probe_params,
            )
            .await;
        },
        Ordering::Greater => {
            probe_mergeable_until = probe_first_incomplete;
            let probe_last_mergeable =
                get_key(probe, probe_params, params, probe_mergeable_until - 1).await;
            build_mergeable_until = binary_search_upper(
                build,
                &probe_last_mergeable,
                ..build_first_incomplete,
                params,
                build_params,
            )
            .await;
        },
    }

If the last completed build key is smaller than the last completed probe key, then every probe row up to that build key has been consumed, so the probe side can be cut at the corresponding position — and that position is found with another binary search. The result is a pair of prefixes, one from each side, that can be merged without ever consulting a row outside the prefix. That is the definition of a mergeable piece.

Once the pieces are known, the build side can be split further so that several pieces are merged in parallel:

rust 
// crates/polars-stream/src/nodes/joins/merge_join.rs:625
fn find_mergeable_partition(
    build: SpillFrameSearchBuffer,
    probe: SpillFrameSearchBuffer,
    fmp: FindMergeableParams<'_>,
) -> PolarsResult<UnitVec<(SpillFrameSearchBuffer, SpillFrameSearchBuffer)>> {
    let morsel_size = get_ideal_morsel_size();

    if fmp.params.preserve_order_probe() || fmp.params.probe_params().emit_unmatched {
        return Ok(UnitVec::from([(build, probe)]));
    }

    let est_out_rows = build.height() * probe.height() + build.height() + probe.height();
    let normal_out_rows = morsel_size.pow(2);
    let partition_count = est_out_rows.div_ceil(normal_out_rows);
    if partition_count <= 1 {
        return Ok(UnitVec::from([(build, probe)]));
    }

The estimate build.height() * probe.height() is the worst case for the number of output rows — the cartesian product. Dividing that by the square of the morsel size gives a partition count that keeps each partition’s output around one morsel. The build side is then cut into chunks and each chunk is paired with the whole probe side; because the build side is sorted and the chunks are contiguous and disjoint, the chunks’ outputs can be concatenated without reordering.

Note the two early-outs: if the join must preserve the probe side’s order, or if unmatched probe rows must be reported, the partition cannot be split, because the output would no longer be in probe order and the unmatched rows would have to be merged back in order.

How the choice is made

The join type is not decided by inspecting the data. It comes from the optimizer’s sortedness analysis, which tracks which expressions are known to be sorted and in which direction:

rust 
// crates/polars-stream/src/physical_plan/lower_ir.rs:1391
            let right_on_sorted = ctx.sortedness.are_keys_sorted_any(
                input_right,
                &right_on,
                expr_arena,
                &input_right_schema,
            );
            let join_keys_sorted_together =
                Option::zip(left_on_sorted.as_ref(), right_on_sorted.as_ref())
                    .is_some_and(|(ls, rs)| ls == rs);
            let mut key_descending = left_on_sorted
                .as_ref()
                .and_then(|v| v.first())
                .and_then(|s| s.descending);
            let key_nulls_last = left_on_sorted
                .as_ref()
                .and_then(|v| v.first())
                .and_then(|s| s.nulls_last);
            let use_streaming_merge_join = args.how.is_equi()
                && join_keys_sorted_together
                && key_descending.is_some()
                && key_nulls_last.is_some();

Four conditions, and each one is load-bearing:

  • args.how.is_equi() — a merge join implements equality only. Range joins and asof joins have their own machinery.
  • join_keys_sorted_together — the left and right key columns must be sorted in the same way. The comparison in the sweep assumes both sides are ordered identically; if the left is ascending and the right is descending, the merge degenerates into a quadratic scan.
  • key_descending.is_some() — the direction must be known, because it decides how the comparison is interpreted.
  • key_nulls_last.is_some() — the placement of nulls must be known, because nulls are either at the front or the back of the sorted array and the sweep has to know where to start.

The sortedness metadata is derived from the plan: a sort node produces sorted output, and that fact propagates through the plan tree. If the user writes df.sort("k").join(other, on="k") where other is also sorted on k, the optimizer knows both sides are sorted and the merge join is chosen. If either side is not known to be sorted, the hash join is used and the sortedness is irrelevant.

For multi-column keys, the streaming engine materialises a single row-encoded key column for each side, because the merge kernel compares one key column at a time:

rust 
// crates/polars-stream/src/physical_plan/lower_ir.rs:1521
                    _ if use_streaming_merge_join => {
                        let keys_are_row_encoded = left_on_names.len() > 1;
                        if keys_are_row_encoded {
                            key_descending = Some(false);
                        }

Row encoding preserves the ordering of the key tuple, so the encoded column is sorted exactly when the original key columns are — but only in ascending order, which is why the direction is forced to ascending here.

The other merge: merge_sorted

A different operation in the same family is worth mentioning, because it uses the same sortedness assumption for a different job. merge_sorted concatenates two frames that are sorted on the same key, instead of joining them:

rust 
// crates/polars-ops/src/frame/join/merge_sorted.rs:249
fn get_merge_indicator<T>(
    mut a_iter: impl ExactSizeIterator<Item = T>,
    mut b_iter: impl ExactSizeIterator<Item = T>,
) -> Vec<bool>
where
    T: PartialOrd + Default + Copy,
{
    const A_INDICATOR: bool = true;
    const B_INDICATOR: bool = false;

    let a_len = a_iter.size_hint().0;
    let b_len = b_iter.size_hint().0;
    if a_len == 0 {
        return vec![B_INDICATOR; b_len];
    };
    if b_len == 0 {
        return vec![A_INDICATOR; a_len];
    }

    let mut current_a = T::default();
    let cap = a_len + b_len;
    let mut out = Vec::with_capacity(cap);

    let mut current_b = b_iter.next().unwrap();

    for a in &mut a_iter {
        current_a = a;
        if a <= current_b {
            out.push(A_INDICATOR);
            continue;
        }
        out.push(B_INDICATOR);

Instead of emitting pairs, it emits a boolean per output row saying which side that row came from. It is the same two-cursor sweep with a different payload: the merge of two sorted runs into one sorted run, which is the last step of a merge sort. The output frame is then built by interleaving the two sides according to that indicator.

The bottom line

The design in one sentence: a sorted equi-join is a merge, and the merge is resumable.

The trade-offs:

  • It requires sortedness that the optimizer can prove. Not “the data happens to be sorted” but “this plan produces sorted output here”. A user who sorts one side but not the other gets the hash join, and the sort on one side is then pure overhead.
  • The kernel handles duplicates, but the output can explode. A many-to-many join on duplicate keys produces the cartesian product of the runs, and the partition estimate in find_mergeable_partition assumes exactly that worst case.
  • Unmatched rows cost extra bookkeeping. Tracking which probe rows were not matched requires the gather_probe_unmatched vector and the probe_first_unmatched watermark, and it disables the parallel partitioning of the build side.
  • The state is two cursors and nothing else. Because the sweep is resumable, the join can be suspended and resumed for free, and the output can be produced in morsels of any size without buffering the whole result. That is the payoff of writing the kernel as a function of offsets rather than as a function of two frames.

PolarsInternalsJoinsPerformanceStreaming