Skip to content

perf: reduce generic hash join memory for low-cardinality builds - #25434

Open
sunchao wants to merge 3 commits into
apache:mainfrom
sunchao:dev/chao/codex/oss-compact-hash-build-25392
Open

sunchao wants to merge 3 commits into
apache:mainfrom
sunchao:dev/chao/codex/oss-compact-hash-build-25392

Conversation

@sunchao

@sunchao sunchao commented Sep 17, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #25392.

Rationale for this change

Generic hash joins reserve lookup-table capacity for every build row, even when many rows share a key. A join over 10,000 customer records with state GA must retain all 10,000 matching rows, but needs only one lookup entry for their hash. Reserving unused buckets can make otherwise small joins fail with ResourcesExhausted.

This PR sizes the lookup table around the build keys while preserving every row in the separate index chain. The regression test joins 65,536 rows containing repeated GA, NULL, and CA keys in a 2 MiB memory pool; the old implementation exhausts the pool, while the compact build preserves all duplicate matches and both NULL-matching modes.

The reported 20,000- and 100,000-key benchmarks exposed another important case. With one million rows, growing directly to the full row count and shrinking afterward allocates an oversized table, then scans it and moves its populated entries into a smaller table. That extra compaction work explains why reducing retained memory could still slow construction. These distributions are now included in the benchmark, and eligible builds estimate their required capacity before making that large allocation.

What changes are included in this PR?

The generic hash table starts with capacity for one bounded chunk. At its first growth, a sample of up to 2,048 row positions across the buffered input can estimate capacity for a single non-null string or binary column. Sampling reads values directly without copying them or evaluating expressions again. When the sample cannot support a useful estimate, or the key representation is unsupported, the builder tries full-row preallocation. Every row still passes through normal hashing and insertion; the estimate affects capacity only.

A smaller sampled allocation is admitted together with the workspace needed to compact it. If that combined reservation cannot fit, the builder tries full-row preallocation, then bounded growth. An underestimated sample can fall back to full preallocation once. If only a replacement table fits, the builder releases its partial table and rebuilds from the buffered rows. Final compaction also handles bounded growth after a denied full fallback, preserving memory needed by subsequent join work.

The compact builder passes its iterator directly to a generic internal insertion helper, removing its per-chunk iterator box and allowing compiler specialization. The existing public boxed helper keeps its signature and delegates to that implementation.

The memory pool accounts for row indices, hash scratch, table overlap, and reserved compaction workspace. build_mem_used reports their tracked reservation peak, which can include workspace reserved before it is allocated. It is not process RSS and does not account for every allocation during concatenation of the buffered payload.

For the one-million-unique-key case, build_mem_used changes from 55,592,952 to 59,937,088 bytes. Most of the increase comes from the pre-existing 4,000,000-byte row-index chain and 65,536-byte hash scratch becoming charged. Growth also adds 278,544 bytes of actual temporary overlap for the initial small table, plus a small fixed bookkeeping allowance in the counter. Both revisions retain the same 39,651,656-byte hash index, including its row-index chain, after construction.

What is the testing strategy for this PR?

Focused tests cover duplicate and NULL semantics, row-chain order, both index widths, empty input, sliced and empty batches, all supported string/binary representations, and dictionary, nested, multi-column, and computed keys. Memory tests compare the reported peak with an independently recording pool and verify cleanup after denied allocations and rebuilds.

Two regressions exercise downstream memory availability: constrained LeftAnti joins that need a visited-row bitmap, and an underestimated sample whose denied full-table fallback previously retained twice the necessary table allocation. Uniform and grouped intermediate-cardinality inputs verify that the construction peak stays below a full-row table allocation.

Final-revision validation:

  • All 19 focused tests passed normally and with both force_hash_collisions and datafusion-common/force_hash_collisions enabled.
  • Extended workspace libraries/tests/binaries, excluding examples, benchmarks, and CLI: 11,932 passed, 8 ignored; all 521 SQL logic test files completed. Core/CLI tests and doctests: 2,479 passed, 57 ignored doctests.
  • cargo fmt --all: passed.
  • cargo clippy --all-targets --all-features -- -D warnings: passed.
  • ./dev/rust_lint.sh: passed, including documentation and dependency checks.
  • Rust 1.98.1 on Linux x86_64. The local registry lacked four locked patch versions, so validation and both benchmark revisions used the same temporary lock (async-compression 0.4.44, compression-codecs 0.4.40, toml 1.1.5, uuid 1.26.0). The original lockfile is restored and is not changed by this PR; all current-head GitHub checks passed with the published dependency versions (38 successful checks, 3 optional skips).

