Skip to content

feat: support reusable prepared hash-join builds - #25491

Open
sunchao wants to merge 3 commits into
apache:mainfrom
sunchao:codex/oss-prepared-hash-join
Open

sunchao wants to merge 3 commits into
apache:mainfrom
sunchao:codex/oss-prepared-hash-join

Conversation

@sunchao

@sunchao sunchao commented Sep 19, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of Comet #6013. Comet #6037 uses this API to share a broadcast build across executor tasks.

The shared row-index accounting fix is split out in #25508 for independent review and backporting; this branch includes the same fix.

Rationale for this change

Independent joins against the same build snapshot currently repeat the work of preparing that snapshot and retain separate hash tables while they run. CollectLeft already shares a build among probe partitions within one join plan, but that sharing does not extend to independently created plans.

For example, a DataFusion service might repeatedly join incoming event batches to a stable users snapshot:

SELECT events.id, users.country
FROM events JOIN users ON events.user_id = users.id;

If users is the build side and eight requests overlap, each request needs its own probe input, but all eight can use the same users lookup. Preparing that lookup once avoids seven redundant builds and lets the requests share one retained hash table. When the users snapshot changes, the service can prepare a replacement while existing consumers finish with the old snapshot. Comet's independent tasks probing a broadcast relation are another instance of the same problem; the API is also useful to DataFusion applications without Comet.

Consumers still probe in parallel. Reuse reduces aggregate build work and duplicated build memory; the effect on request latency depends on build size, probe cost, concurrent execution, and whether preparation has already happened. Keeping a build between requests is an explicit lifetime decision for the application.

What changes are included in this PR?

The join's immutable build data can now live independently of a particular execution. An application prepares a build stream once, then attaches the returned handle to compatible joins with different probe inputs:

let prepared = first_join
    .prepare_build(users_snapshot, durable_pool, config)
    .await?;

let first = first_join.builder()
    .with_prepared_build(Arc::clone(&prepared))
    .build()?;
let second = second_join.builder()
    .with_prepared_build(prepared)
    .build()?;

Both joins share the prepared rows and lookup table. Probe progress, completion, and residual conditions remain local to each consumer. Independently planned consumers supply their own probe plans and dynamic-filter expressions, so one execution can finish or be cancelled without changing another's execution state.

The shared handle also carries the build's memory reservation until its last owner releases it. Preparation admits retained data and temporary working memory to the caller's durable pool; failure or cancellation releases unfinished work. The caller supplies buffers that survive producer-task cleanup and owns snapshot identity, preparation coordination, and eviction. Compatibility checks verify schema, keys, and null equality; they cannot establish that two snapshots contain the same data.

The initial contract is deliberately limited to non-spilling CollectLeft inner joins, direct-column keys with matching types, and fixed-width or UTF-8 build columns. Applications attach the prepared build after physical optimization. The unused build subtree becomes a schema placeholder, and EXPLAIN displays prepared_build=<N> rows in default, verbose, and tree formats. Supported resets and projections retain the handle; incompatible rewrites and serialization of the process-local prepared state fail explicitly.

What is the testing strategy for this PR?

The regressions exercise separate consumers on multiple runtime workers and compare their filtered results with ordinary CollectLeft joins. A bounded differential test uses 24 fixed seeds, both null-equality modes, and two output batch sizes to cover duplicate and skewed keys, nulls, empty inputs, residual conditions, and varied input batch boundaries. Concrete result and ownership tests remain alongside these comparisons.

Memory tests check admission and release across successful preparation, errors, cancellation, aliased or sliced inputs, and concatenation that creates validity buffers. The existing plan-reset regression now checks all three EXPLAIN formats, and the serialization regression checks that a prepared join is rejected as unsupported instead of losing its build rows. The public API example is compiled as a doctest.

The measurements show a benefit for the larger shared builds, including when preparation is timed, but do not show a stable latency benefit for small builds. The ordinary-path controls also vary across runs. Applications should measure their intended workload rather than assume a universal break-even size.

Operator benchmark setup, timings, and memory-accounting samples

The Criterion benchmark covers independent DataFusion INNER joins over one dimension snapshot. It compares rebuilding each join, preparing once inside the timed interval (cold), and reusing an already retained build (warm). Each consumer probes 4,096 rows with 50% matches; input generation is excluded, while plan construction, concurrent execution, and draining outputs are timed. All 24 cases check result counts/checksums and release their pool reservations.

The initial measurements below are mean microseconds for the complete group of one or four consumers, using Rust 1.98.1, release-nonlto, a four-worker runtime, and an AMD EPYC-Milan VM with 32 visible CPUs. Each prepared case uses 20 samples, a 0.5-second warmup, and 2 seconds of measurement; this VM is not a dedicated benchmark host. Negative percentages mean faster than rebuilding.

Key Build rows Consumers Rebuild us Cold us Cold change Warm us Warm change
int64 64 1 239.74 299.79 +25.0% 77.04 -67.9%
int64 64 4 318.18 286.04 -10.1% 462.28 +45.3%
int64 65,536 1 1435.26 1123.49 -21.7% 97.32 -93.2%
int64 65,536 4 5928.36 2051.71 -65.4% 373.74 -93.7%
utf8 64 1 191.26 204.31 +6.8% 177.76 -7.1%
utf8 64 4 393.39 417.90 +6.2% 427.65 +8.7%
utf8 65,536 1 1824.95 1906.01 +4.4% 279.39 -84.7%
utf8 65,536 4 2478.99 2092.53 -15.6% 475.86 -80.8%

