diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 9497d521e6321..829f02707d6f0 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -28,7 +28,7 @@ pub(super) use ordered::GroupValuesOrdered; pub mod primitive; pub mod row_backed; -use std::mem::{self, size_of}; +use std::mem; use std::sync::Arc; use crate::aggregates::group_values::GroupValues; @@ -49,7 +49,7 @@ use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::utils::{has_float_leaf, normalize_float_zero}; use datafusion_common::{Result, not_impl_err}; -use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{EmitTo, GroupSelection}; use datafusion_physical_expr::binary_map::OutputType; @@ -194,9 +194,6 @@ pub struct GroupValuesColumn { /// map: HashTable<(u64, GroupIndexView)>, - /// The size of `map` in bytes - map_size: usize, - /// The lists for group indices with the same hash value /// /// It is possible that hash value collision exists, @@ -290,7 +287,6 @@ impl GroupValuesColumn { group_index_lists: Vec::new(), emit_group_index_list_buffer: Vec::new(), vectorized_operation_buffers: VectorizedOperationBuffers::default(), - map_size: 0, group_values, hashes_buffer: Default::default(), random_state: crate::aggregates::AGGREGATION_HASH_SEED, @@ -431,10 +427,10 @@ impl GroupValuesColumn { } // for hasher function, use precomputed hash value - self.map.insert_accounted( + self.map.insert_unique( + target_hash, (target_hash, GroupIndexView::new_inlined(group_idx as u64)), |(hash, _group_index)| *hash, - &mut self.map_size, ); group_idx } @@ -551,10 +547,10 @@ impl GroupValuesColumn { // Insert the `group index view` and its hash into `map` // for hasher function, use precomputed hash value - self.map.insert_accounted( + self.map.insert_unique( + target_hash, (target_hash, group_index_view), |(hash, _)| *hash, - &mut self.map_size, ); // Add row index to `vectorized_append_row_indices` @@ -1068,7 +1064,12 @@ impl GroupValues for GroupValuesColumn { fn size(&self) -> usize { let group_values_size: usize = self.group_values.iter().map(|v| v.size()).sum(); - group_values_size + self.map_size + self.hashes_buffer.allocated_size() + // `HashTable::allocation_size` reports the complete retained allocation — + // including hashbrown control bytes and trailing layout — which is exactly + // what we want here. This follows the same approach as `ArrowBytesMap::size()`. + group_values_size + + self.map.allocation_size() + + self.hashes_buffer.allocated_size() } fn is_empty(&self) -> bool { @@ -1210,7 +1211,6 @@ impl GroupValues for GroupValuesColumn { .expect("schema previously validated in try_new"); self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared - self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); self.hashes_buffer.clear(); self.hashes_buffer.shrink_to(num_rows); @@ -1258,8 +1258,8 @@ mod tests { compute::{concat_batches, take}, util::pretty::pretty_format_batches, }; - use datafusion_common::utils::proxy::HashTableAllocExt; use datafusion_expr::{EmitTo, GroupSelection}; + use std::mem::size_of; use crate::aggregates::group_values::{ GroupValues, multi_group_by::GroupValuesColumn, @@ -2655,10 +2655,10 @@ mod tests { group_index: u64, ) { let group_index_view = GroupIndexView::new_inlined(group_index); - group_values.map.insert_accounted( + group_values.map.insert_unique( + hash_key, (hash_key, group_index_view), |(hash, _)| *hash, - &mut group_values.map_size, ); } @@ -2670,10 +2670,159 @@ mod tests { let list_offset = group_values.group_index_lists.len(); let group_index_view = GroupIndexView::new_non_inlined(list_offset as u64); group_values.group_index_lists.push(group_indices); - group_values.map.insert_accounted( + group_values.map.insert_unique( + hash_key, (hash_key, group_index_view), |(hash, _)| *hash, - &mut group_values.map_size, + ); + } + + // ----------------------------------------------------------------------- + // Tests for exact hash-table allocation accounting (issue #25736) + // ----------------------------------------------------------------------- + + /// After enough groups are inserted to force a hashbrown table growth, + /// `size()` must include the full `allocation_size()` of the map — control + /// bytes and trailing layout included — not just an entry-capacity estimate. + #[test] + fn map_allocation_size_included_in_size_after_growth() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let mut gv: GroupValuesColumn = + GroupValuesColumn::try_new(Arc::clone(&schema)).unwrap(); + + // Insert enough distinct rows to trigger at least one table resize. + let n: usize = 128; + let keys: Vec = (0..n as i64).collect(); + let col = Arc::new(Int64Array::from(keys)) as ArrayRef; + let mut groups = vec![]; + gv.intern(&[col], &mut groups).unwrap(); + + let reported = gv.size(); + let map_alloc = gv.map.allocation_size(); + + // The reported size must be at least as large as the raw map allocation. + assert!( + reported >= map_alloc, + "size() ({reported}) must be >= map.allocation_size() ({map_alloc})" + ); + + // And the map must actually hold a non-trivial allocation after growth. + assert!( + map_alloc > 0, + "map.allocation_size() should be > 0 after inserting {n} groups" + ); + } + + /// When hash collisions are forced, the collision-chain list allocation + /// (`group_index_lists`) must be accounted for independently of the + /// hash-table term. The two must not be conflated. + #[test] + fn collision_chain_size_is_separate_from_map_allocation() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let mut gv: GroupValuesColumn = + GroupValuesColumn::try_new(Arc::clone(&schema)).unwrap(); + + // Manually insert two entries with the same hash key to force a + // collision chain (non-inlined GroupIndexView). + insert_inline_group_index_view(&mut gv, 42, 0); + insert_non_inline_group_index_view(&mut gv, 42, vec![0, 1]); + + let map_alloc = gv.map.allocation_size(); + let chain_size: usize = gv + .group_index_lists + .iter() + .map(|list| list.capacity() * size_of::()) + .sum(); + + // Both terms must be individually non-zero. + assert!(map_alloc > 0, "map allocation must be > 0"); + assert!(chain_size > 0, "collision-chain allocation must be > 0"); + + // The reported total must include both independently. + let reported = gv.size(); + assert!( + reported >= map_alloc, + "size() must cover the map allocation" + ); + } + + /// After a partial emit (`EmitTo::First`), the map retains its allocated + /// capacity even though logical entries were removed. `size()` must still + /// reflect the retained allocation, not drop to zero. + #[test] + fn size_reflects_retained_capacity_after_partial_emit() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let mut gv: GroupValuesColumn = + GroupValuesColumn::try_new(Arc::clone(&schema)).unwrap(); + + let keys: Vec = (0..64i64).collect(); + let col = Arc::new(Int64Array::from(keys)) as ArrayRef; + let mut groups = vec![]; + gv.intern(&[col], &mut groups).unwrap(); + + let size_before = gv.size(); + let map_alloc_before = gv.map.allocation_size(); + + // Emit only the first 32 groups; the table allocation should be retained. + gv.emit(EmitTo::First(32)).unwrap(); + + let size_after = gv.size(); + let map_alloc_after = gv.map.allocation_size(); + + // Capacity is not released by a partial emit — allocation should persist. + assert!( + map_alloc_after > 0, + "map allocation should be retained after partial emit, got {map_alloc_after}" + ); + assert!( + size_after > 0, + "size() should remain > 0 after partial emit, got {size_after}" + ); + + // The allocation reported before and after should both be real and + // consistent with what the map actually holds. + assert_eq!( + map_alloc_before, map_alloc_after, + "partial emit must not shrink the map allocation (before={map_alloc_before}, after={map_alloc_after})" + ); + let _ = size_before; // used implicitly via the assertions above + } + + /// After a full emit followed by re-use, `size()` must still produce + /// correct group values and the map allocation is reported accurately. + #[test] + fn reuse_after_full_emit_produces_correct_groups() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let mut gv: GroupValuesColumn = + GroupValuesColumn::try_new(Arc::clone(&schema)).unwrap(); + + // First pass: intern 4 distinct keys. + let first_keys: Vec = vec![10, 20, 30, 40]; + let col1 = Arc::new(Int64Array::from(first_keys.clone())) as ArrayRef; + let mut groups = vec![]; + gv.intern(&[col1], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 3]); + + // Full emit clears the map. + let emitted = gv.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 1); + + // Second pass after reuse: intern 3 different keys, numbering restarts from 0. + let second_keys: Vec = vec![100, 200, 300]; + let col2 = Arc::new(Int64Array::from(second_keys)) as ArrayRef; + gv.intern(&[col2], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2]); + + // size() must be consistent with the live map allocation. + let reported = gv.size(); + let map_alloc = gv.map.allocation_size(); + assert!( + reported >= map_alloc, + "size() ({reported}) must cover map.allocation_size() ({map_alloc}) after reuse" ); } } diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 10a9fa70a8f76..f6ef221b21a20 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -70,9 +70,13 @@ use crate::{ metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use arrow::array::{Array, ArrayRef, BooleanBufferBuilder, UInt64Array}; +use arrow::array::{ + Array, ArrayRef, BinaryViewArray, BooleanBufferBuilder, ByteView, + GenericByteViewArray, StringViewArray, UInt64Array, +}; +use arrow::buffer::ScalarBuffer; use arrow::compute::concat_batches; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{ByteViewType, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow::util::bit_util; use arrow_schema::{DataType, Schema}; @@ -2984,6 +2988,8 @@ fn concat_build_batches( }; drop(batches); + let batch = deduplicate_record_batch_view_buffers(&batch); + // The inputs are gone: only hold on to what the concatenated batch retains, // which includes any buffers it still shares with the inputs. let held = inputs_reserved + copy_size; @@ -2999,6 +3005,133 @@ fn concat_build_batches( Ok(batch) } +/// Deduplicates shared data buffer references in a [`GenericByteViewArray`] by pointer identity. +/// +/// When multiple record batches that share underlying buffer allocations are concatenated, +/// Arrow's `concat` kernel appends every batch's `data_buffers` list verbatim, resulting +/// in N × K buffer references for N batches that share K allocations. +/// +/// This function walks the buffer list, identifies duplicates by raw pointer address, +/// and rewrites the 4-byte `buffer_index` inside each non-inline view (length > 12) to +/// point into the deduplicated buffer vector. **No string bytes are copied.** +/// +/// The fast path (0 or 1 data buffers, or no duplicates found) clones the array reference +/// with no allocations. +fn deduplicate_view_array_buffers( + array: &GenericByteViewArray, +) -> GenericByteViewArray { + let data_buffers = array.data_buffers(); + if data_buffers.len() <= 1 { + return array.clone(); + } + + // Use the raw buffer address and length as the deduplication key. Casting to usize is the + // idiomatic way to use pointer values as HashMap keys on stable Rust. + let mut unique_buffers: Vec = + Vec::with_capacity(data_buffers.len()); + let mut pointer_map: HashMap<(usize, usize), u32> = + HashMap::with_capacity(data_buffers.len()); + let mut index_remap: Vec = Vec::with_capacity(data_buffers.len()); + let mut has_duplicates = false; + + for buf in data_buffers.iter() { + let key = (buf.as_ptr() as usize, buf.len()); + if let Some(&new_idx) = pointer_map.get(&key) { + index_remap.push(new_idx); + has_duplicates = true; + } else { + let new_idx = unique_buffers.len() as u32; + pointer_map.insert(key, new_idx); + unique_buffers.push(buf.clone()); + index_remap.push(new_idx); + } + } + + if !has_duplicates { + return array.clone(); + } + + // Rewrite the buffer_index field in the 128-bit view descriptor for every + // non-inline value. Inline values (length <= 12) embed the payload inside + // the descriptor itself and carry no buffer index, so they are left as-is. + let views = array.views(); + let mut new_views: Vec = Vec::with_capacity(views.len()); + for &v in views.iter() { + let mut view = ByteView::from(v); + if view.length > 12 { + view.buffer_index = index_remap[view.buffer_index as usize]; + } + new_views.push(view.as_u128()); + } + + let new_views_buffer = ScalarBuffer::from(new_views); + let nulls = array.nulls().cloned(); + + // SAFETY: `new_views_buffer` contains only valid 128-bit view descriptors + // derived from the source array. Each non-inline view's `buffer_index` has + // been remapped to point at the logically equivalent deduplicated buffer in + // `unique_buffers`, preserving the original byte offsets and lengths. + unsafe { + GenericByteViewArray::::new_unchecked( + new_views_buffer, + unique_buffers.into(), + nulls, + ) + } +} + +/// Deduplicates shared data buffer references across all `Utf8View` and `BinaryView` +/// columns in a [`RecordBatch`], returning a new batch whose view arrays hold at most +/// as many buffer references as there are distinct underlying allocations. +/// +/// Columns of other types are passed through unchanged. If the batch contains no view +/// columns this function returns a cheap clone of the batch reference. +/// Returns a new [`RecordBatch`] with deduplicated view buffer references. +/// +/// This function is infallible: we reconstruct the batch using the original +/// schema and the same set of columns (same lengths, same types). `RecordBatch::try_new` +/// only fails when column lengths or schema mismatches occur — neither can happen here +/// since we only replace view columns with logically equivalent deduplicated versions. +fn deduplicate_record_batch_view_buffers(batch: &RecordBatch) -> RecordBatch { + let has_view_columns = batch + .columns() + .iter() + .any(|col| matches!(col.data_type(), DataType::Utf8View | DataType::BinaryView)); + + if !has_view_columns { + return batch.clone(); + } + + let new_columns: Vec = batch + .columns() + .iter() + .map(|col| match col.data_type() { + DataType::Utf8View => { + let array = col + .as_any() + .downcast_ref::() + .expect("Utf8View column must be StringViewArray"); + Arc::new(deduplicate_view_array_buffers(array)) as ArrayRef + } + DataType::BinaryView => { + let array = col + .as_any() + .downcast_ref::() + .expect("BinaryView column must be BinaryViewArray"); + Arc::new(deduplicate_view_array_buffers(array)) as ArrayRef + } + _ => Arc::clone(col), + }) + .collect(); + + // SAFETY: schema is taken from the original batch unchanged, and every column + // in new_columns has the same length as in the original (deduplication only + // rewrites buffer indices, not the view count). This call cannot fail. + RecordBatch::try_new(batch.schema(), new_columns).expect( + "schema and column-length invariants are preserved by buffer deduplication", + ) +} + /// Collects all batches from the left (build) side stream and creates a hash map for joining. /// /// This function is responsible for: @@ -8144,6 +8277,338 @@ mod tests { Ok(()) } + #[test] + fn concat_build_batches_deduplicates_view_buffers() { + use arrow::array::StringViewBuilder; + + let mut builder = StringViewBuilder::new(); + builder.append_value("this is a long string that exceeds inline size 12"); + builder.append_value("short_inline"); + builder.append_null(); + builder.append_value("another long string that exceeds inline size 12"); + let base_array: StringViewArray = builder.finish(); + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8View, true), + Field::new("id", DataType::Int32, false), + ])); + + let id_array = Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(base_array.clone()), Arc::clone(&id_array)], + ) + .expect("valid batch"); + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(base_array.clone()), Arc::clone(&id_array)], + ) + .expect("valid batch"); + let batch3 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(base_array.clone()), Arc::clone(&id_array)], + ) + .expect("valid batch"); + + // Before deduplication, concat_batches puts 3 duplicate buffer references in data_buffers + let concatenated_raw = + concat_batches(&schema, &[batch1.clone(), batch2.clone(), batch3.clone()]) + .expect("concat"); + let raw_view_arr = concatenated_raw + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(raw_view_arr.data_buffers().len(), 3); + + // After concat_build_batches, buffer references are deduplicated down to 1 + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let batches = vec![batch1, batch2, batch3]; + let (mut reservation, inputs_reserved) = + reserve_inputs(&batches, &pool).expect("reserve"); + + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &mut reservation, + &metrics, + ) + .expect("concat_build"); + + let view_arr = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(view_arr.data_buffers().len(), 1); + assert_eq!(batch.num_rows(), 12); + assert_eq!( + view_arr.value(0), + "this is a long string that exceeds inline size 12" + ); + assert_eq!(view_arr.value(1), "short_inline"); + assert!(view_arr.is_null(2)); + assert_eq!( + view_arr.value(3), + "another long string that exceeds inline size 12" + ); + + let id_col = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(id_col.value(0), 1); + assert_eq!(id_col.value(1), 2); + assert_eq!(id_col.value(2), 3); + assert_eq!(id_col.value(3), 4); + assert_eq!(id_col.value(4), 1); + } + + #[test] + fn concat_build_batches_deduplicates_binary_view_buffers() { + use arrow::array::BinaryViewBuilder; + + let mut builder = BinaryViewBuilder::new(); + builder.append_value(b"this is a long binary that exceeds inline size 12"); + builder.append_value(b"short"); + builder.append_null(); + builder.append_value(b"another long binary that exceeds inline size 12"); + let base_array: BinaryViewArray = builder.finish(); + let schema = Arc::new(Schema::new(vec![ + Field::new("b", DataType::BinaryView, true), + Field::new("id", DataType::Int64, false), + ])); + + let id_array = Arc::new(Int64Array::from(vec![10, 20, 30, 40])) as ArrayRef; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(base_array.clone()), Arc::clone(&id_array)], + ) + .expect("valid batch"); + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(base_array.clone()), Arc::clone(&id_array)], + ) + .expect("valid batch"); + let batch3 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(base_array.clone()), Arc::clone(&id_array)], + ) + .expect("valid batch"); + + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let batches = vec![batch1, batch2, batch3]; + let (mut reservation, inputs_reserved) = + reserve_inputs(&batches, &pool).expect("reserve"); + + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &mut reservation, + &metrics, + ) + .expect("concat_build"); + + let view_arr = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(view_arr.data_buffers().len(), 1); + assert_eq!(batch.num_rows(), 12); + assert_eq!( + view_arr.value(0), + b"this is a long binary that exceeds inline size 12" + ); + assert_eq!(view_arr.value(1), b"short"); + assert!(view_arr.is_null(2)); + assert_eq!( + view_arr.value(3), + b"another long binary that exceeds inline size 12" + ); + + let id_col = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(id_col.value(0), 10); + assert_eq!(id_col.value(4), 10); + } + + /// Regression test: two arrays backed by Buffer objects that share the same + /// base pointer but have *different* declared lengths (produced via + /// `Buffer::slice_with_length`) must NOT be collapsed into one buffer entry, + /// because the longer one covers bytes beyond the shorter one's range. + /// Values that live in the longer buffer must survive concatenation intact. + #[test] + fn concat_build_batches_deduplicates_slices_regression() { + use arrow::array::StringViewBuilder; + use arrow::buffer::Buffer; + + // Build a two-element array so the builder allocates a single contiguous + // data buffer large enough to hold both non-inline strings. + let mut builder = StringViewBuilder::new(); + builder.append_value("this is a long string that exceeds inline size 12"); + builder.append_value("another long string that exceeds inline size 12"); + let base_array: StringViewArray = builder.finish(); + + // Pull out the single data buffer Arrow produced. + let data_buffers = base_array.data_buffers(); + assert_eq!(data_buffers.len(), 1, "expected one backing data buffer"); + let full_buf: &Buffer = &data_buffers[0]; + let full_len = full_buf.len(); + + // Slice the same allocation to two different declared lengths. Both + // handles start at offset 0, so `buf.as_ptr()` is identical, but + // `buf.len()` differs. The deduplication key is (ptr, len), so these + // are distinct keys — the longer buffer must not be discarded. + // + // Derive the split point from the first view's actual byte range so + // the short buffer provably covers only the first string. + let view0 = ByteView::from(base_array.views()[0]); + let first_string_end = view0.offset as usize + view0.length as usize; + assert!( + first_string_end < full_len, + "second string must extend beyond the split point" + ); + let short_buf = full_buf.slice_with_length(0, first_string_end); + let long_buf = full_buf.clone(); // full length — covers both strings + + // Borrow the raw 128-bit view words from `base_array`. Each word + // encodes (length, prefix, buffer_index, offset); buffer_index is 0 in + // both because `base_array` has a single data buffer. + let views = base_array.views(); + let view0_raw: u128 = views[0]; // first string — lives within short_buf + let view1_raw: u128 = views[1]; // second string — lives in long_buf only + + let schema = + Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8View, true)])); + + // Array 1: backed only by `short_buf` (shorter declared length) + let arr1: StringViewArray = { + let views_buf = ScalarBuffer::from(vec![view0_raw]); + // SAFETY: view0 was taken verbatim from `base_array` which is valid, + // and `short_buf` is sized to exactly cover the first string's byte + // range (offset 0..first_string_end), so all offsets referenced by + // view0 are in bounds. + unsafe { + GenericByteViewArray::new_unchecked( + views_buf, + vec![short_buf].into(), + None, + ) + } + }; + + // Array 2: backed only by `long_buf` (full declared length) + let arr2: StringViewArray = { + let views_buf = ScalarBuffer::from(vec![view1_raw]); + // SAFETY: view1 was taken verbatim from `base_array` which is valid, + // and `long_buf` is the full allocation, so all byte ranges are covered. + unsafe { + GenericByteViewArray::new_unchecked( + views_buf, + vec![long_buf].into(), + None, + ) + } + }; + + // Sanity-check the arrays read back correctly before we feed them in. + assert_eq!( + arr1.value(0), + "this is a long string that exceeds inline size 12" + ); + assert_eq!( + arr2.value(0), + "another long string that exceeds inline size 12" + ); + + let batch1 = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(arr1)]) + .expect("valid batch"); + let batch2 = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(arr2)]) + .expect("valid batch"); + + // Before deduplication, 4 input batches with duplicate references to short_buf and long_buf + // result in 4 data buffer references in the concatenated raw batch. + let concatenated_raw = concat_batches( + &schema, + &[ + batch1.clone(), + batch2.clone(), + batch1.clone(), + batch2.clone(), + ], + ) + .expect("concat"); + let raw_view_arr = concatenated_raw + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(raw_view_arr.data_buffers().len(), 4); + + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + // Provide 4 batches with duplicate handles to both short_buf and long_buf + let batches = vec![batch1.clone(), batch2.clone(), batch1, batch2]; + let (mut reservation, inputs_reserved) = + reserve_inputs(&batches, &pool).expect("reserve"); + + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &mut reservation, + &metrics, + ) + .expect("concat_build"); + + let view_arr = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + // 4 buffer references before deduplication are collapsed to 2 unique buffers: + // short_buf and long_buf are kept separate despite sharing the same pointer address. + assert_eq!(view_arr.data_buffers().len(), 2); + assert_eq!(batch.num_rows(), 4); + + // Both rows from both batches must be readable after deduplication — the value that lives + // exclusively in the longer buffer must not have been dropped or truncated. + assert_eq!( + view_arr.value(0), + "this is a long string that exceeds inline size 12", + "value from the shorter-declared-length buffer should survive (row 0)" + ); + assert_eq!( + view_arr.value(1), + "another long string that exceeds inline size 12", + "value from the longer-declared-length buffer should survive (row 1)" + ); + assert_eq!( + view_arr.value(2), + "this is a long string that exceeds inline size 12", + "value from duplicate shorter-declared-length buffer should survive (row 2)" + ); + assert_eq!( + view_arr.value(3), + "another long string that exceeds inline size 12", + "value from duplicate longer-declared-length buffer should survive (row 3)" + ); + } + /// The build side is concatenated into a single batch, and that copy must /// be visible to the memory pool while the input batches are still alive. #[rstest] @@ -11021,4 +11486,355 @@ mod tests { assert!(join.set_dynamic_filter(df).is_err()); Ok(()) } + + // ----------------------------------------------------------------------- + // Unit tests for deduplicate_view_array_buffers / + // deduplicate_record_batch_view_buffers + // ----------------------------------------------------------------------- + + /// Fast path: an array with a single data buffer must be returned as-is + /// (pointer-equal clone, zero allocations). + #[test] + fn test_dedup_view_array_single_buffer_is_noop() { + let array = StringViewArray::from(vec![ + "hello world long string!", + "another long value!", + ]); + let result = deduplicate_view_array_buffers(&array); + // Same number of data buffers — nothing removed. + assert_eq!(result.data_buffers().len(), array.data_buffers().len()); + } + + /// Fast path: an array with zero data buffers (all inline values) must be returned as-is. + #[test] + fn test_dedup_view_array_zero_buffers_is_noop() { + let array = StringViewArray::from(vec!["inline_only"]); + assert_eq!(array.data_buffers().len(), 0); + let result = deduplicate_view_array_buffers(&array); + assert_eq!(result.data_buffers().len(), 0); + assert_eq!(result.value(0), "inline_only"); + } + + /// Fast path: multiple distinct buffers (no duplicates) → returned as-is. + #[test] + fn test_dedup_view_array_no_duplicates_is_noop() { + // Build two separate StringViewArrays so their buffers are distinct. + let a = StringViewArray::from(vec!["first long string value here"]); + let b = StringViewArray::from(vec!["second long string value here"]); + // Concatenate to get an array with two *different* buffers. + let combined = arrow::compute::concat(&[&a as _, &b as _]).unwrap(); + let combined = combined.as_any().downcast_ref::().unwrap(); + let n_before = combined.data_buffers().len(); + let result = deduplicate_view_array_buffers(combined); + // Still the same number of unique buffers — nothing deduplicated. + assert_eq!(result.data_buffers().len(), n_before); + } + + /// Deduplication path: concatenating an array with itself produces N refs to + /// the same K buffers; deduplicate_view_array_buffers collapses them to K. + #[test] + fn test_dedup_view_array_removes_duplicate_buffers() { + let base = StringViewArray::from(vec!["this is a long string value abc"]); + // concat gives us 2 references to the same underlying buffer. + let doubled = arrow::compute::concat(&[&base as _, &base as _]).unwrap(); + let doubled = doubled.as_any().downcast_ref::().unwrap(); + assert_eq!(doubled.data_buffers().len(), 2); + let deduped = deduplicate_view_array_buffers(doubled); + // After deduplication only 1 unique buffer should remain. + assert_eq!(deduped.data_buffers().len(), 1); + // Values must be preserved. + assert_eq!(deduped.value(0), "this is a long string value abc"); + assert_eq!(deduped.value(1), "this is a long string value abc"); + } + + /// Inline-string path: strings with length ≤ 12 are stored inline in the + /// view descriptor and carry no buffer index; they must survive deduplication. + #[test] + fn test_dedup_view_array_inline_strings_preserved() { + // "hi" is 2 bytes — well within the 12-byte inline threshold. + let base = StringViewArray::from(vec!["hi", "short"]); + let doubled = arrow::compute::concat(&[&base as _, &base as _]).unwrap(); + let doubled = doubled.as_any().downcast_ref::().unwrap(); + let deduped = deduplicate_view_array_buffers(doubled); + assert_eq!(deduped.value(0), "hi"); + assert_eq!(deduped.value(1), "short"); + assert_eq!(deduped.value(2), "hi"); + assert_eq!(deduped.value(3), "short"); + } + + /// Mixed inline, long strings, and nulls: ensures the view rewriting loop handles + /// non-inline view remapping, skips inline views (length <= 12), and preserves the null buffer. + #[test] + fn test_dedup_view_array_mixed_inline_long_and_nulls() { + let array = StringViewArray::from(vec![ + Some("this is a long string that exceeds inline size 12"), + Some("inline"), + None, + ]); + let doubled = arrow::compute::concat(&[&array as _, &array as _]).unwrap(); + let doubled = doubled.as_any().downcast_ref::().unwrap(); + assert_eq!(doubled.data_buffers().len(), 2); + let deduped = deduplicate_view_array_buffers(doubled); + assert_eq!(deduped.data_buffers().len(), 1); + assert_eq!( + deduped.value(0), + "this is a long string that exceeds inline size 12" + ); + assert_eq!(deduped.value(1), "inline"); + assert!(deduped.is_null(2)); + assert_eq!( + deduped.value(3), + "this is a long string that exceeds inline size 12" + ); + assert_eq!(deduped.value(4), "inline"); + assert!(deduped.is_null(5)); + } + + /// Direct BinaryView path: exercises deduplicate_view_array_buffers directly on BinaryViewArray + /// with mixed long values, inline values, and nulls. + #[test] + fn test_dedup_view_array_binary_view_direct() { + let base = BinaryViewArray::from_iter(vec![ + Some(&b"this is definitely longer than 12 bytes"[..]), + Some(&b"short"[..]), + None, + ]); + let doubled = arrow::compute::concat(&[&base as _, &base as _]).unwrap(); + let doubled = doubled.as_any().downcast_ref::().unwrap(); + assert_eq!(doubled.data_buffers().len(), 2); + let deduped = deduplicate_view_array_buffers(doubled); + assert_eq!(deduped.data_buffers().len(), 1); + assert_eq!(deduped.value(0), b"this is definitely longer than 12 bytes"); + assert_eq!(deduped.value(1), b"short"); + assert!(deduped.is_null(2)); + assert_eq!(deduped.value(3), b"this is definitely longer than 12 bytes"); + assert_eq!(deduped.value(4), b"short"); + assert!(deduped.is_null(5)); + } + + /// BinaryView path: `deduplicate_record_batch_view_buffers` must also + /// deduplicate `BinaryView` columns (exercises the BinaryView arm). + #[test] + fn test_dedup_record_batch_binary_view_column() { + let schema = Arc::new(Schema::new(vec![Field::new( + "data", + DataType::BinaryView, + false, + )])); + let base = BinaryViewArray::from_iter_values(vec![ + b"this is definitely longer than 12 bytes", + ]); + let doubled = arrow::compute::concat(&[&base as _, &base as _]).unwrap(); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![doubled]).unwrap(); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .data_buffers() + .len(), + 2 + ); + let deduped = deduplicate_record_batch_view_buffers(&batch); + let result = deduped + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.data_buffers().len(), 1); + } + + /// Mixed columns path: `deduplicate_record_batch_view_buffers` must deduplicate + /// both Utf8View and BinaryView columns while passing through non-view columns (Int32, Utf8). + #[test] + fn test_dedup_record_batch_mixed_view_and_non_view_columns() { + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8View, true), + Field::new("b", DataType::BinaryView, true), + Field::new("n", DataType::Int32, false), + Field::new("str", DataType::Utf8, false), + ])); + let str_base = StringViewArray::from(vec!["long string value here abc"]); + let bin_base = + BinaryViewArray::from_iter_values(vec![b"long binary value here abc"]); + let str_doubled = + arrow::compute::concat(&[&str_base as _, &str_base as _]).unwrap(); + let bin_doubled = + arrow::compute::concat(&[&bin_base as _, &bin_base as _]).unwrap(); + let int_col: ArrayRef = Arc::new(Int32Array::from(vec![42, 43])); + let utf8_col: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar"])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + str_doubled, + bin_doubled, + Arc::clone(&int_col), + Arc::clone(&utf8_col), + ], + ) + .unwrap(); + + let deduped = deduplicate_record_batch_view_buffers(&batch); + let deduped_str = deduped + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let deduped_bin = deduped + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(deduped_str.data_buffers().len(), 1); + assert_eq!(deduped_bin.data_buffers().len(), 1); + // Non-view columns passed through unmodified + assert!(Arc::ptr_eq(deduped.column(2), &int_col)); + assert!(Arc::ptr_eq(deduped.column(3), &utf8_col)); + } + + /// Fast path: a batch with no Utf8View or BinaryView columns must be + /// returned as a cheap clone with no allocations. + #[test] + fn test_dedup_record_batch_no_view_columns_is_noop() { + let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap(); + // Should succeed without error and return the same data. + let result = deduplicate_record_batch_view_buffers(&batch); + assert_eq!(result.num_rows(), 3); + } + + /// Reverse concatenation order: exercises concat_build_batches with reverse=true + /// on view columns to verify both reverse order and buffer deduplication work together. + #[test] + fn test_concat_build_batches_reverse_order_deduplication() { + use arrow::array::StringViewBuilder; + + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Utf8View, + false, + )])); + + let mut builder = StringViewBuilder::new(); + builder.append_value("batch1: long string value exceeding twelve bytes"); + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(builder.finish()) as ArrayRef], + ) + .expect("valid batch"); + + let mut builder = StringViewBuilder::new(); + builder.append_value("batch2: long string value exceeding twelve bytes"); + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(builder.finish()) as ArrayRef], + ) + .expect("valid batch"); + + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + // 4 batches with two duplicate references each — raw concat yields 4 buffer handles + let batches = vec![batch1.clone(), batch2.clone(), batch1, batch2]; + let (mut reservation, inputs_reserved) = + reserve_inputs(&batches, &pool).expect("reserve"); + + let batch = concat_build_batches( + &schema, + batches, + true, + inputs_reserved, + &mut reservation, + &metrics, + ) + .expect("concat"); + + let view_arr = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + // After deduplication 4 handles collapse to 2 unique buffers. + assert_eq!(view_arr.data_buffers().len(), 2); + assert_eq!(batch.num_rows(), 4); + // reverse=true: batch2, batch1, batch2, batch1 + assert_eq!( + view_arr.value(0), + "batch2: long string value exceeding twelve bytes" + ); + assert_eq!( + view_arr.value(1), + "batch1: long string value exceeding twelve bytes" + ); + assert_eq!( + view_arr.value(2), + "batch2: long string value exceeding twelve bytes" + ); + assert_eq!( + view_arr.value(3), + "batch1: long string value exceeding twelve bytes" + ); + } + + /// Exercises the `retained > held` growth path in concat_build_batches. + /// + /// When a single batch is provided `copy_size` is 0, so we only pre-reserve the + /// input size. If the returned batch happens to occupy more memory than what we + /// reserved (e.g., due to deduplication creating new buffers or tracking overhead + /// differences), the function must call `reservation.try_grow(retained - held)`. + /// We trigger this by passing `inputs_reserved = 0` directly so that `held == 0` + /// and any non-empty batch forces the grow branch. + #[test] + fn test_concat_build_batches_grow_branch() { + use arrow::array::StringViewBuilder; + + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Utf8View, + false, + )])); + + let mut builder = StringViewBuilder::new(); + builder.append_value("long string that exceeds the twelve byte inline threshold"); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(builder.finish()) as ArrayRef], + ) + .expect("valid batch"); + + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut reservation = MemoryConsumer::new("test").register(&pool); + + // Pass inputs_reserved = 0 so that held == 0 + copy_size == 0 == 0, + // while retained == get_record_batch_memory_size(&batch) > 0 — forcing + // the `retained > held` branch to call reservation.try_grow. + let result = concat_build_batches( + &schema, + vec![batch], + false, + 0, // inputs_reserved deliberately zero + &mut reservation, + &metrics, + ) + .expect("concat"); + + let view_arr = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.num_rows(), 1); + assert_eq!( + view_arr.value(0), + "long string that exceeds the twelve byte inline threshold" + ); + } }