The seven Criterion cases build from one million UTF-8 rows in batches of 8,192 and probe with one absent key. Input generation is outside timing; execution consumes and drops the full join.

Build-key distribution Base ms, passes 1 / 2 Updated PR ms, passes 1 / 2 Elapsed change, passes 1 / 2 Tracked build memory MB, base → updated PR
64 repeated keys 10.34 / 12.66 8.76 / 9.00 -15.4% / -28.9% 47.66 → 16.35
20,000 distinct keys 32.52 / 64.11 13.03 / 13.46 -59.9% / -79.0% 54.54 → 26.58
100,000 distinct keys 37.25 / 74.63 29.33 / 28.79 -21.3% / -61.4% 55.40 → 50.83
900,000 distinct keys 38.87 / 71.91 33.84 / 68.44 -13.0% / -4.8% 55.53 → 59.87
1,000,000 distinct keys 36.33 / 69.36 33.43 / 65.42 -8.0% / -5.7% 55.59 → 59.94
Unique rows first, repeated keys second 22.70 / 56.83 21.46 / 54.96 -5.4% / -3.3% 51.59 → 55.94
Repeated keys first, unique rows second 22.63 / 57.26 21.82 / 55.44 -3.6% / -3.2% 51.59 → 55.94

The last two cases contain the same multiset: 500,000 unique keys and 500,000 rows repeating 64 other keys, in opposite input order. These measurements establish behavior for the tested workloads; sampling remains a capacity heuristic rather than a universal performance guarantee.

Compared exact base 62f039f0d9583054b81d47b2f08b6728a14ac452 with 048610951fbcf5cd7aa566bfd288f5485fd10109 on AMD EPYC Milan/Linux x86_64 using Rust 1.98.1, separate target directories, and the identical submitted seven-case harness. Both use the default release profile (lto=true, one codegen unit), CPU 0, one Tokio worker, 30 samples, 2 s warmup, and 5 s target measurement. Process order was base → PR → PR → base. The table preserves both paired passes because both binaries showed distinct fast/slow timing modes; it does not pool them into a single speedup.

Additional runs with the submitted harness confirmed the mostly-unique, unique, and both input-order cases. A duplicate-first control had one long sample, which is retained in the results; a preselected follow-up with 80 samples and 10 s target measurement improved by 7.1% and 5.4% in its two reversed pairs. Earlier input-order controls overlapped an unrelated build that saturated the host; those observations are preserved but excluded from the performance conclusion.

A separate non-LTO comparison of exact base → reviewed PR head 57791ea2 → updated PR reproduced the old head's middle-cardinality slowdown. At 20k distinct keys, the three means were 66.07 → 80.97 → 30.30 ms; at 100k, 71.54 → 93.11 → 29.06 ms. All seven updated-PR cases improved against the base in this screen. This diagnostic harness selects the requested case before allocating inputs, so its absolute timings are kept separate from the submitted-harness LTO table.

Seven further non-LTO stress cases also improved in both reversed pairs: 6,144/6,145 distinct keys around the initial-table boundary; half unique rows mixed with 8,000 repeated keys in both orders; a 1.5% unique tail mixed with one repeated key in both orders; and grouped 100k keys. The final candidate showed no reproducible timing regression in the seven PR cases or this stress matrix. Reservation peaks are reported separately from timing and retained storage; this is not a claim that every memory counter decreases.

To reproduce each revision after installing the same toolchain/dependencies:

CARGO_TARGET_DIR=/path/to/separate-target cargo bench --locked \
  -p datafusion-physical-plan --features test_utils \
  --bench hash_join_semi_anti --no-run
TOKIO_WORKER_THREADS=1 taskset -c 0 /path/to/benchmark-binary \
  --bench '^hash_join_build/' --sample-size 30 \
  --warm-up-time 2 --measurement-time 5 --noplot

Copy the benchmark file from the updated PR onto the base checkout so the requested middle-cardinality cases exist in both, and reverse the run order for the second pair.

Are there any user-facing changes?

Duplicate-heavy generic hash joins can execute with smaller memory pools while preserving their results. Intermediate-cardinality string and binary builds can avoid full-row lookup-table allocation. No configuration or public API changes are required. The specialized perfect-hash ArrayMap and auxiliary NULL-aware map construction paths are unchanged.

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 17, 2026
@sunchao
sunchao marked this pull request as ready for review September 17, 2026 21:32
@codecov-commenter

codecov-commenter commented Sep 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 88.02589% with 37 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.42%. Comparing base (62f039f) to head (0486109).
⚠️ Report is 49 commits behind head on main.

