jayzhan211 opened a new pull request, #24579:
URL: https://github.com/apache/datafusion/pull/24579
## 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:
```rust
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
