jayzhan211 commented on code in PR #24573:
URL: https://github.com/apache/datafusion/pull/24573#discussion_r3853747653
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -1664,60 +1666,129 @@ impl MaterializingSortMergeJoinStream {
}
// Multiple source batches: map each buffered_batch_idx to a
- // contiguous source index, reserving source 0 for a null sentinel.
- let mut batch_idx_to_source: HashMap<usize, usize> = HashMap::new();
+ // contiguous source index. A null sentinel array is prepended as
+ // source 0 only when some right index is actually null (an
+ // unmatched streamed row inside an otherwise matched chunk);
+ // `interleave` walks a null buffer for *every* output row as soon as
+ // any input is nullable, so an always-present sentinel would tax the
+ // common all-matched case.
+ let needs_null_sentinel = matched_chunks
+ .iter()
+ .any(|(_, _, right)| right.null_count() > 0);
+ let source_offset = usize::from(needs_null_sentinel);
+
+ // A group spans only a handful of buffered batches, so a linear
+ // scan beats hashing here. Measured over 8192 rows in 2048 chunks,
+ // against a `HashMap<usize, usize>` built in one pass and read back
+ // in a second (what this used to do):
+ //
+ // distinct sources | hashmap | linear scan
+ // -----------------+-----------+-------------
+ // 4 | 21.5 us | 5.0 us
+ // 16 | 22.0 us | 9.4 us
+ // 32 | 22.4 us | 13.6 us
+ // 64 | 22.9 us | 23.5 us
+ // 128 | 24.1 us | 44.8 us
+ //
+ // `std::collections::HashMap` hashes with SipHash-1-3, so a single
+ // `usize` lookup costs several ns of serial latency before the probe
+ // begins, while a scan over a handful of `usize` is one L1-resident
+ // cache line with a perfectly predicted trip count. The map is also
+ // purely additive state: `source_batches` has to be built regardless
+ // (`source_data` is gathered from it), so hashing means maintaining
+ // two containers holding the same keys.
+ //
+ // The crossover is ~32 distinct sources. That bound follows from how
+ // pairs accumulate, not from any assumption about key skew:
+ //
+ // 1. `pair_streamed_row_with_group` appends exactly one pair per
+ // buffered row and re-checks `num_unfrozen_pairs() < batch_size`
+ // before each append, so at most `batch_size` pairs accumulate
+ // between two `freeze_streamed()` calls.
+ // 2. `BufferedData::scanning_advance` walks the group's rows in
+ // order, so those pairs cover a *contiguous run* of buffered
+ // rows.
+ // 3. So the distinct `buffered_batch_idx` values seen here are the
+ // batches spanned by at most `batch_size` consecutive buffered
+ // rows: `len(source_batches) <= batch_size / R + 1`, where `R`
+ // is the smallest buffered batch in that run.
+ //
+ // The assumption is therefore not "key groups are narrow" — a group
+ // of any width still only contributes `batch_size` rows per freeze —
+ // but "buffered batches are not tiny relative to `batch_size`".
+ // Exceeding 32 sources needs `R < batch_size / 31`, i.e. under ~264
+ // rows per batch at the default `batch_size` of 8192. The buffered
+ // side of a merge join is sorted input, and every operator that
+ // normally feeds it emits ~`batch_size` batches: `SortExec` chunks
+ // its output with `sort_batch_chunked(.., batch_size)`, and
+ // `FilterExec` and `RepartitionExec` each embed a
+ // `LimitedBatchCoalescer` targeting `batch_size`.
+ //
+ // If something does feed tiny batches, this degrades gradually rather
+ // than falling off a cliff, and never affects correctness: at 4
+ // sources this loop is ~13% of the cost of the `interleave` calls it
+ // feeds (3 columns, 8192 rows), so even the 128-source case above
+ // leaves `interleave` the dominant term.
let mut source_batches: Vec<usize> = Vec::new();
- for (batch_idx, _, _) in matched_chunks {
- batch_idx_to_source.entry(*batch_idx).or_insert_with(|| {
- let idx = source_batches.len() + 1;
- source_batches.push(*batch_idx);
- idx
- });
- }
-
let mut interleave_indices: Vec<(usize, usize)> =
Vec::with_capacity(total_matched_rows);
for (batch_idx, _, right) in matched_chunks {
- let source = batch_idx_to_source[batch_idx];
- for i in 0..right.len() {
- if right.is_null(i) {
- interleave_indices.push((0, 0));
- } else {
- interleave_indices.push((source, right.value(i) as usize));
+ let source = match source_batches.iter().position(|b| b ==
batch_idx) {
Review Comment:
**TL;DR** — You're right, and it's worse than described: the chunk sequence
*cycles* (`scanning_reset` fires per streamed row, not per freeze), so the
chunk count isn't bounded by `S` at all and the scan is `O(chunks × S)` — 8.35
ms in a single freeze, measured. Fixed, but with a direct-addressed table
rather than a `HashMap`: `buffered_batch_idx` is an index into
`buffered_data.batches`, so the keys are dense small integers and hashing them
is pure overhead. Result is **~4× faster and ~6.4× smaller than the map** in
the normal case, ties the linear scan at its best, and needs no crossover
threshold or fallback branch. One honest tradeoff in the wrapped-freeze case,
detailed at the end.
## Confirming the diagnosis
I instrumented `materialize_right_columns` on a real join (6 one-row
buffered batches, 2 streamed rows, `batch_size` 5) and dumped the chunk
sequence per freeze:
```
FREEZE chunks: [0, 1, 2, 3, 4]
FREEZE chunks: [5, 0, 1, 2, 3] <- wrapped
FREEZE chunks: [4, 5] <- never sees batch 0
```
The `batch_size` bound in my comment bounded `S`, but not the number of
*chunks* — and the scan runs once per chunk. `pair_streamed_row_with_group`
calls `scanning_reset()` per **streamed row**, not per **freeze**, so the
sequence cycles and one freeze can contain many passes over the same sources.
Cost is `O(chunks × S)`; with one-row batches that's your 33M comparisons,
which I measured at **8.35 ms in a single freeze**. My "degrades gradually
rather than falling off a cliff" line was simply wrong.
## Why a direct-addressed table rather than the map
The keys here aren't opaque. `buffered_batch_idx` is literally an index into
`buffered_data.batches`, a `VecDeque` — so the key space is dense, small,
non-negative integers bounded by the deque length. Hashing those is pure
overhead. A `Vec` indexed by `batch_idx - min` is the natural map, and it beats
the `HashMap` on **both** axes.
**Speed** — 8192 rows in 2048 chunks; the last row is the one-row-per-batch
shape:
| distinct sources | hashmap | linear scan | direct table |
| ---: | ---: | ---: | ---: |
| 4 | 19.7 µs | 4.5 µs | 4.8 µs |
| 32 | 20.7 µs | 13.0 µs | 5.0 µs |
| 128 | 23.5 µs | 42.7 µs | 5.1 µs |
| 1024 | 48.1 µs | 281.7 µs | 5.8 µs |
| 8192 | 293.3 µs | **8347.6 µs** | 16.9 µs |
**Memory** — peak bytes held by the lookup structure alone, measured with a
tracking allocator (`source_batches` and `interleave_indices` excluded, since
every variant needs them):
| shape | distinct | span | map (presized) | direct |
| --- | ---: | ---: | ---: | ---: |
| dense, 32 sources | 32 | 32 | 1.1 KB | 256 B |
| dense, 1024 sources | 1024 | 1024 | 34.0 KB | 8.0 KB |
| dense, 8192 sources | 8192 | 8192 | 272.0 KB | 64.0 KB |
~6.4× less, because `HashMap<usize, usize>` pays 17 bytes/bucket at a 7/8
load factor rounded up to a power of two, while the table pays 8 bytes/slot
flat. (Without pre-sizing, the map peaks at 408 KB during rehash.) It also ties
the linear scan where the scan is at its best, so there's no crossover to tune
and no fallback branch.
## The two cases
Keys within a freeze are contiguous *only within a streamed-row pass*:
- **Dense freeze** (the norm) — `span == distinct`. The table is optimal on
both time and memory.
- **Wrapped freeze** — the freeze straddles a `scanning_reset`, holding the
tail of one pass plus the head of the next. The window wraps and leaves a gap,
so `span` = the group's batch count while `distinct` = `batch_size`.
## The tradeoff
The wrapped case is where the table gives something up, and I want to be
straight about it. **Memory crosses over at a span of ~35k batches** — beyond
that the table's transient exceeds the map's fixed 272 KB, reaching 4 MB at a
524288-batch group.
Time still favours the table, because `scanning_reset` fires once per pass,
so at most **one** freeze per pass wraps and the `O(span)` cost amortizes
against the `O(group)` of useful work the rest of the pass does:
| group spans | freezes | hashmap | direct |
| ---: | ---: | ---: | ---: |
| 8192 | 2 | 264.3 µs | 26.2 µs |
| 131072 | 17 | 2894.0 µs | 304.1 µs |
| 524288 | 65 | 11252.4 µs | 1168.4 µs |
Both linear, table ~10× ahead. The other tradeoff is one allocation per
multi-source freeze that the scan didn't have — that's the 4.8 vs 4.5 µs at 4
sources. The single-source path returns before reaching any of this, so the
common case is untouched.
I also tried an epoch-stamped table kept as reusable scratch on the stream:
only ~15% faster, and it converts a transient allocation into a permanent one
of the same size. Not worth the extra state.
If you'd like a hard cap on that transient, a one-line guard falling back to
the map when `span > k * matched_chunks.len()` would bound it cheaply — happy
to add it, though I'd argue 4 MB against a buffer already holding 524288
`BufferedBatch`es with retained join-key arrays isn't the binding constraint.
## Also in the PR
A regression test covering all three freeze shapes above (including the
wrapped one and the `min > 0` one), plus a `debug_assert` that every key
addresses the live deque — the property the dense key space rests on.
--
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]
