From 716b46c9ba39066e720e78c8eec23491f1b0a786 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 17:26:24 +0800 Subject: [PATCH 01/24] **feat(transformer, stream): add RecordBatchMemoryCounter for retained batches, deduplicate shared buffers, and update reservation lifecycle** - RecordBatchMemoryCounter now tracks transformer retained batches. - Symmetric stream deduplicates input/transformer shared buffers. - Reservation updates occur on retain and release operations. - Added tests for noop, splitter, shared buffers, and reservation lifecycle. --- .../src/joins/symmetric_hash_join.rs | 102 ++++++++++++++++-- datafusion/physical-plan/src/joins/utils.rs | 39 +++++++ 2 files changed, 135 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 0c6e84b36cc55..837f80882b103 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -68,6 +68,7 @@ use arrow::record_batch::RecordBatch; use datafusion_common::hash_utils::create_hashes; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::bisect; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::{ HashSet, JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, plan_err, @@ -1665,11 +1666,14 @@ impl SymmetricHashJoinStream { } StatefulStreamResult::Ready(Some(batch)) => { self.batch_transformer.set_batch(batch); + self.update_reservation()?; } _ => {} } } Some((batch, _)) => { + // The transformer released this batch before it is emitted. + self.update_reservation()?; return self .metrics .baseline_metrics @@ -1921,13 +1925,23 @@ impl SymmetricHashJoinStream { self.state.clone() } + /// Returns the memory retained by this stream at its reservation boundary. + /// + /// Input buffers and a transformer-held output batch can share Arrow buffers, + /// so count them as one sequence rather than summing individual batch sizes. fn size(&self) -> usize { - let mut size = 0; + let mut batch_memory_counter = RecordBatchMemoryCounter::new(); + batch_memory_counter.count_batch(&self.left.input_buffer); + batch_memory_counter.count_batch(&self.right.input_buffer); + self.batch_transformer + .count_memory(&mut batch_memory_counter); + + let mut size = batch_memory_counter.memory_usage(); size += size_of_val(&self.schema); size += size_of_val(&self.filter); size += size_of_val(&self.join_type); - size += self.left.size(); - size += self.right.size(); + size += self.left.size() - self.left.input_buffer.get_array_memory_size(); + size += self.right.size() - self.right.input_buffer.get_array_memory_size(); size += size_of_val(&self.column_indices); size += self.graph.as_ref().map(|g| g.size()).unwrap_or(0); size += size_of_val(&self.left_sorted_filter_expr); @@ -1938,6 +1952,14 @@ impl SymmetricHashJoinStream { size } + /// Resizes the stream reservation to match all memory retained by the stream. + fn update_reservation(&mut self) -> Result<()> { + let capacity = self.size(); + self.metrics.stream_memory_usage.set(capacity); + self.reservation.try_resize(capacity)?; + Ok(()) + } + /// Performs a join operation for the specified `probe_side` (either left or right). /// This function: /// 1. Determines which side is the probe and which is the build side. @@ -2035,9 +2057,7 @@ impl SymmetricHashJoinStream { // Combine results: let result = combine_two_batches(&self.schema, equal_result, anti_result)?; - let capacity = self.size(); - self.metrics.stream_memory_usage.set(capacity); - self.reservation.try_resize(capacity)?; + self.update_reservation()?; Ok(result) } } @@ -2083,6 +2103,7 @@ mod tests { partitioned_sym_join_with_filter, split_record_batches, }; + use arrow::array::Int32Array; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; use datafusion_common::ScalarValue; @@ -2102,6 +2123,75 @@ mod tests { static TABLE_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new())); + fn create_stream( + batch_transformer: T, + input_schema: SchemaRef, + ) -> SymmetricHashJoinStream { + let context = TaskContext::default(); + let metrics = ExecutionPlanMetricsSet::new(); + SymmetricHashJoinStream { + left_stream: Box::pin(EmptyRecordBatchStream::new(Arc::clone(&input_schema))), + right_stream: Box::pin(EmptyRecordBatchStream::new(Arc::clone( + &input_schema, + ))), + schema: Arc::clone(&input_schema), + filter: None, + join_type: JoinType::Inner, + left: OneSideHashJoiner::new( + JoinSide::Left, + vec![], + Arc::clone(&input_schema), + ), + right: OneSideHashJoiner::new(JoinSide::Right, vec![], input_schema), + column_indices: vec![], + graph: None, + left_sorted_filter_expr: None, + right_sorted_filter_expr: None, + random_state: RandomState::default(), + null_equality: NullEquality::NullEqualsNothing, + metrics: StreamJoinMetrics::new(0, &metrics), + reservation: Arc::new( + MemoryConsumer::new("SymmetricHashJoinStream[test]") + .register(context.memory_pool()), + ), + state: SHJStreamState::PullRight, + batch_transformer, + } + } + + fn assert_stream_accounts_for_transformer(batch_transformer: T) { + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from_iter_values(0..10)) as _, + )]) + .unwrap(); + let expected_size = RecordBatchMemoryCounter::new().count_batch(&batch); + let mut stream = create_stream(batch_transformer, batch.schema()); + + let empty_size = stream.size(); + stream.batch_transformer.set_batch(batch.clone()); + stream.update_reservation().unwrap(); + assert_eq!(stream.size() - empty_size, expected_size); + assert_eq!(stream.reservation.size(), stream.size()); + + while stream.batch_transformer.next().is_some() {} + stream.update_reservation().unwrap(); + assert_eq!(stream.reservation.size(), empty_size); + + stream.left.input_buffer = batch; + let size_with_shared_batch = stream.size(); + stream + .batch_transformer + .set_batch(stream.left.input_buffer.clone()); + assert_eq!(stream.size(), size_with_shared_batch); + } + + #[test] + fn stream_accounts_for_transformer_batches_once() { + assert_stream_accounts_for_transformer(NoopBatchTransformer::new()); + assert_stream_accounts_for_transformer(BatchSplitter::new(3)); + } + fn get_or_create_table( cardinality: (i32, i32), batch_size: usize, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 20467a7ec5e33..dfb346c132e94 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -64,6 +64,7 @@ use datafusion_common::cast::as_boolean_array; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::stats::Precision; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::utils::normalize_float_zero; use datafusion_common::{ DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, @@ -1953,6 +1954,9 @@ pub(crate) trait BatchTransformer: Debug + Clone { /// Returns `None` if all batches have been produced. /// The boolean flag indicates whether the batch is the last one. fn next(&mut self) -> Option<(RecordBatch, bool)>; + + /// Counts buffers retained by this transformer. + fn count_memory(&self, counter: &mut RecordBatchMemoryCounter); } #[derive(Debug, Clone)] @@ -1976,6 +1980,12 @@ impl BatchTransformer for NoopBatchTransformer { fn next(&mut self) -> Option<(RecordBatch, bool)> { self.batch.take().map(|batch| (batch, true)) } + + fn count_memory(&self, counter: &mut RecordBatchMemoryCounter) { + if let Some(batch) = &self.batch { + counter.count_batch(batch); + } + } } #[derive(Debug, Clone)] @@ -2024,6 +2034,12 @@ impl BatchTransformer for BatchSplitter { Some((sliced_batch, last)) } + + fn count_memory(&self, counter: &mut RecordBatchMemoryCounter) { + if let Some(batch) = &self.batch { + counter.count_batch(batch); + } + } } /// When the order of the join inputs are changed, the output order of columns @@ -4380,6 +4396,29 @@ mod tests { } } + #[test] + fn batch_transformers_count_retained_batch_memory() { + let batch = create_test_batch(10); + let expected_size = RecordBatchMemoryCounter::new().count_batch(&batch); + + let mut noop = NoopBatchTransformer::new(); + noop.set_batch(batch.clone()); + let mut noop_counter = RecordBatchMemoryCounter::new(); + noop.count_memory(&mut noop_counter); + assert_eq!(noop_counter.memory_usage(), expected_size); + + let mut splitter = BatchSplitter::new(3); + splitter.set_batch(batch.clone()); + let mut splitter_counter = RecordBatchMemoryCounter::new(); + splitter.count_memory(&mut splitter_counter); + assert_eq!(splitter_counter.memory_usage(), expected_size); + + let mut shared_buffer_counter = RecordBatchMemoryCounter::new(); + shared_buffer_counter.count_batch(&batch); + splitter.count_memory(&mut shared_buffer_counter); + assert_eq!(shared_buffer_counter.memory_usage(), expected_size); + } + #[rstest] #[test] fn test_batch_splitter( From 958cc74b45faad1e5d07fee15b1a32e92cac75eb Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 17:44:30 +0800 Subject: [PATCH 02/24] refactor(transformer): deduplicate retained-batch accounting in private helper (no public API/behavior change) --- datafusion/physical-plan/src/joins/utils.rs | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index dfb346c132e94..3b83232997cb3 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1959,6 +1959,15 @@ pub(crate) trait BatchTransformer: Debug + Clone { fn count_memory(&self, counter: &mut RecordBatchMemoryCounter); } +fn count_retained_batch_memory( + batch: &Option, + counter: &mut RecordBatchMemoryCounter, +) { + if let Some(batch) = batch { + counter.count_batch(batch); + } +} + #[derive(Debug, Clone)] /// A batch transformer that does nothing. pub(crate) struct NoopBatchTransformer { @@ -1982,9 +1991,7 @@ impl BatchTransformer for NoopBatchTransformer { } fn count_memory(&self, counter: &mut RecordBatchMemoryCounter) { - if let Some(batch) = &self.batch { - counter.count_batch(batch); - } + count_retained_batch_memory(&self.batch, counter); } } @@ -2036,9 +2043,7 @@ impl BatchTransformer for BatchSplitter { } fn count_memory(&self, counter: &mut RecordBatchMemoryCounter) { - if let Some(batch) = &self.batch { - counter.count_batch(batch); - } + count_retained_batch_memory(&self.batch, counter); } } From 71eb7e9eca2e2d8ee214a5b8ba280a716caaa4b4 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 18:18:35 +0800 Subject: [PATCH 03/24] fix(symmetric_hash_join): preserve retained batch memory accounting * Preserve RecordBatch array-object/non-buffer overhead while deduplicating Arrow buffer capacities in memory reservations. * Ensure transformer-held batches are fully accounted for in bounded-memory enforcement. * Add wide/nested retained-batch regression coverage for non-buffer allocation overhead. * Add execution-level bounded-memory tests using SymmetricHashJoinExec with a tight GreedyMemoryPool. * Verify NoopBatchTransformer and BatchSplitter retained batches correctly affect reservations and memory-limit behavior. --- datafusion/common/src/utils/memory.rs | 55 +++++++++++++++++- .../src/joins/symmetric_hash_join.rs | 58 ++++++++++++++++++- datafusion/physical-plan/src/joins/utils.rs | 8 +-- 3 files changed, 111 insertions(+), 10 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 21c084119e120..5c8f554a17f42 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -19,10 +19,11 @@ use crate::error::_exec_datafusion_err; use crate::{HashSet, Result}; -use arrow::array::ArrayData; +use arrow::array::{Array, ArrayData}; use arrow::record_batch::RecordBatch; use std::mem::size_of; use std::num::NonZero; +use std::sync::Arc; /// Estimates the memory size required for a hash table prior to allocation. /// @@ -152,7 +153,9 @@ pub struct RecordBatchMemoryCounter { /// Start addresses of `Buffer`s that have already been counted (instead of /// actual used data region's pointer represented by current `Array`) counted_buffers: HashSet>, - /// Total memory of all unique buffers counted so far + /// Array objects already counted by [`Self::count_batch_with_array_overhead`] + counted_arrays: HashSet, + /// Total memory of all counted allocations memory_usage: usize, } @@ -179,7 +182,29 @@ impl RecordBatchMemoryCounter { total_size } - /// Total memory of the unique buffers of all batches counted so far. + /// Counts unique buffers and Array objects retained by `batch`. + /// + /// This is useful for accounting a sequence of batches at an operator + /// boundary. It counts buffers once, and also avoids double-counting a + /// top-level Arrow array shared by multiple batches. + pub fn count_batch_with_array_overhead(&mut self, batch: &RecordBatch) -> usize { + let mut total_size = self.count_batch(batch); + let mut array_overhead = 0; + + for array in batch.columns() { + let array_ptr = Arc::as_ptr(array) as *const () as usize; + if self.counted_arrays.insert(array_ptr) { + array_overhead += + array.get_array_memory_size() - array.get_buffer_memory_size(); + } + } + + total_size += array_overhead; + self.memory_usage += array_overhead; + total_size + } + + /// Total memory of all counted allocations. pub fn memory_usage(&self) -> usize { self.memory_usage } @@ -336,6 +361,30 @@ mod record_batch_tests { assert_eq!(size_origin, size_sliced); } + #[test] + fn test_record_batch_memory_counter_array_overhead_shared_across_batches() { + let schema = Arc::new(Schema::new(vec![ + Field::new("ints", DataType::Int32, false), + Field::new("floats", DataType::Float64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])), + Arc::new(Float64Array::from(vec![1., 2., 3., 4., 5., 6.])), + ], + ) + .unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + assert_eq!( + counter.count_batch_with_array_overhead(&batch), + batch.get_array_memory_size() + ); + assert_eq!(counter.count_batch_with_array_overhead(&batch), 0); + assert_eq!(counter.memory_usage(), batch.get_array_memory_size()); + } + #[test] fn test_record_batch_memory_counter_buffer_shared_across_batches() { let schema = Arc::new(Schema::new(vec![Field::new( diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 837f80882b103..bb219ccae470b 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -1931,8 +1931,8 @@ impl SymmetricHashJoinStream { /// so count them as one sequence rather than summing individual batch sizes. fn size(&self) -> usize { let mut batch_memory_counter = RecordBatchMemoryCounter::new(); - batch_memory_counter.count_batch(&self.left.input_buffer); - batch_memory_counter.count_batch(&self.right.input_buffer); + batch_memory_counter.count_batch_with_array_overhead(&self.left.input_buffer); + batch_memory_counter.count_batch_with_array_overhead(&self.right.input_buffer); self.batch_transformer .count_memory(&mut batch_memory_counter); @@ -2102,12 +2102,14 @@ mod tests { join_expr_tests_fixture_temporal, partitioned_hash_join_with_filter, partitioned_sym_join_with_filter, split_record_batches, }; + use crate::test::TestMemoryExec; use arrow::array::Int32Array; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; use datafusion_common::ScalarValue; use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{Column, binary, col, lit}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; @@ -2165,7 +2167,7 @@ mod tests { Arc::new(Int32Array::from_iter_values(0..10)) as _, )]) .unwrap(); - let expected_size = RecordBatchMemoryCounter::new().count_batch(&batch); + let expected_size = batch.get_array_memory_size(); let mut stream = create_stream(batch_transformer, batch.schema()); let empty_size = stream.size(); @@ -2192,6 +2194,56 @@ mod tests { assert_stream_accounts_for_transformer(BatchSplitter::new(3)); } + #[rstest] + #[tokio::test] + async fn symmetric_hash_join_reserves_transformer_batch( + #[values(false, true)] enforce_batch_size_in_joins: bool, + ) -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(0..10))], + )?; + let left = TestMemoryExec::try_new_exec( + &[vec![batch.clone()]], + Arc::clone(&schema), + None, + )?; + let right = + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + let on = vec![(col("id", &schema)?, col("id", &schema)?)]; + let join = SymmetricHashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?; + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(2_400, 1.0) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(1) + .with_enforce_batch_size_in_joins(enforce_batch_size_in_joins), + ) + .with_runtime(runtime), + ); + + let error = crate::common::collect(join.execute(0, context)?) + .await + .unwrap_err(); + assert!(error.to_string().contains("Additional allocation failed")); + Ok(()) + } + fn get_or_create_table( cardinality: (i32, i32), batch_size: usize, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 3b83232997cb3..98d6999d36667 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1955,7 +1955,7 @@ pub(crate) trait BatchTransformer: Debug + Clone { /// The boolean flag indicates whether the batch is the last one. fn next(&mut self) -> Option<(RecordBatch, bool)>; - /// Counts buffers retained by this transformer. + /// Counts memory retained by this transformer. fn count_memory(&self, counter: &mut RecordBatchMemoryCounter); } @@ -1964,7 +1964,7 @@ fn count_retained_batch_memory( counter: &mut RecordBatchMemoryCounter, ) { if let Some(batch) = batch { - counter.count_batch(batch); + counter.count_batch_with_array_overhead(batch); } } @@ -4404,7 +4404,7 @@ mod tests { #[test] fn batch_transformers_count_retained_batch_memory() { let batch = create_test_batch(10); - let expected_size = RecordBatchMemoryCounter::new().count_batch(&batch); + let expected_size = batch.get_array_memory_size(); let mut noop = NoopBatchTransformer::new(); noop.set_batch(batch.clone()); @@ -4419,7 +4419,7 @@ mod tests { assert_eq!(splitter_counter.memory_usage(), expected_size); let mut shared_buffer_counter = RecordBatchMemoryCounter::new(); - shared_buffer_counter.count_batch(&batch); + shared_buffer_counter.count_batch_with_array_overhead(&batch); splitter.count_memory(&mut shared_buffer_counter); assert_eq!(shared_buffer_counter.memory_usage(), expected_size); } From 4e9965ce1758285fcd837428e3fce07dfc64d249 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 18:36:23 +0800 Subject: [PATCH 04/24] feat: improve recursive unique ArrayRef overhead accounting and coverage for children (Struct, List, Map, Union, Dict, RunEnd), plus regression test for shared nested StructArray child --- datafusion/common/src/utils/memory.rs | 130 ++++++++++++++++++++++++-- 1 file changed, 122 insertions(+), 8 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 5c8f554a17f42..a5ea7b5a708ab 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -19,7 +19,12 @@ use crate::error::_exec_datafusion_err; use crate::{HashSet, Result}; -use arrow::array::{Array, ArrayData}; +use arrow::array::types::{ + Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, + UInt64Type, +}; +use arrow::array::{Array, ArrayData, ArrayRef, AsArray, RunArray}; +use arrow::datatypes::DataType; use arrow::record_batch::RecordBatch; use std::mem::size_of; use std::num::NonZero; @@ -192,11 +197,8 @@ impl RecordBatchMemoryCounter { let mut array_overhead = 0; for array in batch.columns() { - let array_ptr = Arc::as_ptr(array) as *const () as usize; - if self.counted_arrays.insert(array_ptr) { - array_overhead += - array.get_array_memory_size() - array.get_buffer_memory_size(); - } + array_overhead += + count_unique_array_object_memory_size(array, &mut self.counted_arrays); } total_size += array_overhead; @@ -210,6 +212,90 @@ impl RecordBatchMemoryCounter { } } +/// Counts the unique Array object memory retained by `array` and its children. +fn count_unique_array_object_memory_size( + array: &ArrayRef, + counted_arrays: &mut HashSet, +) -> usize { + let array_ptr = Arc::as_ptr(array) as *const () as usize; + if !counted_arrays.insert(array_ptr) { + return 0; + } + + let children = array_children(array); + let children_overhead: usize = children + .iter() + .map(|child| child.get_array_memory_size() - child.get_buffer_memory_size()) + .sum(); + let own_overhead = array.get_array_memory_size() + - array.get_buffer_memory_size() + - children_overhead; + + own_overhead + + children + .into_iter() + .map(|child| count_unique_array_object_memory_size(child, counted_arrays)) + .sum::() +} + +/// Returns the `ArrayRef` children whose object allocations may be shared. +fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> { + match array.data_type() { + DataType::Struct(_) => array.as_struct().columns().iter().collect(), + DataType::List(_) => vec![array.as_list::().values()], + DataType::LargeList(_) => vec![array.as_list::().values()], + DataType::ListView(_) => vec![array.as_list_view::().values()], + DataType::LargeListView(_) => vec![array.as_list_view::().values()], + DataType::FixedSizeList(_, _) => vec![array.as_fixed_size_list().values()], + DataType::Map(_, _) => { + let map = array.as_map(); + vec![map.keys(), map.values()] + } + DataType::Union(_, _) => array + .as_union() + .fields() + .iter() + .map(|(type_id, _)| array.as_union().child(type_id)) + .collect(), + DataType::Dictionary(key_type, _) => match key_type.as_ref() { + DataType::Int8 => vec![array.as_dictionary::().values()], + DataType::Int16 => vec![array.as_dictionary::().values()], + DataType::Int32 => vec![array.as_dictionary::().values()], + DataType::Int64 => vec![array.as_dictionary::().values()], + DataType::UInt8 => vec![array.as_dictionary::().values()], + DataType::UInt16 => vec![array.as_dictionary::().values()], + DataType::UInt32 => vec![array.as_dictionary::().values()], + DataType::UInt64 => vec![array.as_dictionary::().values()], + _ => unreachable!("invalid dictionary key type: {key_type}"), + }, + DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() { + DataType::Int16 => vec![ + array + .as_any() + .downcast_ref::>() + .expect("run-end array data type must match its run-end field") + .values(), + ], + DataType::Int32 => vec![ + array + .as_any() + .downcast_ref::>() + .expect("run-end array data type must match its run-end field") + .values(), + ], + DataType::Int64 => vec![ + array + .as_any() + .downcast_ref::>() + .expect("run-end array data type must match its run-end field") + .values(), + ], + _ => unreachable!("invalid run-end type: {run_ends}"), + }, + _ => vec![], + } +} + /// Count the memory usage of `array_data` and its children recursively. fn count_array_data_memory_size( array_data: &ArrayData, @@ -272,8 +358,8 @@ mod tests { #[cfg(test)] mod record_batch_tests { use super::*; - use arrow::array::{Float64Array, Int32Array, ListArray}; - use arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use arrow::array::{ArrayRef, Float64Array, Int32Array, ListArray, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Int32Type, Schema}; use std::sync::Arc; #[test] @@ -385,6 +471,34 @@ mod record_batch_tests { assert_eq!(counter.memory_usage(), batch.get_array_memory_size()); } + #[test] + fn test_record_batch_memory_counter_deduplicates_shared_nested_array_overhead() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields = + Fields::from(vec![Arc::new(Field::new("value", DataType::Int32, false))]); + let first = Arc::new(StructArray::new( + fields.clone(), + vec![Arc::clone(&shared_child)], + None, + )) as _; + let second = Arc::new(StructArray::new( + fields, + vec![Arc::clone(&shared_child)], + None, + )) as _; + let batch = + RecordBatch::try_from_iter(vec![("first", first), ("second", second)]) + .unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + counter.count_batch_with_array_overhead(&batch); + + assert_eq!( + counter.memory_usage(), + batch.get_array_memory_size() - shared_child.get_array_memory_size() + ); + } + #[test] fn test_record_batch_memory_counter_buffer_shared_across_batches() { let schema = Arc::new(Schema::new(vec![Field::new( From 22f6a6616958275d7edb44e6bb2a9002a0f9edd0 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 18:38:44 +0800 Subject: [PATCH 05/24] fix(stream-boundary): add nested shared-child regression test for NoopBatchTransformer and BatchSplitter --- .../src/joins/symmetric_hash_join.rs | 49 ++++++++++++++++++- 1 file changed, 47 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index bb219ccae470b..f6ff041a6b633 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2104,9 +2104,9 @@ mod tests { }; use crate::test::TestMemoryExec; - use arrow::array::Int32Array; + use arrow::array::{ArrayRef, Int32Array, StructArray}; use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; + use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit}; use datafusion_common::ScalarValue; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; @@ -2194,6 +2194,51 @@ mod tests { assert_stream_accounts_for_transformer(BatchSplitter::new(3)); } + fn assert_stream_deduplicates_nested_transformer_batch( + batch_transformer: T, + ) { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields = + Fields::from(vec![Arc::new(Field::new("value", DataType::Int32, false))]); + let left_batch = RecordBatch::try_from_iter(vec![( + "nested", + Arc::new(StructArray::new( + fields.clone(), + vec![Arc::clone(&shared_child)], + None, + )) as ArrayRef, + )]) + .unwrap(); + let transformer_batch = RecordBatch::try_from_iter(vec![( + "nested", + Arc::new(StructArray::new( + fields, + vec![Arc::clone(&shared_child)], + None, + )) as ArrayRef, + )]) + .unwrap(); + let mut stream = create_stream(batch_transformer, left_batch.schema()); + stream.left.input_buffer = left_batch; + let size_without_transformer = stream.size(); + + stream + .batch_transformer + .set_batch(transformer_batch.clone()); + + assert_eq!( + stream.size() - size_without_transformer, + transformer_batch.get_array_memory_size() + - shared_child.get_array_memory_size() + ); + } + + #[test] + fn stream_deduplicates_nested_transformer_batches() { + assert_stream_deduplicates_nested_transformer_batch(NoopBatchTransformer::new()); + assert_stream_deduplicates_nested_transformer_batch(BatchSplitter::new(3)); + } + #[rstest] #[tokio::test] async fn symmetric_hash_join_reserves_transformer_batch( From 188fea530479f1b88d3458a94d79f659644deaf2 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 18:41:59 +0800 Subject: [PATCH 06/24] docs: document BatchTransformer::count_memory shared counter and dedup contract; add regression-boundary rationale comment for 2_400 test limit --- datafusion/physical-plan/src/joins/symmetric_hash_join.rs | 2 ++ datafusion/physical-plan/src/joins/utils.rs | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index f6ff041a6b633..128de9839fead 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2269,6 +2269,8 @@ mod tests { None, StreamJoinPartitionMode::Partitioned, )?; + // This limit is intentionally at the regression boundary: accounting that + // omits the transformer-held batch succeeds, while corrected accounting fails. let runtime = RuntimeEnvBuilder::new() .with_memory_limit(2_400, 1.0) .build_arc()?; diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 98d6999d36667..321f7cc61117a 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1955,7 +1955,10 @@ pub(crate) trait BatchTransformer: Debug + Clone { /// The boolean flag indicates whether the batch is the last one. fn next(&mut self) -> Option<(RecordBatch, bool)>; - /// Counts memory retained by this transformer. + /// Adds all memory retained by this transformer to `counter`. + /// + /// The stream shares this counter with its other retained batches, so + /// implementations must let it deduplicate shared Arrow buffers and array objects. fn count_memory(&self, counter: &mut RecordBatchMemoryCounter); } From 709a488b941290d6fd827daafe0ca4a542e7eb10 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 19:03:17 +0800 Subject: [PATCH 07/24] feat(memory): add table-driven shared-child cases for List, Map, Union, Dictionary, RunEnd and improve shared allocation delta checks --- datafusion/common/src/utils/memory.rs | 160 +++++++++++++++++++++++++- 1 file changed, 158 insertions(+), 2 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index a5ea7b5a708ab..0fcd9b30d3a7d 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -358,8 +358,12 @@ mod tests { #[cfg(test)] mod record_batch_tests { use super::*; - use arrow::array::{ArrayRef, Float64Array, Int32Array, ListArray, StructArray}; - use arrow::datatypes::{DataType, Field, Fields, Int32Type, Schema}; + use arrow::array::{ + ArrayRef, DictionaryArray, Float64Array, Int32Array, ListArray, MapArray, + RunArray, StructArray, UnionArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field, Fields, Int32Type, Schema, UnionFields}; use std::sync::Arc; #[test] @@ -499,6 +503,158 @@ mod record_batch_tests { ); } + #[test] + fn test_record_batch_memory_counter_deduplicates_recursive_shared_children() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let shared_map_key: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); + let list_field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let map_fields = Fields::from(vec![ + Arc::new(Field::new("key", DataType::Int32, false)), + Arc::new(Field::new("value", DataType::Int32, false)), + ]); + let union_fields: UnionFields = + [(0, Arc::new(Field::new("value", DataType::Int32, false)))] + .into_iter() + .collect(); + + let arrays = vec![ + ( + "list", + Arc::new(ListArray::new( + Arc::clone(&list_field), + OffsetBuffer::new(vec![0, 3].into()), + Arc::clone(&shared_child), + None, + )) as ArrayRef, + Arc::new(ListArray::new( + list_field, + OffsetBuffer::new(vec![0, 3].into()), + Arc::clone(&shared_child), + None, + )) as ArrayRef, + shared_child.get_array_memory_size(), + ), + ( + "map", + Arc::new( + MapArray::try_new( + Arc::new(Field::new( + "entries", + DataType::Struct(map_fields.clone()), + false, + )), + OffsetBuffer::new(vec![0, 3].into()), + StructArray::new( + map_fields.clone(), + vec![Arc::clone(&shared_map_key), Arc::clone(&shared_child)], + None, + ), + None, + false, + ) + .unwrap(), + ) as ArrayRef, + Arc::new( + MapArray::try_new( + Arc::new(Field::new( + "entries", + DataType::Struct(map_fields.clone()), + false, + )), + OffsetBuffer::new(vec![0, 3].into()), + StructArray::new( + map_fields, + vec![Arc::clone(&shared_map_key), Arc::clone(&shared_child)], + None, + ), + None, + false, + ) + .unwrap(), + ) as ArrayRef, + shared_map_key.get_array_memory_size() + + shared_child.get_array_memory_size(), + ), + ( + "union", + Arc::new( + UnionArray::try_new( + union_fields.clone(), + vec![0, 0, 0].into(), + None, + vec![Arc::clone(&shared_child)], + ) + .unwrap(), + ) as ArrayRef, + Arc::new( + UnionArray::try_new( + union_fields, + vec![0, 0, 0].into(), + None, + vec![Arc::clone(&shared_child)], + ) + .unwrap(), + ) as ArrayRef, + shared_child.get_buffer_memory_size(), + ), + ( + "dictionary", + Arc::new( + DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 2]), + Arc::clone(&shared_child), + ) + .unwrap(), + ) as ArrayRef, + Arc::new( + DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 2]), + Arc::clone(&shared_child), + ) + .unwrap(), + ) as ArrayRef, + shared_child.get_array_memory_size(), + ), + ( + "run_end_encoded", + Arc::new( + RunArray::::try_new( + &Int32Array::from(vec![1, 2, 3]), + shared_child.as_ref(), + ) + .unwrap(), + ) as ArrayRef, + Arc::new( + RunArray::::try_new( + &Int32Array::from(vec![1, 2, 3]), + shared_child.as_ref(), + ) + .unwrap(), + ) as ArrayRef, + shared_child.get_buffer_memory_size(), + ), + ]; + + for (name, first, second, shared_memory) in arrays { + let first_batch = + RecordBatch::try_from_iter(vec![(name, first.clone())]).unwrap(); + let second_batch = + RecordBatch::try_from_iter(vec![(name, second.clone())]).unwrap(); + let mut counter = RecordBatchMemoryCounter::new(); + + assert_eq!( + counter.count_batch_with_array_overhead(&first_batch), + first.get_array_memory_size(), + "{name}: first batch" + ); + assert_eq!( + counter.count_batch_with_array_overhead(&second_batch), + second.get_array_memory_size() - shared_memory, + "{name}: shared child" + ); + } + } + #[test] fn test_record_batch_memory_counter_buffer_shared_across_batches() { let schema = Arc::new(Schema::new(vec![Field::new( From fb97c27a716c5f4909a6f9f72a6d44a31daae187 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 19:09:08 +0800 Subject: [PATCH 08/24] =?UTF-8?q?feat:=20add=20test=E2=80=91private=20asse?= =?UTF-8?q?rtion=20helper=20and=20remove=20post=E2=80=91construction=20Arr?= =?UTF-8?q?ayRef=20clones?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Introduces a private assertion helper for testing. - Removes unnecessary ArrayRef clones after construction, improving efficiency. --- datafusion/common/src/utils/memory.rs | 41 ++++++++++++++++----------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 0fcd9b30d3a7d..b9298f0d6d297 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -503,6 +503,30 @@ mod record_batch_tests { ); } + fn assert_recursive_shared_child_memory( + name: &str, + first: ArrayRef, + second: ArrayRef, + shared_memory: usize, + ) { + let first_memory = first.get_array_memory_size(); + let second_memory = second.get_array_memory_size(); + let first_batch = RecordBatch::try_from_iter(vec![(name, first)]).unwrap(); + let second_batch = RecordBatch::try_from_iter(vec![(name, second)]).unwrap(); + let mut counter = RecordBatchMemoryCounter::new(); + + assert_eq!( + counter.count_batch_with_array_overhead(&first_batch), + first_memory, + "{name}: first batch" + ); + assert_eq!( + counter.count_batch_with_array_overhead(&second_batch), + second_memory - shared_memory, + "{name}: shared child" + ); + } + #[test] fn test_record_batch_memory_counter_deduplicates_recursive_shared_children() { let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); @@ -636,22 +660,7 @@ mod record_batch_tests { ]; for (name, first, second, shared_memory) in arrays { - let first_batch = - RecordBatch::try_from_iter(vec![(name, first.clone())]).unwrap(); - let second_batch = - RecordBatch::try_from_iter(vec![(name, second.clone())]).unwrap(); - let mut counter = RecordBatchMemoryCounter::new(); - - assert_eq!( - counter.count_batch_with_array_overhead(&first_batch), - first.get_array_memory_size(), - "{name}: first batch" - ); - assert_eq!( - counter.count_batch_with_array_overhead(&second_batch), - second.get_array_memory_size() - shared_memory, - "{name}: shared child" - ); + assert_recursive_shared_child_memory(name, first, second, shared_memory); } } From 792fa412f0362896f437bbb315f8708f0bf8fbc8 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 14 Aug 2026 19:21:54 +0800 Subject: [PATCH 09/24] test(symmetric_hash_join): add tests for transformer output exhaustion and memory usage - Polls transformer output to exhaustion. - Asserts reservation + stream_memory_usage grow on retain, return baseline on release. - Covers Noop + BatchSplitter. --- .../src/joins/symmetric_hash_join.rs | 23 ++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 128de9839fead..5b0544eeb79ae 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2113,6 +2113,7 @@ mod tests { use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{Column, binary, col, lit}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use std::task::Context as PollContext; use rstest::*; @@ -2161,7 +2162,9 @@ mod tests { } } - fn assert_stream_accounts_for_transformer(batch_transformer: T) { + fn assert_stream_accounts_for_transformer( + batch_transformer: T, + ) { let batch = RecordBatch::try_from_iter(vec![( "a", Arc::new(Int32Array::from_iter_values(0..10)) as _, @@ -2175,10 +2178,24 @@ mod tests { stream.update_reservation().unwrap(); assert_eq!(stream.size() - empty_size, expected_size); assert_eq!(stream.reservation.size(), stream.size()); + assert_eq!(stream.metrics.stream_memory_usage.value(), stream.size()); - while stream.batch_transformer.next().is_some() {} - stream.update_reservation().unwrap(); + loop { + let poll = stream.poll_next_unpin(&mut PollContext::from_waker( + futures::task::noop_waker_ref(), + )); + match poll { + Poll::Ready(Some(Ok(_))) => {} + Poll::Ready(None) => break, + Poll::Ready(Some(Err(error))) => { + panic!("unexpected stream error: {error}") + } + Poll::Pending => panic!("empty input streams must not pend"), + } + } + assert_eq!(stream.size(), empty_size); assert_eq!(stream.reservation.size(), empty_size); + assert_eq!(stream.metrics.stream_memory_usage.value(), empty_size); stream.left.input_buffer = batch; let size_with_shared_batch = stream.size(); From a31dc42dca1261afabb8cc3ee2bd96a2c92a89fe Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 15 Aug 2026 11:17:01 +0800 Subject: [PATCH 10/24] feat(pool): replace fixed pool limit with dynamic peak profiling and tighten reserves - Replace fixed 2_400 pool limit with dynamic peak profiling using PeakRecordingPool. - Tighten pool reserves to stay below transformer batch delta. - Assert join values and introduce typed ResourcesExhausted. - Ensure coverage of Noop and BatchSplitter. --- datafusion/common/src/utils/memory.rs | 249 +++++++++--------- .../src/joins/symmetric_hash_join.rs | 104 ++++++-- 2 files changed, 208 insertions(+), 145 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index b9298f0d6d297..e40a2bd6cab36 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -528,140 +528,143 @@ mod record_batch_tests { } #[test] - fn test_record_batch_memory_counter_deduplicates_recursive_shared_children() { + fn test_record_batch_memory_counter_deduplicates_shared_list_child() { let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let shared_map_key: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); let list_field = Arc::new(Field::new_list_field(DataType::Int32, false)); - let map_fields = Fields::from(vec![ + + let first = Arc::new(ListArray::new( + Arc::clone(&list_field), + OffsetBuffer::new(vec![0, 3].into()), + Arc::clone(&shared_child), + None, + )); + let second = Arc::new(ListArray::new( + list_field, + OffsetBuffer::new(vec![0, 3].into()), + Arc::clone(&shared_child), + None, + )); + + assert_recursive_shared_child_memory( + "list", + first, + second, + shared_child.get_array_memory_size(), + ); + } + + fn map_with_shared_children( + fields: &Fields, + shared_key: &ArrayRef, + shared_value: &ArrayRef, + ) -> ArrayRef { + Arc::new( + MapArray::try_new( + Arc::new(Field::new( + "entries", + DataType::Struct(fields.clone()), + false, + )), + OffsetBuffer::new(vec![0, 3].into()), + StructArray::new( + fields.clone(), + vec![Arc::clone(shared_key), Arc::clone(shared_value)], + None, + ), + None, + false, + ) + .unwrap(), + ) + } + + #[test] + fn test_record_batch_memory_counter_deduplicates_shared_map_children() { + let shared_key: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); + let shared_value: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields = Fields::from(vec![ Arc::new(Field::new("key", DataType::Int32, false)), Arc::new(Field::new("value", DataType::Int32, false)), ]); - let union_fields: UnionFields = + + let first = map_with_shared_children(&fields, &shared_key, &shared_value); + let second = map_with_shared_children(&fields, &shared_key, &shared_value); + + assert_recursive_shared_child_memory( + "map", + first, + second, + shared_key.get_array_memory_size() + shared_value.get_array_memory_size(), + ); + } + + #[test] + fn test_record_batch_memory_counter_deduplicates_shared_union_child() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields: UnionFields = [(0, Arc::new(Field::new("value", DataType::Int32, false)))] .into_iter() .collect(); - - let arrays = vec![ - ( - "list", - Arc::new(ListArray::new( - Arc::clone(&list_field), - OffsetBuffer::new(vec![0, 3].into()), - Arc::clone(&shared_child), + let make_union = || { + Arc::new( + UnionArray::try_new( + fields.clone(), + vec![0, 0, 0].into(), None, - )) as ArrayRef, - Arc::new(ListArray::new( - list_field, - OffsetBuffer::new(vec![0, 3].into()), + vec![Arc::clone(&shared_child)], + ) + .unwrap(), + ) as ArrayRef + }; + + assert_recursive_shared_child_memory( + "union", + make_union(), + make_union(), + shared_child.get_buffer_memory_size(), + ); + } + + #[test] + fn test_record_batch_memory_counter_deduplicates_shared_dictionary_child() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let make_dictionary = || { + Arc::new( + DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 2]), Arc::clone(&shared_child), - None, - )) as ArrayRef, - shared_child.get_array_memory_size(), - ), - ( - "map", - Arc::new( - MapArray::try_new( - Arc::new(Field::new( - "entries", - DataType::Struct(map_fields.clone()), - false, - )), - OffsetBuffer::new(vec![0, 3].into()), - StructArray::new( - map_fields.clone(), - vec![Arc::clone(&shared_map_key), Arc::clone(&shared_child)], - None, - ), - None, - false, - ) - .unwrap(), - ) as ArrayRef, - Arc::new( - MapArray::try_new( - Arc::new(Field::new( - "entries", - DataType::Struct(map_fields.clone()), - false, - )), - OffsetBuffer::new(vec![0, 3].into()), - StructArray::new( - map_fields, - vec![Arc::clone(&shared_map_key), Arc::clone(&shared_child)], - None, - ), - None, - false, - ) - .unwrap(), - ) as ArrayRef, - shared_map_key.get_array_memory_size() - + shared_child.get_array_memory_size(), - ), - ( - "union", - Arc::new( - UnionArray::try_new( - union_fields.clone(), - vec![0, 0, 0].into(), - None, - vec![Arc::clone(&shared_child)], - ) - .unwrap(), - ) as ArrayRef, - Arc::new( - UnionArray::try_new( - union_fields, - vec![0, 0, 0].into(), - None, - vec![Arc::clone(&shared_child)], - ) - .unwrap(), - ) as ArrayRef, - shared_child.get_buffer_memory_size(), - ), - ( - "dictionary", - Arc::new( - DictionaryArray::::try_new( - Int32Array::from(vec![0, 1, 2]), - Arc::clone(&shared_child), - ) - .unwrap(), - ) as ArrayRef, - Arc::new( - DictionaryArray::::try_new( - Int32Array::from(vec![0, 1, 2]), - Arc::clone(&shared_child), - ) - .unwrap(), - ) as ArrayRef, - shared_child.get_array_memory_size(), - ), - ( - "run_end_encoded", - Arc::new( - RunArray::::try_new( - &Int32Array::from(vec![1, 2, 3]), - shared_child.as_ref(), - ) - .unwrap(), - ) as ArrayRef, - Arc::new( - RunArray::::try_new( - &Int32Array::from(vec![1, 2, 3]), - shared_child.as_ref(), - ) - .unwrap(), - ) as ArrayRef, - shared_child.get_buffer_memory_size(), - ), - ]; + ) + .unwrap(), + ) as ArrayRef + }; + + assert_recursive_shared_child_memory( + "dictionary", + make_dictionary(), + make_dictionary(), + shared_child.get_array_memory_size(), + ); + } - for (name, first, second, shared_memory) in arrays { - assert_recursive_shared_child_memory(name, first, second, shared_memory); - } + #[test] + fn test_record_batch_memory_counter_deduplicates_shared_run_end_encoded_child() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let make_run_array = || { + Arc::new( + RunArray::::try_new( + &Int32Array::from(vec![1, 2, 3]), + shared_child.as_ref(), + ) + .unwrap(), + ) as ArrayRef + }; + + assert_recursive_shared_child_memory( + "run_end_encoded", + make_run_array(), + make_run_array(), + shared_child.get_buffer_memory_size(), + ); } #[test] diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 5b0544eeb79ae..abc03db96b6cc 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2104,12 +2104,15 @@ mod tests { }; use crate::test::TestMemoryExec; - use arrow::array::{ArrayRef, Int32Array, StructArray}; + use arrow::array::{ArrayRef, AsArray, Int32Array, StructArray}; use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit}; - use datafusion_common::ScalarValue; + use arrow::datatypes::{DataType, Field, Fields, Int32Type, IntervalUnit, TimeUnit}; + use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_execution::config::SessionConfig; - use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryPool, PeakRecordingPool, UnboundedMemoryPool, + }; + use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{Column, binary, col, lit}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; @@ -2256,11 +2259,7 @@ mod tests { assert_stream_deduplicates_nested_transformer_batch(BatchSplitter::new(3)); } - #[rstest] - #[tokio::test] - async fn symmetric_hash_join_reserves_transformer_batch( - #[values(false, true)] enforce_batch_size_in_joins: bool, - ) -> Result<()> { + fn transformer_memory_test_join() -> Result { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); let batch = RecordBatch::try_new( @@ -2275,7 +2274,7 @@ mod tests { let right = TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; let on = vec![(col("id", &schema)?, col("id", &schema)?)]; - let join = SymmetricHashJoinExec::try_new( + SymmetricHashJoinExec::try_new( left, right, on, @@ -2285,13 +2284,14 @@ mod tests { None, None, StreamJoinPartitionMode::Partitioned, - )?; - // This limit is intentionally at the regression boundary: accounting that - // omits the transformer-held batch succeeds, while corrected accounting fails. - let runtime = RuntimeEnvBuilder::new() - .with_memory_limit(2_400, 1.0) - .build_arc()?; - let context = Arc::new( + ) + } + + fn transformer_memory_test_context( + runtime: Arc, + enforce_batch_size_in_joins: bool, + ) -> Arc { + Arc::new( TaskContext::default() .with_session_config( SessionConfig::new() @@ -2299,12 +2299,72 @@ mod tests { .with_enforce_batch_size_in_joins(enforce_batch_size_in_joins), ) .with_runtime(runtime), - ); + ) + } - let error = crate::common::collect(join.execute(0, context)?) - .await - .unwrap_err(); - assert!(error.to_string().contains("Additional allocation failed")); + fn assert_transformer_memory_test_output(batches: &[RecordBatch]) { + let actual_rows = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .iter() + .zip(batch.column(1).as_primitive::().values()) + .map(|(&left, &right)| (left, right)) + .collect::>() + }) + .collect::>(); + let expected_rows = (0..10).map(|id| (id, id)).collect::>(); + assert_eq!(actual_rows, expected_rows); + } + + #[rstest] + #[tokio::test] + async fn symmetric_hash_join_reserves_transformer_batch( + #[values(false, true)] enforce_batch_size_in_joins: bool, + ) -> Result<()> { + let recording_pool = Arc::new(PeakRecordingPool::new(Arc::new( + UnboundedMemoryPool::default(), + ))); + let pool: Arc = Arc::clone(&recording_pool) as _; + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc()?; + let batches = crate::common::collect(transformer_memory_test_join()?.execute( + 0, + transformer_memory_test_context(runtime, enforce_batch_size_in_joins), + )?) + .await?; + assert_transformer_memory_test_output(&batches); + + let peak_reservation = recording_pool.peak_reserved(); + let transformer_batch_memory = batches + .iter() + .map(RecordBatch::get_array_memory_size) + .max() + .expect("join emits a transformer batch"); + let memory_limit = peak_reservation + .checked_sub(transformer_batch_memory) + .expect("transformer batch is part of the peak reservation") + + 1; + + // The direct stream tests prove this batch is the reservation delta. This + // limit leaves room for the old accounting but not the retained batch. + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))) + .build_arc()?; + let error = crate::common::collect(transformer_memory_test_join()?.execute( + 0, + transformer_memory_test_context(runtime, enforce_batch_size_in_joins), + )?) + .await + .unwrap_err(); + assert!( + matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)), + "expected a memory-pool error, got: {error}" + ); Ok(()) } From b039d6c9be255bb3cb6d55bc10628571473b8a19 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 15 Aug 2026 11:29:29 +0800 Subject: [PATCH 11/24] fix(pool): adjust limit to observed peak - 1 and clarify ownership comment - Updated pool limit calculation to be observed peak - 1. - Added comment clarifying direct vs integration ownership. - All focused tests pass, and formatting/diff checks succeed. --- .../src/joins/symmetric_hash_join.rs | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index abc03db96b6cc..cdf08d8f76576 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2340,18 +2340,13 @@ mod tests { assert_transformer_memory_test_output(&batches); let peak_reservation = recording_pool.peak_reserved(); - let transformer_batch_memory = batches - .iter() - .map(RecordBatch::get_array_memory_size) - .max() - .expect("join emits a transformer batch"); let memory_limit = peak_reservation - .checked_sub(transformer_batch_memory) - .expect("transformer batch is part of the peak reservation") - + 1; + .checked_sub(1) + .expect("join must reserve memory"); - // The direct stream tests prove this batch is the reservation delta. This - // limit leaves room for the old accounting but not the retained batch. + // Direct stream tests prove transformer retain/release accounting and + // shared-buffer deduplication. This test verifies the physical plan + // enforces its observed peak reservation. let runtime = RuntimeEnvBuilder::new() .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))) .build_arc()?; From 3106c2ad5377d15cb35c70df5fbbb1f8063fd2a7 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 15 Aug 2026 11:39:37 +0800 Subject: [PATCH 12/24] docs(common/utils): update documentation to clarify array counting includes shared objects --- datafusion/common/src/utils/memory.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index e40a2bd6cab36..f077b98d80ce8 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -190,8 +190,8 @@ impl RecordBatchMemoryCounter { /// Counts unique buffers and Array objects retained by `batch`. /// /// This is useful for accounting a sequence of batches at an operator - /// boundary. It counts buffers once, and also avoids double-counting a - /// top-level Arrow array shared by multiple batches. + /// boundary. It counts buffers and recursively reachable Arrow array objects + /// once, including objects shared by multiple batches. pub fn count_batch_with_array_overhead(&mut self, batch: &RecordBatch) -> usize { let mut total_size = self.count_batch(batch); let mut array_overhead = 0; From c83dedb87f013f825eed5b292bd8f0deb7647927 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 15 Aug 2026 13:25:53 +0800 Subject: [PATCH 13/24] =?UTF-8?q?test:=20replace=20noisy=20whole=E2=80=91e?= =?UTF-8?q?xec=20OOM=20test,=20add=20exact=20stream=E2=80=91boundary=20bou?= =?UTF-8?q?nded=E2=80=91pool=20regression,=20and=20cover=20Noop=20+=20Batc?= =?UTF-8?q?hSplitter;=20ablation=20confirms=20count=5Fmemory=20is=20requir?= =?UTF-8?q?ed=20(transformer=20batch=20exceeds=20by=201=20byte)**?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/joins/symmetric_hash_join.rs | 161 +++++------------- 1 file changed, 39 insertions(+), 122 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index cdf08d8f76576..d60f9b6ac4afc 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2102,21 +2102,15 @@ mod tests { join_expr_tests_fixture_temporal, partitioned_hash_join_with_filter, partitioned_sym_join_with_filter, split_record_batches, }; - use crate::test::TestMemoryExec; - - use arrow::array::{ArrayRef, AsArray, Int32Array, StructArray}; + use arrow::array::{ArrayRef, Int32Array, StructArray}; use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, Fields, Int32Type, IntervalUnit, TimeUnit}; + use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit}; use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_execution::config::SessionConfig; - use datafusion_execution::memory_pool::{ - GreedyMemoryPool, MemoryPool, PeakRecordingPool, UnboundedMemoryPool, - }; - use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{Column, binary, col, lit}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; - use std::task::Context as PollContext; use rstest::*; @@ -2134,6 +2128,14 @@ mod tests { input_schema: SchemaRef, ) -> SymmetricHashJoinStream { let context = TaskContext::default(); + create_stream_with_context(batch_transformer, input_schema, &context) + } + + fn create_stream_with_context( + batch_transformer: T, + input_schema: SchemaRef, + context: &TaskContext, + ) -> SymmetricHashJoinStream { let metrics = ExecutionPlanMetricsSet::new(); SymmetricHashJoinStream { left_stream: Box::pin(EmptyRecordBatchStream::new(Arc::clone(&input_schema))), @@ -2165,9 +2167,7 @@ mod tests { } } - fn assert_stream_accounts_for_transformer( - batch_transformer: T, - ) { + fn assert_stream_accounts_for_transformer(batch_transformer: T) { let batch = RecordBatch::try_from_iter(vec![( "a", Arc::new(Int32Array::from_iter_values(0..10)) as _, @@ -2181,24 +2181,10 @@ mod tests { stream.update_reservation().unwrap(); assert_eq!(stream.size() - empty_size, expected_size); assert_eq!(stream.reservation.size(), stream.size()); - assert_eq!(stream.metrics.stream_memory_usage.value(), stream.size()); - loop { - let poll = stream.poll_next_unpin(&mut PollContext::from_waker( - futures::task::noop_waker_ref(), - )); - match poll { - Poll::Ready(Some(Ok(_))) => {} - Poll::Ready(None) => break, - Poll::Ready(Some(Err(error))) => { - panic!("unexpected stream error: {error}") - } - Poll::Pending => panic!("empty input streams must not pend"), - } - } - assert_eq!(stream.size(), empty_size); + while stream.batch_transformer.next().is_some() {} + stream.update_reservation().unwrap(); assert_eq!(stream.reservation.size(), empty_size); - assert_eq!(stream.metrics.stream_memory_usage.value(), empty_size); stream.left.input_buffer = batch; let size_with_shared_batch = stream.size(); @@ -2259,103 +2245,27 @@ mod tests { assert_stream_deduplicates_nested_transformer_batch(BatchSplitter::new(3)); } - fn transformer_memory_test_join() -> Result { - let schema = - Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int32Array::from_iter_values(0..10))], - )?; - let left = TestMemoryExec::try_new_exec( - &[vec![batch.clone()]], - Arc::clone(&schema), - None, - )?; - let right = - TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; - let on = vec![(col("id", &schema)?, col("id", &schema)?)]; - SymmetricHashJoinExec::try_new( - left, - right, - on, - None, - &JoinType::Inner, - NullEquality::NullEqualsNothing, - None, - None, - StreamJoinPartitionMode::Partitioned, - ) - } - - fn transformer_memory_test_context( - runtime: Arc, - enforce_batch_size_in_joins: bool, - ) -> Arc { - Arc::new( - TaskContext::default() - .with_session_config( - SessionConfig::new() - .with_batch_size(1) - .with_enforce_batch_size_in_joins(enforce_batch_size_in_joins), - ) - .with_runtime(runtime), - ) - } - - fn assert_transformer_memory_test_output(batches: &[RecordBatch]) { - let actual_rows = batches - .iter() - .flat_map(|batch| { - batch - .column(0) - .as_primitive::() - .values() - .iter() - .zip(batch.column(1).as_primitive::().values()) - .map(|(&left, &right)| (left, right)) - .collect::>() - }) - .collect::>(); - let expected_rows = (0..10).map(|id| (id, id)).collect::>(); - assert_eq!(actual_rows, expected_rows); - } - - #[rstest] - #[tokio::test] - async fn symmetric_hash_join_reserves_transformer_batch( - #[values(false, true)] enforce_batch_size_in_joins: bool, + fn assert_transformer_reservation_exhausts_pool( + batch_transformer: T, ) -> Result<()> { - let recording_pool = Arc::new(PeakRecordingPool::new(Arc::new( - UnboundedMemoryPool::default(), - ))); - let pool: Arc = Arc::clone(&recording_pool) as _; - let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(pool) - .build_arc()?; - let batches = crate::common::collect(transformer_memory_test_join()?.execute( - 0, - transformer_memory_test_context(runtime, enforce_batch_size_in_joins), - )?) - .await?; - assert_transformer_memory_test_output(&batches); - - let peak_reservation = recording_pool.peak_reserved(); - let memory_limit = peak_reservation - .checked_sub(1) - .expect("join must reserve memory"); - - // Direct stream tests prove transformer retain/release accounting and - // shared-buffer deduplication. This test verifies the physical plan - // enforces its observed peak reservation. + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from_iter_values(0..10)) as _, + )])?; + let empty_size = + create_stream(NoopBatchTransformer::new(), batch.schema()).size(); let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))) + .with_memory_limit(empty_size + batch.get_array_memory_size() - 1, 1.0) .build_arc()?; - let error = crate::common::collect(transformer_memory_test_join()?.execute( - 0, - transformer_memory_test_context(runtime, enforce_batch_size_in_joins), - )?) - .await - .unwrap_err(); + let context = TaskContext::default().with_runtime(runtime); + let mut stream = + create_stream_with_context(batch_transformer, batch.schema(), &context); + + // The empty stream fits, but retaining this batch exceeds the exact + // stream reservation boundary by one byte. + stream.update_reservation()?; + stream.batch_transformer.set_batch(batch); + let error = stream.update_reservation().unwrap_err(); assert!( matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)), "expected a memory-pool error, got: {error}" @@ -2363,6 +2273,13 @@ mod tests { Ok(()) } + #[test] + fn transformer_reservation_exhausts_pool() -> Result<()> { + assert_transformer_reservation_exhausts_pool(NoopBatchTransformer::new())?; + assert_transformer_reservation_exhausts_pool(BatchSplitter::new(3))?; + Ok(()) + } + fn get_or_create_table( cardinality: (i32, i32), batch_size: usize, From 4769cb704d679051b6c8aa35cd0cd9367dc605aa Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 15 Aug 2026 14:17:20 +0800 Subject: [PATCH 14/24] =?UTF-8?q?=20Fix=20whole-stream=20OOM=20assertion?= =?UTF-8?q?=20=20-=20Added=20real=20SymmetricHashJoinExec=20poll-lifecycle?= =?UTF-8?q?=20regression.=20=20-=20Verifies=20retain/release=20reservation?= =?UTF-8?q?=20calls:=20=20=20=20=20=20-=20Noop:=20retain=20=E2=86=92=20imm?= =?UTF-8?q?ediate=20release.=20=20=20=20=20=20-=20BatchSplitter:=20retaine?= =?UTF-8?q?d=20after=20non-final=20slice=20=E2=86=92=20released=20final=20?= =?UTF-8?q?=20=20=20=20=20=20=20slice.=20=20-=20Temporary=20ablation=20rem?= =?UTF-8?q?oving=20production=20resize=20calls:=20both=20cases=20fail.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/joins/symmetric_hash_join.rs | 169 +++++++++++++++++- 1 file changed, 168 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index d60f9b6ac4afc..47eb227278724 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2102,12 +2102,16 @@ mod tests { join_expr_tests_fixture_temporal, partitioned_hash_join_with_filter, partitioned_sym_join_with_filter, split_record_batches, }; + use crate::test::TestMemoryExec; use arrow::array::{ArrayRef, Int32Array, StructArray}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit}; use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_execution::config::SessionConfig; - use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_execution::memory_pool::{ + MemoryLimit, MemoryPool, MemoryReservation, UnboundedMemoryPool, + }; + use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{Column, binary, col, lit}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; @@ -2280,6 +2284,169 @@ mod tests { Ok(()) } + #[derive(Debug, Default)] + struct RecordingMemoryPool { + inner: UnboundedMemoryPool, + symmetric_join_changes: Mutex>, + } + + impl fmt::Display for RecordingMemoryPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&self.inner, f) + } + } + + impl RecordingMemoryPool { + fn record(&self, reservation: &MemoryReservation, change: isize) { + if reservation + .consumer() + .name() + .starts_with("SymmetricHashJoinStream") + { + self.symmetric_join_changes + .lock() + .expect("recording pool mutex is not poisoned") + .push(change); + } + } + + fn changes(&self) -> Vec { + self.symmetric_join_changes + .lock() + .expect("recording pool mutex is not poisoned") + .clone() + } + } + + impl MemoryPool for RecordingMemoryPool { + fn name(&self) -> &str { + self.inner.name() + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record(reservation, additional as isize); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + self.record(reservation, -(shrink as isize)); + } + + fn try_grow( + &self, + reservation: &MemoryReservation, + additional: usize, + ) -> Result<()> { + self.inner.try_grow(reservation, additional)?; + self.record(reservation, additional as isize); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } + } + + fn transformer_lifecycle_test_join() -> Result { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(0..10))], + )?; + let left = TestMemoryExec::try_new_exec( + &[vec![batch.clone()]], + Arc::clone(&schema), + None, + )?; + let right = + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + let on = vec![(col("id", &schema)?, col("id", &schema)?)]; + SymmetricHashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + ) + } + + fn transformer_lifecycle_test_context( + runtime: Arc, + enforce_batch_size_in_joins: bool, + ) -> Arc { + Arc::new( + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(3) + .with_enforce_batch_size_in_joins(enforce_batch_size_in_joins), + ) + .with_runtime(runtime), + ) + } + + #[rstest] + #[tokio::test] + async fn symmetric_hash_join_updates_reservation_while_transforming_output( + #[values(false, true)] enforce_batch_size_in_joins: bool, + ) -> Result<()> { + let pool = Arc::new(RecordingMemoryPool::default()); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build_arc()?; + let mut stream = transformer_lifecycle_test_join()?.execute( + 0, + transformer_lifecycle_test_context(runtime, enforce_batch_size_in_joins), + )?; + + let first_batch = stream.next().await.transpose()?.unwrap(); + let retained_batch_size = first_batch.get_array_memory_size() as isize; + let changes_after_first_batch = pool.changes(); + assert!( + changes_after_first_batch.contains(&retained_batch_size), + "expected transformer retain reservation update: {changes_after_first_batch:?}" + ); + if enforce_batch_size_in_joins { + assert!( + !changes_after_first_batch.contains(&-retained_batch_size), + "splitter must retain the output batch after a non-final slice: {changes_after_first_batch:?}" + ); + } else { + assert!( + changes_after_first_batch + .windows(2) + .any(|changes| changes == [retained_batch_size, -retained_batch_size]), + "expected Noop transformer release after emission: {changes_after_first_batch:?}" + ); + } + + let remaining_batches = crate::common::collect(stream).await?; + let changes_after_completion = pool.changes(); + if enforce_batch_size_in_joins { + assert!( + changes_after_completion.contains(&-retained_batch_size), + "expected splitter release after final slice: {changes_after_completion:?}" + ); + } + let output_rows = first_batch.num_rows() + + remaining_batches + .iter() + .map(RecordBatch::num_rows) + .sum::(); + assert_eq!(output_rows, 10); + Ok(()) + } + fn get_or_create_table( cardinality: (i32, i32), batch_size: usize, From bb6e5efd0f312745bb9d86e26cb60a948a5b4f66 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 15 Aug 2026 22:11:18 +0800 Subject: [PATCH 15/24] =?UTF-8?q?fix(datafusion/physical-plan/src/joins/sy?= =?UTF-8?q?mmetric=5Fhash=5Fjoin):=20adjust=20memory=20reservation=20accou?= =?UTF-8?q?nting,=20add=20finite=E2=80=91pool=20regression=20test,=20pin?= =?UTF-8?q?=20retain=20event=20by=20poll=20lifecycle=20position,=20restore?= =?UTF-8?q?=20count=5Fmemory=20handling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/joins/symmetric_hash_join.rs | 83 ++++++++++++++++++- 1 file changed, 79 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 47eb227278724..28b97c2a2c6f9 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2284,10 +2284,16 @@ mod tests { Ok(()) } + #[derive(Clone, Debug)] + struct RecordedReservationChange { + change: isize, + reserved: usize, + } + #[derive(Debug, Default)] struct RecordingMemoryPool { inner: UnboundedMemoryPool, - symmetric_join_changes: Mutex>, + symmetric_join_changes: Mutex>, } impl fmt::Display for RecordingMemoryPool { @@ -2306,11 +2312,14 @@ mod tests { self.symmetric_join_changes .lock() .expect("recording pool mutex is not poisoned") - .push(change); + .push(RecordedReservationChange { + change, + reserved: self.inner.reserved(), + }); } } - fn changes(&self) -> Vec { + fn changes(&self) -> Vec { self.symmetric_join_changes .lock() .expect("recording pool mutex is not poisoned") @@ -2412,6 +2421,10 @@ mod tests { let first_batch = stream.next().await.transpose()?.unwrap(); let retained_batch_size = first_batch.get_array_memory_size() as isize; let changes_after_first_batch = pool.changes(); + let changes_after_first_batch: Vec<_> = changes_after_first_batch + .iter() + .map(|change| change.change) + .collect(); assert!( changes_after_first_batch.contains(&retained_batch_size), "expected transformer retain reservation update: {changes_after_first_batch:?}" @@ -2431,7 +2444,11 @@ mod tests { } let remaining_batches = crate::common::collect(stream).await?; - let changes_after_completion = pool.changes(); + let changes_after_completion: Vec<_> = pool + .changes() + .into_iter() + .map(|change| change.change) + .collect(); if enforce_batch_size_in_joins { assert!( changes_after_completion.contains(&-retained_batch_size), @@ -2447,6 +2464,64 @@ mod tests { Ok(()) } + #[rstest] + #[tokio::test] + async fn symmetric_hash_join_transformer_retention_exhausts_bounded_pool( + #[values(false, true)] enforce_batch_size_in_joins: bool, + ) -> Result<()> { + let calibration_pool = Arc::new(RecordingMemoryPool::default()); + let calibration_runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&calibration_pool) as Arc) + .build_arc()?; + let mut calibration_stream = transformer_lifecycle_test_join()?.execute( + 0, + transformer_lifecycle_test_context( + calibration_runtime, + enforce_batch_size_in_joins, + ), + )?; + + let first_batch = calibration_stream.next().await.transpose()?.unwrap(); + let retained_batch_size = first_batch.get_array_memory_size() as isize; + let changes = calibration_pool.changes(); + // `poll_next_impl` has no reservation operation after transformer `next()`. + // Therefore the retain operation is last for a splitter and penultimate + // for Noop, whose `next()` immediately releases the retained batch. + let retain_index = changes + .len() + .checked_sub(if enforce_batch_size_in_joins { 1 } else { 2 }) + .expect("transformer retain must update the stream reservation"); + let retain_change = &changes[retain_index]; + assert_eq!( + retain_change.change, retained_batch_size, + "expected transformer retain reservation update: {changes:?}" + ); + if !enforce_batch_size_in_joins { + assert_eq!( + changes.last().map(|change| change.change), + Some(-retained_batch_size), + "Noop must release the retained batch before emitting it: {changes:?}" + ); + } + + // This limit is between the reservation immediately before transformer + // retention and the reservation immediately after it. The old accounting + // omitted this growth and would emit the first batch instead of failing. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(retain_change.reserved - 1, 1.0) + .build_arc()?; + let mut stream = transformer_lifecycle_test_join()?.execute( + 0, + transformer_lifecycle_test_context(runtime, enforce_batch_size_in_joins), + )?; + let error = stream.next().await.transpose().unwrap_err(); + assert!( + matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)), + "expected transformer retention to exhaust the pool, got: {error}" + ); + Ok(()) + } + fn get_or_create_table( cardinality: (i32, i32), batch_size: usize, From 48ef0d52381c4d3372bf925cdd54cffc66dd7eaf Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 15 Aug 2026 22:22:37 +0800 Subject: [PATCH 16/24] feat: bound union once, dedup transformer test assertion, add change_deltas test helper, skip risky cache/macro ideas --- datafusion/common/src/utils/memory.rs | 14 ++++++++------ .../src/joins/symmetric_hash_join.rs | 19 +++++++++---------- datafusion/physical-plan/src/joins/utils.rs | 17 +++++++++++------ 3 files changed, 28 insertions(+), 22 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index f077b98d80ce8..f8508808d4b13 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -251,12 +251,14 @@ fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> { let map = array.as_map(); vec![map.keys(), map.values()] } - DataType::Union(_, _) => array - .as_union() - .fields() - .iter() - .map(|(type_id, _)| array.as_union().child(type_id)) - .collect(), + DataType::Union(_, _) => { + let union = array.as_union(); + union + .fields() + .iter() + .map(|(type_id, _)| union.child(type_id)) + .collect() + } DataType::Dictionary(key_type, _) => match key_type.as_ref() { DataType::Int8 => vec![array.as_dictionary::().values()], DataType::Int16 => vec![array.as_dictionary::().values()], diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 28b97c2a2c6f9..e7e2b1eb13d12 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2325,6 +2325,13 @@ mod tests { .expect("recording pool mutex is not poisoned") .clone() } + + fn change_deltas(&self) -> Vec { + self.changes() + .into_iter() + .map(|change| change.change) + .collect() + } } impl MemoryPool for RecordingMemoryPool { @@ -2420,11 +2427,7 @@ mod tests { let first_batch = stream.next().await.transpose()?.unwrap(); let retained_batch_size = first_batch.get_array_memory_size() as isize; - let changes_after_first_batch = pool.changes(); - let changes_after_first_batch: Vec<_> = changes_after_first_batch - .iter() - .map(|change| change.change) - .collect(); + let changes_after_first_batch = pool.change_deltas(); assert!( changes_after_first_batch.contains(&retained_batch_size), "expected transformer retain reservation update: {changes_after_first_batch:?}" @@ -2444,11 +2447,7 @@ mod tests { } let remaining_batches = crate::common::collect(stream).await?; - let changes_after_completion: Vec<_> = pool - .changes() - .into_iter() - .map(|change| change.change) - .collect(); + let changes_after_completion = pool.change_deltas(); if enforce_batch_size_in_joins { assert!( changes_after_completion.contains(&-retained_batch_size), diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 321f7cc61117a..31f5a0cc7c30b 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -4404,6 +4404,15 @@ mod tests { } } + fn assert_transformer_memory( + transformer: &impl BatchTransformer, + expected_size: usize, + ) { + let mut counter = RecordBatchMemoryCounter::new(); + transformer.count_memory(&mut counter); + assert_eq!(counter.memory_usage(), expected_size); + } + #[test] fn batch_transformers_count_retained_batch_memory() { let batch = create_test_batch(10); @@ -4411,15 +4420,11 @@ mod tests { let mut noop = NoopBatchTransformer::new(); noop.set_batch(batch.clone()); - let mut noop_counter = RecordBatchMemoryCounter::new(); - noop.count_memory(&mut noop_counter); - assert_eq!(noop_counter.memory_usage(), expected_size); + assert_transformer_memory(&noop, expected_size); let mut splitter = BatchSplitter::new(3); splitter.set_batch(batch.clone()); - let mut splitter_counter = RecordBatchMemoryCounter::new(); - splitter.count_memory(&mut splitter_counter); - assert_eq!(splitter_counter.memory_usage(), expected_size); + assert_transformer_memory(&splitter, expected_size); let mut shared_buffer_counter = RecordBatchMemoryCounter::new(); shared_buffer_counter.count_batch_with_array_overhead(&batch); From 52e5f0092cfc0f568e1368384ab3c79b79c9b1a5 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Sat, 15 Aug 2026 22:38:38 +0800 Subject: [PATCH 17/24] refactor: change_deltas() now projects under one mutex lock, removing clone - Ensures change_deltas() uses a single mutex lock for thread safety. - Removes unnecessary cloning of data within the function. --- datafusion/physical-plan/src/joins/symmetric_hash_join.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index e7e2b1eb13d12..865bdcfc177c4 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2327,8 +2327,10 @@ mod tests { } fn change_deltas(&self) -> Vec { - self.changes() - .into_iter() + self.symmetric_join_changes + .lock() + .expect("recording pool mutex is not poisoned") + .iter() .map(|change| change.change) .collect() } From eb20ffddf53e8292997576d5c4e07c78f3417b6f Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 18 Aug 2026 11:56:24 +0800 Subject: [PATCH 18/24] refactor(join): remove duplicate tests and simplify recording pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Removed duplicate direct‑stream, transformer, nested, OOM/calibration tests. - Simplified recording pool using `Vec` and removed change structs/snapshots. - Kept a real poll‑lifecycle regression test (Noop + splitter retain/release). --- .../src/joins/symmetric_hash_join.rs | 245 +----------------- datafusion/physical-plan/src/joins/utils.rs | 28 -- 2 files changed, 6 insertions(+), 267 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 865bdcfc177c4..8dd888877a236 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2103,10 +2103,10 @@ mod tests { partitioned_sym_join_with_filter, split_record_batches, }; use crate::test::TestMemoryExec; - use arrow::array::{ArrayRef, Int32Array, StructArray}; + use arrow::array::Int32Array; use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit}; - use datafusion_common::{DataFusionError, ScalarValue}; + use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; + use datafusion_common::ScalarValue; use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::{ MemoryLimit, MemoryPool, MemoryReservation, UnboundedMemoryPool, @@ -2127,173 +2127,10 @@ mod tests { static TABLE_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new())); - fn create_stream( - batch_transformer: T, - input_schema: SchemaRef, - ) -> SymmetricHashJoinStream { - let context = TaskContext::default(); - create_stream_with_context(batch_transformer, input_schema, &context) - } - - fn create_stream_with_context( - batch_transformer: T, - input_schema: SchemaRef, - context: &TaskContext, - ) -> SymmetricHashJoinStream { - let metrics = ExecutionPlanMetricsSet::new(); - SymmetricHashJoinStream { - left_stream: Box::pin(EmptyRecordBatchStream::new(Arc::clone(&input_schema))), - right_stream: Box::pin(EmptyRecordBatchStream::new(Arc::clone( - &input_schema, - ))), - schema: Arc::clone(&input_schema), - filter: None, - join_type: JoinType::Inner, - left: OneSideHashJoiner::new( - JoinSide::Left, - vec![], - Arc::clone(&input_schema), - ), - right: OneSideHashJoiner::new(JoinSide::Right, vec![], input_schema), - column_indices: vec![], - graph: None, - left_sorted_filter_expr: None, - right_sorted_filter_expr: None, - random_state: RandomState::default(), - null_equality: NullEquality::NullEqualsNothing, - metrics: StreamJoinMetrics::new(0, &metrics), - reservation: Arc::new( - MemoryConsumer::new("SymmetricHashJoinStream[test]") - .register(context.memory_pool()), - ), - state: SHJStreamState::PullRight, - batch_transformer, - } - } - - fn assert_stream_accounts_for_transformer(batch_transformer: T) { - let batch = RecordBatch::try_from_iter(vec![( - "a", - Arc::new(Int32Array::from_iter_values(0..10)) as _, - )]) - .unwrap(); - let expected_size = batch.get_array_memory_size(); - let mut stream = create_stream(batch_transformer, batch.schema()); - - let empty_size = stream.size(); - stream.batch_transformer.set_batch(batch.clone()); - stream.update_reservation().unwrap(); - assert_eq!(stream.size() - empty_size, expected_size); - assert_eq!(stream.reservation.size(), stream.size()); - - while stream.batch_transformer.next().is_some() {} - stream.update_reservation().unwrap(); - assert_eq!(stream.reservation.size(), empty_size); - - stream.left.input_buffer = batch; - let size_with_shared_batch = stream.size(); - stream - .batch_transformer - .set_batch(stream.left.input_buffer.clone()); - assert_eq!(stream.size(), size_with_shared_batch); - } - - #[test] - fn stream_accounts_for_transformer_batches_once() { - assert_stream_accounts_for_transformer(NoopBatchTransformer::new()); - assert_stream_accounts_for_transformer(BatchSplitter::new(3)); - } - - fn assert_stream_deduplicates_nested_transformer_batch( - batch_transformer: T, - ) { - let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let fields = - Fields::from(vec![Arc::new(Field::new("value", DataType::Int32, false))]); - let left_batch = RecordBatch::try_from_iter(vec![( - "nested", - Arc::new(StructArray::new( - fields.clone(), - vec![Arc::clone(&shared_child)], - None, - )) as ArrayRef, - )]) - .unwrap(); - let transformer_batch = RecordBatch::try_from_iter(vec![( - "nested", - Arc::new(StructArray::new( - fields, - vec![Arc::clone(&shared_child)], - None, - )) as ArrayRef, - )]) - .unwrap(); - let mut stream = create_stream(batch_transformer, left_batch.schema()); - stream.left.input_buffer = left_batch; - let size_without_transformer = stream.size(); - - stream - .batch_transformer - .set_batch(transformer_batch.clone()); - - assert_eq!( - stream.size() - size_without_transformer, - transformer_batch.get_array_memory_size() - - shared_child.get_array_memory_size() - ); - } - - #[test] - fn stream_deduplicates_nested_transformer_batches() { - assert_stream_deduplicates_nested_transformer_batch(NoopBatchTransformer::new()); - assert_stream_deduplicates_nested_transformer_batch(BatchSplitter::new(3)); - } - - fn assert_transformer_reservation_exhausts_pool( - batch_transformer: T, - ) -> Result<()> { - let batch = RecordBatch::try_from_iter(vec![( - "a", - Arc::new(Int32Array::from_iter_values(0..10)) as _, - )])?; - let empty_size = - create_stream(NoopBatchTransformer::new(), batch.schema()).size(); - let runtime = RuntimeEnvBuilder::new() - .with_memory_limit(empty_size + batch.get_array_memory_size() - 1, 1.0) - .build_arc()?; - let context = TaskContext::default().with_runtime(runtime); - let mut stream = - create_stream_with_context(batch_transformer, batch.schema(), &context); - - // The empty stream fits, but retaining this batch exceeds the exact - // stream reservation boundary by one byte. - stream.update_reservation()?; - stream.batch_transformer.set_batch(batch); - let error = stream.update_reservation().unwrap_err(); - assert!( - matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)), - "expected a memory-pool error, got: {error}" - ); - Ok(()) - } - - #[test] - fn transformer_reservation_exhausts_pool() -> Result<()> { - assert_transformer_reservation_exhausts_pool(NoopBatchTransformer::new())?; - assert_transformer_reservation_exhausts_pool(BatchSplitter::new(3))?; - Ok(()) - } - - #[derive(Clone, Debug)] - struct RecordedReservationChange { - change: isize, - reserved: usize, - } - #[derive(Debug, Default)] struct RecordingMemoryPool { inner: UnboundedMemoryPool, - symmetric_join_changes: Mutex>, + symmetric_join_changes: Mutex>, } impl fmt::Display for RecordingMemoryPool { @@ -2312,27 +2149,15 @@ mod tests { self.symmetric_join_changes .lock() .expect("recording pool mutex is not poisoned") - .push(RecordedReservationChange { - change, - reserved: self.inner.reserved(), - }); + .push(change); } } - fn changes(&self) -> Vec { - self.symmetric_join_changes - .lock() - .expect("recording pool mutex is not poisoned") - .clone() - } - fn change_deltas(&self) -> Vec { self.symmetric_join_changes .lock() .expect("recording pool mutex is not poisoned") - .iter() - .map(|change| change.change) - .collect() + .clone() } } @@ -2465,64 +2290,6 @@ mod tests { Ok(()) } - #[rstest] - #[tokio::test] - async fn symmetric_hash_join_transformer_retention_exhausts_bounded_pool( - #[values(false, true)] enforce_batch_size_in_joins: bool, - ) -> Result<()> { - let calibration_pool = Arc::new(RecordingMemoryPool::default()); - let calibration_runtime = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&calibration_pool) as Arc) - .build_arc()?; - let mut calibration_stream = transformer_lifecycle_test_join()?.execute( - 0, - transformer_lifecycle_test_context( - calibration_runtime, - enforce_batch_size_in_joins, - ), - )?; - - let first_batch = calibration_stream.next().await.transpose()?.unwrap(); - let retained_batch_size = first_batch.get_array_memory_size() as isize; - let changes = calibration_pool.changes(); - // `poll_next_impl` has no reservation operation after transformer `next()`. - // Therefore the retain operation is last for a splitter and penultimate - // for Noop, whose `next()` immediately releases the retained batch. - let retain_index = changes - .len() - .checked_sub(if enforce_batch_size_in_joins { 1 } else { 2 }) - .expect("transformer retain must update the stream reservation"); - let retain_change = &changes[retain_index]; - assert_eq!( - retain_change.change, retained_batch_size, - "expected transformer retain reservation update: {changes:?}" - ); - if !enforce_batch_size_in_joins { - assert_eq!( - changes.last().map(|change| change.change), - Some(-retained_batch_size), - "Noop must release the retained batch before emitting it: {changes:?}" - ); - } - - // This limit is between the reservation immediately before transformer - // retention and the reservation immediately after it. The old accounting - // omitted this growth and would emit the first batch instead of failing. - let runtime = RuntimeEnvBuilder::new() - .with_memory_limit(retain_change.reserved - 1, 1.0) - .build_arc()?; - let mut stream = transformer_lifecycle_test_join()?.execute( - 0, - transformer_lifecycle_test_context(runtime, enforce_batch_size_in_joins), - )?; - let error = stream.next().await.transpose().unwrap_err(); - assert!( - matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)), - "expected transformer retention to exhaust the pool, got: {error}" - ); - Ok(()) - } - fn get_or_create_table( cardinality: (i32, i32), batch_size: usize, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 31f5a0cc7c30b..c1b93c2d8f92c 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -4404,34 +4404,6 @@ mod tests { } } - fn assert_transformer_memory( - transformer: &impl BatchTransformer, - expected_size: usize, - ) { - let mut counter = RecordBatchMemoryCounter::new(); - transformer.count_memory(&mut counter); - assert_eq!(counter.memory_usage(), expected_size); - } - - #[test] - fn batch_transformers_count_retained_batch_memory() { - let batch = create_test_batch(10); - let expected_size = batch.get_array_memory_size(); - - let mut noop = NoopBatchTransformer::new(); - noop.set_batch(batch.clone()); - assert_transformer_memory(&noop, expected_size); - - let mut splitter = BatchSplitter::new(3); - splitter.set_batch(batch.clone()); - assert_transformer_memory(&splitter, expected_size); - - let mut shared_buffer_counter = RecordBatchMemoryCounter::new(); - shared_buffer_counter.count_batch_with_array_overhead(&batch); - splitter.count_memory(&mut shared_buffer_counter); - assert_eq!(shared_buffer_counter.memory_usage(), expected_size); - } - #[rstest] #[test] fn test_batch_splitter( From 00f20ed5afc17e0468036f8964c9d6c717046c80 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 18 Aug 2026 12:21:25 +0800 Subject: [PATCH 19/24] fix(symmetric_hash_join): restore direct stream alias tests and exact bounded stream reservation test; add retain-transition test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Restored direct stream alias tests for Noop + Splitter and shared batch + nested shared child. - Restored exact bounded stream reservation test. - Added bounded SymmetricHashJoinExec retain-transition test that validates independent retained‑batch delta; test fails when `count_memory` is removed. --- .../src/joins/symmetric_hash_join.rs | 220 +++++++++++++++++- 1 file changed, 217 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 8dd888877a236..ca1ebe8ba0099 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2103,10 +2103,10 @@ mod tests { partitioned_sym_join_with_filter, split_record_batches, }; use crate::test::TestMemoryExec; - use arrow::array::Int32Array; + use arrow::array::{ArrayRef, Int32Array, StructArray}; use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; - use datafusion_common::ScalarValue; + use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit}; + use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::{ MemoryLimit, MemoryPool, MemoryReservation, UnboundedMemoryPool, @@ -2127,6 +2127,161 @@ mod tests { static TABLE_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new())); + fn create_stream( + batch_transformer: T, + input_schema: SchemaRef, + ) -> SymmetricHashJoinStream { + let context = TaskContext::default(); + create_stream_with_context(batch_transformer, input_schema, &context) + } + + fn create_stream_with_context( + batch_transformer: T, + input_schema: SchemaRef, + context: &TaskContext, + ) -> SymmetricHashJoinStream { + let metrics = ExecutionPlanMetricsSet::new(); + SymmetricHashJoinStream { + left_stream: Box::pin(EmptyRecordBatchStream::new(Arc::clone(&input_schema))), + right_stream: Box::pin(EmptyRecordBatchStream::new(Arc::clone( + &input_schema, + ))), + schema: Arc::clone(&input_schema), + filter: None, + join_type: JoinType::Inner, + left: OneSideHashJoiner::new( + JoinSide::Left, + vec![], + Arc::clone(&input_schema), + ), + right: OneSideHashJoiner::new(JoinSide::Right, vec![], input_schema), + column_indices: vec![], + graph: None, + left_sorted_filter_expr: None, + right_sorted_filter_expr: None, + random_state: RandomState::default(), + null_equality: NullEquality::NullEqualsNothing, + metrics: StreamJoinMetrics::new(0, &metrics), + reservation: Arc::new( + MemoryConsumer::new("SymmetricHashJoinStream[test]") + .register(context.memory_pool()), + ), + state: SHJStreamState::PullRight, + batch_transformer, + } + } + + fn assert_stream_accounts_for_transformer(batch_transformer: T) { + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from_iter_values(0..10)) as _, + )]) + .unwrap(); + let expected_size = batch.get_array_memory_size(); + let mut stream = create_stream(batch_transformer, batch.schema()); + + let empty_size = stream.size(); + stream.batch_transformer.set_batch(batch.clone()); + stream.update_reservation().unwrap(); + assert_eq!(stream.size() - empty_size, expected_size); + assert_eq!(stream.reservation.size(), stream.size()); + + while stream.batch_transformer.next().is_some() {} + stream.update_reservation().unwrap(); + assert_eq!(stream.reservation.size(), empty_size); + + stream.left.input_buffer = batch; + let size_with_shared_batch = stream.size(); + stream + .batch_transformer + .set_batch(stream.left.input_buffer.clone()); + assert_eq!(stream.size(), size_with_shared_batch); + } + + #[test] + fn stream_accounts_for_transformer_batches_once() { + assert_stream_accounts_for_transformer(NoopBatchTransformer::new()); + assert_stream_accounts_for_transformer(BatchSplitter::new(3)); + } + + fn assert_stream_deduplicates_nested_transformer_batch( + batch_transformer: T, + ) { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields = + Fields::from(vec![Arc::new(Field::new("value", DataType::Int32, false))]); + let left_batch = RecordBatch::try_from_iter(vec![( + "nested", + Arc::new(StructArray::new( + fields.clone(), + vec![Arc::clone(&shared_child)], + None, + )) as ArrayRef, + )]) + .unwrap(); + let transformer_batch = RecordBatch::try_from_iter(vec![( + "nested", + Arc::new(StructArray::new( + fields, + vec![Arc::clone(&shared_child)], + None, + )) as ArrayRef, + )]) + .unwrap(); + let mut stream = create_stream(batch_transformer, left_batch.schema()); + stream.left.input_buffer = left_batch; + let size_without_transformer = stream.size(); + + stream + .batch_transformer + .set_batch(transformer_batch.clone()); + + assert_eq!( + stream.size() - size_without_transformer, + transformer_batch.get_array_memory_size() + - shared_child.get_array_memory_size() + ); + } + + #[test] + fn stream_deduplicates_nested_transformer_batches() { + assert_stream_deduplicates_nested_transformer_batch(NoopBatchTransformer::new()); + assert_stream_deduplicates_nested_transformer_batch(BatchSplitter::new(3)); + } + + fn assert_transformer_reservation_exhausts_pool( + batch_transformer: T, + ) -> Result<()> { + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from_iter_values(0..10)) as _, + )])?; + let empty_size = + create_stream(NoopBatchTransformer::new(), batch.schema()).size(); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(empty_size + batch.get_array_memory_size() - 1, 1.0) + .build_arc()?; + let context = TaskContext::default().with_runtime(runtime); + let mut stream = + create_stream_with_context(batch_transformer, batch.schema(), &context); + + stream.update_reservation()?; + stream.batch_transformer.set_batch(batch); + let error = stream.update_reservation().unwrap_err(); + assert!( + matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)), + "expected a memory-pool error, got: {error}" + ); + Ok(()) + } + + #[test] + fn transformer_reservation_exhausts_pool() -> Result<()> { + assert_transformer_reservation_exhausts_pool(NoopBatchTransformer::new())?; + assert_transformer_reservation_exhausts_pool(BatchSplitter::new(3))?; + Ok(()) + } + #[derive(Debug, Default)] struct RecordingMemoryPool { inner: UnboundedMemoryPool, @@ -2290,6 +2445,65 @@ mod tests { Ok(()) } + #[rstest] + #[tokio::test] + async fn symmetric_hash_join_transformer_retention_exhausts_bounded_pool( + #[values(false, true)] enforce_batch_size_in_joins: bool, + ) -> Result<()> { + let calibration_pool = Arc::new(RecordingMemoryPool::default()); + let calibration_runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&calibration_pool) as Arc) + .build_arc()?; + let mut calibration_stream = transformer_lifecycle_test_join()?.execute( + 0, + transformer_lifecycle_test_context( + calibration_runtime, + enforce_batch_size_in_joins, + ), + )?; + + let first_batch = calibration_stream.next().await.transpose()?.unwrap(); + let retained_size = first_batch.get_array_memory_size() as isize; + let changes = calibration_pool.change_deltas(); + let retain_index = changes + .len() + .checked_sub(if enforce_batch_size_in_joins { 1 } else { 2 }) + .expect("transformer retain must update the stream reservation"); + assert_eq!( + changes[retain_index], retained_size, + "expected transformer retain reservation update: {changes:?}" + ); + if !enforce_batch_size_in_joins { + assert_eq!( + changes.last(), + Some(&-retained_size), + "Noop must release the retained batch before emitting it: {changes:?}" + ); + } + let pre_retain_size: isize = changes[..retain_index].iter().sum(); + assert!( + pre_retain_size > 0, + "expected a positive reservation before retention: {changes:?}" + ); + + // This limit is one byte below the independently measured retained + // output batch added to the reservation immediately before it. + // Without retain-side accounting, the first batch is emitted. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit((pre_retain_size + retained_size) as usize - 1, 1.0) + .build_arc()?; + let mut stream = transformer_lifecycle_test_join()?.execute( + 0, + transformer_lifecycle_test_context(runtime, enforce_batch_size_in_joins), + )?; + let error = stream.next().await.transpose().unwrap_err(); + assert!( + matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)), + "expected transformer retention to exhaust the pool, got: {error}" + ); + Ok(()) + } + fn get_or_create_table( cardinality: (i32, i32), batch_size: usize, From 901b57ac080123b07e4900376d0905b26768df72 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 18 Aug 2026 14:10:30 +0800 Subject: [PATCH 20/24] fix: fix reservation leak and stream memory metric reset, add retained-stream terminal-poll regression test --- .../physical-plan/src/joins/symmetric_hash_join.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index ca1ebe8ba0099..b14a52558a3ff 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -1656,6 +1656,8 @@ impl SymmetricHashJoinStream { final_result: false, } => self.prepare_for_final_results_after_exhaustion(), SHJStreamState::BothExhausted { final_result: true } => { + self.reservation.free(); + self.metrics.stream_memory_usage.set(0); return Poll::Ready(None); } }; @@ -2204,6 +2206,18 @@ mod tests { assert_stream_accounts_for_transformer(BatchSplitter::new(3)); } + #[tokio::test] + async fn symmetric_hash_join_releases_reservation_when_complete() { + let schema = Arc::new(Schema::empty()); + let mut stream = create_stream(NoopBatchTransformer::new(), schema); + stream.update_reservation().unwrap(); + assert_ne!(stream.reservation.size(), 0); + + stream.set_state(SHJStreamState::BothExhausted { final_result: true }); + assert!(stream.next().await.is_none()); + assert_eq!(stream.reservation.size(), 0); + } + fn assert_stream_deduplicates_nested_transformer_batch( batch_transformer: T, ) { From 200662ba7c2d55e65dff30e33a6296c7d86b3ee9 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Thu, 20 Aug 2026 13:23:09 +0800 Subject: [PATCH 21/24] =?UTF-8?q?fix(runend,union):=20remove=20object?= =?UTF-8?q?=E2=80=91identity=20branches/tests,=20add=20size=5Fwithout=5Fin?= =?UTF-8?q?put=5Fbuffer,=20update=20metrics=20after=20successful=20resize,?= =?UTF-8?q?=20add=20failed=E2=80=91resize=20regression=20assertion?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Removes Union/RunEnd object‑identity branches/tests. - Adds size_without_input_buffer(). - Updates metrics only after successful resize. - Adds failed‑resize metric regression assertion. --- datafusion/common/src/utils/memory.rs | 82 +------------------ .../src/joins/symmetric_hash_join.rs | 12 ++- 2 files changed, 10 insertions(+), 84 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index ad273e20c8da7..5af12b6cefe07 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -412,14 +412,6 @@ fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> { let map = array.as_map(); vec![map.keys(), map.values()] } - DataType::Union(_, _) => { - let union = array.as_union(); - union - .fields() - .iter() - .map(|(type_id, _)| union.child(type_id)) - .collect() - } DataType::Dictionary(key_type, _) => match key_type.as_ref() { DataType::Int8 => vec![array.as_dictionary::().values()], DataType::Int16 => vec![array.as_dictionary::().values()], @@ -431,30 +423,6 @@ fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> { DataType::UInt64 => vec![array.as_dictionary::().values()], _ => unreachable!("invalid dictionary key type: {key_type}"), }, - DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() { - DataType::Int16 => vec![ - array - .as_any() - .downcast_ref::>() - .expect("run-end array data type must match its run-end field") - .values(), - ], - DataType::Int32 => vec![ - array - .as_any() - .downcast_ref::>() - .expect("run-end array data type must match its run-end field") - .values(), - ], - DataType::Int64 => vec![ - array - .as_any() - .downcast_ref::>() - .expect("run-end array data type must match its run-end field") - .values(), - ], - _ => unreachable!("invalid run-end type: {run_ends}"), - }, _ => vec![], } } @@ -542,7 +510,7 @@ mod record_batch_tests { use arrow::array::{ ArrayData, ArrayRef, BinaryViewArray, DictionaryArray, Float64Array, Int16Array, Int32Array, Int64Array, LargeListViewArray, ListArray, ListViewArray, MapArray, - RunArray, StringArray, StringViewArray, StructArray, UnionArray, new_null_array, + RunArray, StringArray, StringViewArray, StructArray, new_null_array, }; use arrow::buffer::OffsetBuffer; use arrow::datatypes::{ @@ -816,33 +784,6 @@ mod record_batch_tests { ); } - #[test] - fn test_record_batch_memory_counter_deduplicates_shared_union_child() { - let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let fields: UnionFields = - [(0, Arc::new(Field::new("value", DataType::Int32, false)))] - .into_iter() - .collect(); - let make_union = || { - Arc::new( - UnionArray::try_new( - fields.clone(), - vec![0, 0, 0].into(), - None, - vec![Arc::clone(&shared_child)], - ) - .unwrap(), - ) as ArrayRef - }; - - assert_recursive_shared_child_memory( - "union", - make_union(), - make_union(), - shared_child.get_buffer_memory_size(), - ); - } - #[test] fn test_record_batch_memory_counter_deduplicates_shared_dictionary_child() { let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); @@ -864,27 +805,6 @@ mod record_batch_tests { ); } - #[test] - fn test_record_batch_memory_counter_deduplicates_shared_run_end_encoded_child() { - let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let make_run_array = || { - Arc::new( - RunArray::::try_new( - &Int32Array::from(vec![1, 2, 3]), - shared_child.as_ref(), - ) - .unwrap(), - ) as ArrayRef - }; - - assert_recursive_shared_child_memory( - "run_end_encoded", - make_run_array(), - make_run_array(), - shared_child.get_buffer_memory_size(), - ); - } - #[test] fn test_record_batch_memory_counter_buffer_shared_across_batches() { let schema = Arc::new(Schema::new(vec![Field::new( diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index b14a52558a3ff..0dd39db4fceeb 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -1455,6 +1455,10 @@ impl OneSideHashJoiner { size += size_of_val(&self.deleted_offset); size } + + fn size_without_input_buffer(&self) -> usize { + self.size() - self.input_buffer.get_array_memory_size() + } pub fn new( build_side: JoinSide, on: Vec, @@ -1942,8 +1946,8 @@ impl SymmetricHashJoinStream { size += size_of_val(&self.schema); size += size_of_val(&self.filter); size += size_of_val(&self.join_type); - size += self.left.size() - self.left.input_buffer.get_array_memory_size(); - size += self.right.size() - self.right.input_buffer.get_array_memory_size(); + size += self.left.size_without_input_buffer(); + size += self.right.size_without_input_buffer(); size += size_of_val(&self.column_indices); size += self.graph.as_ref().map(|g| g.size()).unwrap_or(0); size += size_of_val(&self.left_sorted_filter_expr); @@ -1957,8 +1961,8 @@ impl SymmetricHashJoinStream { /// Resizes the stream reservation to match all memory retained by the stream. fn update_reservation(&mut self) -> Result<()> { let capacity = self.size(); - self.metrics.stream_memory_usage.set(capacity); self.reservation.try_resize(capacity)?; + self.metrics.stream_memory_usage.set(capacity); Ok(()) } @@ -2280,12 +2284,14 @@ mod tests { create_stream_with_context(batch_transformer, batch.schema(), &context); stream.update_reservation()?; + let reserved_size = stream.reservation.size(); stream.batch_transformer.set_batch(batch); let error = stream.update_reservation().unwrap_err(); assert!( matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)), "expected a memory-pool error, got: {error}" ); + assert_eq!(stream.metrics.stream_memory_usage.value(), reserved_size); Ok(()) } From a1127188603e04c5a39d29cbddda90396d8c5ee1 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 26 Aug 2026 19:22:18 +0800 Subject: [PATCH 22/24] =?UTF-8?q?fix(stream,=20traversal):=20preserve=20re?= =?UTF-8?q?servation/metric=20in=20terminal=20streams,=20restore=20Union/R?= =?UTF-8?q?unEnd=20child=20traversal,=20and=20add=20slice=E2=80=91sharing?= =?UTF-8?q?=20regression=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Ensure that terminal streams retain their reservation and associated metric information until the stream is explicitly dropped, preventing premature release or loss of accounting data. - Restore correct child traversal behavior for `Union` and `RunEnd` nodes, aligning the implementation with the documented identity guarantees and fixing the regression that caused incorrect traversal order. - Add comprehensive regression tests for slice‑sharing scenarios to catch future breakage and verify that shared slices maintain proper reservation and metric tracking. --- datafusion/common/src/utils/memory.rs | 110 +++++++++++++++++- .../src/joins/symmetric_hash_join.rs | 24 ++-- 2 files changed, 124 insertions(+), 10 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 5af12b6cefe07..90b627e4f705f 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -195,8 +195,10 @@ impl RecordBatchMemoryCounter { /// Counts unique buffers and Array objects retained by `batch`. /// /// This is useful for accounting a sequence of batches at an operator - /// boundary. It counts buffers and recursively reachable Arrow array objects - /// once, including objects shared by multiple batches. + /// boundary. It counts buffers once, including buffers shared by multiple + /// batches. Array-object overhead is deduplicated recursively when a shared + /// child has the same `ArrayRef` identity; Arrow constructors that rebuild + /// child array objects may conservatively count their object overhead again. pub fn count_batch_with_array_overhead(&mut self, batch: &RecordBatch) -> usize { let mut total_size = self.count_batch(batch); let mut array_overhead = 0; @@ -412,6 +414,14 @@ fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> { let map = array.as_map(); vec![map.keys(), map.values()] } + DataType::Union(_, _) => { + let union = array.as_union(); + union + .fields() + .iter() + .map(|(type_id, _)| union.child(type_id)) + .collect() + } DataType::Dictionary(key_type, _) => match key_type.as_ref() { DataType::Int8 => vec![array.as_dictionary::().values()], DataType::Int16 => vec![array.as_dictionary::().values()], @@ -423,6 +433,24 @@ fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> { DataType::UInt64 => vec![array.as_dictionary::().values()], _ => unreachable!("invalid dictionary key type: {key_type}"), }, + DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() { + DataType::Int16 => array + .as_any() + .downcast_ref::>() + .map(|array| vec![array.values()]) + .unwrap_or_default(), + DataType::Int32 => array + .as_any() + .downcast_ref::>() + .map(|array| vec![array.values()]) + .unwrap_or_default(), + DataType::Int64 => array + .as_any() + .downcast_ref::>() + .map(|array| vec![array.values()]) + .unwrap_or_default(), + _ => vec![], + }, _ => vec![], } } @@ -510,7 +538,7 @@ mod record_batch_tests { use arrow::array::{ ArrayData, ArrayRef, BinaryViewArray, DictionaryArray, Float64Array, Int16Array, Int32Array, Int64Array, LargeListViewArray, ListArray, ListViewArray, MapArray, - RunArray, StringArray, StringViewArray, StructArray, new_null_array, + RunArray, StringArray, StringViewArray, StructArray, UnionArray, new_null_array, }; use arrow::buffer::OffsetBuffer; use arrow::datatypes::{ @@ -784,6 +812,82 @@ mod record_batch_tests { ); } + #[test] + fn test_record_batch_memory_counter_deduplicates_union_slice_child_overhead() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields: UnionFields = + [(0, Arc::new(Field::new("value", DataType::Int32, false)))] + .into_iter() + .collect(); + let first = Arc::new( + UnionArray::try_new( + fields, + vec![0, 0, 0].into(), + Some(vec![0, 1, 2].into()), + vec![Arc::clone(&shared_child)], + ) + .unwrap(), + ) as ArrayRef; + let second = Arc::new(first.as_union().slice(0, first.len())) as ArrayRef; + let first_batch = + RecordBatch::try_from_iter(vec![("union", Arc::clone(&first))]).unwrap(); + let second_batch = + RecordBatch::try_from_iter(vec![("union", Arc::clone(&second))]).unwrap(); + let child_overhead = + shared_child.get_array_memory_size() - shared_child.get_buffer_memory_size(); + let second_parent_overhead = second.get_array_memory_size() + - second.get_buffer_memory_size() + - child_overhead; + let mut counter = RecordBatchMemoryCounter::new(); + + assert_eq!( + counter.count_batch_with_array_overhead(&first_batch), + first.get_array_memory_size() + ); + assert_eq!( + counter.count_batch_with_array_overhead(&second_batch), + second_parent_overhead + ); + } + + #[test] + fn test_record_batch_memory_counter_deduplicates_run_end_slice_child_overhead() { + let values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let first = Arc::new( + RunArray::::try_new( + &Int32Array::from(vec![1, 2, 3]), + values.as_ref(), + ) + .unwrap(), + ) as ArrayRef; + let second = Arc::new( + first + .as_any() + .downcast_ref::>() + .unwrap() + .slice(0, first.len()), + ) as ArrayRef; + let first_batch = + RecordBatch::try_from_iter(vec![("run", Arc::clone(&first))]).unwrap(); + let second_batch = + RecordBatch::try_from_iter(vec![("run", Arc::clone(&second))]).unwrap(); + let child_overhead = + values.get_array_memory_size() - values.get_buffer_memory_size(); + let second_parent_overhead = second.get_array_memory_size() + - second.get_buffer_memory_size() + - child_overhead; + let mut counter = RecordBatchMemoryCounter::new(); + + assert_eq!( + counter.count_batch_with_array_overhead(&first_batch), + first.get_array_memory_size() + ); + assert_eq!( + counter.count_batch_with_array_overhead(&second_batch), + second_parent_overhead + ); + } + #[test] fn test_record_batch_memory_counter_deduplicates_shared_dictionary_child() { let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 0dd39db4fceeb..2484f28987c59 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -1660,8 +1660,6 @@ impl SymmetricHashJoinStream { final_result: false, } => self.prepare_for_final_results_after_exhaustion(), SHJStreamState::BothExhausted { final_result: true } => { - self.reservation.free(); - self.metrics.stream_memory_usage.set(0); return Poll::Ready(None); } }; @@ -2211,15 +2209,27 @@ mod tests { } #[tokio::test] - async fn symmetric_hash_join_releases_reservation_when_complete() { + async fn symmetric_hash_join_retains_reservation_until_dropped() -> Result<()> { + let pool = Arc::new(RecordingMemoryPool::default()); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build_arc()?; + let context = TaskContext::default().with_runtime(runtime); let schema = Arc::new(Schema::empty()); - let mut stream = create_stream(NoopBatchTransformer::new(), schema); - stream.update_reservation().unwrap(); - assert_ne!(stream.reservation.size(), 0); + let mut stream = + create_stream_with_context(NoopBatchTransformer::new(), schema, &context); + stream.update_reservation()?; + let reserved_size = stream.reservation.size(); + assert_ne!(reserved_size, 0); stream.set_state(SHJStreamState::BothExhausted { final_result: true }); assert!(stream.next().await.is_none()); - assert_eq!(stream.reservation.size(), 0); + assert_eq!(stream.reservation.size(), reserved_size); + assert_eq!(stream.metrics.stream_memory_usage.value(), reserved_size); + + drop(stream); + assert_eq!(pool.reserved(), 0); + Ok(()) } fn assert_stream_deduplicates_nested_transformer_batch( From db7a6f98b8804550568da2414388e032f9039af2 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 26 Aug 2026 19:33:59 +0800 Subject: [PATCH 23/24] fix: resolve test gaps - Terminal test now drives normal exhaustion/finalization. - Union + RunEnd use nonzero partial slices. --- datafusion/common/src/utils/memory.rs | 4 ++-- datafusion/physical-plan/src/joins/symmetric_hash_join.rs | 7 ++++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 90b627e4f705f..30cedeef39b4a 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -828,7 +828,7 @@ mod record_batch_tests { ) .unwrap(), ) as ArrayRef; - let second = Arc::new(first.as_union().slice(0, first.len())) as ArrayRef; + let second = Arc::new(first.as_union().slice(1, 2)) as ArrayRef; let first_batch = RecordBatch::try_from_iter(vec![("union", Arc::clone(&first))]).unwrap(); let second_batch = @@ -865,7 +865,7 @@ mod record_batch_tests { .as_any() .downcast_ref::>() .unwrap() - .slice(0, first.len()), + .slice(1, 2), ) as ArrayRef; let first_batch = RecordBatch::try_from_iter(vec![("run", Arc::clone(&first))]).unwrap(); diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 2484f28987c59..da28256e6a3bf 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -2222,8 +2222,13 @@ mod tests { let reserved_size = stream.reservation.size(); assert_ne!(reserved_size, 0); - stream.set_state(SHJStreamState::BothExhausted { final_result: true }); + // Empty input streams drive the normal exhaustion/finalization path: + // PullRight -> RightExhausted -> BothExhausted(false) -> final_result. assert!(stream.next().await.is_none()); + assert!(matches!( + stream.state(), + SHJStreamState::BothExhausted { final_result: true } + )); assert_eq!(stream.reservation.size(), reserved_size); assert_eq!(stream.metrics.stream_memory_usage.value(), reserved_size); From d922d3df672d0da0160da55d59b90bc66bfe07c6 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 26 Aug 2026 22:09:37 +0800 Subject: [PATCH 24/24] fix(common/utils/memory): use std::iter::once for iter_on_single_items - Fixed the `iter_on_single_items` function in `datafusion/common/src/utils/memory.rs`. - Replaced the custom iterator logic with `std::iter::once(...)`. --- datafusion/common/src/utils/memory.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 30cedeef39b4a..c2209afbdde45 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -816,8 +816,7 @@ mod record_batch_tests { fn test_record_batch_memory_counter_deduplicates_union_slice_child_overhead() { let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); let fields: UnionFields = - [(0, Arc::new(Field::new("value", DataType::Int32, false)))] - .into_iter() + std::iter::once((0, Arc::new(Field::new("value", DataType::Int32, false)))) .collect(); let first = Arc::new( UnionArray::try_new(