The two large/four-consumer workloads were also measured in reverse mode order. Int64 measured 5.45 ms rebuild / 1.45 ms cold / 0.393 ms warm; UTF-8 measured 5.83 / 1.81 / 0.493 ms. Both still benefit from reuse, but absolute times varied. The small int64/four-consumer control measured 0.597 / 0.682 / 0.614 ms in reverse order, so it does not show a stable benefit; several small/cold cases above are slower.

The unchanged ordinary RightSemi benchmark was compared against base a522cd5 in base/candidate/candidate/base order. The initial perfect-map control was -13.1% and the general-hash control +6.0%. A longer reverse-order repeat of the general-hash control was +0.8%, with overlapping intervals and mixed adjacent comparisons. Both controls use 100,000 build rows and 1,000,000 matching probe rows. The first comparison used 40 samples and 5 seconds per case; the reverse comparison used 60 samples and 10 seconds. These controls do not establish a repeatable ordinary-path improvement or regression.

The benchmark records DataFusion pool reservations separately from timing. Those are accounting samples, not allocation/RSS measurements: shared input buffers can be charged once per ordinary join, concatenation/scratch accounting differs, and concurrent build lifetimes depend on scheduling. These are synthetic DataFusion execution-API results; they do not measure SQL end-to-end performance, CPU utilization, or Comet decoding, cache lookup, and single-flight waiting.

Ordinary control means in milliseconds (A is base, B is this branch):

Ordinary case Base A1 ms Candidate B1 ms Candidate B2 ms Base A2 ms Candidate change
right_semi_d100_h100 14.751 13.261 12.927 15.386 -13.09%
right_semi_d10_h100 26.697 26.292 28.302 24.815 +5.98%

The longer general-hash confirmation ran B/A/A/B, with means 23.266 / 22.174 / 22.906 / 22.169 ms.

Pool reservation samples in bytes:

Key Build rows Consumers Rebuild peak Cold peak Warm peak / retained
int64 64 1 3,512 4,024 3,512 / 3,512
int64 64 4 3,512 4,024 3,512 / 3,512
int64 65,536 1 3,539,000 4,671,544 3,539,000 / 3,539,000
int64 65,536 4 3,539,000 4,671,544 3,539,000 / 3,539,000
utf8 64 1 4,344 4,856 4,344 / 4,344
utf8 64 4 13,032 4,856 4,344 / 4,344
utf8 65,536 1 4,325,496 5,982,872 4,063,292 / 4,063,292
utf8 65,536 4 12,976,488 5,982,872 4,063,292 / 4,063,292

The warm column includes the retained build baseline. Cold and rebuild release their build after the group finishes. The Int64 four-consumer sample had the same peak as one consumer, illustrating that spawned consumers do not guarantee overlapping build lifetimes. These samples cannot establish a fourfold physical-memory saving.

Run the new benchmark with cargo bench -p datafusion-physical-plan --features test_utils --profile release-nonlto --bench prepared_hash_join.

Validation of 65f6a1d passed: formatting, full workspace Clippy with all targets/features and -D warnings, the repository extended-test command (11,995 Rust tests and 521 SQL test files; eight tests ignored), all 550 focused hash-join tests, the serialization regression, and the compiled API example. All 24 benchmark cases also passed their row-count and payload checks.

The local registry mirror lacked several requested releases. Validation used temporary source overrides for official Arrow 60.0.0 (ef1fa157), object_store 0.14.2 (279572ea), sqlparser 0.63.0 (85b1a6f2), and the published compression-codecs 0.4.42 crate, with Rust 1.98.1. These overrides and generated lockfile changes are not included in the PR. This replaces the earlier branch-55-only validation; CI on the newly pushed revision remains an independent check.

Are there any user-facing changes?

New opt-in Rust APIs for applications that construct physical plans, plus an EXPLAIN marker for attached builds. Ordinary SQL planning and join selection are unchanged; applications explicitly prepare and attach a build to reuse it.

@github-actions github-actions Bot added proto Related to proto crate physical-plan Changes to the physical-plan crate labels Sep 19, 2026
@codecov-commenter

codecov-commenter commented Sep 19, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 86.15752% with 58 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.37%. Comparing base (a522cd5) to head (65f6a1d).
⚠️ Report is 6 commits behind head on main.

Files with missing lines Patch % Lines
...tafusion/physical-plan/src/joins/hash_join/exec.rs 84.16% 13 Missing and 22 partials ⚠️
...physical-plan/src/joins/hash_join/exec/prepared.rs 88.26% 17 Missing and 6 partials ⚠️
Additional details and impacted files
@@           Coverage Diff            @@
##             main   #25491    +/-   ##
========================================
  Coverage   82.37%   82.37%            
========================================
  Files        1138     1139     +1     
  Lines      433506   434422   +916     
  Branches   433506   434422   +916     
========================================
+ Hits       357102   357874   +772     
- Misses      54850    54901    +51     
- Partials    21554    21647    +93     

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

@sunchao
sunchao marked this pull request as ready for review September 19, 2026 05:30
@sunchao

sunchao commented Sep 19, 2026

Copy link
Copy Markdown
Member Author

cc @andygrove @alamb @comphead @viirya @jayzhan211 could you take a look at this PR? thanks!

@comphead comphead 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 given the PR size, starting with AI review to shape the PR feedback

@comphead comphead 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.

Review

