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-25496-c09448f8423344665c8b3ee2a85daa6563919828 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 53ccb70d8201e67bed312e69fd1ed370dbcd2aab Author: Jay Zhan <[email protected]> AuthorDate: Mon Sep 21 16:08:59 2026 +0000 fix: reserve the concatenated build batch and computed join keys in HashJoinExec (#25496) ## Which issue does this PR close? - Closes #25269. ## Rationale for this change A hash join reserves every build-side batch as it arrives, then copies all of them into a single batch without reserving the copy, while the input batches stay alive until the build finishes. Near the memory limit the process therefore holds about twice the build side while the memory pool sees it once, so instead of a `ResourcesExhausted` error the process can be killed by the OS. It also means nothing that reacts to pool pressure can rely on what a hash join reports. Measured on a 1.25 GB build side (TPC-H SF1 `lineitem`, 16 columns, as the build side): the pool was charged ~1.25 GB while peak RSS was ~2.46 GB. ## What changes are included in this PR? - `concat_build_batches`: reserves the copy before `concat_batches`, drops the input batches, then trims the reservation to what the concatenated batch retains (`get_record_batch_memory_size`), growing instead if the estimate was low. `build_mem_used` follows every grow and shrink. - Nothing extra is reserved when there is nothing to copy: concatenating a single batch is zero-copy, and for `Utf8View`/`BinaryView` columns only the views are copied while the data buffers stay shared with the inputs. - `try_create_array_map` becomes `array_map_key_range`: it only decides whether the perfect hash join applies and reserves the `ArrayMap`. Both the hash map and the `ArrayMap` paths now hand the batches to the helper by value, so the inputs are gone before the visited bitmap and the null-aware scope maps are allocated. - Join key arrays that do not share the buffers of the build batch (any on-expression that is not a plain column) are kept for the whole join and are now reserved too. This uses a new public `RecordBatchMemoryCounter::count_array`. This does not reduce memory usage: the build still peaks at about twice the build side, it is now counted. A memory-limited query that only fit because the copy was not counted can now fail with `ResourcesExhausted`. ## What is the testing strategy for this PR? - `join_build_concat_is_reserved` (CollectLeft and Partitioned, hash map and `ArrayMap`): fails with `ResourcesExhausted` at map + 1.5x inputs, passes at map + 3x inputs. It fails on `main` in all four variants. - `join_build_key_arrays_are_reserved` (hash map and `ArrayMap`): with a single build batch and an `a1 + 1` join key, the join fails at inputs + map + half the key array and passes with room for it; a plain column key passes at the lower limit. It fails when the reservation is removed. - Unit tests of the helper: output equals `concat_batches` in both orders, the copy is reserved and the reservation then matches the retained batch, a single batch is not reserved twice, view data buffers are not reserved twice. - Existing `single_partition_join_overallocation` / `partitioned_join_overallocation` are unchanged. ## Are there any user-facing changes? One public addition: `RecordBatchMemoryCounter::count_array`. Hash joins running close to their memory limit may now report `ResourcesExhausted` where they previously exceeded the limit unnoticed. --- datafusion/common/src/utils/memory.rs | 20 + .../physical-plan/src/joins/hash_join/exec.rs | 557 +++++++++++++++++++-- 2 files changed, 523 insertions(+), 54 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index bb006295f8..ac5e15301a 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -193,6 +193,14 @@ impl RecordBatchMemoryCounter { self.memory_usage - previous_memory_usage } + /// Count `array`, returning the memory used by its buffers that have not + /// been counted before. + pub fn count_array(&mut self, array: &dyn Array) -> usize { + let previous_memory_usage = self.memory_usage; + self.count_array_memory_size(array); + self.memory_usage - previous_memory_usage + } + /// Counts unique buffers and Array objects retained by `batch`. /// /// This is useful for accounting a sequence of batches at an operator @@ -581,6 +589,18 @@ mod record_batch_tests { assert_eq!(counter.memory_usage(), array_data_memory_size(array)); } + #[test] + fn test_count_array_counts_shared_buffers_once() { + let array = Int32Array::from(vec![1, 2, 3, 4, 5]); + let size = array_data_memory_size(&array); + + let mut counter = RecordBatchMemoryCounter::new(); + assert_eq!(counter.count_array(&array), size); + // A slice shares the buffer that is already counted + assert_eq!(counter.count_array(&array.slice(1, 2)), 0); + assert_eq!(counter.memory_usage(), size); + } + #[test] fn test_get_record_batch_memory_size() { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 9858a4c06c..3f1d92653b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -77,7 +77,9 @@ use arrow::util::bit_util; use arrow_schema::{DataType, Schema}; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; -use datafusion_common::utils::memory::{RecordBatchMemoryCounter, estimate_memory_size}; +use datafusion_common::utils::memory::{ + RecordBatchMemoryCounter, estimate_memory_size, get_record_batch_memory_size, +}; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, assert_or_internal_err, internal_err, plan_err, project_schema, @@ -107,8 +109,14 @@ pub(crate) const HASH_JOIN_SEED: SeededRandomState = const ARRAY_MAP_CREATED_COUNT_METRIC_NAME: &str = "array_map_created_count"; +/// Decides whether the build side should be joined with an [`ArrayMap`] +/// (perfect hash join), returning the `(min, max)` key range to build it over. +/// +/// On `Some`, the memory of the [`ArrayMap`] has already been added to +/// `reservation`; the caller concatenates the build batches and creates the map +/// with [`ArrayMap::try_new`]. #[expect(clippy::too_many_arguments)] -fn try_create_array_map( +fn array_map_key_range( bounds: Option<&PartitionBounds>, schema: &SchemaRef, batches: &[RecordBatch], @@ -117,7 +125,7 @@ fn try_create_array_map( perfect_hash_join_small_build_threshold: usize, perfect_hash_join_min_key_density: f64, null_equality: NullEquality, -) -> Result<Option<(ArrayMap, RecordBatch, Vec<ArrayRef>)>> { +) -> Result<Option<(u64, u64)>> { // `bounds` are also collected for dynamic filters, on any key type, so the // key type cannot be inferred from their presence. if !is_perfect_hash_join_candidate(on_left, schema)? { @@ -184,12 +192,7 @@ fn try_create_array_map( let mem_size = ArrayMap::estimate_memory_size(min_val, max_val, num_row); reservation.try_grow(mem_size)?; - let batch = concat_batches(schema, batches)?; - let left_values = evaluate_expressions_to_arrays(on_left, &batch)?; - - let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?; - - Ok(Some((array_map, batch, left_values))) + Ok(Some((min_val, max_val))) } /// Correlation-scope hash map over only the build rows whose scalar `NOT IN` @@ -678,7 +681,7 @@ impl From<&HashJoinExec> for HashJoinExecBuilder { /// 4. build_side.num_rows() < u32::MAX /// 5. NullEqualsNothing || (NullEqualsNull && build side doesn't contain null) /// -/// See [`try_create_array_map`] for more details. +/// See [`array_map_key_range`] for more details. /// /// Note that when using [`PartitionMode::Partitioned`], the build side is split into multiple /// partitions. This can cause a dense build side to become sparse within each partition, @@ -2773,7 +2776,7 @@ impl BuildSideState { /// (perfect hash join): a single join key of a supported integer type. /// /// Only a candidate: the final decision also depends on the observed key -/// range and density, see [`try_create_array_map`]. +/// range and density, see [`array_map_key_range`]. fn is_perfect_hash_join_candidate( on_left: &[PhysicalExprRef], schema: &SchemaRef, @@ -2824,6 +2827,75 @@ fn new_join_hashmap( } } +/// Estimates the memory newly allocated for `array`'s share of a concatenated +/// array. +/// +/// Concatenating view arrays copies only the views and keeps sharing the data +/// buffers of the inputs, which are already accounted for. +fn estimate_concat_allocation(array: &dyn Array) -> Result<usize> { + match array.data_type() { + DataType::Utf8View | DataType::BinaryView => { + let nulls = array + .nulls() + .map(|_| bit_util::ceil(array.len(), 8)) + .unwrap_or_default(); + Ok(array.len() * size_of::<u128>() + nulls) + } + _ => Ok(array.to_data().get_slice_memory_size()?), + } +} + +/// Concatenates the build side `batches` into a single batch, in reverse order +/// if `reverse` is set, keeping `reservation` ahead of the actual memory usage. +/// +/// `inputs_reserved` is the memory `reservation` already holds for `batches`. +/// The copy is reserved before it is made, and once `batches` are dropped the +/// reservation is trimmed to the memory retained by the returned batch. +fn concat_build_batches( + schema: &SchemaRef, + batches: Vec<RecordBatch>, + reverse: bool, + inputs_reserved: usize, + reservation: &mut MemoryReservation, + metrics: &BuildProbeJoinMetrics, +) -> Result<RecordBatch> { + // Concatenating a single batch is zero-copy + let copy_size = if batches.len() > 1 { + let mut copy_size = 0; + for batch in &batches { + for array in batch.columns() { + copy_size += estimate_concat_allocation(array.as_ref())?; + } + } + copy_size + } else { + 0 + }; + reservation.try_grow(copy_size)?; + metrics.build_mem_used.add(copy_size); + + let batch = if reverse { + concat_batches(schema, batches.iter().rev())? + } else { + concat_batches(schema, batches.iter())? + }; + drop(batches); + + // The inputs are gone: only hold on to what the concatenated batch retains, + // which includes any buffers it still shares with the inputs. + let held = inputs_reserved + copy_size; + let retained = get_record_batch_memory_size(&batch); + if retained > held { + reservation.try_grow(retained - held)?; + metrics.build_mem_used.add(retained - held); + } else { + reservation.shrink(held - retained); + metrics.build_mem_used.sub(held - retained); + } + + Ok(batch) +} + /// Collects all batches from the left (build) side stream and creates a hash map for joining. /// /// This function is responsible for: @@ -2844,6 +2916,13 @@ fn new_join_hashmap( /// * `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 /// +/// # Memory Accounting +/// Build batches are added to `reservation` as they arrive. They are then copied +/// into a single batch, see [`concat_build_batches`]: the copy is reserved +/// before it is made, and the reservation is trimmed to what the single batch +/// retains once the input batches are dropped. Join key arrays that do not +/// share the buffers of that batch are reserved as well. +/// /// # Dynamic Filter Coordination /// When `should_compute_dynamic_filters` is true, this function computes the min/max bounds /// for each join key column but does NOT update the dynamic filter. Instead, the @@ -2920,8 +2999,9 @@ async fn collect_left_input( metrics, mut reservation, bounds_accumulators, - memory_counter: _, + memory_counter, } = state; + let inputs_reserved = memory_counter.memory_usage(); // Compute bounds let mut bounds = match bounds_accumulators { @@ -2935,8 +3015,8 @@ async fn collect_left_input( _ => None, }; - let (join_hash_map, batch, left_values) = - if let Some((array_map, batch, left_value)) = try_create_array_map( + let (join_hash_map, batch, left_values) = if let Some((min_val, max_val)) = + array_map_key_range( bounds.as_ref(), &schema, &batches, @@ -2946,48 +3026,75 @@ async fn collect_left_input( config.execution.perfect_hash_join_min_key_density, null_equality, )? { - array_map_created_count.add(1); - metrics.build_mem_used.add(array_map.size()); + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &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)?; - (Map::ArrayMap(array_map), batch, left_value) - } else { - // Estimation of memory size, required for hashtable, prior to allocation. - // Final result can be verified using `RawTable.allocation_info()` - // 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 - let mut hashmap = new_join_hashmap(num_rows, &mut reservation, &metrics)?; - - let mut hashes_buffer = Vec::new(); - let mut offset = 0; - - let batches_iter = batches.iter().rev(); - - // Updating hashmap starting from the last batch - for batch in batches_iter.clone() { - hashes_buffer.clear(); - hashes_buffer.resize(batch.num_rows(), 0); - update_hash( - &on_left, - batch, - &mut *hashmap, - offset, - &random_state, - &mut hashes_buffer, - 0, - true, - null_equality, - )?; - offset += batch.num_rows(); - } + array_map_created_count.add(1); + metrics.build_mem_used.add(array_map.size()); - // Merge all batches into a single batch, so we can directly index into the arrays - let batch = concat_batches(&schema, batches_iter.clone())?; + (Map::ArrayMap(array_map), batch, left_values) + } else { + // Estimation of memory size, required for hashtable, prior to allocation. + // Final result can be verified using `RawTable.allocation_info()` + // 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 + let mut hashmap = new_join_hashmap(num_rows, &mut reservation, &metrics)?; + + let mut hashes_buffer = Vec::new(); + let mut offset = 0; + + // Updating hashmap starting from the last batch + for batch in batches.iter().rev() { + hashes_buffer.clear(); + hashes_buffer.resize(batch.num_rows(), 0); + update_hash( + &on_left, + batch, + &mut *hashmap, + offset, + &random_state, + &mut hashes_buffer, + 0, + true, + null_equality, + )?; + offset += batch.num_rows(); + } - let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + // 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, + )?; - (Map::HashMap(hashmap), batch, left_values) - }; + let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + + (Map::HashMap(hashmap), batch, left_values) + }; + + // Join keys that are plain columns share the buffers of `batch`, any other + // expression evaluates to new arrays that are kept for the whole join. + let mut key_counter = RecordBatchMemoryCounter::new(); + key_counter.count_batch(&batch); + let keys_size = left_values + .iter() + .map(|values| key_counter.count_array(values.as_ref())) + .sum::<usize>(); + reservation.try_grow(keys_size)?; + metrics.build_mem_used.add(keys_size); let allocate_bitmap = || -> Result<BooleanBufferBuilder> { let bitmap_size = bit_util::ceil(batch.num_rows(), 8); @@ -3207,8 +3314,9 @@ mod tests { }; use arrow::array::{ - Array, ArrayRef, AsArray, Date32Array, DictionaryArray, Int32Array, Int64Array, - StructArray, UInt32Array, UInt64Array, + Array, ArrayRef, AsArray, BinaryViewArray, Date32Array, DictionaryArray, + Int32Array, Int64Array, StringArray, StringViewArray, StructArray, UInt32Array, + UInt64Array, }; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, Field, Int32Type}; @@ -3219,6 +3327,9 @@ mod tests { exec_err, internal_err, }; use datafusion_execution::config::SessionConfig; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, + }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; @@ -7185,6 +7296,344 @@ mod tests { Ok(()) } + /// Build side batches of `num_batches` x `num_rows` rows with distinct + /// buffers: an Int32 key, a Utf8, a Utf8View and a BinaryView payload. + fn concat_test_batches(num_batches: usize, num_rows: usize) -> Vec<RecordBatch> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("s", DataType::Utf8, true), + Field::new("v", DataType::Utf8View, true), + Field::new("bv", DataType::BinaryView, true), + ])); + (0..num_batches) + .map(|b| { + let start = (b * num_rows) as i32; + let keys = Int32Array::from_iter_values(start..start + num_rows as i32); + let strings = (0..num_rows) + .map(|i| (i % 7 != 0).then(|| format!("string-{b}-{i}"))) + .collect::<StringArray>(); + let views = (0..num_rows) + .map(|i| { + (i % 5 != 0).then(|| format!("a long string view value {b}-{i}")) + }) + .collect::<StringViewArray>(); + let binary_views = (0..num_rows) + .map(|i| { + (i % 3 != 0).then(|| format!("a long binary view value {b}-{i}")) + }) + .collect::<BinaryViewArray>(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(keys), + Arc::new(strings), + Arc::new(views), + Arc::new(binary_views), + ], + ) + .unwrap() + }) + .collect() + } + + /// Reserves `batches` the way `collect_left_input` does. + fn reserve_inputs( + batches: &[RecordBatch], + pool: &Arc<dyn MemoryPool>, + ) -> Result<(MemoryReservation, usize)> { + let reservation = MemoryConsumer::new("HashJoinInput").register(pool); + let mut counter = RecordBatchMemoryCounter::new(); + for batch in batches { + reservation.try_grow(counter.count_batch(batch))?; + } + Ok((reservation, counter.memory_usage())) + } + + #[test] + fn concat_build_batches_matches_concat_batches() -> Result<()> { + let batches = concat_test_batches(4, 100); + let schema = batches[0].schema(); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + + for reverse in [false, true] { + let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default()); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + let expected = if reverse { + concat_batches(&schema, batches.iter().rev())? + } else { + concat_batches(&schema, batches.iter())? + }; + let batch = concat_build_batches( + &schema, + batches.clone(), + reverse, + inputs_reserved, + &mut reservation, + &metrics, + )?; + assert_eq!(batch, expected); + } + Ok(()) + } + + #[test] + fn concat_build_batches_reserves_copy() -> Result<()> { + let batches = concat_test_batches(4, 1000); + let schema = batches[0].schema(); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let inputs: usize = batches.iter().map(get_record_batch_memory_size).sum(); + + // The inputs fit, but not together with their concatenated copy + let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(inputs * 5 / 4)); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + assert_eq!(inputs_reserved, inputs); + let err = concat_build_batches( + &schema, + batches.clone(), + false, + inputs_reserved, + &mut reservation, + &metrics, + ) + .unwrap_err(); + assert_contains!(err.to_string(), "Resources exhausted"); + drop(reservation); + + // With room for the copy, the reservation ends up at what is retained + let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(inputs * 2)); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &mut reservation, + &metrics, + )?; + assert_eq!(reservation.size(), get_record_batch_memory_size(&batch)); + assert_eq!(pool.reserved(), reservation.size()); + Ok(()) + } + + /// Concatenating a single batch is zero-copy, so nothing more is reserved. + #[test] + fn concat_build_batches_single_batch_not_reserved_twice() -> Result<()> { + let batches = concat_test_batches(1, 1000); + let schema = batches[0].schema(); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let inputs = get_record_batch_memory_size(&batches[0]); + + let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(inputs)); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + let batch = concat_build_batches( + &schema, + batches, + true, + inputs_reserved, + &mut reservation, + &metrics, + )?; + assert_eq!(batch.num_rows(), 1000); + assert_eq!(reservation.size(), inputs); + Ok(()) + } + + /// Only the views of a view array are copied, its data buffers stay shared. + #[rstest] + fn concat_build_batches_view_data_not_reserved_twice( + #[values("v", "bv")] column: &str, + ) -> Result<()> { + let batches = concat_test_batches(4, 1000); + let column = batches[0].schema().index_of(column)?; + let batches = batches + .into_iter() + .map(|batch| batch.project(&[column])) + .collect::<Result<Vec<_>, _>>()?; + let schema = batches[0].schema(); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let inputs: usize = batches.iter().map(get_record_batch_memory_size).sum(); + let views = 4 * 1000 * size_of::<u128>(); + assert!(inputs > 2 * views); + + let pool: Arc<dyn MemoryPool> = + Arc::new(GreedyMemoryPool::new(inputs + views * 5 / 4)); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &mut reservation, + &metrics, + )?; + assert_eq!(reservation.size(), get_record_batch_memory_size(&batch)); + Ok(()) + } + + /// The build side is concatenated into a single batch, and that copy must + /// be visible to the memory pool while the input batches are still alive. + #[rstest] + #[tokio::test] + async fn join_build_concat_is_reserved( + #[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)] + mode: PartitionMode, + #[values(false, true)] use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let num_rows = 4000; + let values = (0..num_rows).collect::<Vec<i32>>(); + let batches = (0..4) + .map(|_| build_table_i32(("a1", &values), ("b1", &values), ("c1", &values))) + .collect::<Vec<_>>(); + let inputs: usize = batches.iter().map(get_record_batch_memory_size).sum(); + let left_schema = batches[0].schema(); + let right = build_table( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left_schema)?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + // Reserved on top of the inputs and their copy + let build_rows = batches.len() * num_rows as usize; + let map = if use_perfect_hash_join_as_possible { + ArrayMap::estimate_memory_size(0, num_rows as u64 - 1, build_rows) + } else { + estimate_memory_size::<(u32, u64)>(build_rows, size_of::<JoinHashMapU32>())? + + build_rows * size_of::<u32>() + }; + + for (limit, fits) in [(map + inputs * 3 / 2, false), (map + inputs * 3, true)] { + let left = TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&left_schema), + None, + )?; + let join = HashJoinExec::try_new( + left, + Arc::clone(&right), + on.clone(), + None, + &JoinType::Inner, + None, + mode, + NullEquality::NullEqualsNothing, + false, + )?; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .build_arc()?; + let task_ctx = prepare_task_ctx(8192, use_perfect_hash_join_as_possible); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(task_ctx.session_config().clone()) + .with_runtime(runtime), + ); + + let result = common::collect(join.execute(0, task_ctx)?).await; + if fits { + result?; + } else { + assert_contains!( + result.unwrap_err().to_string(), + "Resources exhausted: Additional allocation failed for HashJoinInput" + ); + } + } + Ok(()) + } + + /// Join keys that are not plain columns evaluate to new arrays, which are + /// kept for the whole join and must be reserved. + #[rstest] + #[tokio::test] + async fn join_build_key_arrays_are_reserved( + #[values(false, true)] use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let num_rows = 16000; + let values = (0..num_rows).collect::<Vec<i32>>(); + // A single build batch, so that there is no concatenated copy + let batch = build_table_i32(("a1", &values), ("b1", &values), ("c1", &values)); + let inputs = get_record_batch_memory_size(&batch); + let keys = inputs / 3; + let left_schema = batch.schema(); + let right = build_table( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ); + + let map = if use_perfect_hash_join_as_possible { + ArrayMap::estimate_memory_size(1, num_rows as u64, num_rows as usize) + } else { + estimate_memory_size::<(u32, u64)>( + num_rows as usize, + size_of::<JoinHashMapU32>(), + )? + num_rows as usize * size_of::<u32>() + }; + + for (computed_key, limit, fits) in [ + (false, inputs + map + keys / 2, true), + (true, inputs + map + keys / 2, false), + (true, inputs + map + keys * 2, true), + ] { + let column = Arc::new(Column::new_with_schema("a1", &left_schema)?) as _; + let left_key: PhysicalExprRef = if computed_key { + Arc::new(BinaryExpr::new( + column, + Operator::Plus, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )) + } else { + column + }; + let on = vec![( + left_key, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let left = TestMemoryExec::try_new_exec( + &[vec![batch.clone()]], + Arc::clone(&left_schema), + None, + )?; + let join = HashJoinExec::try_new( + left, + Arc::clone(&right), + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .build_arc()?; + let task_ctx = prepare_task_ctx(8192, use_perfect_hash_join_as_possible); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(task_ctx.session_config().clone()) + .with_runtime(runtime), + ); + + let result = common::collect(join.execute(0, task_ctx)?).await; + if fits { + result?; + } else { + assert_contains!( + result.unwrap_err().to_string(), + "Resources exhausted: Additional allocation failed for HashJoinInput" + ); + } + } + Ok(()) + } + fn build_table_struct( struct_name: &str, field_name_and_values: (&str, &Vec<Option<i32>>), --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