Files with missing lines Patch % Lines
...sical-plan/src/joins/hash_join/compact_hash_map.rs 90.57% 15 Missing and 11 partials ⚠️
...tafusion/physical-plan/src/joins/hash_join/exec.rs 52.17% 11 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #25434      +/-   ##
==========================================
+ Coverage   82.33%   82.42%   +0.09%     
==========================================
  Files        1137     1139       +2     
  Lines      431887   435702    +3815     
  Branches   431887   435702    +3815     
==========================================
+ Hits       355580   359125    +3545     
- Misses      54811    54863      +52     
- Partials    21496    21714     +218     

☔ 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.

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @sunchao , 2 suggestions

// Keep the old allocation charged during rehash. The fixed-size
// allowance covers control-group padding; at least eight elements
// also covers hashbrown's minimum bucket sizes.
let minimum = (table.len() + additional).max(chunk_rows);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Second growth always jumps to num_rows, and it triggers once len + 8192 > 14336, so any build with more than ~6,144 distinct hashes gets the same table as main. Measured with 1M Utf8 rows: 6,000 distinct → 278,536 B table; 6,200 distinct → 35,651,592 B; 10K and 100K distinct → 35,651,592 B.

Cheap fix: shrink after the build. With this, 6,200 → 139 KB, 10K → 278 KB, 100K → 2.2 MB, ≥500K unchanged, and all 6 new tests still pass. Peak is still row-count-sized when the pool admits it

         break;
     }
+    // The one-shot row-count preallocation overshoots when only a fraction
+    // of the rows carry distinct hashes; return the unused buckets.
+    if table.capacity() / 4 > table.len() {
+        let old_bytes = table.allocation_size();
+        table.shrink_to(table.len(), |&(hash, _)| hash);
+        reservation.shrink(old_bytes - table.allocation_size());
+    }
     drop(hashes);

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 57791ea. After releasing the hash scratch buffer, tables with capacity greater than four times their entry count are compacted. The replacement allocation is admitted while the old table remains charged, and the overlap is included in peak memory. If admission fails, compaction is skipped and the completed join still succeeds. As you noted, this reduces retained memory; the earlier row-count preallocation peak can still occur.

use crate::test::TestMemoryExec;

#[tokio::test]
async fn compact_hash_build_with_duplicates_and_nulls() -> Result<()> {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No case with HASH_BUILD_CHUNK_ROWS < distinct << rows, so nothing pins table size in the mid-cardinality range. Add one that bounds the final allocation.

let rows = 1_000_000;
let distinct = 10_000;
let batches: Vec<_> = (0..rows)
    .step_by(HASH_BUILD_CHUNK_ROWS)
    .map(|start| {
        let end = (start + HASH_BUILD_CHUNK_ROWS).min(rows);
        RecordBatch::try_from_iter([(
            "key",
            Arc::new(Int32Array::from_iter_values(
                (start..end).map(|i| (i % distinct) as i32),
            )) as ArrayRef,
        )])
    })
    .collect::<Result<_, _>>()?;
let on = vec![Arc::new(Column::new("key", 0)) as PhysicalExprRef];
let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 30));
let reservation = MemoryConsumer::new("mid cardinality").register(&pool);
let (table, _next) = build_compact_hash_map::<u32>(
    &batches,
    &on,
    rows,
    HASH_JOIN_SEED.random_state(),
    NullEquality::NullEqualsNothing,
    &reservation,
    &mut 0,
)?;
assert!(
    table.allocation_size()
        <= estimate_memory_size::<(u64, u32)>(4 * distinct, size_of::<JoinHashMapU32>())?
);

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added the 1,000,000-row / 10,000-key final-allocation regression in 57791ea. It fails against the original PR head and passes with compaction, for both u32 and u64 row indices. It also checks every row chain and exact retained accounting. I extended the growth test with a constrained 100,000-key case that skips compaction while preserving the same heads/chains, and compare reported peak memory with PeakRecordingPool. All seven focused tests pass both normally and with forced hash collisions.

@jayzhan211

jayzhan211 commented Sep 19, 2026

Copy link
Copy Markdown
Contributor

Speculate-to-num_rows then shrink_to_fit regresses mid-cardinality builds; bench_hash_join_build has no such shape. 1M Utf8 rows, base 62f039f0d vs this branch, separate target dirs, interleaved:

distinct base PR Δ peak build_mem_used
20k 4.61 ms 4.83 ms +4.8% 54.5 → 59.1 MB
100k 5.63 ms 6.42 ms +14–16% 55.4 → 61.6 MB

