sunchao commented on code in PR #25491: URL: https://github.com/apache/datafusion/pull/25491#discussion_r4054126489
########## datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs: ########## @@ -0,0 +1,787 @@ +// 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, + empty::EmptyExec, + filter::FilterExec, + memory::MemoryStream, + sorts::sort::SortExec, + statistics::{ChildStats, StatisticsArgs, StatisticsContext}, + stream::RecordBatchStreamAdapter, + test::TestMemoryExec, +}; +use arrow::array::{Array, Int64Array, LargeStringArray, StringViewArray}; +use arrow::compute::kernels::sort::SortOptions; +use arrow_schema::Field; +use datafusion_common::assert_batches_eq; +use datafusion_common::utils::memory::get_record_batch_memory_size; +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()?; + 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)?); + assert_eq!(base.child_stats_requests(None)[0], ChildStats::At(None)); + assert_eq!(attached.child_stats_requests(None)[0], ChildStats::Skip); + assert_eq!(attached.child_stats_requests(Some(0))[0], ChildStats::Skip); + let stats = StatisticsContext::new().compute(&attached, &StatisticsArgs::new())?; + assert_ne!(stats.num_rows.get_value(), Some(&0)); + let output = run(&attached).await?; + assert_batches_eq!( + [ + "+-----+---------+-----+", + "| key | payload | key |", + "+-----+---------+-----+", + "| 1 | 2 | 1 |", + "| 1 | 1 | 1 |", + "+-----+---------+-----+", + ], + &output + ); + // 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] +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(); + 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 task = with_probe_filter(task)? + .builder() + .with_prepared_build(Arc::clone(&prepared)) + .build()?; + plans.push(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 (first, second) = tokio::join!(run(&plans[0]), run(&plans[1])); + let output = [first?, second?]; + for (index, plan) in plans.iter().enumerate() { + assert_batches_eq!( + [ + "+-----+---------+-----+---------+", + "| key | payload | key | payload |", + "+-----+---------+-----+---------+", + [ + "| 1 | 11 | 1 | 10 |", + "| 1 | 10 | 1 | 11 |" + ][index], + "+-----+---------+-----+---------+", + ], + &output[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 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_accounts_for_hash_map_row_indices() -> Result<()> { Review Comment: Shared the repeated admission-rejection and reservation-release assertions. The cases retain their distinct checks: retained bytes excluding scratch, composite/null-key results, and UTF-8 alias/copy admission. In particular, the UTF-8 case rejects at concat admission and uses a separately funded success run; it is not another exact-peak success case. Keeping those checks local avoids combining different contracts into one large parameterized fixture. ########## datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs: ########## @@ -0,0 +1,787 @@ +// 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, + empty::EmptyExec, + filter::FilterExec, + memory::MemoryStream, + sorts::sort::SortExec, + statistics::{ChildStats, StatisticsArgs, StatisticsContext}, + stream::RecordBatchStreamAdapter, + test::TestMemoryExec, +}; +use arrow::array::{Array, Int64Array, LargeStringArray, StringViewArray}; +use arrow::compute::kernels::sort::SortOptions; +use arrow_schema::Field; +use datafusion_common::assert_batches_eq; +use datafusion_common::utils::memory::get_record_batch_memory_size; +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()?; + 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)?); + assert_eq!(base.child_stats_requests(None)[0], ChildStats::At(None)); + assert_eq!(attached.child_stats_requests(None)[0], ChildStats::Skip); + assert_eq!(attached.child_stats_requests(Some(0))[0], ChildStats::Skip); + let stats = StatisticsContext::new().compute(&attached, &StatisticsArgs::new())?; + assert_ne!(stats.num_rows.get_value(), Some(&0)); + let output = run(&attached).await?; + assert_batches_eq!( + [ + "+-----+---------+-----+", + "| key | payload | key |", + "+-----+---------+-----+", + "| 1 | 2 | 1 |", + "| 1 | 1 | 1 |", + "+-----+---------+-----+", + ], + &output + ); + // 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] +async fn prepared_build_reuses_data_with_independent_dynamic_filters() -> Result<()> { Review Comment: The existing independent-consumer test now uses a two-worker runtime, separately spawned `SpawnedTask` consumers, and a start barrier. Each result is compared with an ordinary CollectLeft join using the same inputs and residual filter. Changing runtime flavor alone would not parallelize `tokio::join!`. Added 24 fixed random seeds, each compared under both null-equality modes and output batch sizes 1 and 7: 96 differential comparisons covering duplicates/skew, empty inputs, residual filters, and varying input batch boundaries, using the existing rand dependency. ########## 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: Moved the scratch and copy sizing into `prepared.rs` helpers, accumulated copy bytes/max batch rows while ingesting, and removed the prepared-only row-chain charge in favor of the shared constructor fix. The ordinary and prepared paths now preallocate the hash buffer from the observed maximum batch size. Validation and the before/after ordinary-path benchmark are summarized in the updated PR description. ########## datafusion/physical-plan/src/joins/hash_join/exec.rs: ########## @@ -3104,18 +3250,37 @@ async fn collect_left_input( && !left_values.is_empty() && left_values[0].logical_null_count() > 0; + if prepared { + drop(batches); + let retained = RecordBatchMemoryCounter::new().count_batch(&batch); + let allowance = input_bytes.checked_add(copy_bytes).ok_or_else(|| { + datafusion_common::exec_datafusion_err!( + "Prepared hash-join payload size overflow" + ) + })?; + if retained > allowance { Review Comment: I kept the invariant error here. The API promises admission before allocating the copy; growing the reservation afterward cannot cover the earlier peak overlap with the original batches. The post-check prevents publishing an undercharged build, but is not a substitute for a correct preflight bound. The sizing refactor now reuses Arrow's slice measurement and has an allocation-capacity regression for mixed physical validity. ########## datafusion/physical-plan/src/joins/hash_join/exec.rs: ########## @@ -1625,6 +1660,9 @@ impl ExecutionPlan for HashJoinExec { partition: usize, context: Arc<TaskContext>, ) -> Result<SendableRecordBatchStream> { + if let Some(prepared) = &self.prepared_build { Review Comment: Added that explanation and retained the execute-time check. `on` is public, so callers can change it after builder validation; the existing direct-mutation test continues to cover this. ########## 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: Addressed the collector's readability with named sizing helpers, ingestion-time totals, and a small `BuildMode` enum. This removes the long inline arithmetic and the prepared-only row-chain top-up without adding an admission object with its own lifecycle. ########## 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 { Review Comment: I strengthened the method's caller contract and example, and kept snapshot identity/invalidation in the embedding executor for this revision. An opaque token could catch cache-selection mistakes if a consumer separately supplies its expected identity, but the caller would still be responsible for assigning that identity to the right stream. The current join has no independently known snapshot identity to compare against. The docs now lead with this responsibility and explicitly state that schema/key compatibility cannot establish input identity. ########## 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(), Review Comment: Added a short explanation that the reservation stays empty for the supported INNER joins. I did not add an unused consumer-pool parameter now; extending the API to outer joins will need to put mutable probe allocations in the consuming task's pool as part of that extension. ########## 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. Review Comment: Reworked these docs to lead with caller-owned snapshot identity and the supplied-stream behavior. They now also state that numeric IN-list publication allocates per-row literals outside the prepared reservation. The configured threshold measures input-array bytes, not the resulting expression heap, so it should not be read as a hard heap limit. ########## 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> { + let rows = batch.num_rows(); + batch.columns().iter().try_fold(0usize, |total, array| { + let values = match array.data_type() { + DataType::Utf8 => rows + .checked_add(1) + .and_then(|len| len.checked_mul(4)) + .and_then(|offsets| utf8_value_span(array.as_ref()).checked_add(offsets)) + // UTF-8 has a third allocation for offsets. + .and_then(|bytes| bytes.checked_add(64)), + DataType::Null => Some(0), + DataType::Boolean => Some(rows.div_ceil(8)), + DataType::FixedSizeBinary(width) => usize::try_from(*width) + .ok() + .and_then(|width| width.checked_mul(rows)), + ty => ty + .primitive_width() + .and_then(|width| width.checked_mul(rows)), + }; + values + .and_then(|bytes| bytes.checked_add(rows.div_ceil(8))) + .and_then(|bytes| bytes.checked_add(2 * 64)) + .and_then(|bytes| total.checked_add(bytes)) + .ok_or_else(|| exec_datafusion_err!("Prepared hash-join copy size overflow")) + }) +} + +/// Include values hidden by nulls but exclude bytes outside a sliced array. +fn utf8_value_span(array: &dyn Array) -> usize { + let offsets = array.as_string::<i32>().value_offsets(); + (offsets[offsets.len() - 1] - offsets[0]) as usize +} + +/// Add `batch`'s byte-column spans to `totals`, one entry per schema column, before +/// an eventual single-batch concat. Reject offset overflow before allocating +/// the value buffer; fixed-width columns leave their totals unchanged. +pub(super) fn check_byte_concat_sizes( + batch: &RecordBatch, + totals: &mut [usize], +) -> Result<()> { + for (array, total) in batch.columns().iter().zip(totals) { + if matches!(array.data_type(), DataType::Utf8) { + *total = total.checked_add(utf8_value_span(array.as_ref())) + .filter(|&sum| i32::try_from(sum).is_ok()) + .ok_or_else(|| exec_datafusion_err!( + "Prepared hash-join UTF-8 column exceeds its offset limit; a compact build is required" + ))?; + } + } + Ok(()) +} + +impl HashJoinExecBuilder { + /// Attach a fully prepared build to a fresh, compatible join execution. + /// + /// [`Self::build`] validates compatibility. Attaching resets the build future + /// and execution metrics. A previously attached task-local dynamic filter + /// keeps its expression handle (also referenced by the probe plan), while + /// its build-report accumulator is reset. The caller must provide a fresh + /// filter expression/probe plan for each independent task. + /// Residual filters also remain consumer-local, so compatible INNER joins + /// may use different predicates with the same prepared data. + /// The resulting join plan retains the build lease. A caller retaining a + /// dynamic-filter expression beyond that plan must retain a prepared-build + /// lease alongside it, because membership filters can reference build data. + /// Attach after child-rewriting physical optimizations. The unused left + /// subtree is replaced with an empty schema placeholder so plan resets + /// preserve it. Probe-only rewrites must retain the attached join's `left()`. + /// Replacing that child or changing to incompatible join keys fails; other + /// incompatible join-mode/type changes fail validation in `build`. Review Comment: Separated the attachment obligations into a `Caller contract` section, including fresh consumer filters, snapshot choice, attachment after build-child rewrites, and retaining a prepared-build lease when an external filter outlives the consuming plan. The lease preserves the corresponding memory charge; the filter's map/array Arcs preserve allocation lifetime on their own. Probe-only rewrites must preserve the attached join's build placeholder. -- 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]