The cut between JoinBuildData (immutable, owns the reservation) and JoinLeftData (per-execution) is the right one, and the narrow allowlist plus explicit plan_err! is the right way to ship a first version of this API. Most of my comments are about reuse and traversal, not correctness. One memory-accounting finding is a pre-existing bug the PR fixes for itself only.

Verification I ran

cargo build -p datafusion-physical-plan --all-features succeeds on this head against current main (Arrow 60), and all 15 new prepared tests pass. The note in the description that local compilation is blocked on Arrow 60 does not reproduce here, so the draft can probably be un-drafted once CI runs. Clippy did not finish in my environment, so I make no claim there.

Benchmarks

None. datafusion/physical-plan/benches/hash_join_semi_anti.rs is the pattern to extend. Two numbers are worth having:

  1. Prepared attach versus rebuild-per-plan. This is the feature's entire justification and nothing in the PR measures it.
  2. The ordinary CollectLeft path before and after. collect_left_input's ingest loop was rewritten from try_fold to while let and gained per-batch branches, so the unchanged path needs a no-regression number.

Not applicable: .slt

These are opt-in Rust APIs with no planner, SQL or config surface, so sqllogictest cannot reach them. The Rust tests are in the right place.

Minor

exec.rs:1876: Arc::unwrap_or_clone(Arc::clone(&input_stats[1])) always deep-clones because the refcount is at least 2. It is input_stats[1].as_ref().clone() written the long way. Pre-existing, but the line moved in this PR.

Positive

  • Reservation lifetime is right, and the tests prove it: cancellation mid-preparation releases, and the pool returns to zero only on last-lease drop.
  • Rejecting serialization loudly instead of silently encoding an EmptyExec placeholder that would lose the build rows.
  • The internal_err! post-check on the concat bound turns an accounting breach into an error rather than an unaccounted allocation, which covers the case where Arrow stops short-circuiting single-array concat.

