Skip to content

perf: replace PiecewiseMergeJoin's visited bitmap with a suffix watermark (~1.8x faster Left join) - #24579

Open
jayzhan211 wants to merge 2 commits into
apache:mainfrom
jayzhan211:pwmj-suffix-watermark
Open

perf: replace PiecewiseMergeJoin's visited bitmap with a suffix watermark (~1.8x faster Left join)#24579
jayzhan211 wants to merge 2 commits into
apache:mainfrom
jayzhan211:pwmj-suffix-watermark

Conversation

@jayzhan211

Copy link
Copy Markdown
Contributor

Rationale for this change

PiecewiseMergeJoin (PWMJ) needs to know which buffered (left) rows were matched, so that
Left and Full joins can emit the unmatched ones in a final pass. Today it tracks this
with a bitmap of one bit per buffered row, guarded by a mutex.

That bitmap is more machinery than the operator needs. PWMJ's whole trick is that both
sides are sorted, so when a streamed row finds its first match at buffered index k,
every buffered row from k to the end matches — and the operator emits that entire
range at once. Marking is therefore never scattered: it is always the suffix
[k, buffered_len).

A union of suffixes that all end at the last row is just another suffix:

⋃ᵢ [kᵢ, len)  =  [min kᵢ, len)

So the entire bitmap only ever encodes a single number — the smallest k seen. Keeping a
bit per row costs, for every matched streamed row, a mutex acquisition plus one set_bit
per emitted row, and the final pass then has to walk the bitmap to build an index array and
run take over it.

The existence side of this operator (LeftSemi/LeftAnti) already exploits exactly this
property — it stores one AtomicUsize watermark and no bitmap. In fact the old
build_visited_indices_map doc comment spelled the property out, but only for the existence
path. This PR carries the same observation across to the classic joins.

On a LEFT JOIN with a 1M-row buffered side and a 4k-row streamed side (release build,
datafusion-cli):

before after
run 1 4.789 s 2.662 s
run 2 4.819 s 2.645 s

~1.8x faster.

What changes are included in this PR?

  • BufferedSideData drops visited_indices_bitmap; its existing existence-join field
    existence_min_marked is generalized to min_marked and is now maintained by both
    streams. The bitmap allocation, its memory reservation, the Mutex, and the
    BooleanBufferBuilder / SharedBitmapBuilder / bit_util imports all go away.
  • Marking (classic_join.rs): a mutex lock plus one set_bit per emitted row becomes a
    single fetch_min. Because buffer_idx only moves forward within a stream batch, a
    batch_min_marked field in BatchProcessState means only the batch's first match
    touches the shared atomic at all.
  • Final pass: get_final_indices_from_shared_bitmap + take_record_batch becomes
    buffered_batch.slice(0, min_marked). The unmatched buffered rows are exactly the
    complementary prefix, so the final pass is now zero-copy.
  • Dead code: build_visited_indices_map() is removed — every arm other than
    Full/Left named a join type that try_new rejects. A small unused row-count
    accumulator in build_buffered_data's try_fold goes too.

Net: 136 insertions, 113 deletions across 5 files.

Why this is safe

The suffix property is syntactic, not a consequence of the merge algorithm. The bitmap had
exactly one writer, and both of its call sites passed the same range:

let count = buffered_len - buffer_idx;   // always runs to the end of the array((buffer_idx, count), (row_idx, count),)

So every write was [buffer_idx, buffered_len) by construction, and the union of those is
[min buffer_idx, buffered_len) — regardless of sortedness, operator direction, NULL
placement, or how partitions interleave. The old and new encodings are therefore
unconditionally equal, not equal-under-an-assumption.

Output ordering is unchanged as well: the old reader returned ascending indices whose bit was
false, and the complement of [min, len) is [0, min), also ascending.

Are these changes tested?

Yes.

The two existing PWMJ fuzz entry points are merged into one differential test against a
NestedLoopJoin oracle, fuzz_pwmj_matches_nested_loop, now covering all six supported
join types (Inner/Left/Right/Full/LeftSemi/LeftAnti) rather than the existence
pair alone — 60 seeds x 4 operators, with NULL keys, duplicate keys, 1–3 streamed partitions
executed concurrently, and batch_size = 3. The previous existence-only test and its
collector are folded in, so this adds coverage while shrinking the file's duplication.

Why an added test was warranted, stated precisely: pwmj.slt runs partitions=1
throughout, and every classic unit test builds the exec with a single-partition streamed side
at the default batch size. Neither reaches several partitions racing to run the final pass,
nor the mid-scan resume path. I confirmed this by mutation:

mutation 23 unit tests pwmj.slt this test
fetch_min(k)fetch_min(k + 1) fails ✅ fails ✅ fails ✅
drop the min(num_rows) clamp (min_marked is usize::MAX when nothing matched) passes passes fails ✅

The second row is the gap this closes — the existing suite already catches coarse breakage,
but not the multi-partition / never-matched cases the new encoding introduces.

Also verified unchanged: the 23 PWMJ unit tests, pwmj.slt, the full joins:: suite, and
clippy -D warnings. As a further cross-check, old and new release binaries produce
byte-identical results for inner / left / right / full / semi / anti / <= over a 20k x 3k
join with NULLs on both sides at target_partitions = 4.

Are there any user-facing changes?

No. This is an internal change to how matched buffered rows are recorded — query results,
output ordering, and public APIs are unchanged. PiecewiseMergeJoin uses slightly less
memory, since the per-buffered-row bitmap is no longer allocated.

@github-actions github-actions Bot added core Core DataFusion crate physical-plan Changes to the physical-plan crate labels Aug 22, 2026
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.93939% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.37%. Comparing base (32badcb) to head (cbdcc72).
⚠️ Report is 2 commits behind head on main.

Files with missing lines Patch % Lines
...lan/src/joins/piecewise_merge_join/classic_join.rs 95.23% 0 Missing and 1 partial ⚠️
...ysical-plan/src/joins/piecewise_merge_join/exec.rs 85.71% 0 Missing and 1 partial ⚠️
Additional details and impacted files
@@           Coverage Diff            @@
##             main   #24579    +/-   ##
========================================
  Coverage   81.37%   81.37%            
========================================
  Files        1116     1116            
  Lines      397509   397645   +136     
  Branches   397509   397645   +136     
========================================
+ Hits       323461   323581   +120     
- Misses      55110    55121    +11     
- Partials    18938    18943     +5     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants