You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/26 15:15:24 UTC
(datafusion) branch gh-readonly-queue/main/pr-25491-6a792c671376c6c91f778530220358eb70a91271 created (now acf5c89e64)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-25491-6a792c671376c6c91f778530220358eb70a91271
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at acf5c89e64 feat: support reusable prepared hash-join builds (#25491)
This branch includes the following new commits:
new acf5c89e64 feat: support reusable prepared hash-join builds (#25491)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: feat: support reusable prepared hash-join builds (#25491)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25491-6a792c671376c6c91f778530220358eb70a91271
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit acf5c89e6454469c763c0fb5539e5a294d3439ed
Author: Chao Sun <su...@apple.com>
AuthorDate: Sat Sep 26 15:15:00 2026 +0000
feat: support reusable prepared hash-join builds (#25491)
## Which issue does this PR close?
Part of [Comet
#6013](https://github.com/apache/datafusion-comet/issues/6013). [Comet
#6037](https://github.com/apache/datafusion-comet/pull/6037) uses this
API. The shared row-index accounting fix in
[#25508](https://github.com/apache/datafusion/pull/25508) is included in
the base branch.
## Rationale for this change
Independent joins against the same build snapshot repeat preparation and
retain separate hash tables. `CollectLeft` shares a build among
partitions of one join, but cannot share it across independently created
plans.
For example, eight concurrent requests joining events to the same users
snapshot can share one users lookup while probing their own event
batches. Preparing that lookup once avoids seven redundant builds.
Applications can replace the snapshot while existing consumers finish
with the old one. This supports both Comet broadcast tasks and other
DataFusion applications.
Comet creates independent native joins for its executor tasks, so their
existing builds can overlap across task slots. The intended benefit is
sharing one retained build batch and hash table and avoiding repeated
materialization and hashing work. Removing repeated builds does not
imply a proportional latency improvement: cold consumers wait for
preparation, then probe independently. Latency depends on scheduling,
contention, build size, coordination, and whether a prepared build is
already retained. This PR does not establish end-to-end Comet latency,
process CPU-time, or RSS savings.
## What changes are included in this PR?
`HashJoinExec::prepare_build` returns an immutable
`PreparedHashJoinBuild` that compatible joins attach explicitly:
```rust
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()?;
```
Consumers share build rows, lookup storage, and the retained memory
reservation. Probe progress, residual predicates, and dynamic filters
remain independent. Preparation charges retained data and temporary
working memory; errors and cancellation release unfinished work.
Concatenation alignment is charged once per output buffer, so small
input batches do not multiply that reservation.
Supported builds use non-spilling `CollectLeft` inner joins,
direct-column keys with matching types, and fixed-width or UTF-8 build
columns. Attachment and execution check schema, keys, and null equality,
replaces the unused build subtree with a schema placeholder, and
preserves the handle through supported resets and projections. Prepared
state is process-local and cannot be serialized.
The caller owns snapshot identity, concurrent preparation, and eviction.
Input buffers must survive producer cleanup, and the supplied pool must
cover every consumer. Attach after physical optimization, supply fresh
probe plans and dynamic-filter expressions, and retain a build handle
while any associated filter outlives its plan. Compatibility checks
cannot verify snapshot contents.
## What is the testing strategy for this PR?
Fifteen focused tests cover equivalence with ordinary joins, both
null-equality modes, empty and duplicate keys, batch boundaries,
concurrent consumers with independent filters and predicates, plan
rewrites, public key changes after attachment, admission failures,
small-batch concatenation under a fixed memory limit, cancellation, and
buffer/reservation lifetimes. The latest regressions exercise
under-admission without shrinking the reservation, checked copy/scratch
arithmetic, and IN-list buffer sharing across one/two keys and one/two
input batches. Existing serialization coverage rejects prepared state;
the API example is a doctest. The benchmark checks outputs across 24
cases.
Prepared copy and scratch sizing retain checked arithmetic. A concat
output that exceeds its admitted bound returns an error in every build
profile before shrinking the reservation. Numeric IN-list membership
shares the retained batch's key buffers; charging its full array size
again would duplicate that payload charge. General array/schema metadata
and consumer-local filter expressions are outside the buffer
reservation.
Validation of
[8a964a88f4](https://github.com/apache/datafusion/commit/8a964a88f4a656a6c745b8bfd8c8e543a670ad51)
passed: all 15 focused prepared-build tests in the normal `ci` profile;
`cargo fmt --all`; full workspace Clippy with all targets/features and
`-D warnings`; the complete repository lint suite, including Rust and
HTML documentation builds; and the extended workspace suite (12,159 Rust
tests and all 524 SQL logic files, with 8 Rust tests ignored). All 15
prepared-build tests also passed in the optimized `release-nonlto`
profile.
The public-key-mutation regression was verified to fail before restoring
execution-time validation: the ordinary join returned its two expected
rows while the stale prepared lookup returned none. With the guard
restored, the mutated descriptor is rejected with a planning error.
Local validation used Rust 1.98.1 and temporary source overrides for the
official Arrow 60.0.0 (`ef1fa157`), object_store 0.14.2 (`279572ea`),
and sqlparser 0.63.0 (`85b1a6f2`) releases. Dependency overrides and
generated lockfile changes are not included.
The rebased head has not been benchmarked.
```bash
cargo bench -p datafusion-physical-plan --features test_utils --profile release-nonlto --bench prepared_hash_join
```
## Are there any user-facing changes?
Opt-in Rust APIs and an `EXPLAIN` marker for prepared builds. Ordinary
SQL planning and join selection are unchanged.
---
datafusion/physical-plan/Cargo.toml | 5 +
.../physical-plan/benches/prepared_hash_join.rs | 272 +++++++
.../physical-plan/src/joins/hash_join/exec.rs | 311 ++++---
.../src/joins/hash_join/exec/prepared.rs | 367 +++++++++
.../src/joins/hash_join/exec/prepared/tests.rs | 905 +++++++++++++++++++++
.../physical-plan/src/joins/hash_join/mod.rs | 2 +-
.../src/joins/hash_join/partitioned_hash_eval.rs | 17 +-
.../physical-plan/src/joins/hash_join/stream.rs | 4 +-
datafusion/physical-plan/src/joins/mod.rs | 3 +-
datafusion/proto/tests/cases/plans/joins.rs | 46 ++
10 files changed, 1817 insertions(+), 115 deletions(-)
diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml
index 85c64fe527..17e98a4764 100644
--- a/datafusion/physical-plan/Cargo.toml
+++ b/datafusion/physical-plan/Cargo.toml
@@ -177,6 +177,11 @@ harness = false
name = "hash_join_semi_anti"
required-features = ["test_utils"]
+[[bench]]
+harness = false
+name = "prepared_hash_join"
+required-features = ["test_utils"]
+
[[bench]]
harness = false
name = "multi_group_by"
diff --git a/datafusion/physical-plan/benches/prepared_hash_join.rs b/datafusion/physical-plan/benches/prepared_hash_join.rs
new file mode 100644
index 0000000000..113a2cf75b
--- /dev/null
+++ b/datafusion/physical-plan/benches/prepared_hash_join.rs
@@ -0,0 +1,272 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Independent INNER joins over the same in-memory build snapshot.
+//!
+//! SQL planning does not expose prepared builds, so these benchmarks exercise
+//! the public execution API. Cases vary build size, integer/string keys, and
+//! concurrent consumers. Every consumer probes 4096 rows with 50% matches.
+//! `cold` includes preparation; `warm` retains a build across iterations.
+//! Input generation is excluded, but plan construction and draining outputs are
+//! timed. This does not measure Comet decoding, its cache, or its single-flight
+//! wait. Sparse integer keys avoid perfect-hash selection.
+
+use std::hint::black_box;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray};
+use arrow::datatypes::{DataType, Field, Schema};
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_common::{JoinType, Result};
+use datafusion_common_runtime::SpawnedTask;
+use datafusion_execution::TaskContext;
+use datafusion_physical_expr::expressions::col;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::joins::{
+ HashJoinExec, HashJoinExecBuilder, PartitionMode, PreparedHashJoinBuild,
+};
+use datafusion_physical_plan::memory::MemoryStream;
+use datafusion_physical_plan::test::TestMemoryExec;
+use futures::TryStreamExt;
+use tokio::runtime::Builder;
+use tokio::sync::Barrier;
+
+const PROBE_ROWS: usize = 4096;
+
+struct Workload {
+ build: Vec<RecordBatch>,
+ probe: RecordBatch,
+ expected_sum: i64,
+}
+
+impl Workload {
+ fn new(build_rows: usize, strings: bool) -> Self {
+ let batch = |keys: Vec<usize>| {
+ let array: ArrayRef = if strings {
+ Arc::new(StringArray::from_iter_values(
+ keys.iter().map(|key| format!("key-{key:08}")),
+ ))
+ } else {
+ Arc::new(Int64Array::from_iter_values(
+ keys.iter().map(|key| (key * 17) as i64),
+ ))
+ };
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![
+ Field::new("key", array.data_type().clone(), false),
+ Field::new("payload", DataType::Int64, false),
+ ])),
+ vec![
+ array,
+ Arc::new(Int64Array::from_iter_values(
+ keys.iter().map(|key| *key as i64),
+ )),
+ ],
+ )
+ .unwrap()
+ };
+ let build = batch((0..build_rows).collect());
+ let probe = batch(
+ (0..PROBE_ROWS)
+ .map(|i| (i / 2) % build_rows + (i % 2) * build_rows)
+ .collect(),
+ );
+ Self {
+ build: (0..build_rows)
+ .step_by(8192)
+ .map(|offset| build.slice(offset, (build_rows - offset).min(8192)))
+ .collect(),
+ probe,
+ expected_sum: (0..PROBE_ROWS / 2).map(|i| (i % build_rows) as i64).sum(),
+ }
+ }
+
+ fn join(&self) -> Result<HashJoinExec> {
+ let schema = self.probe.schema();
+ HashJoinExecBuilder::new(
+ TestMemoryExec::try_new_exec(
+ std::slice::from_ref(&self.build),
+ Arc::clone(&schema),
+ None,
+ )?,
+ TestMemoryExec::try_new_exec(
+ &[vec![self.probe.clone()]],
+ Arc::clone(&schema),
+ None,
+ )?,
+ vec![(col("key", &schema)?, col("key", &schema)?)],
+ JoinType::Inner,
+ )
+ .with_partition_mode(PartitionMode::CollectLeft)
+ .build()
+ }
+
+ async fn prepare(
+ &self,
+ join: &HashJoinExec,
+ context: &TaskContext,
+ ) -> Result<Arc<PreparedHashJoinBuild>> {
+ join.prepare_build(
+ Box::pin(MemoryStream::try_new(
+ self.build.clone(),
+ self.probe.schema(),
+ None,
+ )?),
+ Arc::clone(context.memory_pool()),
+ Arc::clone(context.session_config().options()),
+ )
+ .await
+ }
+}
+
+#[derive(Clone, Copy)]
+enum Mode {
+ Rebuild,
+ Cold,
+ Warm,
+}
+
+impl Mode {
+ fn name(self) -> &'static str {
+ match self {
+ Self::Rebuild => "rebuild",
+ Self::Cold => "cold",
+ Self::Warm => "warm",
+ }
+ }
+}
+
+async fn run(
+ workload: &Workload,
+ consumers: usize,
+ mode: Mode,
+ warm: Option<Arc<PreparedHashJoinBuild>>,
+ context: Arc<TaskContext>,
+) -> Result<(usize, i64)> {
+ let joins = (0..consumers)
+ .map(|_| workload.join())
+ .collect::<Result<Vec<_>>>()?;
+ let prepared = match mode {
+ Mode::Rebuild => None,
+ Mode::Cold => Some(workload.prepare(&joins[0], &context).await?),
+ Mode::Warm => warm,
+ };
+ let start = Arc::new(Barrier::new(consumers));
+ let tasks = joins
+ .into_iter()
+ .map(|join| {
+ let join = match &prepared {
+ Some(build) => join
+ .builder()
+ .with_prepared_build(Arc::clone(build))
+ .build()?,
+ None => join,
+ };
+ let context = Arc::clone(&context);
+ let start = Arc::clone(&start);
+ Ok(SpawnedTask::spawn(async move {
+ start.wait().await;
+ let mut stream = join.execute(0, context)?;
+ let mut result = (0, 0);
+ while let Some(batch) = stream.try_next().await? {
+ result.0 += batch.num_rows();
+ let payload = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ result.1 += payload.values().iter().sum::<i64>();
+ }
+ Ok::<_, datafusion_common::DataFusionError>(result)
+ }))
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let mut result = (0, 0);
+ for task in tasks {
+ let (rows, sum) = task.await.unwrap()?;
+ result.0 += rows;
+ result.1 += sum;
+ }
+ Ok(result)
+}
+
+fn benchmark(c: &mut Criterion) {
+ let runtime = Builder::new_multi_thread()
+ .worker_threads(4)
+ .enable_all()
+ .build()
+ .unwrap();
+ let mut group = c.benchmark_group("prepared_hash_join");
+ group.sample_size(20);
+ group.warm_up_time(Duration::from_millis(500));
+ group.measurement_time(Duration::from_secs(2));
+
+ for (key_type, strings) in [("int64", false), ("utf8", true)] {
+ for rows in [64, 65536] {
+ let workload = Workload::new(rows, strings);
+ for consumers in [1, 4] {
+ for mode in [Mode::Rebuild, Mode::Cold, Mode::Warm] {
+ let context = Arc::new(TaskContext::default());
+ let warm = matches!(mode, Mode::Warm).then(|| {
+ runtime
+ .block_on(
+ workload.prepare(&workload.join().unwrap(), &context),
+ )
+ .unwrap()
+ });
+ let output = runtime
+ .block_on(run(
+ &workload,
+ consumers,
+ mode,
+ warm.clone(),
+ Arc::clone(&context),
+ ))
+ .unwrap();
+ assert_eq!(
+ output,
+ (
+ consumers * PROBE_ROWS / 2,
+ consumers as i64 * workload.expected_sum
+ )
+ );
+ let case = format!("{key_type}/{rows}/{consumers}_consumers");
+ group.bench_function(BenchmarkId::new(mode.name(), case), |b| {
+ b.iter(|| {
+ black_box(
+ runtime
+ .block_on(run(
+ &workload,
+ consumers,
+ mode,
+ warm.clone(),
+ Arc::clone(&context),
+ ))
+ .unwrap(),
+ )
+ });
+ });
+ }
+ }
+ }
+ }
+ group.finish();
+}
+
+criterion_group!(benches, benchmark);
+criterion_main!(benches);
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index c98e7b9b37..1bffdb1ab2 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -76,6 +76,7 @@ use arrow::record_batch::RecordBatch;
use arrow::util::bit_util;
use arrow_schema::{DataType, Schema};
use datafusion_common::config::ConfigOptions;
+use datafusion_common::stats::Precision;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
use datafusion_common::utils::memory::{
RecordBatchMemoryCounter, estimate_memory_size, get_record_batch_memory_size,
@@ -103,6 +104,9 @@ use parking_lot::Mutex;
use super::partitioned_hash_eval::SeededRandomState;
+mod prepared;
+pub use prepared::PreparedHashJoinBuild;
+
/// Hard-coded seed to ensure hash values from the hash join differ from `RepartitionExec`, avoiding collisions.
pub(crate) const HASH_JOIN_SEED: SeededRandomState =
SeededRandomState::with_seed(12210250226015887276);
@@ -288,11 +292,24 @@ impl NullAwareMode {
}
}
-/// HashTable and input data for the left (build side) of a join
+/// Immutable build buffers and their durable reservation.
+struct JoinBuildData {
+ /// Hash table with row indices into `batch`, also shared with dynamic filters.
+ map: Arc<Map>,
+ /// The input rows for the build side.
+ batch: RecordBatch,
+ /// Evaluated build-side key expressions.
+ values: Vec<ArrayRef>,
+ /// Bounds computed from the build side; absent for an empty partition.
+ bounds: Option<PartitionBounds>,
+ /// IN-list values or a hash-table reference used for filter pushdown.
+ membership: PushdownStrategy,
+ // Keep the reservation after the allocations it accounts for.
+ reservation: MemoryReservation,
+}
+
+/// A build lease and mutable bookkeeping for one join execution.
pub(super) struct JoinLeftData {
- /// The hash table with indices into `batch`
- /// Arc is used to allow sharing with SharedBuildAccumulator for hash map pushdown
- pub(super) map: Arc<Map>,
/// Hash table over correlated scope keys for correlated null-aware joins.
///
/// Key 0 is the scalar `NOT IN` value key and keys 1..N are correlated
@@ -303,37 +320,23 @@ pub(super) struct JoinLeftData {
/// The build rows whose value key is NULL (see [`NullValueBuildRows`]).
/// `None` when the build side has no NULL value keys.
null_value_build_rows: Option<NullValueBuildRows>,
- /// The input rows for the build side
- batch: RecordBatch,
- /// The build side on expressions values
- values: Vec<ArrayRef>,
- /// Shared bitmap builder for visited left indices
+ /// Shared bitmap for visited left indices within this consuming join.
visited_indices_bitmap: SharedBitmapBuilder,
- /// Shared bitmap builder for null marks
+ /// Shared bitmap for NULL marks within this consuming join.
null_indices_bitmap: SharedBitmapBuilder,
- /// Tracks which probe partition finishes last and what the partitions
- /// collectively saw. See [`ProbeCompletion`] for the invariant it upholds.
+ /// Tracks the final probe partition and what every partition observed.
+ /// See [`ProbeCompletion`] for the completion invariant.
probe_completion: ProbeCompletion,
- /// We need to keep this field to maintain accurate memory accounting, even though we don't directly use it.
- /// Without holding onto this reservation, the recorded memory usage would become inconsistent with actual usage.
- /// This could hide potential out-of-memory issues, especially when upstream operators increase their memory consumption.
- /// The MemoryReservation ensures proper tracking of memory resources throughout the join operation's lifecycle.
- _reservation: MemoryReservation,
- /// Bounds computed from the build side for dynamic filter pushdown.
- /// If the partition is empty (no rows) this will be None.
- /// If the partition has some rows this will be Some with the bounds for each join key column.
- pub(super) bounds: Option<PartitionBounds>,
- /// Membership testing strategy for filter pushdown
- /// Contains either InList values for small build sides or hash table reference for large build sides
- pub(super) membership: PushdownStrategy,
- // For RightAnti joins, where the build side is a smaller subquery, truthy if has null for the single join key
+ /// Whether the smaller RightAnti subquery contains a NULL key.
pub(super) build_side_has_null: bool,
+ // Drop the reservation after the mutable state it also accounts for.
+ build: Arc<JoinBuildData>,
}
impl JoinLeftData {
/// return a reference to the map
pub(super) fn map(&self) -> &Map {
- &self.map
+ &self.build.map
}
pub(super) fn null_aware_scope_map(&self) -> Option<&dyn JoinHashMapType> {
@@ -346,7 +349,7 @@ impl JoinLeftData {
/// returns a reference to the build side batch
pub(super) fn batch(&self) -> &RecordBatch {
- &self.batch
+ &self.build.batch
}
/// Returns `true` if the build side physically contains rows.
@@ -368,7 +371,7 @@ impl JoinLeftData {
/// returns a reference to the build side expressions values
pub(super) fn values(&self) -> &[ArrayRef] {
- &self.values
+ &self.build.values
}
/// returns a reference to the visited indices bitmap
@@ -382,7 +385,12 @@ impl JoinLeftData {
/// returns a reference to the InList values for filter pushdown
pub(super) fn membership(&self) -> &PushdownStrategy {
- &self.membership
+ &self.build.membership
+ }
+
+ /// Borrow bounds without sharing a consuming join's dynamic-filter state.
+ pub(super) fn bounds(&self) -> Option<&PartitionBounds> {
+ self.build.bounds.as_ref()
}
/// Records what a probe partition saw in one batch, for the null-aware
@@ -444,6 +452,7 @@ impl HashJoinExecBuilder {
filter: None,
join_type,
left_fut: Default::default(),
+ prepared_build: None,
random_state: HASH_JOIN_SEED,
mode: PartitionMode::Auto,
fetch: None,
@@ -535,13 +544,20 @@ impl HashJoinExecBuilder {
children.len() == 2,
"wrong number of children passed into `HashJoinExecBuilder`"
);
+ if self.exec.prepared_build.is_some()
+ && !Arc::ptr_eq(&self.exec.left, &children[0])
+ {
+ return plan_err!(
+ "Cannot replace the build child after attaching a prepared hash-join build"
+ );
+ }
self.preserve_properties &= has_same_children_properties(&self.exec, &children)?;
self.exec.right = children.swap_remove(1);
self.exec.left = children.swap_remove(0);
Ok(self)
}
- /// Reset runtime state.
+ /// Reset task-local runtime state while retaining the immutable prepared build.
pub fn reset_state(mut self) -> Self {
self.exec.left_fut = Default::default();
self.exec.dynamic_filter = None;
@@ -557,10 +573,18 @@ impl HashJoinExecBuilder {
/// Build resulting execution plan.
pub fn build(self) -> Result<HashJoinExec> {
let Self {
- exec,
+ mut exec,
preserve_properties,
} = self;
+ if let Some(prepared) = &exec.prepared_build {
+ prepared.validate(&exec)?;
+ exec.left_fut = Default::default();
+ exec.metrics = ExecutionPlanMetricsSet::new();
+ if let Some(filter) = &mut exec.dynamic_filter {
+ filter.build_accumulator = OnceLock::new();
+ }
+ }
// Validate null_aware flag
exec.null_aware_mode()?;
@@ -575,6 +599,7 @@ impl HashJoinExecBuilder {
filter,
join_type,
left_fut,
+ prepared_build,
random_state,
mode,
metrics,
@@ -622,6 +647,7 @@ impl HashJoinExecBuilder {
join_type,
join_schema,
left_fut,
+ prepared_build,
random_state,
mode,
metrics,
@@ -652,6 +678,7 @@ impl From<&HashJoinExec> for HashJoinExecBuilder {
join_type: exec.join_type,
join_schema: Arc::clone(&exec.join_schema),
left_fut: Arc::clone(&exec.left_fut),
+ prepared_build: exec.prepared_build.clone(),
random_state: exec.random_state.clone(),
mode: exec.mode,
metrics: exec.metrics.clone(),
@@ -881,6 +908,8 @@ pub struct HashJoinExec {
/// Each output stream waits on the `OnceAsync` to signal the completion of
/// the hash table creation.
left_fut: Arc<OnceAsync<JoinLeftData>>,
+ /// Immutable build attached by an embedding executor.
+ prepared_build: Option<Arc<PreparedHashJoinBuild>>,
/// Shared the `SeededRandomState` for the hashing algorithm (seeds preserved for serialization)
random_state: SeededRandomState,
/// Partitioning mode to use
@@ -939,6 +968,7 @@ impl fmt::Debug for HashJoinExec {
.field("join_type", &self.join_type)
.field("join_schema", &self.join_schema)
.field("left_fut", &self.left_fut)
+ .field("prepared_build", &self.prepared_build)
.field("random_state", &self.random_state)
.field("mode", &self.mode)
.field("null_aware", &self.null_aware)
@@ -1457,6 +1487,12 @@ impl DisplayAs for HashJoinExec {
.map_or_else(String::new, |f| format!(", fetch={f}"));
let display_null_aware =
if self.null_aware { ", null_aware" } else { "" };
+ let display_prepared = self
+ .prepared_build
+ .as_ref()
+ .map_or_else(String::new, |prepared| {
+ format!(", prepared_build={} rows", prepared.num_rows())
+ });
let on = self
.on
.iter()
@@ -1465,7 +1501,7 @@ impl DisplayAs for HashJoinExec {
.join(", ");
write!(
f,
- "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}{}",
+ "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}{}{}",
self.mode,
self.join_type,
on,
@@ -1474,6 +1510,7 @@ impl DisplayAs for HashJoinExec {
display_null_equality,
display_fetch,
display_null_aware,
+ display_prepared,
)
}
DisplayFormatType::TreeRender => {
@@ -1492,6 +1529,10 @@ impl DisplayAs for HashJoinExec {
writeln!(f, "on={on}")?;
+ if let Some(prepared) = &self.prepared_build {
+ writeln!(f, "prepared_build={} rows", prepared.num_rows())?;
+ }
+
if self.null_equality() == NullEquality::NullEqualsNull {
writeln!(f, "NullsEqual: true")?;
}
@@ -1708,8 +1749,16 @@ impl ExecutionPlan for HashJoinExec {
let null_aware = self.null_aware_mode()?;
- let left_fut = match self.mode {
- PartitionMode::CollectLeft => self.left_fut.try_once(|| {
+ let left_fut = match (&self.prepared_build, self.mode) {
+ (Some(prepared), _) => {
+ // Public join fields can change after builder validation.
+ prepared.validate(self)?;
+ let prepared = Arc::clone(prepared);
+ self.left_fut.try_once(|| {
+ Ok(async move { Ok(prepared.probe_data(right_partitions)) })
+ })?
+ }
+ (None, PartitionMode::CollectLeft) => self.left_fut.try_once(|| {
let left_stream = self.left.execute(0, Arc::clone(&context))?;
let reservation =
@@ -1722,15 +1771,16 @@ impl ExecutionPlan for HashJoinExec {
join_metrics.clone(),
reservation,
need_produce_result_in_final(self.join_type),
- self.right().output_partitioning().partition_count(),
+ right_partitions,
enable_dynamic_filter_pushdown,
Arc::clone(context.session_config().options()),
self.null_equality,
null_aware,
array_map_created_count,
+ BuildMode::Ordinary,
))
})?,
- PartitionMode::Partitioned => {
+ (None, PartitionMode::Partitioned) => {
let left_stream = self.left.execute(partition, Arc::clone(&context))?;
let reservation =
@@ -1749,9 +1799,10 @@ impl ExecutionPlan for HashJoinExec {
self.null_equality,
null_aware,
array_map_created_count,
+ BuildMode::Ordinary,
))
}
- PartitionMode::Auto => {
+ (None, PartitionMode::Auto) => {
return plan_err!(
"Invalid HashJoinExec, unsupported PartitionMode {:?} in execute()",
PartitionMode::Auto
@@ -1834,11 +1885,17 @@ impl ExecutionPlan for HashJoinExec {
input_stats: &[Arc<Statistics>],
_args: &StatisticsArgs,
) -> Result<Arc<Statistics>> {
- let left_stats = Arc::clone(&input_stats[0]);
- let right_stats = Arc::clone(&input_stats[1]);
+ let left_stats = if let Some(prepared) = &self.prepared_build {
+ // The declared left child supplies the schema but is never executed.
+ Statistics::new_unknown(self.left.schema().as_ref())
+ .with_num_rows(Precision::Exact(prepared.num_rows()))
+ } else {
+ input_stats[0].as_ref().clone()
+ };
+ let right_stats = input_stats[1].as_ref().clone();
let stats = estimate_join_statistics(
- Arc::unwrap_or_clone(left_stats),
- Arc::unwrap_or_clone(right_stats),
+ left_stats,
+ right_stats,
&self.on,
self.null_equality,
&self.join_type,
@@ -1861,6 +1918,9 @@ impl ExecutionPlan for HashJoinExec {
if self.contains_projection() {
return Ok(None);
}
+ if self.prepared_build.is_some() {
+ return try_embed_projection(projection, self);
+ }
let schema = self.schema();
if let Some(JoinData {
@@ -2085,6 +2145,7 @@ impl ExecutionPlan for HashJoinExec {
join_schema: _,
// runtime build-side state, not part of the plan
left_fut: _,
+ prepared_build,
// the fixed `HASH_JOIN_SEED` constant, set identically by the
// builder on decode
random_state: _,
@@ -2096,6 +2157,11 @@ impl ExecutionPlan for HashJoinExec {
cache: _,
} = self;
+ if prepared_build.is_some() {
+ return datafusion_common::not_impl_err!(
+ "HashJoinExec with a prepared build cannot be serialized"
+ );
+ }
let left = ctx.encode_child(left)?;
let right = ctx.encode_child(right)?;
@@ -2750,6 +2816,12 @@ impl CollectLeftAccumulator {
}
}
+#[derive(PartialEq, Eq)]
+enum BuildMode {
+ Ordinary,
+ Prepared,
+}
+
/// State for collecting the build-side data during hash join
struct BuildSideState {
batches: Vec<RecordBatch>,
@@ -2931,8 +3003,8 @@ fn concat_build_batches(
/// * `with_visited_indices_bitmap` - Whether to track visited indices (for outer joins)
/// * `probe_threads_count` - Number of threads that will probe this hash table
/// * `should_compute_dynamic_filters` - Whether to compute min/max bounds for dynamic filtering
-/// * `with_null_aware_mark_state` - Whether to build the per-build-row null-indices bitmap
-/// and correlation-scope maps used by correlated null-aware `LeftMark` joins
+/// * `null_aware` - The null-aware join mode, including whether per-build-row
+/// bitmaps and correlation-scope maps are required
///
/// # Memory Accounting
/// Build batches are added to `reservation` as they arrive. They are then copied
@@ -2954,7 +3026,7 @@ fn concat_build_batches(
#[expect(clippy::too_many_arguments)]
async fn collect_left_input(
random_state: RandomState,
- left_stream: SendableRecordBatchStream,
+ mut left_stream: SendableRecordBatchStream,
on_left: Vec<PhysicalExprRef>,
metrics: BuildProbeJoinMetrics,
reservation: MemoryReservation,
@@ -2965,8 +3037,10 @@ async fn collect_left_input(
null_equality: NullEquality,
null_aware: Option<NullAwareMode>,
array_map_created_count: Count,
+ mode: BuildMode,
) -> Result<JoinLeftData> {
let schema = left_stream.schema();
+ let prepared = mode == BuildMode::Prepared;
// The extra scope maps + null bitmap are only built for correlated
// null-aware joins (see `NullAwareMode`).
@@ -2974,7 +3048,7 @@ async fn collect_left_input(
let is_phj_candidate = is_perfect_hash_join_candidate(&on_left, &schema)?;
- let initial = BuildSideState::try_new(
+ let mut state = BuildSideState::try_new(
metrics,
reservation,
on_left.clone(),
@@ -2982,41 +3056,52 @@ async fn collect_left_input(
should_compute_dynamic_filters || is_phj_candidate,
)?;
- let state = left_stream
- .try_fold(initial, |mut state, batch| async move {
- // Update accumulators if computing bounds
- if let Some(ref mut accumulators) = state.bounds_accumulators {
- for accumulator in accumulators {
- accumulator.update_batch(&batch)?;
- }
+ let mut max_batch_rows = 0;
+ while let Some(batch) = left_stream.try_next().await? {
+ max_batch_rows = max_batch_rows.max(batch.num_rows());
+ if let Some(accumulators) = &mut state.bounds_accumulators {
+ for accumulator in accumulators {
+ accumulator.update_batch(&batch)?;
}
-
- // Decide if we spill or not
- let batch_size = state.memory_counter.count_batch(&batch);
- // Reserve memory for incoming batch
- state.reservation.try_grow(batch_size)?;
- // Update metrics
- state.metrics.build_mem_used.add(batch_size);
- state.metrics.build_input_batches.add(1);
- state.metrics.build_input_rows.add(batch.num_rows());
- // Update row count
- state.num_rows += batch.num_rows();
- // Push batch to output
- state.batches.push(batch);
- Ok(state)
- })
- .await?;
-
- // Extract fields from state
+ }
+ let batch_size = state.memory_counter.count_batch(&batch);
+ state.reservation.try_grow(batch_size)?;
+ state.metrics.build_mem_used.add(batch_size);
+ state.metrics.build_input_batches.add(1);
+ state.metrics.build_input_rows.add(batch.num_rows());
+ state.num_rows += batch.num_rows();
+ state.batches.push(batch);
+ }
+ drop(left_stream);
+ if prepared && state.batches.is_empty() {
+ // Even empty UTF-8 output retains an offset buffer.
+ let empty = RecordBatch::new_empty(Arc::clone(&schema));
+ state
+ .reservation
+ .try_grow(state.memory_counter.count_batch(&empty))?;
+ state.batches.push(empty);
+ }
+ let input_bytes = state.memory_counter.memory_usage();
+
+ // Bind the reservation first so error paths release allocations before their charge.
let BuildSideState {
- batches,
+ mut reservation,
+ mut batches,
num_rows,
metrics,
- mut reservation,
bounds_accumulators,
- memory_counter,
+ memory_counter: _,
} = state;
- let inputs_reserved = memory_counter.memory_usage();
+
+ // Admit concatenation copies while the original batches are retained.
+ // Arrow keeps a single batch as an inexpensive slice.
+ let copy_bytes = if prepared && batches.len() > 1 {
+ let copy_bytes = prepared::prepared_copy_bytes(&batches)?;
+ reservation.try_grow(copy_bytes)?;
+ copy_bytes
+ } else {
+ 0
+ };
// Compute bounds
let mut bounds = match bounds_accumulators {
@@ -3041,14 +3126,18 @@ async fn collect_left_input(
config.execution.perfect_hash_join_min_key_density,
null_equality,
)? {
- let batch = concat_build_batches(
- &schema,
- batches,
- false,
- inputs_reserved,
- &mut reservation,
- &metrics,
- )?;
+ let batch = if prepared {
+ concat_batches(&schema, batches.iter())?
+ } else {
+ concat_build_batches(
+ &schema,
+ std::mem::take(&mut batches),
+ false,
+ input_bytes,
+ &mut reservation,
+ &metrics,
+ )?
+ };
let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?;
@@ -3064,7 +3153,16 @@ async fn collect_left_input(
// Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown
let mut hashmap = new_join_hashmap(num_rows, &mut reservation, &metrics)?;
- let mut hashes_buffer = Vec::new();
+ let scratch_reservation = reservation.new_empty();
+ if prepared {
+ scratch_reservation.try_grow(prepared::prepared_scratch_bytes(
+ max_batch_rows,
+ on_left.len(),
+ null_equality,
+ )?)?;
+ }
+ // The maximum is known: avoid geometric growth and its excess capacity.
+ let mut hashes_buffer = Vec::with_capacity(max_batch_rows);
let mut offset = 0;
// Updating hashmap starting from the last batch
@@ -3086,14 +3184,18 @@ async fn collect_left_input(
}
// Merge all batches into a single batch, so we can directly index into the arrays
- let batch = concat_build_batches(
- &schema,
- batches,
- true,
- inputs_reserved,
- &mut reservation,
- &metrics,
- )?;
+ let batch = if prepared {
+ concat_batches(&schema, batches.iter().rev())?
+ } else {
+ concat_build_batches(
+ &schema,
+ std::mem::take(&mut batches),
+ true,
+ input_bytes,
+ &mut reservation,
+ &metrics,
+ )?
+ };
let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
@@ -3245,18 +3347,33 @@ async fn collect_left_input(
&& !left_values.is_empty()
&& left_values[0].logical_null_count() > 0;
+ if prepared {
+ drop(batches);
+ // Prepared keys are direct columns. IN-list arrays share these batch
+ // buffers, including through multi-key StructArray children.
+ let retained = RecordBatchMemoryCounter::new().count_batch(&batch);
+ prepared::reconcile_prepared_copy_reservation(
+ &reservation,
+ input_bytes,
+ copy_bytes,
+ retained,
+ )?;
+ }
+
let data = JoinLeftData {
- map,
+ build: Arc::new(JoinBuildData {
+ map,
+ batch,
+ values: left_values,
+ bounds,
+ membership,
+ reservation,
+ }),
null_aware_scope_map,
null_value_build_rows,
- batch,
- values: left_values,
visited_indices_bitmap: Mutex::new(visited_indices_bitmap),
null_indices_bitmap: Mutex::new(null_indices_bitmap),
probe_completion: ProbeCompletion::new(probe_threads_count),
- _reservation: reservation,
- bounds,
- membership,
build_side_has_null: build_has_null,
};
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs
new file mode 100644
index 0000000000..fb46cedf96
--- /dev/null
+++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs
@@ -0,0 +1,367 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Explicit immutable build reuse for embedding executors.
+
+use super::*;
+use crate::spill::spill_manager::GetSlicedSize;
+use datafusion_common::exec_datafusion_err;
+use datafusion_execution::memory_pool::MemoryPool;
+
+/// An immutable, fully prepared broadcast build, independent of any probe task.
+///
+/// Created by [`HashJoinExec::prepare_build`]. The embedding executor owns cache
+/// identity, admission, single-flight coordination, cancellation and eviction.
+/// This object retains its input buffers and memory reservation until its last
+/// lease is dropped; it never retains an input stream or task context. Prepared
+/// builds support fixed-width and UTF-8 build columns, with direct-column keys
+/// and non-spilling INNER joins. Residual conditions belong to each consuming
+/// join; null-aware joins remain unsupported.
+///
+/// Hash-join gathers copy supported build columns into output buffers,
+/// including contiguous selections. Output batches can therefore outlive this
+/// 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 {
+ build: Arc<JoinBuildData>,
+ keys: Vec<usize>,
+ null_equality: NullEquality,
+}
+
+impl fmt::Debug for PreparedHashJoinBuild {
+ /// Describe immutable metadata without dumping table contents.
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("PreparedHashJoinBuild")
+ .field("schema", &self.build.batch.schema())
+ .field("keys", &self.keys)
+ .field("rows", &self.num_rows())
+ .field("reserved_bytes", &self.reserved_bytes())
+ .finish()
+ }
+}
+
+impl PreparedHashJoinBuild {
+ /// Return the retained build reservation, excluding all per-probe state.
+ pub fn reserved_bytes(&self) -> usize {
+ self.build.reservation.size()
+ }
+
+ /// Return the complete build row count, including duplicate and null keys.
+ pub fn num_rows(&self) -> usize {
+ self.build.batch.num_rows()
+ }
+
+ /// Create independent mutable state for one consuming join.
+ pub(super) fn probe_data(&self, probe_threads: usize) -> JoinLeftData {
+ JoinLeftData {
+ build: Arc::clone(&self.build),
+ null_aware_scope_map: None,
+ null_value_build_rows: None,
+ visited_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)),
+ null_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)),
+ probe_completion: ProbeCompletion::new(probe_threads),
+ build_side_has_null: false,
+ }
+ }
+
+ /// Validate the build descriptor and current execution restrictions without
+ /// consuming input or modifying either plan. Cache identity is caller-owned.
+ pub(super) fn validate(&self, join: &HashJoinExec) -> Result<()> {
+ let keys = prepared_key_indices(join)?;
+ if join.left.schema() != self.build.batch.schema()
+ || keys != self.keys
+ || join.null_equality != self.null_equality
+ {
+ return plan_err!(
+ "Prepared hash-join build does not match schema, keys or null equality"
+ );
+ }
+ if let Some(filter) = &join.dynamic_filter {
+ let probe_keys = join.on.iter().map(|(_, right)| right);
+ if !filter.filter.children().into_iter().eq(probe_keys) {
+ return plan_err!(
+ "Prepared hash-join dynamic filter keys do not match probe keys"
+ );
+ }
+ }
+ Ok(())
+ }
+}
+
+impl HashJoinExec {
+ /// Prepare the supplied build snapshot using an embedding executor's durable pool.
+ ///
+ /// The caller owns snapshot identity: this consumes `input`, never `self.left`.
+ /// Schema and key compatibility cannot establish that two inputs contain the
+ /// same data. Attaching the wrong snapshot can silently change query results.
+ ///
+ /// # Caller contract
+ ///
+ /// * Input buffers must remain valid independently of producer task cleanup.
+ /// Keep producer allocations charged until the stream releases them.
+ /// * Supply a pool whose lifetime covers every consumer. It accounts for
+ /// retained input, hash storage, and preparation's copy/scratch buffers.
+ /// * Coordinate concurrent preparation, publication, and invalidation in the
+ /// embedding executor. This method does not provide a cache.
+ ///
+ /// Errors and cancellation release unfinished work. `config` controls
+ /// perfect-map and dynamic-filter choices. Byte keys use hash membership
+ /// to avoid copying payloads into range bounds and IN-list literals.
+ /// Consumer-local dynamic-filter allocations are outside this reservation.
+ ///
+ /// # Example
+ ///
+ /// Given two independently planned, compatible joins over different probe
+ /// inputs, a service can prepare one dimension snapshot for both consumers.
+ ///
+ /// ```no_run
+ /// # use std::sync::Arc;
+ /// # use datafusion_common::{Result, config::ConfigOptions};
+ /// # use datafusion_execution::memory_pool::MemoryPool;
+ /// # use datafusion_physical_plan::SendableRecordBatchStream;
+ /// # use datafusion_physical_plan::joins::HashJoinExec;
+ /// # async fn example(first_join: HashJoinExec, second_join: HashJoinExec,
+ /// # dimension_snapshot: SendableRecordBatchStream,
+ /// # durable_pool: Arc<dyn MemoryPool>, config: Arc<ConfigOptions>) -> Result<()> {
+ /// let prepared = first_join
+ /// .prepare_build(dimension_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()?;
+ /// # let _ = (first, second);
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub async fn prepare_build(
+ &self,
+ input: SendableRecordBatchStream,
+ pool: Arc<dyn MemoryPool>,
+ config: Arc<ConfigOptions>,
+ ) -> Result<Arc<PreparedHashJoinBuild>> {
+ let keys = prepared_key_indices(self)?;
+ let schema = self.left.schema();
+ if input.schema() != schema {
+ return plan_err!(
+ "Prepared hash-join input schema does not match build schema"
+ );
+ }
+ let byte_keys = keys.iter().any(|&key| {
+ matches!(
+ schema.field(key).data_type(),
+ DataType::Utf8 | DataType::FixedSizeBinary(_)
+ )
+ });
+ // Range accumulation and IN-list publication materialize ScalarValue
+ // copies of byte keys. Hash membership borrows the admitted table instead.
+ let config = if byte_keys {
+ let mut config = config.as_ref().clone();
+ config.optimizer.hash_join_inlist_pushdown_max_size = 0;
+ Arc::new(config)
+ } else {
+ config
+ };
+ let metrics_set = ExecutionPlanMetricsSet::new();
+ let metrics = BuildProbeJoinMetrics::new(0, &metrics_set);
+ let count = MetricBuilder::new(&metrics_set)
+ .counter(ARRAY_MAP_CREATED_COUNT_METRIC_NAME, 0);
+ let reservation = MemoryConsumer::new("PreparedHashJoinBuild").register(&pool);
+ let data = collect_left_input(
+ self.random_state.random_state().clone(),
+ input,
+ self.on.iter().map(|(left, _)| Arc::clone(left)).collect(),
+ metrics,
+ reservation,
+ false,
+ 0,
+ !byte_keys,
+ config,
+ self.null_equality,
+ None,
+ count,
+ BuildMode::Prepared,
+ )
+ .await?;
+ Ok(Arc::new(PreparedHashJoinBuild {
+ build: data.build,
+ keys,
+ null_equality: self.null_equality,
+ }))
+ }
+}
+
+/// Bound copy allocations, including validity, offsets and alignment. Aliased
+/// columns count separately because concatenation materializes each column.
+pub(super) fn prepared_copy_bytes(batches: &[RecordBatch]) -> Result<usize> {
+ let overflow = || exec_datafusion_err!("Prepared hash-join copy size overflow");
+ let mut bytes = 0usize;
+ for batch in batches {
+ // Concat can materialize validity for an all-valid input when another input
+ // contains nulls. Arrow's slice measurement only counts existing bitmaps.
+ let missing_validity = batch
+ .columns()
+ .iter()
+ .filter(|array| array.nulls().is_none())
+ .count();
+ let validity_bytes = batch
+ .num_rows()
+ .div_ceil(8)
+ .checked_mul(missing_validity)
+ .ok_or_else(overflow)?;
+ bytes = bytes
+ .checked_add(batch.get_sliced_size()?)
+ .and_then(|bytes| bytes.checked_add(validity_bytes))
+ .ok_or_else(overflow)?;
+ }
+ // These are flat arrays: allow one alignment unit per output buffer,
+ // including potential validity. Concat allocates each buffer only once.
+ let buffers = batches[0]
+ .columns()
+ .iter()
+ .try_fold(0usize, |buffers, array| {
+ buffers
+ .checked_add(array.to_data().buffers().len())
+ .and_then(|buffers| buffers.checked_add(1))
+ .ok_or_else(overflow)
+ })?;
+ buffers
+ .checked_mul(64)
+ .and_then(|padding| bytes.checked_add(padding))
+ .ok_or_else(overflow)
+}
+
+/// Bound hashes and flat-key null masks for the largest input batch.
+pub(super) fn prepared_scratch_bytes(
+ max_batch_rows: usize,
+ num_keys: usize,
+ null_equality: NullEquality,
+) -> Result<usize> {
+ let overflow = || exec_datafusion_err!("Prepared hash-join scratch size overflow");
+ let hash_bytes = max_batch_rows
+ .checked_mul(size_of::<u64>())
+ .ok_or_else(overflow)?;
+ if null_equality == NullEquality::NullEqualsNull {
+ return Ok(hash_bytes);
+ }
+ // Allow one logical null mask per key plus the combined mask.
+ let masks = num_keys.checked_add(1).ok_or_else(overflow)?;
+ let mask_bytes = max_batch_rows
+ .div_ceil(8)
+ .checked_add(64)
+ .and_then(|bytes| bytes.checked_mul(masks))
+ .ok_or_else(overflow)?;
+ hash_bytes.checked_add(mask_bytes).ok_or_else(overflow)
+}
+
+/// Release the unused copy allowance only after validating the retained size.
+pub(super) fn reconcile_prepared_copy_reservation(
+ reservation: &MemoryReservation,
+ input_bytes: usize,
+ copy_bytes: usize,
+ retained: usize,
+) -> Result<()> {
+ let allowance = input_bytes.checked_add(copy_bytes).ok_or_else(|| {
+ exec_datafusion_err!("Prepared hash-join copy allowance size overflow")
+ })?;
+ if retained > allowance {
+ return internal_err!(
+ "Prepared hash-join concat exceeded its admitted copy bound"
+ );
+ }
+ reservation.shrink(allowance - retained);
+ Ok(())
+}
+
+impl HashJoinExecBuilder {
+ /// Attach a fully prepared build to a fresh, compatible join execution.
+ ///
+ /// [`Self::build`] checks schema, keys, null equality, and supported join modes.
+ /// Attaching resets execution state and replaces the unused left subtree with
+ /// an empty schema placeholder. Residual predicates remain consumer-local.
+ ///
+ /// # Caller contract
+ ///
+ /// * Choose the correct build snapshot; compatibility checks do not verify
+ /// input identity. See [`HashJoinExec::prepare_build`].
+ /// * Supply fresh dynamic-filter expressions and probe plans for independent
+ /// executions. An existing expression handle is preserved for its probe
+ /// consumers, while the build-report accumulator is reset.
+ /// * Retain a prepared-build lease if a dynamic-filter expression outlives
+ /// this plan. Its array/map references preserve allocation lifetime, but
+ /// the lease is what preserves the corresponding memory charge.
+ /// * Attach after physical optimizations that rewrite the build child.
+ /// Probe-only rewrites must preserve the attached join's `left()`; replacing
+ /// it or making incompatible key, mode, or join-type changes fails.
+ pub fn with_prepared_build(mut self, prepared: Arc<PreparedHashJoinBuild>) -> Self {
+ // This removes the ignored child's ordering/equivalences and preserves
+ // its identity through plan resets.
+ self.exec.left = Arc::new(crate::empty::EmptyExec::new(self.exec.left.schema()));
+ self.exec.prepared_build = Some(prepared);
+ self.preserve_properties = false;
+ self
+ }
+}
+
+/// Validate eligibility and return build-key indices in join-key order.
+/// INNER joins need no shared build-match bitmap.
+fn prepared_key_indices(join: &HashJoinExec) -> Result<Vec<usize>> {
+ if join.join_type != JoinType::Inner
+ || join.mode != PartitionMode::CollectLeft
+ || join.null_aware
+ {
+ return plan_err!("Prepared hash-join builds require a CollectLeft INNER join");
+ }
+ let left = join.left.schema();
+ let right = join.right.schema();
+ if left.fields().iter().any(|field| {
+ let ty = field.data_type();
+ ty.primitive_width().is_none()
+ && !matches!(
+ ty,
+ DataType::Null
+ | DataType::Boolean
+ | DataType::FixedSizeBinary(_)
+ | DataType::Utf8
+ )
+ }) {
+ return plan_err!(
+ "Prepared hash-join builds require fixed-width or UTF-8 build columns"
+ );
+ }
+ join.on
+ .iter()
+ .map(|(l, r)| {
+ let Some(l) = l.downcast_ref::<Column>() else {
+ return plan_err!("Prepared hash-join builds require direct column keys");
+ };
+ let Some(r) = r.downcast_ref::<Column>() else {
+ return plan_err!("Prepared hash-join builds require direct column keys");
+ };
+ if l.data_type(&left)? != r.data_type(&right)? {
+ return plan_err!("Prepared hash-join key types must match");
+ }
+ Ok(l.index())
+ })
+ .collect()
+}
+
+#[cfg(test)]
+mod tests;
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs
new file mode 100644
index 0000000000..e83f953a4c
--- /dev/null
+++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs
@@ -0,0 +1,905 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Regressions for immutable preparation, task isolation and buffer ownership.
+
+use super::*;
+use crate::{
+ common,
+ display::DisplayableExecutionPlan,
+ empty::EmptyExec,
+ filter::FilterExec,
+ memory::MemoryStream,
+ sorts::sort::SortExec,
+ statistics::{StatisticsArgs, StatisticsContext},
+ stream::RecordBatchStreamAdapter,
+ test::TestMemoryExec,
+};
+use arrow::array::{
+ BooleanArray, FixedSizeBinaryArray, Int64Array, NullArray, StringArray,
+ StringViewArray, StructArray,
+};
+use arrow::buffer::Buffer;
+use arrow::compute::kernels::sort::SortOptions;
+use arrow_schema::Field;
+use datafusion_common::test_util::batches_to_sort_string;
+use datafusion_common::{DataFusionError, assert_batches_eq};
+use datafusion_common_runtime::SpawnedTask;
+use datafusion_execution::memory_pool::GreedyMemoryPool;
+use datafusion_expr::Operator;
+use datafusion_physical_expr::expressions::BinaryExpr;
+use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
+use futures::{StreamExt, stream};
+
+fn batch(keys: Vec<Option<i64>>) -> RecordBatch {
+ let rows = keys.len();
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Int64, true),
+ Field::new("payload", DataType::Int64, false),
+ ])),
+ vec![
+ Arc::new(Int64Array::from(keys)),
+ Arc::new(Int64Array::from_iter_values(
+ (0..rows).map(|i| i as i64 + 10),
+ )),
+ ],
+ )
+ .unwrap()
+}
+
+fn join(build_schema: SchemaRef, probe: RecordBatch) -> Result<HashJoinExec> {
+ let right =
+ TestMemoryExec::try_new_exec(&[vec![probe.clone()]], probe.schema(), None)?;
+ HashJoinExecBuilder::new(
+ Arc::new(EmptyExec::new(build_schema)),
+ right,
+ vec![(
+ Arc::new(Column::new("key", 0)),
+ Arc::new(Column::new("key", 0)),
+ )],
+ JoinType::Inner,
+ )
+ .with_partition_mode(PartitionMode::CollectLeft)
+ .build()
+}
+
+fn input(batches: Vec<RecordBatch>, schema: SchemaRef) -> SendableRecordBatchStream {
+ Box::pin(MemoryStream::try_new(batches, schema, None).unwrap())
+}
+
+async fn prepare(
+ join: &HashJoinExec,
+ batches: Vec<RecordBatch>,
+ pool: Arc<dyn MemoryPool>,
+) -> Result<Arc<PreparedHashJoinBuild>> {
+ join.prepare_build(
+ input(batches, join.left().schema()),
+ pool,
+ Arc::new(ConfigOptions::default()),
+ )
+ .await
+}
+
+fn buffer_addresses(arrays: &[ArrayRef]) -> Vec<usize> {
+ arrays
+ .iter()
+ .flat_map(|array| {
+ let data = array.to_data();
+ data.buffers()
+ .iter()
+ .chain(data.nulls().map(|nulls| nulls.inner().inner()))
+ .filter(|buffer| buffer.capacity() != 0)
+ .map(|buffer| buffer.data_ptr().as_ptr() as usize)
+ .collect::<Vec<_>>()
+ })
+ .collect()
+}
+
+async fn run(join: &HashJoinExec) -> Result<Vec<RecordBatch>> {
+ let partitions = (0..join.properties().output_partitioning().partition_count()).map(
+ |partition| async move {
+ common::collect(join.execute(partition, Arc::new(TaskContext::default()))?)
+ .await
+ },
+ );
+ Ok(futures::future::try_join_all(partitions)
+ .await?
+ .into_iter()
+ .flatten()
+ .collect())
+}
+
+fn with_probe_filter(join: HashJoinExec) -> Result<HashJoinExec> {
+ let filter = HashJoinExec::create_dynamic_filter(&join.on);
+ let probe = Arc::new(FilterExec::try_new(
+ Arc::clone(&filter) as _,
+ Arc::clone(join.right()),
+ )?);
+ join.builder()
+ .with_new_children(vec![Arc::clone(join.left()), probe])?
+ .with_dynamic_filter(Some(HashJoinExecDynamicFilter {
+ filter,
+ build_accumulator: OnceLock::new(),
+ }))
+ .build()
+}
+
+fn payload_filter(op: Operator) -> JoinFilter {
+ JoinFilter::new(
+ Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("build_payload", 0)),
+ op,
+ Arc::new(Column::new("probe_payload", 1)),
+ )),
+ vec![
+ ColumnIndex {
+ index: 1,
+ side: JoinSide::Left,
+ },
+ ColumnIndex {
+ index: 1,
+ side: JoinSide::Right,
+ },
+ ],
+ Arc::new(Schema::new(vec![
+ Field::new("build_payload", DataType::Int64, true),
+ Field::new("probe_payload", DataType::Int64, true),
+ ])),
+ )
+}
+
+#[tokio::test]
+async fn prepared_build_placeholder_properties_and_reset() -> Result<()> {
+ let ascending = |name, index| PhysicalSortExpr {
+ expr: Arc::new(Column::new(name, index)),
+ options: SortOptions::default(),
+ };
+ let source_schema =
+ Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false)]));
+ let placeholder = ProjectionExec::try_new(
+ vec![
+ (
+ Arc::new(Column::new("key", 0)) as PhysicalExprRef,
+ "key".into(),
+ ),
+ (
+ Arc::new(Column::new("key", 0)) as PhysicalExprRef,
+ "payload".into(),
+ ),
+ ],
+ Arc::new(EmptyExec::new(source_schema)),
+ )?;
+ let left: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(
+ LexOrdering::new(vec![ascending("payload", 1)]).unwrap(),
+ Arc::new(placeholder),
+ ));
+ let build = RecordBatch::try_new(
+ left.schema(),
+ vec![
+ Arc::new(Int64Array::from(vec![1, 1])),
+ Arc::new(Int64Array::from(vec![2, 1])),
+ ],
+ )?;
+ let probe_schema =
+ Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false)]));
+ let probe = RecordBatch::try_new(
+ Arc::clone(&probe_schema),
+ vec![Arc::new(Int64Array::from(vec![1]))],
+ )?;
+ let right: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(
+ LexOrdering::new(vec![ascending("key", 0)]).unwrap(),
+ TestMemoryExec::try_new_exec(&[vec![probe]], probe_schema, None)?,
+ ));
+ let base = HashJoinExecBuilder::new(
+ left,
+ right,
+ vec![(
+ Arc::new(Column::new("key", 0)),
+ Arc::new(Column::new("key", 0)),
+ )],
+ JoinType::Inner,
+ )
+ .with_partition_mode(PartitionMode::CollectLeft)
+ .build()?;
+ let build_key: PhysicalExprRef = Arc::new(Column::new("key", 0));
+ let build_payload: PhysicalExprRef = Arc::new(Column::new("payload", 1));
+ let asserted_order = [ascending("key", 2), ascending("payload", 1)];
+ let base_properties = base.properties().equivalence_properties();
+ assert!(
+ base_properties
+ .eq_group()
+ .exprs_equal(&build_key, &build_payload)
+ );
+ assert!(base_properties.ordering_satisfy(asserted_order.clone())?);
+
+ let prepared =
+ prepare(&base, vec![build], Arc::new(GreedyMemoryPool::new(1 << 20))).await?;
+ let attached = base.builder().with_prepared_build(prepared).build()?;
+ assert!(
+ DisplayableExecutionPlan::new(&attached)
+ .indent(false)
+ .to_string()
+ .contains("prepared_build=2 rows")
+ );
+ let properties = attached.properties().equivalence_properties();
+ assert!(
+ !properties
+ .eq_group()
+ .exprs_equal(&build_key, &build_payload)
+ );
+ assert!(properties.ordering_satisfy([ascending("key", 2)])?);
+ assert!(!properties.ordering_satisfy(asserted_order)?);
+ let stats = StatisticsContext::new().compute(&attached, &StatisticsArgs::new())?;
+ assert_ne!(stats.num_rows.get_value(), Some(&0));
+ let output = run(&attached).await?;
+ assert_eq!(output.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
+ // Resetting the unused projection/sort must retain the prepared rows.
+ let reset = crate::execution_plan::reset_plan_states(Arc::new(attached))?;
+ assert_eq!(
+ common::collect(reset.execute(0, Arc::new(TaskContext::default()))?).await?,
+ output
+ );
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_build_projection_pushdown_preserves_rows() -> Result<()> {
+ let build = batch(vec![Some(1)]);
+ let base = join(build.schema(), build.clone())?;
+ let prepared =
+ prepare(&base, vec![build], Arc::new(GreedyMemoryPool::new(1 << 20))).await?;
+ let task = Arc::new(base.builder().with_prepared_build(prepared).build()?);
+ let projection = ProjectionExec::try_new(
+ vec![
+ (
+ Arc::new(Column::new("key", 0)) as PhysicalExprRef,
+ "build_key".to_owned(),
+ ),
+ (
+ Arc::new(Column::new("key", 2)) as PhysicalExprRef,
+ "probe_key".to_owned(),
+ ),
+ ],
+ Arc::clone(&task) as _,
+ )?;
+ let transformed = task
+ .try_swapping_with_projection(&projection)?
+ .unwrap_or_else(|| Arc::new(projection));
+ let output =
+ common::collect(transformed.execute(0, Arc::new(TaskContext::default()))?)
+ .await?;
+ assert_batches_eq!(
+ [
+ "+-----------+-----------+",
+ "| build_key | probe_key |",
+ "+-----------+-----------+",
+ "| 1 | 1 |",
+ "+-----------+-----------+",
+ ],
+ &output
+ );
+ Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn prepared_build_reuses_data_with_independent_dynamic_filters() -> Result<()> {
+ let build = batch(vec![Some(1), Some(1), None, Some(-3)]);
+ let base = join(build.schema(), batch(vec![Some(1)]))?;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 20));
+ let prepared = prepare(&base, vec![build.clone()], Arc::clone(&pool)).await?;
+ let bytes = pool.reserved();
+ assert_eq!(prepared.num_rows(), 4);
+ assert_eq!(bytes, prepared.reserved_bytes());
+ assert!(bytes > 0);
+
+ let mut plans = Vec::new();
+ let mut expected = Vec::new();
+ for (op, probe) in [
+ (Operator::Gt, batch(vec![Some(1), Some(2), None])),
+ (Operator::Lt, batch(vec![Some(1), Some(1)])),
+ ] {
+ let right = TestMemoryExec::try_new_exec(
+ &[
+ vec![probe.slice(0, 1)],
+ vec![probe.slice(1, probe.num_rows() - 1)],
+ ],
+ probe.schema(),
+ None,
+ )?;
+ let task = base
+ .builder()
+ .with_new_children(vec![Arc::clone(base.left()), right])?
+ .with_filter(Some(payload_filter(op)))
+ .build()?;
+ let ordinary = task
+ .builder()
+ .reset_state()
+ .with_new_children(vec![
+ TestMemoryExec::try_new_exec(
+ &[vec![build.clone()]],
+ build.schema(),
+ None,
+ )?,
+ Arc::clone(task.right()),
+ ])?
+ .build()?;
+ expected.push(run(&ordinary).await?);
+ let task = with_probe_filter(task)?
+ .builder()
+ .with_prepared_build(Arc::clone(&prepared))
+ .build()?;
+ plans.push(Arc::new(task));
+ }
+ let cancelled = base
+ .builder()
+ .with_prepared_build(Arc::clone(&prepared))
+ .build()?;
+ drop(cancelled.execute(0, Arc::new(TaskContext::default()))?);
+ drop(cancelled);
+ assert_eq!(pool.reserved(), bytes);
+ let start = Arc::new(tokio::sync::Barrier::new(plans.len()));
+ let tasks = plans.iter().map(|plan| {
+ let plan = Arc::clone(plan);
+ let start = Arc::clone(&start);
+ SpawnedTask::spawn(async move {
+ start.wait().await;
+ run(&plan).await
+ })
+ });
+ let output = futures::future::join_all(tasks)
+ .await
+ .into_iter()
+ .map(|task| task.expect("probe task panicked"))
+ .collect::<Result<Vec<_>>>()?;
+ for (index, plan) in plans.iter().enumerate() {
+ assert_eq!(
+ batches_to_sort_string(&output[index]),
+ batches_to_sort_string(&expected[index])
+ );
+ let metrics = plan.metrics().unwrap();
+ assert_eq!(
+ metrics.sum_by_name("build_input_rows").unwrap().as_usize(),
+ 0
+ );
+ assert_eq!(metrics.output_rows(), Some(1));
+ let dynamic = plan.dynamic_filter.as_ref().unwrap();
+ assert!(futures::poll!(Box::pin(dynamic.filter.wait_complete())).is_ready());
+ }
+ assert_eq!(pool.reserved(), bytes);
+ drop(plans);
+ drop(prepared);
+ assert_eq!(pool.reserved(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_build_errors_and_cancellation_release_reservations() -> Result<()> {
+ let build = batch(vec![Some(1), Some(2)]);
+ let base = join(build.schema(), build.clone())?;
+ let exhausted: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(0));
+ assert!(matches!(
+ prepare(&base, vec![build.clone()], Arc::clone(&exhausted)).await,
+ Err(DataFusionError::ResourcesExhausted(_))
+ ));
+ assert_eq!(exhausted.reserved(), 0);
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 20));
+ let failing = stream::iter(vec![
+ Ok(build.clone()),
+ datafusion_common::exec_err!("producer failed"),
+ ]);
+ let result = base
+ .prepare_build(
+ Box::pin(RecordBatchStreamAdapter::new(build.schema(), failing)),
+ Arc::clone(&pool),
+ Arc::new(ConfigOptions::default()),
+ )
+ .await;
+ assert!(result.is_err());
+ assert_eq!(pool.reserved(), 0);
+ let pending = stream::iter(vec![Ok(build.clone())]).chain(stream::pending());
+ let mut preparation = Box::pin(base.prepare_build(
+ Box::pin(RecordBatchStreamAdapter::new(build.schema(), pending)),
+ Arc::clone(&pool),
+ Arc::new(ConfigOptions::default()),
+ ));
+ assert!(futures::poll!(&mut preparation).is_pending());
+ assert!(pool.reserved() > 0);
+ drop(preparation);
+ assert_eq!(pool.reserved(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_build_shares_single_batch_but_owns_output() -> Result<()> {
+ for array in [
+ Arc::new(Int64Array::from(vec![
+ Some(0),
+ Some(1),
+ None,
+ Some(2),
+ Some(2),
+ Some(3),
+ ])) as ArrayRef,
+ Arc::new(StringArray::from(vec![
+ Some("prefix"),
+ Some("a"),
+ None,
+ Some("b"),
+ Some("b"),
+ Some("suffix"),
+ ])),
+ ] {
+ let build = RecordBatch::try_from_iter([("key", array.slice(1, 4))])?;
+ let probe = RecordBatch::try_from_iter([("key", array.slice(1, 3))])?;
+ let base = join(build.schema(), probe)?
+ .builder()
+ .with_null_equality(NullEquality::NullEqualsNull)
+ .build()?;
+ for sources in [
+ vec![build.clone()],
+ vec![build.slice(0, 2), build.slice(2, 2)],
+ ] {
+ let single_batch = sources.len() == 1;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 20));
+ let prepared = prepare(&base, sources, Arc::clone(&pool)).await?;
+ let addresses = buffer_addresses(prepared.build.batch.columns());
+ if single_batch {
+ assert_eq!(addresses, buffer_addresses(build.columns()));
+ }
+ let bytes = prepared.reserved_bytes();
+ let task = base.builder().with_prepared_build(prepared).build()?;
+ let output = run(&task).await?;
+ assert_eq!(output.iter().map(RecordBatch::num_rows).sum::<usize>(), 4);
+ assert!(
+ output
+ .iter()
+ .flat_map(|batch| buffer_addresses(&batch.columns()[..1]))
+ .all(|address| !addresses.contains(&address))
+ );
+ assert_eq!(pool.reserved(), bytes);
+ drop(task);
+ assert_eq!(pool.reserved(), 0);
+ for batch in output {
+ assert_eq!(batch.column(0).to_data(), batch.column(1).to_data());
+ }
+ }
+ }
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_build_validates_descriptor_and_builder_changes() -> Result<()> {
+ let build = batch(vec![Some(1), Some(2)]);
+ let base = with_probe_filter(join(build.schema(), build.clone())?)?;
+ let prepared = prepare(
+ &base,
+ vec![build.clone()],
+ Arc::new(GreedyMemoryPool::new(1 << 20)),
+ )
+ .await?;
+ let keys: JoinOn = vec![(
+ Arc::new(Column::new("payload", 1)),
+ Arc::new(Column::new("payload", 1)),
+ )];
+ for builder in [
+ base.builder()
+ .with_null_equality(NullEquality::NullEqualsNull),
+ base.builder().with_type(JoinType::Right),
+ base.builder()
+ .with_partition_mode(PartitionMode::Partitioned),
+ base.builder().with_on(keys),
+ ] {
+ assert!(
+ builder
+ .with_prepared_build(Arc::clone(&prepared))
+ .build()
+ .is_err()
+ );
+ }
+ let attached = Arc::new(base.builder().with_prepared_build(prepared).build()?);
+ // Changing only the probe key must not reuse its existing dynamic filter.
+ assert!(
+ attached
+ .builder()
+ .with_on(vec![(
+ Arc::new(Column::new("key", 0)),
+ Arc::new(Column::new("payload", 1)),
+ )])
+ .build()
+ .is_err()
+ );
+ let replacement = Arc::new(EmptyExec::new(build.schema())) as Arc<dyn ExecutionPlan>;
+ assert!(
+ Arc::clone(&attached)
+ .replace_children(
+ vec![replacement, Arc::clone(attached.right())],
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ .is_err()
+ );
+ let right = TestMemoryExec::try_new_exec(
+ &[vec![batch(vec![Some(1), Some(3)])]],
+ build.schema(),
+ None,
+ )?;
+ let new_probe = attached
+ .builder()
+ .with_new_children(vec![Arc::clone(attached.left()), right])?
+ .build()?;
+ assert_eq!(
+ run(&new_probe)
+ .await?
+ .iter()
+ .map(RecordBatch::num_rows)
+ .sum::<usize>(),
+ 1
+ );
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_build_revalidates_public_keys_at_execution() -> Result<()> {
+ let build = batch(vec![Some(1), Some(2)]);
+ let base = join(build.schema(), build.clone())?;
+ let keys: JoinOn = vec![(
+ Arc::new(Column::new("payload", 1)),
+ Arc::new(Column::new("payload", 1)),
+ )];
+ let ordinary = base
+ .builder()
+ .with_on(keys.clone())
+ .with_new_children(vec![
+ TestMemoryExec::try_new_exec(&[vec![build.clone()]], build.schema(), None)?,
+ Arc::clone(base.right()),
+ ])?
+ .build()?;
+ assert_batches_eq!(
+ [
+ "+-----+---------+-----+---------+",
+ "| key | payload | key | payload |",
+ "+-----+---------+-----+---------+",
+ "| 1 | 10 | 1 | 10 |",
+ "| 2 | 11 | 2 | 11 |",
+ "+-----+---------+-----+---------+",
+ ],
+ &run(&ordinary).await?
+ );
+
+ let prepared =
+ prepare(&base, vec![build], Arc::new(GreedyMemoryPool::new(1 << 20))).await?;
+ let mut attached = base.builder().with_prepared_build(prepared).build()?;
+ attached.on = keys;
+ let error = run(&attached)
+ .await
+ .expect_err("changed public keys must not probe the old prepared table");
+ assert!(matches!(error, DataFusionError::Plan(_)), "{error}");
+ assert!(
+ error
+ .to_string()
+ .contains("does not match schema, keys or null equality")
+ );
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_build_rejects_view_payload_before_consumption() -> Result<()> {
+ let build = RecordBatch::try_from_iter([
+ ("key", Arc::new(Int64Array::from(vec![1])) as ArrayRef),
+ ("payload", Arc::new(StringViewArray::from(vec!["x"]))),
+ ])?;
+ let base = join(build.schema(), batch(vec![Some(1)]))?;
+ let never_polled =
+ stream::poll_fn(|_| -> std::task::Poll<Option<Result<RecordBatch>>> {
+ panic!("unsupported input must not be polled")
+ });
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 20));
+ assert!(
+ base.prepare_build(
+ Box::pin(RecordBatchStreamAdapter::new(build.schema(), never_polled)),
+ Arc::clone(&pool),
+ Arc::new(ConfigOptions::default()),
+ )
+ .await
+ .is_err()
+ );
+ assert_eq!(pool.reserved(), 0);
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_build_matches_collect_left() -> Result<()> {
+ let split = |batch: &RecordBatch| {
+ (0..batch.num_rows())
+ .step_by(2)
+ .map(|offset| batch.slice(offset, 2.min(batch.num_rows() - offset)))
+ .collect::<Vec<_>>()
+ };
+ for (build_keys, probe_keys) in [
+ (vec![], vec![Some(1)]),
+ (vec![Some(1)], vec![]),
+ (vec![None, None], vec![None, Some(1)]),
+ (
+ vec![None, Some(0), Some(0), Some(-4096), Some(4096)],
+ vec![Some(0), None, Some(4096), Some(2)],
+ ),
+ ] {
+ let build = batch(build_keys);
+ let probe = batch(probe_keys);
+ let sources = split(&build);
+ for nulls in [
+ NullEquality::NullEqualsNothing,
+ NullEquality::NullEqualsNull,
+ ] {
+ let base = join(build.schema(), probe.clone())?
+ .builder()
+ .with_new_children(vec![
+ TestMemoryExec::try_new_exec(
+ std::slice::from_ref(&sources),
+ build.schema(),
+ None,
+ )?,
+ TestMemoryExec::try_new_exec(&[split(&probe)], probe.schema(), None)?,
+ ])?
+ .with_null_equality(nulls)
+ .build()?;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 20));
+ let prepared = prepare(&base, sources.clone(), Arc::clone(&pool)).await?;
+ for batch_size in [1, 7] {
+ let context = Arc::new(
+ TaskContext::default().with_session_config(
+ datafusion_execution::config::SessionConfig::new()
+ .with_batch_size(batch_size),
+ ),
+ );
+ let ordinary = base.builder().reset_state().build()?;
+ let attached = base
+ .builder()
+ .with_prepared_build(Arc::clone(&prepared))
+ .build()?;
+ let mut expected =
+ common::collect(ordinary.execute(0, Arc::clone(&context))?).await?;
+ let mut actual = common::collect(attached.execute(0, context)?).await?;
+ assert!(actual.iter().all(|batch| batch.num_rows() <= batch_size));
+ expected.retain(|batch| batch.num_rows() != 0);
+ actual.retain(|batch| batch.num_rows() != 0);
+ assert_eq!(
+ batches_to_sort_string(&actual),
+ batches_to_sort_string(&expected)
+ );
+ }
+ drop(prepared);
+ assert_eq!(pool.reserved(), 0);
+ }
+ }
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_concat_admission_with_small_batches() -> Result<()> {
+ let rows = 65_536;
+ let build = batch((0..rows).map(|i| Some(i as i64 * 32)).collect());
+ let base = join(build.schema(), batch(vec![Some(0)]))?;
+ for batch_rows in [32, 64, 8192] {
+ let batches = (0..rows)
+ .step_by(batch_rows)
+ .map(|offset| build.slice(offset, batch_rows))
+ .collect();
+ // Small input batches should not multiply concatenation's alignment charge.
+ let pool: Arc<dyn MemoryPool> =
+ Arc::new(GreedyMemoryPool::new(9 * 1024 * 1024 / 2));
+ let prepared = prepare(&base, batches, Arc::clone(&pool)).await?;
+ assert_eq!(prepared.num_rows(), rows);
+ assert_eq!(pool.reserved(), prepared.reserved_bytes());
+ drop(prepared);
+ assert_eq!(pool.reserved(), 0);
+ }
+ Ok(())
+}
+
+#[test]
+fn prepared_concat_admits_validity_for_non_nullable_input_arrays() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "key",
+ DataType::Boolean,
+ true,
+ )]));
+ // A single null in the second batch creates validity for every row in the first.
+ // Boolean values make that extra allocation as large as the values themselves.
+ let batches = [
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(BooleanArray::from(vec![true; 100_000]))],
+ )?,
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(BooleanArray::from(vec![None]))],
+ )?,
+ ];
+ assert!(batches[0].column(0).nulls().is_none());
+ assert!(batches[1].column(0).nulls().is_some());
+ let admitted = prepared_copy_bytes(&batches)?;
+ let copied = concat_batches(&schema, &batches)?;
+ assert_eq!(copied.column(0).null_count(), 1);
+ let allocated = copied.column(0).get_buffer_memory_size();
+ assert!(
+ admitted >= allocated,
+ "copy admission {admitted} is smaller than allocated capacity {allocated}"
+ );
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> {
+ let build = batch(vec![Some(1), Some(1_000_000), None, Some(1)]);
+ for keys in [1, 2] {
+ let on = (0..keys)
+ .map(|index| {
+ let name = build.schema().field(index).name().clone();
+ let key: PhysicalExprRef = Arc::new(Column::new(&name, index));
+ (Arc::clone(&key), key)
+ })
+ .collect();
+ let base = join(build.schema(), build.clone())?
+ .builder()
+ .with_on(on)
+ .build()?;
+ for batches in [
+ vec![build.clone()],
+ vec![build.slice(0, 2), build.slice(2, 2)],
+ ] {
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 20));
+ let prepared = prepare(&base, batches.clone(), Arc::clone(&pool)).await?;
+ let PushdownStrategy::InList(values) = &prepared.build.membership else {
+ panic!("expected IN-list membership");
+ };
+ let columns = prepared.build.batch.columns();
+ if keys == 1 {
+ assert!(Arc::ptr_eq(values, &columns[0]));
+ } else {
+ let values = values.as_any().downcast_ref::<StructArray>().unwrap();
+ for (value, column) in values.columns().iter().zip(columns) {
+ assert!(Arc::ptr_eq(value, column));
+ }
+ }
+ let bytes = prepared.reserved_bytes();
+ let mut counter = RecordBatchMemoryCounter::new();
+ let batch_bytes = counter.count_batch(&prepared.build.batch);
+ assert!(bytes > batch_bytes, "hash storage must also be reserved");
+ let membership_batch =
+ RecordBatch::try_from_iter([("membership", Arc::clone(values))])?;
+ assert_eq!(counter.count_batch(&membership_batch), 0);
+ drop(membership_batch);
+
+ // Disabling IN-list membership changes no retained key buffers.
+ let mut config = ConfigOptions::default();
+ config.optimizer.hash_join_inlist_pushdown_max_size = 0;
+ let map = base
+ .prepare_build(
+ input(batches, build.schema()),
+ Arc::clone(&pool),
+ Arc::new(config),
+ )
+ .await?;
+ assert!(matches!(map.build.membership, PushdownStrategy::Map(_)));
+ assert_eq!(map.reserved_bytes(), bytes);
+ assert_eq!(pool.reserved(), 2 * bytes);
+ drop(map);
+ let lease = Arc::clone(&prepared);
+ drop(prepared);
+ assert_eq!(pool.reserved(), bytes);
+ drop(lease);
+ assert_eq!(pool.reserved(), 0);
+ }
+ }
+ Ok(())
+}
+
+#[test]
+fn prepared_concat_rejects_under_admission_without_shrinking() -> Result<()> {
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024));
+ let reservation = MemoryConsumer::new("prepared concat test").register(&pool);
+ // 128 bytes of map storage plus 128 bytes of input/copy allowance.
+ reservation.try_grow(256)?;
+ let error =
+ reconcile_prepared_copy_reservation(&reservation, 64, 64, 129).unwrap_err();
+ assert!(matches!(error, DataFusionError::Internal(_)));
+ assert!(
+ error
+ .to_string()
+ .contains("exceeded its admitted copy bound")
+ );
+ assert_eq!(reservation.size(), 256);
+ assert!(reconcile_prepared_copy_reservation(&reservation, usize::MAX, 1, 0).is_err());
+ assert_eq!(reservation.size(), 256);
+ reconcile_prepared_copy_reservation(&reservation, 64, 64, 128)?;
+ assert_eq!(reservation.size(), 256);
+ reconcile_prepared_copy_reservation(&reservation, 64, 64, 80)?;
+ assert_eq!(reservation.size(), 208);
+ drop(reservation);
+ assert_eq!(pool.reserved(), 0);
+ Ok(())
+}
+
+#[test]
+fn prepared_sizing_rejects_overflow() -> Result<()> {
+ use NullEquality::{NullEqualsNothing, NullEqualsNull};
+ for (rows, keys, equality) in [
+ (usize::MAX / 8 + 1, 1, NullEqualsNull),
+ (0, usize::MAX, NullEqualsNothing),
+ (0, usize::MAX / 64, NullEqualsNothing),
+ (usize::MAX / 8, 0, NullEqualsNothing),
+ ] {
+ let error = prepared_scratch_bytes(rows, keys, equality).unwrap_err();
+ assert!(error.to_string().contains("scratch size overflow"));
+ }
+ assert!(prepared_scratch_bytes(usize::MAX / 8, 1, NullEqualsNull).is_ok());
+
+ // Null arrays can describe huge row counts without allocating data buffers.
+ // Exercise both per-batch validity multiplication and cumulative addition.
+ let array: ArrayRef = Arc::new(NullArray::new(usize::MAX));
+ for (columns, batches) in [(8, 1), (4, 2)] {
+ let batch = RecordBatch::try_from_iter(
+ (0..columns).map(|i| (format!("key{i}"), Arc::clone(&array))),
+ )?;
+ let error = prepared_copy_bytes(&vec![batch; batches]).unwrap_err();
+ assert!(error.to_string().contains("copy size overflow"));
+ }
+ Ok(())
+}
+
+#[tokio::test]
+async fn prepared_byte_keys_use_hash_membership() -> Result<()> {
+ for array in [
+ Arc::new(StringArray::from(vec!["aaaa", "bbbb"])) as ArrayRef,
+ Arc::new(FixedSizeBinaryArray::try_new(
+ 4,
+ Buffer::from_slice_ref(b"aaaabbbb"),
+ None,
+ )?),
+ ] {
+ let build = RecordBatch::try_from_iter([("key", array.slice(0, 1))])?;
+ let probe = RecordBatch::try_from_iter([("key", array)])?;
+ let base = join(build.schema(), probe)?;
+ let mut options = ConfigOptions::default();
+ options.optimizer.hash_join_inlist_pushdown_max_size = 1024;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 20));
+ let prepared = base
+ .prepare_build(
+ input(vec![build.clone()], build.schema()),
+ Arc::clone(&pool),
+ Arc::new(options),
+ )
+ .await?;
+ assert!(prepared.build.bounds.is_none());
+ assert!(matches!(
+ &prepared.build.membership,
+ PushdownStrategy::Map(_)
+ ));
+ let task = with_probe_filter(base)?
+ .builder()
+ .with_prepared_build(prepared)
+ .build()?;
+ let output = run(&task).await?;
+ assert_eq!(output.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
+ let filter = &task.dynamic_filter.as_ref().unwrap().filter;
+ assert!(futures::poll!(Box::pin(filter.wait_complete())).is_ready());
+ drop(task);
+ assert_eq!(pool.reserved(), 0);
+ }
+ Ok(())
+}
diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs
index 7c1b0d76c2..7b50435e18 100644
--- a/datafusion/physical-plan/src/joins/hash_join/mod.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs
@@ -17,7 +17,7 @@
//! [`HashJoinExec`] Partitioned Hash Join Operator
-pub use exec::{HashJoinExec, HashJoinExecBuilder};
+pub use exec::{HashJoinExec, HashJoinExecBuilder, PreparedHashJoinBuild};
pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState};
mod exec;
diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs
index 82863c0802..0fa5615db3 100644
--- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs
@@ -330,26 +330,15 @@ impl Hash for HashTableLookupExpr {
self.on_columns.dyn_hash(state);
self.description.hash(state);
self.random_state.seed().hash(state);
- // Note that we compare hash_map by pointer equality.
- // Actually comparing the contents of the hash maps would be expensive.
- // The way these hash maps are used in actuality is that 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.
- // Theoretically this is a public API and users could create identical hash maps,
- // but that seems unlikely and not worth paying the cost of deep comparison all the time.
+ // Hash immutable map identity, consistently with pointer equality below.
Arc::as_ptr(&self.map).hash(state);
}
}
impl PartialEq for HashTableLookupExpr {
fn eq(&self, other: &Self) -> bool {
- // Note that we compare hash_map by pointer equality.
- // Actually comparing the contents of the hash maps would be expensive.
- // The way these hash maps are used in actuality is that 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.
- // Theoretically this is a public API and users could create identical hash maps,
- // but that seems unlikely and not worth paying the cost of deep comparison all the time.
+ // Sharing an immutable map implies equal contents, including across
+ // prepared-build consumers. Distinct maps need no expensive deep comparison.
self.on_columns == other.on_columns
&& self.description == other.description
&& self.random_state.seed() == other.random_state.seed()
diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs
index 56c34534fb..03387ea574 100644
--- a/datafusion/physical-plan/src/joins/hash_join/stream.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs
@@ -632,8 +632,8 @@ impl HashJoinStream {
let pushdown = left_data.membership().clone();
let bounds = left_data
- .bounds
- .clone()
+ .bounds()
+ .cloned()
.unwrap_or_else(|| PartitionBounds::new(vec![]));
// Use the logical null count: a dictionary key whose entry points at a
// NULL dictionary value is a NULL key even though the key bitmap has no
diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs
index 10e5793e0f..0d2480c131 100644
--- a/datafusion/physical-plan/src/joins/mod.rs
+++ b/datafusion/physical-plan/src/joins/mod.rs
@@ -25,7 +25,8 @@ pub use asof_join::{AsOfJoinExec, AsOfMatchExpr};
pub use cross_join::CrossJoinExec;
use datafusion_physical_expr::PhysicalExprRef;
pub use hash_join::{
- HashExpr, HashJoinExec, HashJoinExecBuilder, HashTableLookupExpr, SeededRandomState,
+ HashExpr, HashJoinExec, HashJoinExecBuilder, HashTableLookupExpr,
+ PreparedHashJoinBuild, SeededRandomState,
};
pub use nested_loop_join::{NestedLoopJoinExec, NestedLoopJoinExecBuilder};
use parking_lot::Mutex;
diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs
index 6ff4422791..1845789eb8 100644
--- a/datafusion/proto/tests/cases/plans/joins.rs
+++ b/datafusion/proto/tests/cases/plans/joins.rs
@@ -1059,3 +1059,49 @@ async fn roundtrip_planned_piecewise_merge_join() -> Result<()> {
}
Ok(())
}
+
+#[tokio::test]
+async fn prepared_hash_join_cannot_be_serialized_without_its_build() -> Result<()> {
+ use datafusion::arrow::array::Int64Array;
+ use datafusion::arrow::record_batch::RecordBatch;
+ use datafusion::physical_plan::memory::MemoryStream;
+ use datafusion_common::config::ConfigOptions;
+ use datafusion_execution::memory_pool::GreedyMemoryPool;
+
+ let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)]));
+ let join = HashJoinExec::try_new(
+ Arc::new(EmptyExec::new(Arc::clone(&schema))),
+ Arc::new(EmptyExec::new(Arc::clone(&schema))),
+ vec![(
+ Arc::new(Column::new("col", 0)),
+ Arc::new(Column::new("col", 0)),
+ )],
+ None,
+ &JoinType::Inner,
+ None,
+ PartitionMode::CollectLeft,
+ NullEquality::NullEqualsNothing,
+ false,
+ )?;
+ roundtrip_test(Arc::new(join.builder().build()?))?;
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int64Array::from(vec![1_i64]))],
+ )?;
+ let prepared = join
+ .prepare_build(
+ Box::pin(MemoryStream::try_new(vec![batch], schema, None)?),
+ Arc::new(GreedyMemoryPool::new(1 << 20)),
+ Arc::new(ConfigOptions::default()),
+ )
+ .await?;
+ let attached = join.builder().with_prepared_build(prepared).build()?;
+ let error = roundtrip_test(Arc::new(attached))
+ .expect_err("serializing the empty build placeholder would lose prepared rows");
+ assert!(matches!(
+ error,
+ datafusion_common::DataFusionError::NotImplemented(_)
+ ));
+ assert!(error.to_string().contains("prepared build"), "{error}");
+ Ok(())
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org