// `u64` indice variant
// Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown
if prepared {
// new_join_hashmap accounts for buckets but not its row-index chain.

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.

This is fixing a real accounting bug, but only for the prepared path. It should go into new_join_hashmap so every hash join gets it.

JoinHashMapU32::with_capacity(cap) allocates next: vec![0; cap] (joins/join_hash_map.rs:157-162, and JoinHashMapU64 at :237-242). new_join_hashmap (exec.rs:2856) only reserves estimate_memory_size::<(u32, u64)>(num_rows, size_of::<JoinHashMapU32>()), and that helper covers hashbrown buckets plus the struct itself and nothing else (datafusion/common/src/utils/memory.rs:99-120). There is no later try_grow for the chain: the only other map accounting in this file is metrics.build_mem_used.add(array_map.size()) on the ArrayMap branch.

So every ordinary CollectLeft and Partitioned build under-reserves num_rows * 4 bytes today, num_rows * 8 above u32::MAX. A 100M-row build side is 400 MB the pool never sees.

Moving the try_grow next to the existing bucket try_grow inside new_join_hashmap fixes it for all joins, deletes this if prepared block, and turns the hand-computed row_indices term in the tests into an assertion about shared code rather than about a prepared-only special case.

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.

Moved the row-index charge and build-memory metric into new_join_hashmap, so ordinary and prepared joins both account for the chain before allocation. I split the fix into #25508 for independent review/backporting; this PR includes the same change pending that prerequisite. The regression checks rejection one byte below the complete reservation, success at the limit, and release afterward.


/// Bound copy allocations, including validity, offsets and alignment. Aliased
/// columns count separately because concatenation materializes each column.
pub(super) fn prepared_copy_bytes(batch: &RecordBatch) -> Result<usize> {

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.

This reimplements ArrayData::get_slice_memory_size(), which is already wrapped for RecordBatch in this crate as pub(crate) trait GetSlicedSize (physical-plan/src/spill/spill_manager.rs:223-230), documented as exactly "bytes needed if we materialized exactly this slice into fresh buffers". It is already used by sorts/sort.rs:941, sorts/multi_level_merge.rs:700, spill/in_progress_spill_file.rs:111 and aggregates/grouped_hash_stream.rs:1271.

Arrow computes the same thing arm for arm (arrow-data 60, data.rs:587-656): Utf8 gives (len+1)*4 offsets plus offsets[len]-offsets[0] values plus ceil(len/8) nulls, Boolean gives ceil(len/8), fixed-width gives len*width, and NullArray gives 0 via BufferSpec::AlwaysNull. That is this match, line for line.

Reusing it collapses the body to batch.get_sliced_size()? + 64 * buffer_count for the alignment padding, and drops utf8_value_span. It also decouples the size function from the type allowlist in prepared_key_indices: when that later admits Binary or LargeUtf8, nothing here needs to change.

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.

Changed the estimator to reuse GetSlicedSize, the existing wrapper around Arrow's slice-size measurement. It adds alignment and validity bytes for inputs that have no physical null buffer, since another input can make concat materialize validity across all rows. Added a mixed-validity Boolean regression that compares the copy allowance with actual concatenated buffer capacity. The UTF-8 span helper remains for the separate per-column offset-overflow check.

for accumulator in accumulators {
accumulator.update_batch(&batch)?;
}
let mut concat_values = if prepared {

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.

The prepared path walks the batches three times and throws away the result of the first walk.

  1. check_byte_concat_sizes fills concat_values per batch during ingest (prepared.rs:226-240), and the totals are never read. The vector exists only for the i32 overflow side effect.
  2. exec.rs:2989 re-walks every batch and recomputes the identical utf8_value_span inside prepared_copy_bytes. By construction sum_over_batches(utf8_value_span(col)) == concat_values[col], so the byte spans are computed twice.
  3. exec.rs:3056 walks the batches a third time for max(num_rows).

All three fold into the ingest loop as a running max and a running per-column byte total, which yields both the overflow check and the copy bound from one traversal.

Also worth renaming: concat_values holds byte counts, not values.

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.

Copy-size totals and maximum batch rows are now tracked during ingestion, removing the later batch scans, and concat_values is renamed to concat_value_bytes. The cumulative byte totals still enforce each UTF-8 column's offset limit before concat. That check remains separate from Arrow's slice-size measurement, so the two checks still inspect UTF-8 offsets independently.

null_equality: NullEquality,
null_aware: Option<NullAwareMode>,
array_map_created_count: Count,
prepared: bool,

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.

A positional bool on a function that already carries #[expect(clippy::too_many_arguments)], now at 13 parameters, with both ordinary call sites passing a bare false (exec.rs:1750, exec.rs:1773). It gates seven separate branches in the body (2945, 2947, 2954, 2989, 3038, 3055, 3253).

Once the row-index accounting moves into new_join_hashmap and the scans collapse, what is left is the empty-batch placeholder, the copy admission and the scratch admission. Passing Option<&PreparedAdmission> or a small BuildMode enum instead names the intent, keeps the ordinary call sites from reading as a mystery false, and matches C-BOOL-ARG.

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.

Replaced the positional boolean with a private BuildMode and moved sizing arithmetic into helpers in prepared.rs. This keeps the call sites explicit without adding a separate admission object or lifecycle API.

let rows = batches.iter().map(RecordBatch::num_rows).max().unwrap_or(0);
// Combining nullable keys can hold an old and a new validity
// bitmap at once. NullArray also materializes logical validity.
let mask_count = if null_equality == NullEquality::NullEqualsNothing {

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.

This hardcodes update_hash's internal temporaries from the outside. It models matchable_join_keys (joins/utils.rs:2272-2290): among the allowed types only NullArray materializes in logical_nulls(), and NullBuffer::union_many holds at most two intermediates while folding. That is correct today, but nothing connects the two, and update_hash can grow a temporary without anything here failing.

The tests restate the same arithmetic (tests.rs:764: 2 * (1001usize.div_ceil(8) + 64), and tests.rs:796), so they confirm the formula was applied rather than that it is right.

Either drop a comment in matchable_join_keys pointing back here, or give up the precision and admit a flat on_left.len() + 1 masks. The saved bytes do not look worth the coupling to another function's body.

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.

Moved the scratch calculation into hash_scratch_bytes and added the cross-reference in matchable_join_keys. I retained the conservative bound rather than tying the estimate more closely to the current in-place bitmap union optimization.

/// prepared object is returned. Concurrent preparation/cache publication is
/// the caller's responsibility. Bounds and membership are prepared once, but
/// each consuming join publishes them into its own dynamic filter.
pub async fn prepare_build(

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.

prepare_build takes &self but the prepared data depends only on the left schema, the left key columns, null_equality, random_state and config. self.left is validated and never executed.

The result is that callers construct a throwaway join just to reach this method, then with_prepared_build replaces left with a second EmptyExec (prepared.rs:263). Both test helpers do exactly that (tests.rs:91-105, and proto/tests/cases/plans/joins.rs).

Not blocking, since the &self form gets validate compatibility checking for free. But the required call sequence (build a placeholder join, prepare, then attach to the real one) is not obvious from the signature and is not in the rustdoc. Either document it, or expose PreparedHashJoinBuild::try_new(schema, on, null_equality, config, pool, stream) and leave validate as the compatibility gate.

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 a prepare/attach rustdoc example using two independently planned compatible joins, and made the supplied-stream contract prominent. A throwaway join is not required: an existing join can prepare the stream and attach through its builder, as Comet does. I kept the current construction API for this revision.

} = self;

if prepared_build.is_some() {
return plan_err!("HashJoinExec with a prepared build cannot be serialized");

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.

not_impl_err! fits better than plan_err! here. The plan is valid, it just cannot be encoded, which is the "missing feature" case the repo reserves not_impl_err! for. plan_err! reads as though the user built an invalid plan.

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.

Changed this to not_impl_err! and strengthened the existing serialization regression to assert NotImplemented, in addition to the message.

}

#[tokio::test]
async fn prepared_build_accounts_for_hash_map_row_indices() -> 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.

This test, prepared_composite_keys_admit_hash_and_null_mask_scratch (:689), prepared_null_keys_admit_materialized_validity (:752) and the admission half of plain_bytes.rs are one test with four inputs: compute peak, assert peak - 1 yields ResourcesExhausted with reserved() == 0, assert peak succeeds with reserved() == peak.

Collapsing them into one table-driven case list, the way prepared_build_empty_and_all_null_inputs (:652) already loops, keeps the coverage and drops roughly 100 lines.

Separately, each of them recomputes the production formula (retained + buckets + row_indices + scratch) in the test body. A test that re-derives the implementation's arithmetic can only detect that the formula was applied, not that it is correct. Moving the row-index term into new_join_hashmap removes it from all four.

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.

Shared the repeated admission-rejection and reservation-release assertions. The cases retain their distinct checks: retained bytes excluding scratch, composite/null-key results, and UTF-8 alias/copy admission. In particular, the UTF-8 case rejects at concat admission and uses a separately funded success run; it is not another exact-peak success case. Keeping those checks local avoids combining different contracts into one large parameterized fixture.

}

#[tokio::test]
async fn prepared_build_reuses_data_with_independent_dynamic_filters() -> 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.

Three coverage gaps around the central claim of the PR, that a shared immutable build yields the same answers as a private one.

Equivalence. No test asserts that a prepared build produces the same output as an ordinary CollectLeft build over the same data. Every test here checks a row count or a hand-written expected batch. Running the same on, filter and probe through both paths and comparing sorted output is the one test that would catch a mis-shared JoinLeftData, and it is cheap.

Concurrency. This test uses tokio::join! under the default current-thread runtime, so the two plans never actually run in parallel. The production scenario is N executor tasks probing one build simultaneously. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] with several plans would exercise it.

Fuzz. datafusion/core/tests/fuzz_cases/join_fuzz.rs already exists. A prepared-versus-normal property over random build and probe data (nulls, duplicates, skewed keys, empty build, empty probe, varying batch sizes) is exactly what that infrastructure is for, and joins are on the repo's list of operators that warrant it.

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.

The existing independent-consumer test now uses a two-worker runtime, separately spawned SpawnedTask consumers, and a start barrier. Each result is compared with an ordinary CollectLeft join using the same inputs and residual filter. Changing runtime flavor alone would not parallelize tokio::join!. Added 24 fixed random seeds, each compared under both null-equality modes and output batch sizes 1 and 7: 96 differential comparisons covering duplicates/skew, empty inputs, residual filters, and varying input batch boundaries, using the existing rand dependency.

@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 , one suggestion for you

null_equality: NullEquality,
null_aware: Option<NullAwareMode>,
array_map_created_count: Count,
prepared: bool,

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.

prepared: bool forks collect_left_input into two accounting regimes: concat copy, row-index chain and hash/validity scratch are charged only when prepared, though the normal path allocates the same memory. ~110 gated lines inline in a 13-arg function. Move the sizing into prepared.rs helpers so the shared path keeps one if prepared per charge

-            let mut hashes_buffer = if prepared {
-                let rows = batches.iter().map(RecordBatch::num_rows).max().unwrap_or(0);
-                // ... mask/validity arithmetic ...
-                scratch_reservation.try_grow(bytes)?;
-                Vec::with_capacity(rows)
-            } else {
-                Vec::new()
-            };
+            let rows = batches.iter().map(RecordBatch::num_rows).max().unwrap_or(0);
+            if prepared {
+                scratch_reservation.try_grow(prepared::hash_scratch_bytes(
+                    rows,
+                    &on_left,
+                    &schema,
+                    null_equality,
+                )?)?;
+            }
+            let mut hashes_buffer = Vec::with_capacity(rows);

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.

Moved the scratch and copy sizing into prepared.rs helpers, accumulated copy bytes/max batch rows while ingesting, and removed the prepared-only row-chain charge in favor of the shared constructor fix. The ordinary and prepared paths now preallocate the hash buffer from the observed maximum batch size. Validation and the before/after ordinary-path benchmark are summarized in the updated PR description.

@viirya viirya left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks @sunchao — the core refactor here is clean. Splitting JoinLeftData into an immutable JoinBuildData plus per-execution state draws the line in exactly the right place: everything the probe phase writes (visited_indices_bitmap, probe_completion) stays per-execution, so there is no shared mutable state to get wrong. The restrictions are conservative and each one carries weight rather than being arbitrary — INNER avoids the match bitmap entirely, direct-column keys keep values a pointer clone so nothing escapes the memory accounting, and the type allowlist is what makes prepared_copy_bytes computable. Incompatible cases fail loudly instead of degrading silently, including serialization. Test coverage is well above average: asserting on buffer addresses for output ownership, and testing the memory pool at plus/minus one byte, are both the right instincts.

Two design questions that I think need answers before this lands, then smaller things inline.

1. Nothing ties a prepared build to the data it came from

prepare_build consumes a caller-supplied stream; self.left is only read for its schema and never executed. The stored fingerprint is schema + key indices + null equality. None of those can distinguish the intended snapshot from any other stream that happens to share a schema, so this compiles and runs clean:

let join = /* left = scan(users) */;
let prepared = join.prepare_build(stream_of_unrelated_data, pool, config).await?;

The result is a wrong answer with no error anywhere.

I understand this is deliberate — "the caller decides which inputs identify the same snapshot" — and I agree DataFusion should not try to prove two plans equivalent; that is expensive and would be wrong anyway, since the same scan at two points in time is legitimately two snapshots. But the asymmetry bothers me: this PR rejects a replaced build child, rejects serialization, and rejects a key-type mismatch, then leaves the one dimension where a mistake is silent completely unguarded.

Two suggestions, either of which would help:

  • Let the caller pass an opaque identity token (a u64 or Arc<str> that DataFusion stores and compares but never interprets). with_prepared_build then rejects a mismatched snapshot instead of returning stale rows. Cost is near zero and it converts the worst failure mode into a loud one.
  • Consider moving construction off HashJoinExec, e.g. PreparedHashJoinBuild::try_new(descriptor, stream, pool, config). join.prepare_build(stream, ..) reads as "prepare this join's build side", which is not what it does; a free constructor would not mislead.

At minimum the caller's responsibility for snapshot identity belongs at the top of the prepare_build docs rather than only in the PR description.

2. What does this actually save — latency or memory?

The motivating case is several Comet tasks probing one broadcast relation. Those tasks run in parallel today, so each builds its own table but the builds overlap in wall-clock. With a prepared build the caller has to serialize: prepare_build has no single-flight of its own ("concurrent preparation is the caller's responsibility"), so N tasks wait on the first build before any of them can probe.

That trades max(build_1..build_N) for build_1 plus N probes. If the build is single-threaded, wall-clock is roughly unchanged — what actually drops is CPU and, more importantly, the N-way duplication of the hash table in one executor's memory. That seems like the real win, and a good one for Spark-style broadcasts where duplicated build tables are a genuine OOM source. But the rationale reads as though this is about avoiding redundant work in general, which invites the wrong expectation.

Concretely:

  • Is the intended benefit memory, latency, or both? If memory, I would rewrite the rationale around that.
  • Is Comet's broadcast build single-threaded? If those builds are already parallel, the serialization a shared build introduces is a wall-clock regression unless memory pressure is the binding constraint.
  • Roughly where is the break-even build size? There is per-consumer overhead — each attached plan still goes through try_once, allocating a BoxFuture/Shared and a fresh JoinLeftData for a value that is already ready. For small builds that could plausibly cost more than just rebuilding the table, and the docs should steer people away from that case.

"No end-to-end performance result is claimed" is fine for a draft, but for a performance-only feature I do not think we can evaluate whether the design meets its goal without at least one number — even a micro-benchmark of N sequential joins over one build, prepared vs not, with the single-flight wait included.

3. EXPLAIN does not show that a build is attached

fmt_as is untouched, so a plan with a prepared build renders as:

HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(key@0, key@0)]
  EmptyExec
  <probe subtree>

Nothing distinguishes that from a join whose build side is genuinely empty — and an empty build on an INNER join means "no output rows", which is the opposite of what is happening. For a feature whose entire mechanism is "the build data lives outside the plan", the explain output has to say so. This is the one item here I would treat as required.

PreparedHashJoinBuild::num_rows() is already available, so something like:

let display_prepared = self
    .prepared_build
    .as_ref()
    .map_or_else(String::new, |p| format!(", prepared_build={} rows", p.num_rows()));

appended to the Default/Verbose write!, plus a line in TreeRender.

4. A stale comment outside the diff

HashTableLookupExpr's Hash/PartialEq in partitioned_hash_eval.rs (untouched here) justifies pointer equality with:

"HashJoinExec creates one per partition per query execution, thus it is never possible for two different hash maps to have the same content in practice."

This PR's whole point is sharing one Arc<Map> across independent joins and plans. Behaviour is still correct — pointer equality now implies content equality, which is the safe direction — but the stated premise is no longer true, and leaving it invites someone to rely on "different join implies different map". Worth a one-line update.


Nothing inline is a correctness blocker; I found no path that produces wrong results within the supported surface.

On CI: I see all 42 checks green on ff8fe3a, so the validation you were waiting on has landed. I could not build locally myself — my registry mirror does not carry Arrow 60, which is a restriction on my side and not a comment on this PR — so everything above is from reading the code against upstream/main @ a522cd5, not from running it.

"Prepared hash-join payload size overflow"
)
})?;
if retained > allowance {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

allowance is hand-computed from Arrow's current buffer layout and 64-byte alignment rounding. I worked through the multi-batch case and the margin looks sufficient (roughly 192 bytes of padding budget per batch, and 4*(rows_i+1) summed over two or more batches dominates 4*(total+1)), so I do not think this fires today.

But it binds us to Arrow's allocation behaviour. If a future arrow-rs changes capacity rounding — say to power-of-two growth — users hit a hard Internal error on a perfectly valid input, with a message asking them to report a bug.

Would you consider degrading instead of failing?

if retained > allowance {
    debug_assert!(false, "prepared concat exceeded admitted bound");
    reservation.try_grow(retained - allowance)?;   // let the pool decide
} else {
    reservation.shrink(allowance - retained);
}

That keeps the invariant loud in debug builds and in CI, while in release an estimate that drifts becomes "we reserved a bit more memory" rather than a failed query.

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.

I kept the invariant error here. The API promises admission before allocating the copy; growing the reservation afterward cannot cover the earlier peak overlap with the original batches. The post-check prevents publishing an undercharged build, but is not a substitute for a correct preflight bound. The sizing refactor now reuses Arrow's slice measurement and has an allocation-capacity regression for mixed physical validity.

partition: usize,
context: Arc<TaskContext>,
) -> Result<SendableRecordBatchStream> {
if let Some(prepared) = &self.prepared_build {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This looks redundant next to the validate in build(), and I expect someone will eventually delete it as dead weight. It is not redundant — it catches direct field mutation that bypasses the builder, which your own test exercises:

let mut mutated = attached.builder().build()?;
mutated.on = keys;

Worth a comment saying so, e.g. "fields such as on are writable within the crate, so re-check here: build() cannot be the only gate".

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 that explanation and retained the execute-time check. on is public, so callers can change it after builder validation; the existing direct-mutation test continues to cover this.

null_equality: NullEquality,
null_aware: Option<NullAwareMode>,
array_map_created_count: Count,
prepared: bool,

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

collect_left_input already carries #[expect(clippy::too_many_arguments)]; this makes it a 13th parameter and adds six if prepared sites scattered through the hottest function in the hash join. Some of them — the validity-mask scratch sizing, the row-index chain top-up — are subtle enough that I had to read them twice to confirm they do not affect the normal path.

Would you consider pulling the admission arithmetic into its own type, e.g. a PreparedAdmission holding concat_values / input_bytes / copy_bytes with observe_batch / admit_copies / settle, and passing Option<PreparedAdmission>? Six if prepared collapse to three if let Some(..), and the accounting math ends up somewhere it can be unit-tested on its own — which, given how carefully tuned it is, seems worth having.

Not a blocker, but this function is on every hash join's path and the review cost of the current shape is real.

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 the collector's readability with named sizing helpers, ingestion-time totals, and a small BuildMode enum. This removes the long inline arithmetic and the prepared-only row-chain top-up without adding an admission object with its own lifecycle.

/// object without retaining unaccounted cached payload. View, dictionary and
/// nested build columns remain unsupported. UTF-8 and fixed-size binary keys
/// use hash-table membership filters instead of copying range or IN-list values.
pub struct PreparedHashJoinBuild {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This is where an identity token would live if you take up point 1 in the main review — something like snapshot: Option<Arc<str>>, compared in validate and never interpreted. Flagging here so the thread hangs off the struct rather than only in the summary.

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.

I strengthened the method's caller contract and example, and kept snapshot identity/invalidation in the embedding executor for this revision. An opaque token could catch cache-selection mistakes if a consumer separately supplies its expected identity, but the caller would still be responsible for assigning that identity to the right stream. The current join has no independently known snapshot identity to compare against. The docs now lead with this responsibility and explicitly state that schema/key compatibility cannot establish input identity.

null_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)),
probe_completion: ProbeCompletion::new(probe_threads),
build_side_has_null: false,
_probe_reservation: self.build.reservation.new_empty(),

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

self.build.reservation.new_empty() inherits the pool passed to prepare_build — the caller's long-lived cache pool — not the consuming task's pool.

Harmless today: INNER joins allocate no bitmap, so this reservation stays at zero. But probe_data is exactly where outer-join support would be added, and at that point the visited bitmap would be charged to the cache pool instead of the task pool, which is the wrong lifetime. Either a comment marking the hazard, or having probe_data take the consumer's pool now (ignored while it is always zero), would keep the interface honest.

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 a short explanation that the reservation stays empty for the supported INNER joins. I did not add an unused consumer-pool parameter now; extending the API to outer joins will need to put mutable probe allocations in the consuming task's pool as part of that extension.

/// future cancellation, all work and reservations are dropped; no partially
/// prepared object is returned. Concurrent preparation/cache publication is
/// the caller's responsibility. Bounds and membership are prepared once, but
/// each consuming join publishes them into its own dynamic filter.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Two things about this doc block:

  1. The most important caller obligation — that snapshot identity is theirs to get right, and that getting it wrong yields silently wrong results — is not stated here at all. It only appears in the PR description. Could it lead the doc block?

  2. "range bounds and IN-list literals would allocate unaccounted key copies" reads as though only byte keys copy. InListExpr::try_new_from_array materialises one ScalarValue + lit() per row (not per distinct value) for every type, including Int64 — it is just bounded there by hash_join_inlist_pushdown_max_size (128 KB). The real reason to exclude byte keys is that their copy volume is unbounded, not that they are the only types that copy. Worth stating precisely, or a reader will assume the numeric path has no unaccounted allocation.

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.

Reworked these docs to lead with caller-owned snapshot identity and the supplied-stream behavior. They now also state that numeric IN-list publication allocates per-row literals outside the prepared reservation. The configured threshold measures input-array bytes, not the resulting expression heap, so it should not be read as a hard heap limit.

/// subtree is replaced with an empty schema placeholder so plan resets
/// preserve it. Probe-only rewrites must retain the attached join's `left()`.
/// Replacing that child or changing to incompatible join keys fails; other
/// incompatible join-mode/type changes fail validation in `build`.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Fourteen lines of prose carrying six distinct constraints. One of them is a correctness obligation:

A caller retaining a dynamic-filter expression beyond that plan must retain a prepared-build lease alongside it, because membership filters can reference build data.

That is real — PushdownStrategy::InList holds an Arc to the build key array and PushdownStrategy::Map holds the whole hash map — and it is currently line 10 of a paragraph. Could this be restructured, e.g. a # Caller contract section with the obligations as bullets, separated from the descriptive text? As written it is hard for a reader to confirm they have not missed one.

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.

Separated the attachment obligations into a Caller contract section, including fresh consumer filters, snapshot choice, attachment after build-child rewrites, and retaining a prepared-build lease when an external filter outlives the consuming plan. The lease preserves the corresponding memory charge; the filter's map/array Arcs preserve allocation lifetime on their own. Probe-only rewrites must preserve the attached join's build placeholder.

return plan_err!(
"Prepared hash-join builds require fixed-width or UTF-8 build columns"
);
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This allowlist and the match arms in prepared_copy_bytes cover exactly the same set — I checked each against DataType::primitive_width. Nothing enforces that, though. Adding LargeUtf8 here without touching prepared_copy_bytes would silently undercount memory rather than fail to compile. A cross-reference comment in both directions would make the coupling visible.

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.

The copy estimator now reuses Arrow's slice-size API, which removes the duplicated type match. The allowlist remains the eligibility gate. One detail about the earlier implementation: a newly admitted unsupported width would have returned an error, rather than silently undercounted.

@viirya

viirya commented Sep 19, 2026

Copy link
Copy Markdown
Member

Catching up on @comphead's and @jayzhan211's reviews, which landed while I was reading. I checked the two claims that bear on what I wrote, and both hold — with one caveat on the GetSlicedSize migration that I think matters.

new_join_hashmap under-reserves on main today (@comphead's note on exec.rs:3039) — confirmed against upstream/main @ a522cd5. JoinHashMapU32::with_capacity allocates next: vec![0; cap] (join_hash_map.rs:160), and estimate_memory_size covers only size_of::<T>() * buckets + buckets + fixed_size (common/src/utils/memory.rs:105-116) — no chain term, and no later try_grow for it. So every CollectLeft and Partitioned build under-reserves num_rows * 4 bytes, * 8 above u32::MAX, on main right now.

That reframes my B2 comment. I read the if prepared row-index try_grow as prepared-only bookkeeping and suggested softening the internal_err! around it. That was the wrong altitude: this is a pre-existing accounting bug that the prepared path happens to fix for itself. Moving it into new_join_hashmap fixes every hash join and deletes the special case. Worth splitting into its own PR so it can be backported independently of this feature.

prepared_copy_bytes duplicating ArrayData::get_slice_memory_size (prepared.rs:190) — also confirmed. arrow-data 59 data.rs:510-561 computes the same arms, and GetSlicedSize already wraps it for RecordBatch in this crate (spill/spill_manager.rs:222-240), used in four places. Agreed this should be reused.

One correction to the proposed substitution, though. batch.get_sliced_size()? + 64 * buffer_count is not equivalent to the current function on null handling:

// Arrow: validity counted only when the array actually has a null buffer
if self.nulls().is_some() {
    result += bit_util::ceil(self.len, 8);
}

// PR: validity counted unconditionally, for every column
.and_then(|bytes| bytes.checked_add(rows.div_ceil(8)))

The PR is deliberately conservative here, and it has to be: concat of a non-nullable-typed array with a nullable one materialises a validity buffer that none of the inputs had. Measuring the inputs with Arrow's rule would then under-admit the output. Any migration needs to keep the unconditional ceil(rows/8) term (or add it back per column), not just swap in get_sliced_size.

The current tests would not catch this: plain_bytes.rs:36-37 makes both byte columns nullable, and tests.rs:46-47 is Int64, where the values term dominates and the padding slack absorbs the difference. A mixed-nullability fixture would be worth adding alongside the refactor.

On the prepared: bool parameter — three of us landed on this independently, so I will not pile on further. @jayzhan211's sketch of pushing the sizing into prepared.rs helpers and keeping one if prepared per charge looks like the right shape, especially combined with @comphead's point that the row-index term leaves entirely once it moves into new_join_hashmap.

@comphead's equivalence-test gap (tests.rs:301) is the one I most want to endorse. No test currently asserts that a prepared build and an ordinary CollectLeft build produce the same output over the same data — every test checks a row count or a hand-written batch. That single property test is what would catch a mis-shared JoinLeftData, and it is directly the claim the PR is making. Same for running the plans under multi_thread; tokio::join! on the default current-thread runtime does not exercise the concurrent-probe scenario this exists for.

None of this changes my two design questions above — snapshot identity, and whether the win is latency or memory. Those are still the ones I think need answers before this lands.

@sunchao

sunchao commented Sep 19, 2026

Copy link
Copy Markdown
Member Author

Thanks @comphead, @jayzhan211, and @viirya. I pushed 65f6a1d and replied to the inline threads. The shared row-index accounting bug is now a separate prerequisite, #25508, and this branch includes the same fix. The collector now uses named sizing helpers and ingestion-time totals; EXPLAIN identifies attached builds, and the public docs lead with the snapshot and lifetime contracts.

The motivation is sharing a retained build and avoiding repeated build work across independent plans. In the companion Comet integration, one producer materializes and hashes the shared build while matching consumers wait; the consumers then probe independently and can run in parallel. It does not serialize all probes. Ordinary joins already use OnceFut/Shared and JoinLeftData, although attaching a prepared build still has setup costs. The new benchmark includes cold preparation before concurrent probes and warm reuse with an explicit retained lease. The large four-consumer cases still favored reuse in reverse-order runs, while small builds remained mixed. An initial +6.0% ordinary general-hash result became +0.8% in a longer reverse-order comparison with overlapping intervals; I am not treating those controls as evidence of a repeatable ordinary-path improvement or regression.

The updated description contains the full timing matrix and the ordinary-path before/after controls, including the small-build cases. These are DataFusion execution-API measurements, not Comet end-to-end measurements: decoding, cache coordination, waiter behavior, and actual process CPU/RSS are outside this benchmark. Pool reservations are reported with their scope and are not presented as physical-memory ratios or a universal break-even threshold.

I kept the existing construction API with a compiled example and a prominent supplied-stream/snapshot contract. I also kept the concat invariant error: reserving after a copy has already been allocated cannot admit the earlier peak. The mixed-validity regression checks actual Arrow buffer capacity, including the case where a null-containing input creates validity for rows from an input without a null bitmap.

Validation now passes locally against the requested Arrow 60 release sources: full workspace Clippy with all targets/features and warnings denied, the extended workspace run (11,995 Rust tests plus 521 SQL test files; eight tests ignored), 550 focused hash-join tests, the serialization regression, and the API doctest. Temporary local dependency-source overrides are documented in the description and are not in the patch. I removed the stale draft/Arrow-blocker text; CI on the new revision is separate from these local results.

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 proto Related to proto crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants