A Memory Ordering Bug in DataFusion's Hash Join, Found by Model Checking Rather Than Testing

TL;DR: The probe partitions of a partitioned hash join have to agree on two facts: which one finishes last (it emits the build-side rows), and what they collectively saw (did anyone see a NULL join key?). Reading the second without ordering it against the first produced wrong results. The fix is a single AcqRel counter decrement, and the bug it fixed could not be found by a concurrent test — it was found by loom, a model checker that enumerates every interleaving the memory model allows.

Concurrency bugs in query engines are usually found the hard way: a query that returns wrong answers once a week, in production, on a machine with 64 cores. DataFusion has taken the other approach for one particular protocol: model check it.

The protocol

datafusion/physical-plan/src/joins/hash_join/probe_completion.rs is a small file with a very carefully written module doc. The setup:

The probe partitions of one HashJoinExec run concurrently over a shared build side. Two facts have to cross partition boundaries:
which partition emits the build-side rows in the final stage, namely the last one to finish, and
what the partitions collectively saw, which the null-aware NOT IN rules need: did any partition see a row, and did any see a NULL join key.

So there are N probe tasks, each streaming a slice of the probe side, all sharing one build side. When a probe row’s key is NULL, the null-aware NOT IN semantics say the whole result is NULL — but only if some other partition saw a NULL. And exactly one partition — the last to finish — has to emit the build-side rows for the LeftMark case.

The state is three atomics:

rust
pub(super) struct ProbeCompletion {
    /// Probe partitions that have not finished yet.
    running: Counter,
    /// Set once any probe partition has seen a row.
    saw_row: Flag,
    /// Set once any probe partition has seen a NULL join key.
    saw_null_key: Flag,
}

The bug

Here is the module doc’s account of what went wrong:

[ProbeCompletion] owns both, because reading the second without ordering it against the first produced wrong results. A partition that read the NULL flag before decrementing the counter could observe false, have a sibling record a NULL and finish, then become the last partition itself and emit build rows that NOT IN must suppress.

Read that sequence carefully, because it is a textbook time-of-check / time-of-use across threads:

Partition A reads saw_null_key → false. (It read the flag before decrementing the counter.)
Partition B records a NULL key → saw_null_key = true.
Partition B finishes, decrements the counter.
Partition A decrements the counter — and is now the last partition.
Partition A emits build-side rows, believing no NULL was seen.

But a NULL was seen, so NOT IN should have suppressed those rows. The query returns wrong results.

And it is a narrow window: partition A must read the flag, then partition B must record and finish, then A must decrement and find itself last. That is a few instructions wide, and it needs a specific interleaving on a machine with many cores.

The invariant

The fix is to state the property precisely and then design the type so that the property is hard to violate:

The invariant is therefore: whichever caller observes itself to be last sees every fact recorded by every other partition. It rests on two properties, and the type is shaped to keep both.
Order. The facts are reachable only through the value returned by [ProbeCompletion::report_completed], so no caller can read them before its own decrement.
Visibility. That decrement is AcqRel, so the last partition synchronizes with the release of every partition that finished before it. Were the counter Relaxed, nothing would order a sibling’s store before the final load even when the decrement happened first.

Property 1 is structural: the summary is not a field you can read. It comes back from report_completed, and the only way to get it is to have decremented. You cannot read the flags early because there is no API to read them early.

rust
/// What every probe partition together saw, as observed by the last one to
/// finish.
///
/// Obtainable only from [`ProbeCompletion::report_completed`], which is what
/// stops the final stage from reading the shared state too early.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct ProbeSideSummary {
    /// Some probe partition saw a row. An entirely empty probe side is not the
    /// same as one that matched nothing: `NULL NOT IN (empty)` is TRUE.
    pub(super) non_empty: bool,
    /// Some probe partition saw a NULL join key.
    pub(super) has_null: bool,
}

Property 2 is the memory ordering:

rust
pub(super) fn report_completed(&self) -> Option<ProbeSideSummary> {
    // `fetch_sub` returns the value from *before* the decrement,
    // so `1` means this call is the one that took it to zero.
    let was_last = self.running.fetch_sub(1, Ordering::AcqRel) == 1;
    was_last.then(|| ProbeSideSummary {
        non_empty: self.saw_row.load(Ordering::Relaxed),
        has_null: self.saw_null_key.load(Ordering::Relaxed),
    })
}

The AcqRel decrement does double duty. It releases this partition’s own stores (so other partitions that finish later can see them), and it acquires the stores of the partitions that finished earlier (so the last partition sees their facts). The subsequent Relaxed loads are then safe, because the acquire side of the decrement has already established the happens-before relationship.

The Relaxed stores during the run are fine too:

rust
/// Records what one probe partition saw in one batch.
///
/// `Relaxed` suffices: the `AcqRel` decrement in
/// [`Self::report_completed`] publishes these stores to the last
/// partition.
pub(super) fn record_batch(&self, non_empty: bool, has_null: bool) {
    if non_empty {
        self.saw_row.store(true, Ordering::Relaxed);
    }
    if has_null {
        self.saw_null_key.store(true, Ordering::Relaxed);
    }
}

