diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index f6ff292556c24..a79173fdf9912 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -70,10 +70,13 @@ use crate::{ metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, UInt64Array}; -use arrow::buffer::NullBuffer; +use arrow::array::{ + Array, ArrayRef, BinaryViewArray, BooleanArray, BooleanBufferBuilder, ByteView, + GenericByteViewArray, StringViewArray, UInt64Array, +}; +use arrow::buffer::{NullBuffer, 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}; @@ -3129,6 +3132,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; @@ -3144,6 +3149,134 @@ 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 their `(pointer, length)` pair, +/// 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) avoids rewriting the views +/// by cloning the original array reference. +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: hashbrown::HashMap<(usize, usize), u32> = + hashbrown::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. `RecordBatch::try_new` only fails on structural +/// mismatches (like column count or length) or invalid nullability, none of which +/// can occur here because the original valid batch's schema, column count, lengths, types, +/// and null validity are all perfectly preserved when deduplicating the view columns. +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: @@ -8290,6 +8423,341 @@ 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); + assert_eq!(reservation.size(), get_record_batch_memory_size(&batch)); + } + + #[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); + assert_eq!(reservation.size(), get_record_batch_memory_size(&batch)); + } + + /// 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)" + ); + assert_eq!(reservation.size(), get_record_batch_memory_size(&batch)); + } + /// 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] @@ -11200,4 +11668,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" + ); + } }