diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index fd405e06a262e..c2209afbdde45 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -19,9 +19,12 @@ use crate::error::_exec_datafusion_err; use crate::{HashSet, Result}; -use arrow::array::types::{ByteArrayType, ByteViewType, RunEndIndexType}; +use arrow::array::types::{ + ByteArrayType, ByteViewType, Int8Type, Int16Type, Int32Type, Int64Type, + RunEndIndexType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, +}; use arrow::array::{ - Array, AsArray, GenericByteArray, GenericByteViewArray, GenericListArray, + Array, ArrayRef, AsArray, GenericByteArray, GenericByteViewArray, GenericListArray, GenericListViewArray, RunArray, }; use arrow::buffer::Buffer; @@ -30,6 +33,7 @@ use arrow::downcast_primitive_array; use arrow::record_batch::RecordBatch; use std::mem::size_of; use std::num::NonZero; +use std::sync::Arc; /// Maximum number of distinct buffer IDs retained inline before promotion to /// a [`HashSet`]. Sixteen keeps small buffer sets allocation-free while @@ -165,7 +169,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: BufferIdSet, - /// 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, } @@ -186,7 +192,28 @@ impl RecordBatchMemoryCounter { self.memory_usage - previous_memory_usage } - /// 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, 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; + + for array in batch.columns() { + array_overhead += + count_unique_array_object_memory_size(array, &mut self.counted_arrays); + } + + 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 } @@ -273,21 +300,9 @@ impl RecordBatchMemoryCounter { self.count_array_memory_size(array.entries()); } DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() { - DataType::Int16 => { - self.count_run_array_memory_size::( - array, - ); - } - DataType::Int32 => { - self.count_run_array_memory_size::( - array, - ); - } - DataType::Int64 => { - self.count_run_array_memory_size::( - array, - ); - } + DataType::Int16 => self.count_run_array_memory_size::(array), + DataType::Int32 => self.count_run_array_memory_size::(array), + DataType::Int64 => self.count_run_array_memory_size::(array), // Arrow only permits Int16, Int32, and Int64 run-end indexes. A // custom Array implementation may still expose malformed data; // retain correct accounting for it without panicking. @@ -360,6 +375,86 @@ 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(_, _) => { + 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()], + 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 => 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![], + } +} + /// Tracks a small number of buffers inline, avoiding a heap allocation for /// typical batches, and promotes to a hash set when more buffers are seen. #[derive(Debug)] @@ -441,12 +536,14 @@ mod tests { mod record_batch_tests { use super::*; use arrow::array::{ - ArrayData, ArrayRef, BinaryViewArray, Float64Array, Int16Array, Int32Array, - Int64Array, LargeListViewArray, ListArray, ListViewArray, RunArray, StringArray, - StringViewArray, new_null_array, + ArrayData, ArrayRef, BinaryViewArray, DictionaryArray, Float64Array, Int16Array, + Int32Array, Int64Array, LargeListViewArray, ListArray, ListViewArray, MapArray, + RunArray, StringArray, StringViewArray, StructArray, UnionArray, new_null_array, }; + use arrow::buffer::OffsetBuffer; use arrow::datatypes::{ - DataType, Field, Int16Type, Int32Type, Int64Type, Schema, UnionFields, UnionMode, + DataType, Field, Fields, Int16Type, Int32Type, Int64Type, Schema, UnionFields, + UnionMode, }; use std::sync::Arc; @@ -568,6 +665,249 @@ 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_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() + ); + } + + 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_shared_list_child() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let list_field = Arc::new(Field::new_list_field(DataType::Int32, false)); + + 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 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_union_slice_child_overhead() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields: UnionFields = + std::iter::once((0, Arc::new(Field::new("value", DataType::Int32, false)))) + .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(1, 2)) 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(1, 2), + ) 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])); + let make_dictionary = || { + Arc::new( + DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 2]), + Arc::clone(&shared_child), + ) + .unwrap(), + ) as ArrayRef + }; + + assert_recursive_shared_child_memory( + "dictionary", + make_dictionary(), + make_dictionary(), + 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( diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 0c6e84b36cc55..da28256e6a3bf 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, @@ -1454,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, @@ -1665,11 +1670,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 +1929,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_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); + + 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_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); @@ -1938,6 +1956,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.reservation.try_resize(capacity)?; + self.metrics.stream_memory_usage.set(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 +2061,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) } } @@ -2082,11 +2106,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, 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, + }; + 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; @@ -2102,6 +2131,414 @@ 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)); + } + + #[tokio::test] + 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_with_context(NoopBatchTransformer::new(), schema, &context); + stream.update_reservation()?; + let reserved_size = stream.reservation.size(); + assert_ne!(reserved_size, 0); + + // 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); + + drop(stream); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + + 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()?; + 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(()) + } + + #[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, + 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 change_deltas(&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.change_deltas(); + 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.change_deltas(); + 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(()) + } + + #[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, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index ecf056560e015..9e6a5c3a96c02 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, @@ -1956,6 +1957,21 @@ 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)>; + + /// 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); +} + +fn count_retained_batch_memory( + batch: &Option, + counter: &mut RecordBatchMemoryCounter, +) { + if let Some(batch) = batch { + counter.count_batch_with_array_overhead(batch); + } } #[derive(Debug, Clone)] @@ -1979,6 +1995,10 @@ 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) { + count_retained_batch_memory(&self.batch, counter); + } } #[derive(Debug, Clone)] @@ -2027,6 +2047,10 @@ impl BatchTransformer for BatchSplitter { Some((sliced_batch, last)) } + + fn count_memory(&self, counter: &mut RecordBatchMemoryCounter) { + count_retained_batch_memory(&self.batch, counter); + } } /// When the order of the join inputs are changed, the output order of columns