Relaxed stores plus one AcqRel fence at the end is the standard “accumulate in a relaxed manner, publish once” pattern, and it is much cheaper than an acquire-release store on every batch.

The hint that may lag

There is also a hint interface, and it is explicitly allowed to be stale:

rust
/// Whether some partition has already recorded a NULL join key.
///
/// A hint for skipping work during the probe phase, which may lag
/// behind sibling partitions. It must never decide what the final
/// stage emits; use the [`ProbeSideSummary`] from
/// [`Self::report_completed`] for that.
pub(super) fn saw_null_key_hint(&self) -> bool {
    self.saw_null_key.load(Ordering::Relaxed)
}

A hint is allowed to be wrong in one direction (it may miss a NULL that was just recorded) but not the other (it may never report a NULL that was not recorded). The test says exactly that:

rust
/// The hint may lag, but it must never report a NULL that was not
/// recorded.

So the type system and the naming (hint vs summary) encode the difference between “an optimisation that may be conservative” and “the fact you must use for correctness”.

The model check

Now the part that makes this a great story. The module doc:

Neither property can be pinned down by an ordinary concurrent test: the window is a few instructions wide, and on x86 a Relaxed decrement lowers to the same instruction as an AcqRel one. Both are model-checked in loom_tests instead.

That last point is the one that makes model checking necessary rather than merely nice. On x86, fetch_sub with Relaxed and with AcqRel compile to the same machine instruction (lock xadd). The difference is only visible to the memory model — to the compiler’s and CPU’s permission to reorder and to the hardware’s cache-coherence visibility.

Which means:

A test on x86 might not catch it, because the hardware is strong enough to make the bug invisible.
A test on ARM might catch it by accident, but only sometimes.
A model checker catches it deterministically, because it simulates the weak-memory model rather than running on real hardware.

The loom model checker enumerates every interleaving and every store visibility the memory model allows, and fails if the protocol violates the invariant. And the tests fail if the decrement is weakened to Relaxed:

These instantiate the protocol over loom’s atomics (see the module docs) and let loom enumerate the interleavings and store visibility the memory model allows. They fail if the decrement in report_completed is weakened to Relaxed.

One copy of the protocol, two instantiations

The way the code makes the model check meaningful is a macro:

rust
/// Defines [`ProbeCompletion`] over the given atomic types.
///
/// [loom] can only explore its own atomics, so the protocol is written once
/// here and instantiated twice: over [`std::sync::atomic`] for the join, and
/// over `loom::sync::atomic` for the model in `loom_tests`. Both share this
/// single copy of the orderings and the call sequence, which is what lets
/// the model speak for the real thing. Keep it that way: a second copy of these
/// orderings would not be covered.
macro_rules! define_probe_completion {
    (counter: $counter:ty, flag: $flag:ty) => { ... }
}

This is the crucial bit. If the production code and the model were two separate implementations, the model would be checking a different program and proving nothing. By making the protocol a macro parameterised only by the atomic types, the orderings — the actual subject of the model — are literally the same source text in both.

The doc is emphatic: “Keep it that way: a second copy of these orderings would not be covered.”

Running it

The loom tests are behind a cfg flag, because loom is slow — it explores an exponential space:

rust

/// RUSTFLAGS="--cfg datafusion_loom" cargo test -p datafusion-physical-plan --lib loom_tests

So the model check is not in the normal test run. It is a separate invocation, run deliberately.

The catch

Loom is exponential. It explores the full interleaving space, so the test is limited to a small number of participants. It proves the protocol correct for a small case, not the whole join.
The model is per-atomic-type. Loom only understands its own atomics, which is why the macro exists. Any state the protocol touches that is not an atomic — a Mutex, say — is outside the model.
The Relaxed stores are only safe because of the AcqRel decrement. Change the decrement to Relaxed and the model check fails. Change the stores to Relaxed (they already are) and nothing breaks. The ordering discipline is a single point of failure, which is why the model check exists.
The hint interface is a permanent footgun. It is one field access away from being used for something it must not decide. The naming and the doc comment are the only defences.

Why it matters

This file is a good template for how to handle a subtle concurrency protocol in a codebase that people actually contribute to:

State the invariant in one sentence, in the module doc, in bold.
Make the wrong thing impossible to express — the summary is only obtainable through the operation that establishes the ordering.
Use the cheapest memory ordering that is correct, and explain why the more expensive one is not needed.
Model check it, and make the model check the same code as production.

Point 4 is the one that generalises. “Add a concurrency test” is the usual advice and it is often wrong: a test can only find the interleaving it happens to hit, and on a strong memory model the bug may not exist at the hardware level at all. A model checker explores the space the language permits, which is the space you actually have to be correct in.