Disabling this block restores base time (5.85 ms vs 6.03 ms), so the cost is the compaction rehash. Acceptable as a trade for ~33 MB less retained during probe, but please cover it in the bench and note it in the description.

     for distribution in [
         "duplicates",
+        "mid_20k",
+        "mid_100k",
         "mostly_unique",
                             "duplicates" => row % 64,
+                            "mid_20k" => row % 20_000,
+                            "mid_100k" => row % 100_000,
                             "mostly_unique" => row % 900_000,

@sunchao I think we should try to improve mid_20k and mid 100k case 🤔

┌──────────────────────┬─────────┬─────────┬─────────┬─────────────────────┬──────────┐
│     distribution     │  base   │   PR    │ Δ time  │ base build_mem_used │    PR    │
├──────────────────────┼─────────┼─────────┼─────────┼─────────────────────┼──────────┤
│ duplicates (64 keys) │ 3.12 ms │ 3.03 ms │     −3% │            47.66 MB │ 16.35 MB │
├──────────────────────┼─────────┼─────────┼─────────┼─────────────────────┼──────────┤
│ mid_20k              │ 4.61 ms │ 4.83 ms │   +4.8% │            54.54 MB │ 59.10 MB │
├──────────────────────┼─────────┼─────────┼─────────┼─────────────────────┼──────────┤
│ mid_100k             │ 5.63 ms │ 6.42 ms │ +14–16% │            55.40 MB │ 61.62 MB │
├──────────────────────┼─────────┼─────────┼─────────┼─────────────────────┼──────────┤
│ mostly_unique        │ 6.07 ms │ 5.73 ms │     −5% │            55.53 MB │ 59.87 MB │
├──────────────────────┼─────────┼─────────┼─────────┼─────────────────────┼──────────┤
│ unique               │ 6.01 ms │ 5.62 ms │     −6% │            55.59 MB │ 59.94 MB │
├──────────────────────┼─────────┼─────────┼─────────┼─────────────────────┼──────────┤
│ unique_prefix        │ 4.72 ms │ 4.52 ms │     −4% │            51.59 MB │ 55.94 MB │
├──────────────────────┼─────────┼─────────┼─────────┼─────────────────────┼──────────┤
│ duplicate_prefix     │ 4.74 ms │ 4.54 ms │     −4% │            51.59 MB │ 55.94 MB │
└──────────────────────┴─────────┴─────────┴─────────┴─────────────────────┴──────────┘

@sunchao

sunchao commented Sep 20, 2026

Copy link
Copy Markdown
Member Author

Addressed in commit 048610951, including the requested mid_20k and mid_100k benchmark cases.

The slowdown came from allocating a table for all one million rows, then moving its entries into a smaller table. At the first growth, the builder now samples up to 2,048 positions across the buffered string/binary column and uses that only to choose capacity. Every row is still inserted normally. Uncertain estimates retain full preallocation; underestimated estimates can grow again. Smaller sampled allocations reserve their final compaction workspace upfront. The insertion loop also uses a generic internal iterator while preserving the public boxed API.

Against exact upstream base 62f039f0, both default-release/LTO passes improved all seven benchmark cases. The 20k case went from 32.52 → 13.03 ms and 64.11 → 13.46 ms; the 100k case went from 37.25 → 29.33 ms and 74.63 → 28.79 ms. These are separate paired runs, not pooled averages: both binaries showed fast/slow timing modes. Runs used the submitted benchmark file, separate target directories, CPU 0, one Tokio worker, 30 samples, 2 s warmup, 5 s measurement, and reversed process order. Extra quiet runs confirmed the unique-key and input-order controls. A separate non-LTO comparison also improved all seven cases and reproduced the old PR head's 20k/100k slowdown.

Seven additional stress cases also improved in both paired runs: the initial-table size boundary, skewed inputs with unique rows at either end, and grouped 100k keys. No timing regression reproduced in the final PR benchmark runs or these stress cases.

The new tests also cover two failures found during development: a tight-memory LeftAnti query that needs a visited-row bitmap, and an underestimated sample whose denied full-growth attempt previously retained an oversized table. Both regressions pass and verify reservation cleanup; separate focused tests validate row chains and FIFO order. All 19 focused tests pass normally and with forced hash collisions. Extended workspace tests, all 521 SQL logic files, core/CLI tests, strict Clippy, and the full repository lint suite pass locally.

The unique-key build_mem_used counter still includes charges missing from the base: its existing 4 MB row-index vector and 64 KiB scratch. Growth also temporarily holds a 278,544-byte initial table beside the replacement. Retained unique-key index storage is unchanged; the counter measures reservations, not RSS. The PR description records this distinction and the identical temporary dependency lock used for local comparisons.

All GitHub checks for 048610951 are complete: 38 successful, 3 optional skips, no failures. The Rust CI run includes the published lockfile.

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

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Reduce generic hash join bucket allocation for low-cardinality build keys

3 participants