sunchao commented on code in PR #25491:
URL: https://github.com/apache/datafusion/pull/25491#discussion_r4054125868
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2942,9 +3035,61 @@ async fn collect_left_input(
// Use `u32` indices for the JoinHashMap when num_rows ≤ u32::MAX,
otherwise use the
// `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.
Review Comment:
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](https://github.com/apache/datafusion/pull/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.
##########
datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs:
##########
@@ -0,0 +1,326 @@
+// 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 arrow::array::{Array, AsArray};
+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_mark_scope_map: None,
+ null_value_scope_map: 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,
+ _probe_reservation: self.build.reservation.new_empty(),
+ }
+ }
+
+ /// 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 filter_keys = filter.filter.children();
+ if filter_keys.len() != join.on.len()
+ || filter_keys
+ .iter()
+ .zip(&join.on)
+ .any(|(filter_key, (_, probe_key))| {
+ filter_key.as_ref() != probe_key.as_ref()
+ })
+ {
+ return plan_err!(
+ "Prepared hash-join dynamic filter keys do not match probe
keys"
+ );
+ }
+ }
+ Ok(())
+ }
+}
+
+impl HashJoinExec {
+ /// Prepare one immutable build using an embedding executor's durable pool.
+ ///
+ /// The supplied stream must own its native buffers independently of
producer
+ /// task cleanup. This method consumes only that stream, never `self.left`,
+ /// and reserves retained data, hash buckets and row-index chains against
+ /// `pool`. The caller must keep original producer allocations charged
until
+ /// its stream releases them. `config` controls ordinary perfect-map and
+ /// dynamic-filter choices. UTF-8 and fixed-size binary keys
+ /// retain hash membership only:
+ /// range bounds and IN-list literals would allocate unaccounted key
copies.
+ ///
+ /// Validates eligibility and the stream schema before polling. On error or
+ /// 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.
+ 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,
+ true,
+ )
+ .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(batch: &RecordBatch) -> Result<usize> {
Review Comment:
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.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2866,49 +2933,75 @@ 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(),
&schema,
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 concat_values = if prepared {
Review Comment:
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.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2854,6 +2920,7 @@ async fn collect_left_input(
null_equality: NullEquality,
null_aware: Option<NullAwareMode>,
array_map_created_count: Count,
+ prepared: bool,
Review Comment:
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.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2942,9 +3035,61 @@ async fn collect_left_input(
// Use `u32` indices for the JoinHashMap when num_rows ≤ u32::MAX,
otherwise use the
// `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.
+ let index_width = if num_rows > u32::MAX as usize {
+ size_of::<u64>()
+ } else {
+ size_of::<u32>()
+ };
+ let bytes = num_rows.checked_mul(index_width).ok_or_else(|| {
+ datafusion_common::exec_datafusion_err!(
+ "Prepared hash-join row-index size overflow"
+ )
+ })?;
+ reservation.try_grow(bytes)?;
+ }
let mut hashmap = new_join_hashmap(num_rows, &mut reservation,
&metrics)?;
- let mut hashes_buffer = Vec::new();
+ let scratch_reservation = reservation.new_empty();
+ let mut hashes_buffer = if prepared {
+ 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 {
Review Comment:
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.
##########
datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs:
##########
@@ -0,0 +1,326 @@
+// 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 arrow::array::{Array, AsArray};
+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_mark_scope_map: None,
+ null_value_scope_map: 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,
+ _probe_reservation: self.build.reservation.new_empty(),
+ }
+ }
+
+ /// 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 filter_keys = filter.filter.children();
+ if filter_keys.len() != join.on.len()
+ || filter_keys
+ .iter()
+ .zip(&join.on)
+ .any(|(filter_key, (_, probe_key))| {
+ filter_key.as_ref() != probe_key.as_ref()
+ })
+ {
+ return plan_err!(
+ "Prepared hash-join dynamic filter keys do not match probe
keys"
+ );
+ }
+ }
+ Ok(())
+ }
+}
+
+impl HashJoinExec {
+ /// Prepare one immutable build using an embedding executor's durable pool.
+ ///
+ /// The supplied stream must own its native buffers independently of
producer
+ /// task cleanup. This method consumes only that stream, never `self.left`,
+ /// and reserves retained data, hash buckets and row-index chains against
+ /// `pool`. The caller must keep original producer allocations charged
until
+ /// its stream releases them. `config` controls ordinary perfect-map and
+ /// dynamic-filter choices. UTF-8 and fixed-size binary keys
+ /// retain hash membership only:
+ /// range bounds and IN-list literals would allocate unaccounted key
copies.
+ ///
+ /// Validates eligibility and the stream schema before polling. On error or
+ /// 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.
+ pub async fn prepare_build(
Review Comment:
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.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2075,6 +2138,9 @@ impl ExecutionPlan for HashJoinExec {
cache: _,
} = self;
+ if prepared_build.is_some() {
+ return plan_err!("HashJoinExec with a prepared build cannot be
serialized");
Review Comment:
Changed this to `not_impl_err!` and strengthened the existing serialization
regression to assert `NotImplemented`, in addition to the message.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
