From d7cf40c4aa8a7c04f467adf79adf16fc5fba42ab Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 29 Sep 2026 04:58:46 +0000 Subject: [PATCH 1/7] feat(arrow-select): add compact byte-view interleave --- arrow-select/src/interleave.rs | 88 ++++++++- .../tests/interleave_byte_view_compact.rs | 179 ++++++++++++++++++ 2 files changed, 266 insertions(+), 1 deletion(-) create mode 100644 arrow-select/tests/interleave_byte_view_compact.rs diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 74e29defe835..723b7fffa3b9 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -20,7 +20,7 @@ use crate::concat::concat; use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values}; use arrow_array::ListLikeArray; -use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder}; +use arrow_array::builder::{BooleanBufferBuilder, GenericByteViewBuilder, PrimitiveBuilder}; use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; @@ -74,6 +74,9 @@ macro_rules! dict_helper { /// ``` /// /// For selecting values by index from a single array see [`crate::take`] +/// +/// To copy selected byte view values into owned buffers, see +/// [`interleave_byte_view_compact`]. pub fn interleave( values: &[&dyn Array], indices: &[(usize, usize)], @@ -128,6 +131,89 @@ pub fn interleave( } } +/// Interleaves byte view arrays into independently owned, compact buffers. +/// +/// Each pair in `indices` selects an array in `values` and a row in that array, +/// as in [`interleave`]. This supports both [`StringViewArray`] and +/// [`BinaryViewArray`]. Only selected, non-null values are copied; inline values +/// need no data buffer. The result does not retain any input buffers. +/// +/// Use this when selected rows must release their references to input storage, +/// for example when partitioning a large batch for spilling. Unlike calling +/// [`interleave`] followed by [`GenericByteViewArray::gc`], this does not first +/// construct an intermediate array referencing the input data buffers. +/// +/// Copying can increase total memory usage while the inputs remain alive. +/// Repeated selections copy their payload each time; values are not deduplicated. +/// Call [`interleave`] to share input data buffers instead. +/// +/// # Errors +/// +/// Returns an error if `values` is empty or the selected payload exceeds the +/// supported size limits. +/// +/// # Panics +/// +/// Panics if an array or row index is out of bounds. +/// +/// # Example +/// +/// ``` +/// use arrow_array::StringViewArray; +/// use arrow_select::interleave::interleave_byte_view_compact; +/// +/// let a = StringViewArray::from(vec![Some("a long selected value"), None]); +/// let b = StringViewArray::from(vec!["another selected value"]); +/// let result = interleave_byte_view_compact(&[&a, &b], &[(1, 0), (0, 1), (0, 0)])?; +/// assert_eq!(result, StringViewArray::from(vec![ +/// Some("another selected value"), None, Some("a long selected value") +/// ])); +/// # Ok::<(), arrow_schema::ArrowError>(()) +/// ``` +pub fn interleave_byte_view_compact( + values: &[&GenericByteViewArray], + indices: &[(usize, usize)], +) -> Result, ArrowError> { + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + + let bytes = indices.iter().try_fold(0usize, |bytes, &(source, row)| { + let array = values[source]; + if array.is_null(row) { + return Ok(bytes); + } + let len = array.views()[row] as u32 as usize; + if len >= u32::MAX as usize { + return Err(ArrowError::OffsetOverflowError(len)); + } + if len <= arrow_data::MAX_INLINE_VIEW_LEN as usize { + return Ok(bytes); + } + bytes + .checked_add(len) + .ok_or(ArrowError::OffsetOverflowError(usize::MAX)) + })?; + + let mut builder = GenericByteViewBuilder::::with_capacity(indices.len()); + // Avoid the default minimum block size for small selections. Larger + // outputs use the builder's normal growth across multiple blocks. + if bytes < u32::MAX as usize { + builder = builder.with_fixed_block_size(bytes.max(1) as u32); + } + for &(source, row) in indices { + let array = values[source]; + if array.is_null(row) { + builder.append_null(); + } else { + builder.try_append_value(array.value(row))?; + } + } + Ok(builder.finish()) +} + /// Common functionality for interleaving arrays /// /// T is the concrete Array type diff --git a/arrow-select/tests/interleave_byte_view_compact.rs b/arrow-select/tests/interleave_byte_view_compact.rs new file mode 100644 index 000000000000..e1589f397d8b --- /dev/null +++ b/arrow-select/tests/interleave_byte_view_compact.rs @@ -0,0 +1,179 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::builder::{BinaryViewBuilder, StringViewBuilder}; +use arrow_array::cast::AsArray; +use arrow_array::types::StringViewType; +use arrow_array::{Array, StringViewArray}; +use arrow_buffer::NullBuffer; +use arrow_schema::ArrowError; +use arrow_select::interleave::{interleave, interleave_byte_view_compact}; + +#[test] +fn sliced_multibuffer_string_sources() { + let mut first = StringViewBuilder::new().with_fixed_block_size(32); + first.append_value("discard before slice"); + first.append_value("first selected long string"); + first.append_null(); + first.append_value("second selected long string"); + first.append_value("discard after slice"); + let first = first.finish().slice(1, 3); + + let mut second = StringViewBuilder::new().with_fixed_block_size(32); + second.append_value("discard before slice"); + second.append_value("third selected long string"); + second.append_value("short"); + second.append_value("fourth selected long string"); + let second = second.finish().slice(1, 3); + assert!(first.data_buffers().len() > 1); + assert!(second.data_buffers().len() > 1); + + let indices = [(1, 2), (0, 1), (0, 2), (1, 1), (0, 0), (1, 2)]; + let compact = interleave_byte_view_compact(&[&first, &second], &indices).unwrap(); + let ordinary = interleave(&[&first, &second], &indices).unwrap(); + assert_eq!( + compact.iter().collect::>(), + ordinary.as_string_view().iter().collect::>() + ); + assert_eq!( + compact.iter().collect::>(), + vec![ + Some("fourth selected long string"), + None, + Some("second selected long string"), + Some("short"), + Some("first selected long string"), + Some("fourth selected long string"), + ] + ); +} + +#[test] +fn binary_values_across_inline_boundary() { + let inline = [0xff; 12]; + let external = [0xfe; 13]; + let mut input = BinaryViewBuilder::new(); + input.append_value(inline); + input.append_value(external); + input.append_null(); + input.append_value([]); + let input = input.finish(); + + let indices = [(0, 1), (0, 0), (0, 2), (0, 3), (0, 1)]; + let compact = interleave_byte_view_compact(&[&input], &indices).unwrap(); + let ordinary = interleave(&[&input], &indices).unwrap(); + assert_eq!( + compact.iter().collect::>(), + ordinary.as_binary_view().iter().collect::>() + ); + assert_eq!(compact.value(0), &external); + assert_eq!(compact.value(1), &inline); + assert!(compact.is_null(2)); + assert_eq!(compact.value(3), b""); + assert_eq!( + compact + .data_buffers() + .iter() + .map(|b| b.len()) + .sum::(), + 2 * external.len() + ); +} + +#[test] +fn copies_only_selected_nonnull_payload() { + let unselected = "u".repeat(32 * 1024); + let hidden = "n".repeat(32 * 1024); + let selected = "selected long string"; + let input = StringViewArray::from(vec![unselected.as_str(), selected, hidden.as_str()]); + // Keep a valid, nonempty view underneath the null bit. + let input = StringViewArray::try_new( + input.views().clone(), + input.data_buffers().clone(), + Some(NullBuffer::from(vec![true, true, false])), + ) + .unwrap(); + let indices = [(0, 1), (0, 2), (0, 1)]; + let compact = interleave_byte_view_compact(&[&input], &indices).unwrap(); + let ordinary = interleave(&[&input], &indices).unwrap(); + assert_eq!( + compact.iter().collect::>(), + ordinary.as_string_view().iter().collect::>() + ); + assert_eq!( + compact + .data_buffers() + .iter() + .map(|b| b.len()) + .sum::(), + 2 * selected.len() + ); + assert!( + compact + .data_buffers() + .iter() + .map(|b| b.capacity()) + .sum::() + < 1024 + ); + for output in compact.data_buffers().iter() { + for source in input.data_buffers().iter() { + // Compare allocation starts, including when a buffer is a slice. + assert_ne!( + output.as_ptr() as usize - output.ptr_offset(), + source.as_ptr() as usize - source.ptr_offset() + ); + } + } + assert_ne!(compact.views().as_ptr(), input.views().as_ptr()); + assert_ne!( + compact.nulls().unwrap().buffer().as_ptr(), + input.nulls().unwrap().buffer().as_ptr() + ); + drop(ordinary); + drop(input); + assert_eq!( + compact.iter().collect::>(), + vec![Some(selected), None, Some(selected)] + ); +} + +#[test] +fn empty_null_and_inline_selections_need_no_data_buffers() { + let input = StringViewArray::from(vec![Some("unselected long string"), None, Some("inline")]); + for indices in [vec![], vec![(0, 1), (0, 1)], vec![(0, 2), (0, 1)]] { + let compact = interleave_byte_view_compact(&[&input], &indices).unwrap(); + let ordinary = interleave(&[&input], &indices).unwrap(); + assert_eq!( + compact.iter().collect::>(), + ordinary.as_string_view().iter().collect::>() + ); + assert!(compact.data_buffers().is_empty()); + } + let empty = StringViewArray::from(Vec::<&str>::new()); + assert!( + interleave_byte_view_compact(&[&empty], &[]) + .unwrap() + .is_empty() + ); +} + +#[test] +fn missing_input_is_an_error_even_for_empty_selection() { + let error = interleave_byte_view_compact::(&[], &[]).unwrap_err(); + assert!(matches!(error, ArrowError::InvalidArgumentError(_))); +} From f9b5aa02a9a597b3e807c2a7717d37fc22791a12 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 2 Oct 2026 23:21:37 +0000 Subject: [PATCH 2/7] Revise compact interleave with configurable copying and range sharing --- arrow-select/src/interleave.rs | 766 +++++++++++++++--- .../tests/interleave_byte_view_compact.rs | 179 ---- .../benches/interleave_byte_view_cases/mod.rs | 197 +++++ arrow/benches/interleave_kernels.rs | 62 +- 4 files changed, 901 insertions(+), 303 deletions(-) delete mode 100644 arrow-select/tests/interleave_byte_view_compact.rs create mode 100644 arrow/benches/interleave_byte_view_cases/mod.rs diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 723b7fffa3b9..658fe1f37269 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -20,19 +20,20 @@ use crate::concat::concat; use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values}; use arrow_array::ListLikeArray; -use arrow_array::builder::{BooleanBufferBuilder, GenericByteViewBuilder, PrimitiveBuilder}; +use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder}; use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; use arrow_buffer::bit_mask::set_bits; use arrow_buffer::bit_util; use arrow_buffer::{ - ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer, + ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, NullBufferBuilder, + OffsetBuffer, }; -use arrow_data::ByteView; use arrow_data::transform::MutableArrayData; +use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN}; use arrow_schema::{ArrowError, DataType, FieldRef, Fields}; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; macro_rules! primitive_helper { @@ -76,142 +77,290 @@ macro_rules! dict_helper { /// For selecting values by index from a single array see [`crate::take`] /// /// To copy selected byte view values into owned buffers, see -/// [`interleave_byte_view_compact`]. +/// [`Interleaver::with_compact_byte_views`]. pub fn interleave( values: &[&dyn Array], indices: &[(usize, usize)], ) -> Result { - if values.is_empty() { - return Err(ArrowError::InvalidArgumentError( - "interleave requires input of at least one array".to_string(), - )); - } - let data_type = values[0].data_type(); + Interleaver::default().interleave(values, indices) +} + +/// Configurable interleaving of elements from multiple arrays. +/// +/// The default configuration has the same behavior as [`interleave`]. +#[derive(Debug, Default, Clone)] +pub struct Interleaver { + compact_byte_views: bool, + preserve_byte_view_sharing: bool, +} + +impl Interleaver { + /// Creates an interleaver with the same behavior as [`interleave`]. + pub fn new() -> Self { + Self::default() + } + + /// Copies selected top-level [`StringViewArray`] and [`BinaryViewArray`] + /// values into independent buffers when enabled. + /// + /// Only selected, non-null, non-inline values need payload storage. No input + /// buffers are retained. Repeated values are copied for each selected row + /// unless [`Self::with_preserve_byte_view_sharing`] is enabled. + /// Other array types, including byte views nested inside other arrays, keep + /// the behavior of [`interleave`]. This option is disabled by default. + /// + /// This is useful when selected rows need independent lifetimes, such as + /// partitions of a spilling operator. Unlike [`interleave`] followed by + /// [`GenericByteViewArray::gc`], it avoids an intermediate shared array and + /// source-buffer remapping tables. + /// Copying can increase total live memory while inputs remain retained. + /// + /// ``` + /// use arrow_array::{Array, StringViewArray}; + /// use arrow_select::interleave::Interleaver; + /// + /// let a = StringViewArray::from(vec![Some("a long selected value"), None]); + /// let b = StringViewArray::from(vec!["another selected value"]); + /// let result = Interleaver::new() + /// .with_compact_byte_views(true) + /// .interleave(&[&a, &b], &[(1, 0), (0, 1), (0, 0)])?; + /// assert_eq!(result.as_ref(), &StringViewArray::from(vec![ + /// Some("another selected value"), None, Some("a long selected value") + /// ]) as &dyn Array); + /// # Ok::<(), arrow_schema::ArrowError>(()) + /// ``` + pub fn with_compact_byte_views(mut self, compact: bool) -> Self { + self.compact_byte_views = compact; + self + } + + /// Preserves existing source-range sharing when compacting byte views. + /// + /// Repeated references to the same source byte range share one copy in the + /// output, including references through different input arrays. Equal values + /// stored at different addresses are not deduplicated. The output still + /// retains no input buffers. + /// + /// This option only affects [`Self::with_compact_byte_views`]. It is disabled + /// by default: tracking ranges adds a lookup per non-inline value and + /// temporary memory proportional to the number of distinct selected ranges. + /// Enable it for selections with shared long values to avoid copying their + /// payload repeatedly, for example after a join or dictionary decoding. + pub fn with_preserve_byte_view_sharing(mut self, preserve: bool) -> Self { + self.preserve_byte_view_sharing = preserve; + self + } + + /// Selects elements using `(array index, row index)` pairs, as in [`interleave`]. + /// + /// # Errors + /// + /// Returns an error for empty inputs, differing input types, or unsupported + /// sizes. Compact byte views support values and output buffer indices up to + /// [`i32::MAX`], splitting larger total payloads across multiple buffers. + /// + /// # Panics + /// + /// Panics if an array or row index is out of bounds. + pub fn interleave( + &self, + values: &[&dyn Array], + indices: &[(usize, usize)], + ) -> Result { + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + let data_type = values[0].data_type(); + + for array in values.iter().skip(1) { + if array.data_type() != data_type { + return Err(ArrowError::InvalidArgumentError(format!( + "It is not possible to interleave arrays of different data types ({} and {})", + data_type, + array.data_type() + ))); + } + } + + if indices.is_empty() { + return Ok(new_empty_array(data_type)); + } + + if self.compact_byte_views { + match data_type { + DataType::Utf8View => { + return self.interleave_views_compact::( + values, + indices, + i32::MAX as usize, + ); + } + DataType::BinaryView => { + return self.interleave_views_compact::( + values, + indices, + i32::MAX as usize, + ); + } + _ => {} + } + } - for array in values.iter().skip(1) { - if array.data_type() != data_type { - return Err(ArrowError::InvalidArgumentError(format!( - "It is not possible to interleave arrays of different data types ({} and {})", - data_type, - array.data_type() - ))); + downcast_primitive! { + data_type => (primitive_helper, values, indices, data_type), + DataType::Utf8 => interleave_bytes::(values, indices), + DataType::LargeUtf8 => interleave_bytes::(values, indices), + DataType::Binary => interleave_bytes::(values, indices), + DataType::LargeBinary => interleave_bytes::(values, indices), + DataType::BinaryView => interleave_views::(values, indices), + DataType::Utf8View => interleave_views::(values, indices), + DataType::Dictionary(k, _) => downcast_integer! { + k.as_ref() => (dict_helper, values, indices), + _ => unreachable!("illegal dictionary key type {k}") + }, + DataType::Struct(fields) => interleave_struct(fields, values, indices), + DataType::List(field) => interleave_list::(values, indices, field), + DataType::LargeList(field) => interleave_list::(values, indices, field), + DataType::FixedSizeList(field, size) => interleave_fixed_size_list(values, indices, field, *size), + DataType::Map(field, ordered) => interleave_map(values, indices, field, *ordered), + DataType::RunEndEncoded(r, _) => match r.data_type() { + DataType::Int16 => interleave_run_end::(values, indices), + DataType::Int32 => interleave_run_end::(values, indices), + DataType::Int64 => interleave_run_end::(values, indices), + t => unreachable!("illegal run-end type {t}"), + }, + DataType::ListView(field) => interleave_list_view::(values, indices, field), + DataType::LargeListView(field) => interleave_list_view::(values, indices, field), + _ => interleave_fallback(values, indices) } } - if indices.is_empty() { - return Ok(new_empty_array(data_type)); - } - - downcast_primitive! { - data_type => (primitive_helper, values, indices, data_type), - DataType::Utf8 => interleave_bytes::(values, indices), - DataType::LargeUtf8 => interleave_bytes::(values, indices), - DataType::Binary => interleave_bytes::(values, indices), - DataType::LargeBinary => interleave_bytes::(values, indices), - DataType::BinaryView => interleave_views::(values, indices), - DataType::Utf8View => interleave_views::(values, indices), - DataType::Dictionary(k, _) => downcast_integer! { - k.as_ref() => (dict_helper, values, indices), - _ => unreachable!("illegal dictionary key type {k}") - }, - DataType::Struct(fields) => interleave_struct(fields, values, indices), - DataType::List(field) => interleave_list::(values, indices, field), - DataType::LargeList(field) => interleave_list::(values, indices, field), - DataType::FixedSizeList(field, size) => interleave_fixed_size_list(values, indices, field, *size), - DataType::Map(field, ordered) => interleave_map(values, indices, field, *ordered), - DataType::RunEndEncoded(r, _) => match r.data_type() { - DataType::Int16 => interleave_run_end::(values, indices), - DataType::Int32 => interleave_run_end::(values, indices), - DataType::Int64 => interleave_run_end::(values, indices), - t => unreachable!("illegal run-end type {t}"), - }, - DataType::ListView(field) => interleave_list_view::(values, indices, field), - DataType::LargeListView(field) => interleave_list_view::(values, indices, field), - _ => interleave_fallback(values, indices) + fn interleave_views_compact( + &self, + values: &[&dyn Array], + indices: &[(usize, usize)], + max_buffer_size: usize, + ) -> Result { + if self.preserve_byte_view_sharing { + interleave_views_compact::(values, indices, max_buffer_size) + } else { + interleave_views_compact::(values, indices, max_buffer_size) + } } } -/// Interleaves byte view arrays into independently owned, compact buffers. -/// -/// Each pair in `indices` selects an array in `values` and a row in that array, -/// as in [`interleave`]. This supports both [`StringViewArray`] and -/// [`BinaryViewArray`]. Only selected, non-null values are copied; inline values -/// need no data buffer. The result does not retain any input buffers. -/// -/// Use this when selected rows must release their references to input storage, -/// for example when partitioning a large batch for spilling. Unlike calling -/// [`interleave`] followed by [`GenericByteViewArray::gc`], this does not first -/// construct an intermediate array referencing the input data buffers. -/// -/// Copying can increase total memory usage while the inputs remain alive. -/// Repeated selections copy their payload each time; values are not deduplicated. -/// Call [`interleave`] to share input data buffers instead. -/// -/// # Errors -/// -/// Returns an error if `values` is empty or the selected payload exceeds the -/// supported size limits. -/// -/// # Panics -/// -/// Panics if an array or row index is out of bounds. -/// -/// # Example -/// -/// ``` -/// use arrow_array::StringViewArray; -/// use arrow_select::interleave::interleave_byte_view_compact; -/// -/// let a = StringViewArray::from(vec![Some("a long selected value"), None]); -/// let b = StringViewArray::from(vec!["another selected value"]); -/// let result = interleave_byte_view_compact(&[&a, &b], &[(1, 0), (0, 1), (0, 0)])?; -/// assert_eq!(result, StringViewArray::from(vec![ -/// Some("another selected value"), None, Some("a long selected value") -/// ])); -/// # Ok::<(), arrow_schema::ArrowError>(()) -/// ``` -pub fn interleave_byte_view_compact( - values: &[&GenericByteViewArray], +fn interleave_views_compact( + values: &[&dyn Array], indices: &[(usize, usize)], -) -> Result, ArrowError> { - if values.is_empty() { - return Err(ArrowError::InvalidArgumentError( - "interleave requires input of at least one array".to_string(), - )); - } - - let bytes = indices.iter().try_fold(0usize, |bytes, &(source, row)| { - let array = values[source]; + max_buffer_size: usize, +) -> Result { + let arrays: Vec<_> = values + .iter() + .map(|array| array.as_byte_view::()) + .collect(); + let mut views = Vec::with_capacity(indices.len()); + let mut nulls = NullBufferBuilder::new(indices.len()); + let mut block_sizes: Vec = Vec::new(); + let mut copied = + HashMap::<(usize, u32), usize, _>::with_hasher(ahash::RandomState::with_seeds(0, 0, 0, 0)); + for &(source, row) in indices { + let array = arrays[source]; + let raw = array.views()[row]; if array.is_null(row) { - return Ok(bytes); + views.push(0); + nulls.append_null(); + continue; + } + nulls.append_non_null(); + let mut view = ByteView::from(raw); + if view.length <= MAX_INLINE_VIEW_LEN { + views.push(raw); + continue; } - let len = array.views()[row] as u32 as usize; - if len >= u32::MAX as usize { + let len = view.length as usize; + if len > i32::MAX as usize { return Err(ArrowError::OffsetOverflowError(len)); } - if len <= arrow_data::MAX_INLINE_VIEW_LEN as usize { - return Ok(bytes); + if PRESERVE_SHARING { + let buffer = &array.data_buffers()[view.buffer_index as usize]; + // Effective addresses identify aliases through differently sliced buffers. + let address = buffer.as_ptr().wrapping_add(view.offset as usize) as usize; + match copied.entry((address, view.length)) { + std::collections::hash_map::Entry::Occupied(entry) => { + views.push(views[*entry.get()]); + continue; + } + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(views.len()); + } + } } - bytes - .checked_add(len) - .ok_or(ArrowError::OffsetOverflowError(usize::MAX)) - })?; - - let mut builder = GenericByteViewBuilder::::with_capacity(indices.len()); - // Avoid the default minimum block size for small selections. Larger - // outputs use the builder's normal growth across multiple blocks. - if bytes < u32::MAX as usize { - builder = builder.with_fixed_block_size(bytes.max(1) as u32); + if block_sizes + .last() + .is_none_or(|bytes| bytes + len > max_buffer_size) + { + if block_sizes.len() > i32::MAX as usize { + return Err(ArrowError::OffsetOverflowError(block_sizes.len())); + } + block_sizes.push(0); + } + let buffer_index = (block_sizes.len() - 1) as u32; + let bytes = block_sizes.last_mut().unwrap(); + if PRESERVE_SHARING { + view.buffer_index = buffer_index; + view.offset = *bytes as u32; + } + *bytes += len; + views.push(view.as_u128()); } - for &(source, row) in indices { - let array = values[source]; - if array.is_null(row) { - builder.append_null(); - } else { - builder.try_append_value(array.value(row))?; + drop(copied); + + let mut buffers: Vec> = block_sizes + .iter() + .map(|&size| Vec::with_capacity(size)) + .collect(); + let mut current_buffer = 0; + for (raw, &(source, row)) in views.iter_mut().zip(indices) { + let mut view = ByteView::from(*raw); + if view.length <= MAX_INLINE_VIEW_LEN { + continue; } + let array = arrays[source]; + let buffer = if PRESERVE_SHARING { + let buffer = &mut buffers[view.buffer_index as usize]; + // First occurrences fill consecutive offsets. Repeated ranges point + // behind the write cursor and have already been copied. + if view.offset as usize != buffer.len() { + continue; + } + view = ByteView::from(array.views()[row]); + buffer + } else { + // Keep source views until this pass, avoiding a second random read + // of the input views. Rewrite the final views in place as we copy. + if buffers[current_buffer].len() == block_sizes[current_buffer] { + current_buffer += 1; + } + let buffer = &mut buffers[current_buffer]; + *raw = view + .with_buffer_index(current_buffer as u32) + .with_offset(buffer.len() as u32) + .as_u128(); + buffer + }; + let start = view.offset as usize; + buffer.extend_from_slice( + &array.data_buffers()[view.buffer_index as usize][start..start + view.length as usize], + ); } - Ok(builder.finish()) + let buffers: Vec<_> = buffers.into_iter().map(Buffer::from_vec).collect(); + // SAFETY: inline views are unchanged, null views are zero, and every other + // view addresses its copied source range with the original length and prefix. + Ok(Arc::new(unsafe { + GenericByteViewArray::::new_unchecked(views.into(), buffers.into(), nulls.finish()) + })) } /// Common functionality for interleaving arrays @@ -1055,12 +1204,387 @@ mod tests { use super::*; use arrow_array::Int32RunArray; use arrow_array::builder::{ - GenericListBuilder, Int32Builder, PrimitiveBuilder, PrimitiveRunBuilder, + BinaryViewBuilder, GenericListBuilder, Int32Builder, PrimitiveBuilder, PrimitiveRunBuilder, + StringViewBuilder, }; use arrow_array::types::{Decimal128Type, Int8Type, TimestampMicrosecondType}; use arrow_buffer::ScalarBuffer; use arrow_schema::{Field, TimeUnit}; + #[test] + fn test_compact_sliced_multibuffer_sources() { + let mut first = StringViewBuilder::new().with_fixed_block_size(32); + first.append_value("discard before slice"); + first.append_value("first selected long string"); + first.append_null(); + first.append_value("second selected long string"); + first.append_value("discard after slice"); + let first = first.finish().slice(1, 3); + + let mut second = StringViewBuilder::new().with_fixed_block_size(32); + second.append_value("discard before slice"); + second.append_value("third selected long string"); + second.append_value("short"); + second.append_value("fourth selected long string"); + let second = second.finish().slice(1, 3); + assert!(first.data_buffers().len() > 1); + assert!(second.data_buffers().len() > 1); + + let indices = [(1, 2), (0, 1), (0, 2), (1, 1), (0, 0), (1, 2)]; + let ordinary = interleave(&[&first, &second], &indices).unwrap(); + for preserve in [false, true] { + let compact = Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(preserve) + .interleave(&[&first, &second], &indices) + .unwrap(); + assert_eq!(compact.as_ref(), ordinary.as_ref()); + let compact = compact.as_string_view(); + assert_eq!( + compact.iter().collect::>(), + vec![ + Some("fourth selected long string"), + None, + Some("second selected long string"), + Some("short"), + Some("first selected long string"), + Some("fourth selected long string"), + ] + ); + assert_eq!(compact.views()[0] == compact.views()[5], preserve); + } + } + + #[test] + fn test_compact_binary_inline_boundary() { + let inline = [0xff; 12]; + let external = [0xfe; 13]; + let mut input = BinaryViewBuilder::new(); + input.append_value(inline); + input.append_value(external); + input.append_null(); + input.append_value([]); + let input = input.finish(); + + let indices = [(0, 1), (0, 0), (0, 2), (0, 3), (0, 1)]; + for preserve in [false, true] { + let compact = Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(preserve) + .interleave(&[&input], &indices) + .unwrap(); + assert_eq!( + compact.as_ref(), + interleave(&[&input], &indices).unwrap().as_ref() + ); + let compact = compact.as_binary_view(); + assert_eq!(compact.value(0), &external); + assert_eq!(compact.value(1), &inline); + assert!(compact.is_null(2)); + assert_eq!(compact.value(3), b""); + assert_eq!(compact.data_buffers().len(), 1); + let copies = if preserve { 1 } else { 2 }; + assert_eq!(compact.data_buffers()[0].len(), copies * external.len()); + assert_eq!(compact.views()[0] == compact.views()[4], preserve); + } + } + + #[test] + fn test_compact_skips_hidden_null_payload_and_owns_buffers() { + for preserve in [false, true] { + let unselected = "u".repeat(32 * 1024); + let hidden = "n".repeat(32 * 1024); + let selected = "selected long string"; + let input = StringViewArray::from(vec![unselected.as_str(), selected, hidden.as_str()]); + let input = StringViewArray::try_new( + input.views().clone(), + input.data_buffers().clone(), + Some(NullBuffer::from(vec![true, true, false])), + ) + .unwrap(); + let compact = Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(preserve) + .interleave(&[&input], &[(0, 1), (0, 2), (0, 1)]) + .unwrap(); + let compact = compact.as_string_view(); + assert_eq!(compact.data_buffers().len(), 1); + let copies = if preserve { 1 } else { 2 }; + assert_eq!(compact.data_buffers()[0].len(), copies * selected.len()); + assert_eq!( + compact.data_buffers()[0].capacity(), + copies * selected.len() + ); + assert_eq!(compact.views()[1], 0); + for output in compact.data_buffers().iter() { + for source in input.data_buffers().iter() { + assert_ne!( + output.as_ptr() as usize - output.ptr_offset(), + source.as_ptr() as usize - source.ptr_offset() + ); + } + } + assert_ne!(compact.views().as_ptr(), input.views().as_ptr()); + assert_ne!( + compact.nulls().unwrap().buffer().as_ptr(), + input.nulls().unwrap().buffer().as_ptr() + ); + drop(input); + assert_eq!( + compact.iter().collect::>(), + vec![Some(selected), None, Some(selected)] + ); + } + } + + #[test] + fn test_compact_repeated_range_copies_payload_once() { + let value = "x".repeat(1024); + let input = StringViewArray::from(vec![value.as_str()]); + let compact = Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(true) + .interleave(&[&input], &vec![(0, 0); 10_000]) + .unwrap(); + let compact = compact.as_string_view(); + assert_eq!(compact.len(), 10_000); + assert_eq!(compact.data_buffers().len(), 1); + assert_eq!(compact.data_buffers()[0].len(), 1024); + assert_eq!(compact.data_buffers()[0].capacity(), 1024); + assert_ne!( + compact.data_buffers()[0].as_ptr(), + input.data_buffers()[0].as_ptr() + ); + drop(input); + assert!(compact.iter().all(|actual| actual == Some(value.as_str()))); + assert!( + compact + .views() + .iter() + .all(|view| *view == compact.views()[0]) + ); + compact.to_data().validate_full().unwrap(); + } + + #[test] + fn test_compact_sharing_is_opt_in() { + let input = StringViewArray::from(vec!["long repeated payload"]); + for interleaver in [ + Interleaver::new().with_compact_byte_views(true), + Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(true) + .with_preserve_byte_view_sharing(false), + ] { + let output = interleaver + .interleave(&[&input], &[(0, 0), (0, 0), (0, 0)]) + .unwrap(); + let output = output.as_string_view(); + assert_eq!(output.data_buffers()[0].len(), 3 * input.value(0).len()); + assert_eq!( + output.data_buffers()[0].capacity(), + 3 * input.value(0).len() + ); + assert_ne!(output.views()[0], output.views()[1]); + assert_ne!(output.views()[1], output.views()[2]); + assert!(output.iter().all(|value| value == Some(input.value(0)))); + } + } + + #[test] + fn test_compact_aliases_across_sliced_buffers() { + fn view(buffer: Buffer, offset: u32, len: u32) -> StringViewArray { + let mut builder = StringViewBuilder::new(); + let block = builder.append_block(buffer); + builder.try_append_view(block, offset, len).unwrap(); + builder.finish() + } + + let backing = Buffer::from_vec( + b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ".to_vec(), + ); + let a = view(backing.clone(), 7, 20); + let b = view(backing.slice(7), 0, 20); + let c = view(backing.slice(7), 0, 21); + let d = view(Buffer::from_vec(backing[7..27].to_vec()), 0, 20); + let overlap = view(backing.slice(8), 0, 20); + let compact = Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(true) + .interleave( + &[&a, &b, &c, &d, &overlap], + &[(0, 0), (1, 0), (2, 0), (3, 0), (4, 0)], + ) + .unwrap(); + let compact = compact.as_string_view(); + assert_eq!(compact.views()[0], compact.views()[1]); + assert_ne!(compact.views()[0], compact.views()[2]); + assert_ne!(compact.views()[0], compact.views()[3]); + assert_ne!(compact.views()[0], compact.views()[4]); + assert_eq!(compact.value(0), compact.value(3)); + assert_eq!(compact.data_buffers()[0].len(), 20 + 21 + 20 + 20); + assert_eq!(compact.data_buffers()[0].capacity(), 20 + 21 + 20 + 20); + assert_ne!(compact.data_buffers()[0].as_ptr(), backing.as_ptr()); + assert_eq!( + compact.iter().collect::>(), + vec![ + Some(a.value(0)), + Some(b.value(0)), + Some(c.value(0)), + Some(d.value(0)), + Some(overlap.value(0)) + ] + ); + compact.to_data().validate_full().unwrap(); + } + + #[test] + fn test_compact_splits_buffers_and_reuses_earlier_ranges() { + let first = "a".repeat(20); + let second = "b".repeat(13); + let third = "c".repeat(18); + let input = StringViewArray::from(vec![ + Some(first.as_str()), + Some(second.as_str()), + None, + Some(third.as_str()), + Some("inline"), + ]); + let indices = [(0, 0), (0, 1), (0, 3), (0, 0), (0, 2), (0, 4), (0, 1)]; + for preserve in [false, true] { + let compact = if preserve { + interleave_views_compact::(&[&input], &indices, 32) + } else { + interleave_views_compact::(&[&input], &indices, 32) + } + .unwrap(); + assert_eq!( + compact.as_ref(), + interleave(&[&input], &indices).unwrap().as_ref() + ); + let compact = compact.as_string_view(); + assert_eq!(compact.data_buffers().len(), if preserve { 2 } else { 4 }); + assert_eq!(compact.data_buffers()[0].len(), 20); + assert_eq!(compact.data_buffers()[1].len(), 31); + for buffer in compact.data_buffers().iter() { + assert!(buffer.len() <= 32); + assert_eq!(buffer.len(), buffer.capacity()); + } + assert_eq!(ByteView::from(compact.views()[0]).buffer_index, 0); + assert_eq!(ByteView::from(compact.views()[1]).buffer_index, 1); + assert_eq!(ByteView::from(compact.views()[2]).offset, 13); + assert_eq!(compact.views()[0] == compact.views()[3], preserve); + assert_eq!(compact.views()[1] == compact.views()[6], preserve); + assert_eq!(compact.views()[4], 0); + compact.to_data().validate_full().unwrap(); + } + } + + #[test] + fn test_compact_empty_null_and_inline_selections() { + let input = + StringViewArray::from(vec![Some("unselected long string"), None, Some("inline")]); + for preserve in [false, true] { + for indices in [vec![], vec![(0, 1), (0, 1)], vec![(0, 2), (0, 1)]] { + let compact = Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(preserve) + .interleave(&[&input], &indices) + .unwrap(); + assert_eq!( + compact.as_ref(), + interleave(&[&input], &indices).unwrap().as_ref() + ); + assert!(compact.as_string_view().data_buffers().is_empty()); + } + let empty = StringViewArray::from(Vec::<&str>::new()); + assert!( + Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(preserve) + .interleave(&[&empty], &[]) + .unwrap() + .is_empty() + ); + } + } + + #[test] + fn test_compact_configuration_and_type_errors() { + let input = StringViewArray::from(vec!["first long view value", "second long view value"]); + let indices = [(0, 1), (0, 0), (0, 1)]; + let ordinary = interleave(&[&input], &indices).unwrap(); + for interleaver in [ + Interleaver::default(), + Interleaver::new().with_preserve_byte_view_sharing(true), + Interleaver::new() + .with_compact_byte_views(true) + .with_compact_byte_views(false), + ] { + let output = interleaver.interleave(&[&input], &indices).unwrap(); + assert_eq!(output.as_ref(), ordinary.as_ref()); + assert_eq!( + output.as_string_view().data_buffers()[0].as_ptr(), + input.data_buffers()[0].as_ptr() + ); + } + let ints = Int32Array::from(vec![1, 2]); + for preserve in [false, true] { + for compact in [false, true] { + let interleaver = Interleaver::new() + .with_compact_byte_views(compact) + .with_preserve_byte_view_sharing(preserve); + let output = interleaver.interleave(&[&ints], &indices).unwrap(); + assert_eq!( + output.as_ref(), + interleave(&[&ints], &indices).unwrap().as_ref() + ); + assert!(matches!( + interleaver.interleave(&[], &[]), + Err(ArrowError::InvalidArgumentError(_)) + )); + assert!(matches!( + interleaver.interleave(&[&input, &ints], &[]), + Err(ArrowError::InvalidArgumentError(_)) + )); + assert!(matches!( + interleaver.interleave(&[&input, &ints], &[(0, 0)]), + Err(ArrowError::InvalidArgumentError(_)) + )); + } + } + } + + #[test] + fn test_compact_option_leaves_nested_views_unchanged() { + let values = StringViewArray::from(vec![ + "first nested long string", + "second nested long string", + ]); + let pointer = values.data_buffers()[0].as_ptr(); + let input = ListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8View, false)), + OffsetBuffer::new(vec![0_i32, 1, 2].into()), + Arc::new(values), + None, + ); + let compact = Interleaver::new() + .with_compact_byte_views(true) + .interleave(&[&input], &[(0, 1)]) + .unwrap(); + let ordinary = interleave(&[&input], &[(0, 1)]).unwrap(); + assert_eq!(compact.as_ref(), ordinary.as_ref()); + assert_eq!( + compact + .as_list::() + .values() + .as_string_view() + .data_buffers()[0] + .as_ptr(), + pointer + ); + } + #[test] fn test_primitive() { let a = Int32Array::from_iter_values([1, 2, 3, 4]); diff --git a/arrow-select/tests/interleave_byte_view_compact.rs b/arrow-select/tests/interleave_byte_view_compact.rs deleted file mode 100644 index e1589f397d8b..000000000000 --- a/arrow-select/tests/interleave_byte_view_compact.rs +++ /dev/null @@ -1,179 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use arrow_array::builder::{BinaryViewBuilder, StringViewBuilder}; -use arrow_array::cast::AsArray; -use arrow_array::types::StringViewType; -use arrow_array::{Array, StringViewArray}; -use arrow_buffer::NullBuffer; -use arrow_schema::ArrowError; -use arrow_select::interleave::{interleave, interleave_byte_view_compact}; - -#[test] -fn sliced_multibuffer_string_sources() { - let mut first = StringViewBuilder::new().with_fixed_block_size(32); - first.append_value("discard before slice"); - first.append_value("first selected long string"); - first.append_null(); - first.append_value("second selected long string"); - first.append_value("discard after slice"); - let first = first.finish().slice(1, 3); - - let mut second = StringViewBuilder::new().with_fixed_block_size(32); - second.append_value("discard before slice"); - second.append_value("third selected long string"); - second.append_value("short"); - second.append_value("fourth selected long string"); - let second = second.finish().slice(1, 3); - assert!(first.data_buffers().len() > 1); - assert!(second.data_buffers().len() > 1); - - let indices = [(1, 2), (0, 1), (0, 2), (1, 1), (0, 0), (1, 2)]; - let compact = interleave_byte_view_compact(&[&first, &second], &indices).unwrap(); - let ordinary = interleave(&[&first, &second], &indices).unwrap(); - assert_eq!( - compact.iter().collect::>(), - ordinary.as_string_view().iter().collect::>() - ); - assert_eq!( - compact.iter().collect::>(), - vec![ - Some("fourth selected long string"), - None, - Some("second selected long string"), - Some("short"), - Some("first selected long string"), - Some("fourth selected long string"), - ] - ); -} - -#[test] -fn binary_values_across_inline_boundary() { - let inline = [0xff; 12]; - let external = [0xfe; 13]; - let mut input = BinaryViewBuilder::new(); - input.append_value(inline); - input.append_value(external); - input.append_null(); - input.append_value([]); - let input = input.finish(); - - let indices = [(0, 1), (0, 0), (0, 2), (0, 3), (0, 1)]; - let compact = interleave_byte_view_compact(&[&input], &indices).unwrap(); - let ordinary = interleave(&[&input], &indices).unwrap(); - assert_eq!( - compact.iter().collect::>(), - ordinary.as_binary_view().iter().collect::>() - ); - assert_eq!(compact.value(0), &external); - assert_eq!(compact.value(1), &inline); - assert!(compact.is_null(2)); - assert_eq!(compact.value(3), b""); - assert_eq!( - compact - .data_buffers() - .iter() - .map(|b| b.len()) - .sum::(), - 2 * external.len() - ); -} - -#[test] -fn copies_only_selected_nonnull_payload() { - let unselected = "u".repeat(32 * 1024); - let hidden = "n".repeat(32 * 1024); - let selected = "selected long string"; - let input = StringViewArray::from(vec![unselected.as_str(), selected, hidden.as_str()]); - // Keep a valid, nonempty view underneath the null bit. - let input = StringViewArray::try_new( - input.views().clone(), - input.data_buffers().clone(), - Some(NullBuffer::from(vec![true, true, false])), - ) - .unwrap(); - let indices = [(0, 1), (0, 2), (0, 1)]; - let compact = interleave_byte_view_compact(&[&input], &indices).unwrap(); - let ordinary = interleave(&[&input], &indices).unwrap(); - assert_eq!( - compact.iter().collect::>(), - ordinary.as_string_view().iter().collect::>() - ); - assert_eq!( - compact - .data_buffers() - .iter() - .map(|b| b.len()) - .sum::(), - 2 * selected.len() - ); - assert!( - compact - .data_buffers() - .iter() - .map(|b| b.capacity()) - .sum::() - < 1024 - ); - for output in compact.data_buffers().iter() { - for source in input.data_buffers().iter() { - // Compare allocation starts, including when a buffer is a slice. - assert_ne!( - output.as_ptr() as usize - output.ptr_offset(), - source.as_ptr() as usize - source.ptr_offset() - ); - } - } - assert_ne!(compact.views().as_ptr(), input.views().as_ptr()); - assert_ne!( - compact.nulls().unwrap().buffer().as_ptr(), - input.nulls().unwrap().buffer().as_ptr() - ); - drop(ordinary); - drop(input); - assert_eq!( - compact.iter().collect::>(), - vec![Some(selected), None, Some(selected)] - ); -} - -#[test] -fn empty_null_and_inline_selections_need_no_data_buffers() { - let input = StringViewArray::from(vec![Some("unselected long string"), None, Some("inline")]); - for indices in [vec![], vec![(0, 1), (0, 1)], vec![(0, 2), (0, 1)]] { - let compact = interleave_byte_view_compact(&[&input], &indices).unwrap(); - let ordinary = interleave(&[&input], &indices).unwrap(); - assert_eq!( - compact.iter().collect::>(), - ordinary.as_string_view().iter().collect::>() - ); - assert!(compact.data_buffers().is_empty()); - } - let empty = StringViewArray::from(Vec::<&str>::new()); - assert!( - interleave_byte_view_compact(&[&empty], &[]) - .unwrap() - .is_empty() - ); -} - -#[test] -fn missing_input_is_an_error_even_for_empty_selection() { - let error = interleave_byte_view_compact::(&[], &[]).unwrap_err(); - assert!(matches!(error, ArrowError::InvalidArgumentError(_))); -} diff --git a/arrow/benches/interleave_byte_view_cases/mod.rs b/arrow/benches/interleave_byte_view_cases/mod.rs new file mode 100644 index 000000000000..0f33961b720c --- /dev/null +++ b/arrow/benches/interleave_byte_view_cases/mod.rs @@ -0,0 +1,197 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{StringViewArray, StringViewBuilder}; +use arrow_buffer::Buffer; +use rand::{RngExt, SeedableRng, rngs::StdRng}; + +const SOURCE_COUNT: usize = 4; +const ROWS_PER_SOURCE: usize = 8192; +const ALIASED_VALUES: usize = 64; + +pub struct ByteViewCase { + pub name: String, + pub arrays: Vec, + pub indices: Vec<(usize, usize)>, +} + +#[derive(Clone, Copy)] +enum Layout { + Unique, + Random, + Sparse, + RepeatedIndices, + SharedRanges, + SlicedNulls, +} + +impl Layout { + fn name(self) -> &'static str { + match self { + Self::Unique => "unique", + Self::Random => "random", + Self::Sparse => "sparse", + Self::RepeatedIndices => "repeated_indices", + Self::SharedRanges => "shared_ranges", + Self::SlicedNulls => "sliced_nulls", + } + } +} + +fn value(id: usize, min_len: usize, max_len: usize) -> String { + let len = min_len + id.wrapping_mul(4051) % (max_len - min_len + 1); + // Include the row identity even in the shortest values, so the unique + // controls do not accidentally contain equal strings. + let mut value = format!("{id:08x}:"); + value.extend((value.len()..len).map(|i| (b'a' + ((id + i * 17) % 26) as u8) as char)); + value +} + +fn make_case(layout: Layout, selected: usize, min_len: usize, max_len: usize) -> ByteViewCase { + let arrays = if matches!(layout, Layout::SharedRanges) { + let mut bytes = Vec::new(); + let mut ranges = Vec::new(); + for id in 0..ALIASED_VALUES { + let value = value(id, min_len, max_len); + ranges.push((bytes.len() as u32, value.len() as u32)); + bytes.extend_from_slice(value.as_bytes()); + } + let buffer = Buffer::from_vec(bytes); + (0..SOURCE_COUNT) + .map(|source| { + let mut builder = StringViewBuilder::with_capacity(ROWS_PER_SOURCE); + let block = builder.append_block(buffer.clone()); + for row in 0..ROWS_PER_SOURCE { + let (offset, len) = ranges[(row + source * 17) % ALIASED_VALUES]; + builder.try_append_view(block, offset, len).unwrap(); + } + builder.finish() + }) + .collect() + } else { + (0..SOURCE_COUNT) + .map(|source| { + let sliced = matches!(layout, Layout::SlicedNulls); + let extra = if sliced { 37 } else { 0 }; + let block_size = if matches!(layout, Layout::Sparse) { + 1024 + } else { + 64 * 1024 + }; + let mut builder = StringViewBuilder::with_capacity(ROWS_PER_SOURCE + extra) + .with_fixed_block_size(block_size); + for row in 0..ROWS_PER_SOURCE + extra { + if sliced && row % 7 == 0 { + builder.append_null(); + } else if sliced && row % 11 == 0 { + builder.append_value("inline"); + } else { + builder.append_value(value( + source * (ROWS_PER_SOURCE + extra) + row, + min_len, + max_len, + )); + } + } + let array = builder.finish(); + if sliced { + array.slice(17, ROWS_PER_SOURCE) + } else { + array + } + }) + .collect() + }; + let mut rng = StdRng::seed_from_u64(0); + let indices = (0..selected) + .map(|i| { + if matches!(layout, Layout::Random) { + return ( + rng.random_range(0..SOURCE_COUNT), + rng.random_range(0..ROWS_PER_SOURCE), + ); + } + let source = i % SOURCE_COUNT; + let row = i / SOURCE_COUNT; + let row = if matches!(layout, Layout::RepeatedIndices) { + row % 16 + } else { + row + }; + // An odd multiplier permutes the power-of-two source length, + // distributing selections across its buffers without replacement. + (source, (row * 4051) % ROWS_PER_SOURCE) + }) + .collect(); + ByteViewCase { + name: format!( + "{}_len{min_len}-{max_len}_selected{selected}", + layout.name() + ), + arrays, + indices, + } +} + +fn source_buffer_case(buffer_count: usize, selected: usize) -> ByteViewCase { + let rows = SOURCE_COUNT * ROWS_PER_SOURCE / 2; + let buffers_per_source = buffer_count / 2; + let block_size = (rows / buffers_per_source * 16) as u32; + let arrays: Vec<_> = (0..2) + .map(|source| { + let mut builder = + StringViewBuilder::with_capacity(rows).with_fixed_block_size(block_size); + for row in 0..rows { + builder.append_value(value(source * rows + row, 16, 16)); + } + let array = builder.finish(); + assert_eq!(array.data_buffers().len(), buffers_per_source); + array + }) + .collect(); + ByteViewCase { + name: format!("source_buffers{buffer_count}_selected{selected}"), + arrays, + indices: (0..selected) + .map(|i| (i % 2, ((i / 2) * 4051) % rows)) + .collect(), + } +} + +pub fn for_each_case(mut visit: impl FnMut(ByteViewCase)) { + for (min_len, max_len) in [(13, 20), (20, 60), (100, 400)] { + for selected in [512, 8192] { + for layout in [ + Layout::Unique, + Layout::Random, + Layout::Sparse, + Layout::RepeatedIndices, + Layout::SharedRanges, + Layout::SlicedNulls, + ] { + visit(make_case(layout, selected, min_len, max_len)); + } + } + } + // Keep the same source rows, payload, and selections while independently + // varying the number of source buffers that interleave must track. + for buffers in [4, 4096, 32768] { + for selected in [3, 32] { + visit(source_buffer_case(buffers, selected)); + } + } +} diff --git a/arrow/benches/interleave_kernels.rs b/arrow/benches/interleave_kernels.rs index ee5bfb13d980..cfb6f19448dd 100644 --- a/arrow/benches/interleave_kernels.rs +++ b/arrow/benches/interleave_kernels.rs @@ -18,7 +18,9 @@ #[macro_use] extern crate criterion; -use criterion::Criterion; +mod interleave_byte_view_cases; + +use criterion::{BenchmarkId, Criterion, Throughput}; use std::ops::Range; use rand::RngExt; @@ -27,7 +29,7 @@ use arrow::datatypes::*; use arrow::util::test_util::seedable_rng; use arrow::{array::*, util::bench_util::*}; use arrow_buffer::ScalarBuffer; -use arrow_select::interleave::interleave; +use arrow_select::interleave::{Interleaver, interleave}; use std::hint; use std::sync::Arc; @@ -66,6 +68,60 @@ fn bench_values(c: &mut Criterion, name: &str, len: usize, values: &[&dyn Array] }); } +fn bench_byte_view_compaction(c: &mut Criterion) { + let mut group = c.benchmark_group("interleave_byte_view_compaction"); + interleave_byte_view_cases::for_each_case(|case| { + let values: Vec<&dyn Array> = case.arrays.iter().map(|a| a as &dyn Array).collect(); + group.throughput(Throughput::Elements(case.indices.len() as u64)); + group.bench_with_input( + BenchmarkId::new("interleave", &case.name), + &case.indices, + |b, indices| b.iter(|| hint::black_box(interleave(&values, indices).unwrap())), + ); + group.bench_with_input( + BenchmarkId::new("interleave_gc", &case.name), + &case.indices, + |b, indices| { + b.iter(|| { + let selected = interleave(&values, indices).unwrap(); + let selected = selected.as_any().downcast_ref::().unwrap(); + hint::black_box(selected.gc()) + }) + }, + ); + group.bench_with_input( + BenchmarkId::new("compact", &case.name), + &case.indices, + |b, indices| { + b.iter(|| { + hint::black_box( + Interleaver::new() + .with_compact_byte_views(true) + .interleave(&values, indices) + .unwrap(), + ) + }) + }, + ); + group.bench_with_input( + BenchmarkId::new("compact_shared", &case.name), + &case.indices, + |b, indices| { + b.iter(|| { + hint::black_box( + Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(true) + .interleave(&values, indices) + .unwrap(), + ) + }) + }, + ); + }); + group.finish(); +} + fn add_benchmark(c: &mut Criterion) { let i32 = create_primitive_array::(1024, 0.); let i32_opt = create_primitive_array::(1024, 0.5); @@ -218,5 +274,5 @@ fn add_benchmark(c: &mut Criterion) { } } -criterion_group!(benches, add_benchmark); +criterion_group!(benches, add_benchmark, bench_byte_view_compaction); criterion_main!(benches); From a61dff4dbc4ef49b651cfd00f7fa228f361a94f8 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 2 Oct 2026 23:28:05 +0000 Subject: [PATCH 3/7] Optimize single-buffer compact copies and validate index panics --- arrow-select/src/interleave.rs | 116 +++++++++++++++++++++------- arrow/benches/interleave_kernels.rs | 2 +- 2 files changed, 89 insertions(+), 29 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 658fe1f37269..051c67ecef8c 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -321,39 +321,60 @@ fn interleave_views_compact( .iter() .map(|&size| Vec::with_capacity(size)) .collect(); - let mut current_buffer = 0; - for (raw, &(source, row)) in views.iter_mut().zip(indices) { - let mut view = ByteView::from(*raw); - if view.length <= MAX_INLINE_VIEW_LEN { - continue; - } - let array = arrays[source]; - let buffer = if PRESERVE_SHARING { - let buffer = &mut buffers[view.buffer_index as usize]; - // First occurrences fill consecutive offsets. Repeated ranges point - // behind the write cursor and have already been copied. - if view.offset as usize != buffer.len() { + if !PRESERVE_SHARING && buffers.len() == 1 { + let buffer = &mut buffers[0]; + for (raw, &(source, _)) in views.iter_mut().zip(indices) { + let view = ByteView::from(*raw); + if view.length <= MAX_INLINE_VIEW_LEN { continue; } - view = ByteView::from(array.views()[row]); - buffer - } else { - // Keep source views until this pass, avoiding a second random read - // of the input views. Rewrite the final views in place as we copy. - if buffers[current_buffer].len() == block_sizes[current_buffer] { - current_buffer += 1; - } - let buffer = &mut buffers[current_buffer]; *raw = view - .with_buffer_index(current_buffer as u32) + .with_buffer_index(0) .with_offset(buffer.len() as u32) .as_u128(); - buffer - }; - let start = view.offset as usize; - buffer.extend_from_slice( - &array.data_buffers()[view.buffer_index as usize][start..start + view.length as usize], - ); + // SAFETY: the first pass checked every source index and retained + // this non-null source view, whose payload is valid for its array. + unsafe { + copy_view_payload(arrays.get_unchecked(source), view, buffer); + } + } + } else { + let mut current_buffer = 0; + for (raw, &(source, row)) in views.iter_mut().zip(indices) { + let mut view = ByteView::from(*raw); + if view.length <= MAX_INLINE_VIEW_LEN { + continue; + } + // SAFETY: the first pass checked every source index. + let array = unsafe { arrays.get_unchecked(source) }; + let buffer = if PRESERVE_SHARING { + let buffer = &mut buffers[view.buffer_index as usize]; + // First occurrences fill consecutive offsets. Repeated ranges point + // behind the write cursor and have already been copied. + if view.offset as usize != buffer.len() { + continue; + } + view = ByteView::from(array.views()[row]); + buffer + } else { + // Keep source views until this pass, avoiding a second random read + // of the input views. Rewrite the final views in place as we copy. + if buffers[current_buffer].len() == block_sizes[current_buffer] { + current_buffer += 1; + } + let buffer = &mut buffers[current_buffer]; + *raw = view + .with_buffer_index(current_buffer as u32) + .with_offset(buffer.len() as u32) + .as_u128(); + buffer + }; + // SAFETY: view is a non-null source view from array, either retained + // from the first pass or read above, so its payload range is valid. + unsafe { + copy_view_payload(array, view, buffer); + } + } } let buffers: Vec<_> = buffers.into_iter().map(Buffer::from_vec).collect(); // SAFETY: inline views are unchanged, null views are zero, and every other @@ -363,6 +384,29 @@ fn interleave_views_compact( })) } +/// Copies the payload addressed by a non-inline source view. +/// +/// # Safety +/// +/// The view must reference a valid payload range in `array`. +#[inline] +unsafe fn copy_view_payload( + array: &GenericByteViewArray, + view: ByteView, + output: &mut Vec, +) { + let start = view.offset as usize; + // SAFETY: the caller guarantees that the buffer index and payload range + // address valid bytes in array. + let bytes = unsafe { + array + .data_buffers() + .get_unchecked(view.buffer_index as usize) + .get_unchecked(start..start + view.length as usize) + }; + output.extend_from_slice(bytes); +} + /// Common functionality for interleaving arrays /// /// T is the concrete Array type @@ -1555,6 +1599,22 @@ mod tests { } } + #[test] + fn test_compact_invalid_indices_panic_before_copying() { + let input = StringViewArray::from(vec!["valid selected long string"]); + for preserve in [false, true] { + for invalid in [(1, 0), (usize::MAX, 0), (0, 1), (0, usize::MAX)] { + let interleaver = Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(preserve); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + interleaver.interleave(&[&input], &[(0, 0), invalid]) + })); + assert!(result.is_err(), "{invalid:?}, preserve={preserve}"); + } + } + } + #[test] fn test_compact_option_leaves_nested_views_unchanged() { let values = StringViewArray::from(vec![ diff --git a/arrow/benches/interleave_kernels.rs b/arrow/benches/interleave_kernels.rs index cfb6f19448dd..3c8f1481b0a6 100644 --- a/arrow/benches/interleave_kernels.rs +++ b/arrow/benches/interleave_kernels.rs @@ -85,7 +85,7 @@ fn bench_byte_view_compaction(c: &mut Criterion) { b.iter(|| { let selected = interleave(&values, indices).unwrap(); let selected = selected.as_any().downcast_ref::().unwrap(); - hint::black_box(selected.gc()) + hint::black_box(Arc::new(selected.gc()) as ArrayRef) }) }, ); From 33c0f7d03fba239a8d48d36d74ccdb2cfc19b0e3 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 2 Oct 2026 23:33:10 +0000 Subject: [PATCH 4/7] Reuse interleave validity and streamline compact buffer sizing --- arrow-select/src/interleave.rs | 97 ++++++++++++++++++++++++++-------- 1 file changed, 74 insertions(+), 23 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 051c67ecef8c..12815a2dde0e 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -27,8 +27,7 @@ use arrow_array::*; use arrow_buffer::bit_mask::set_bits; use arrow_buffer::bit_util; use arrow_buffer::{ - ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, NullBufferBuilder, - OffsetBuffer, + ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer, }; use arrow_data::transform::MutableArrayData; use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN}; @@ -256,24 +255,24 @@ fn interleave_views_compact( indices: &[(usize, usize)], max_buffer_size: usize, ) -> Result { - let arrays: Vec<_> = values - .iter() - .map(|array| array.as_byte_view::()) - .collect(); + let interleaved = Interleave::>::new(values, indices); + let arrays = &interleaved.arrays; let mut views = Vec::with_capacity(indices.len()); - let mut nulls = NullBufferBuilder::new(indices.len()); let mut block_sizes: Vec = Vec::new(); + let mut current_size = 0; let mut copied = HashMap::<(usize, u32), usize, _>::with_hasher(ahash::RandomState::with_seeds(0, 0, 0, 0)); for &(source, row) in indices { let array = arrays[source]; let raw = array.views()[row]; - if array.is_null(row) { + if interleaved + .nulls + .as_ref() + .is_some_and(|nulls| nulls.is_null(views.len())) + { views.push(0); - nulls.append_null(); continue; } - nulls.append_non_null(); let mut view = ByteView::from(raw); if view.length <= MAX_INLINE_VIEW_LEN { views.push(raw); @@ -297,25 +296,25 @@ fn interleave_views_compact( } } } - if block_sizes - .last() - .is_none_or(|bytes| bytes + len > max_buffer_size) - { - if block_sizes.len() > i32::MAX as usize { - return Err(ArrowError::OffsetOverflowError(block_sizes.len())); + if current_size != 0 && current_size + len > max_buffer_size { + let next_buffer_index = block_sizes.len() + 1; + if next_buffer_index > i32::MAX as usize { + return Err(ArrowError::OffsetOverflowError(next_buffer_index)); } - block_sizes.push(0); + block_sizes.push(current_size); + current_size = 0; } - let buffer_index = (block_sizes.len() - 1) as u32; - let bytes = block_sizes.last_mut().unwrap(); if PRESERVE_SHARING { - view.buffer_index = buffer_index; - view.offset = *bytes as u32; + view.buffer_index = block_sizes.len() as u32; + view.offset = current_size as u32; } - *bytes += len; + current_size += len; views.push(view.as_u128()); } drop(copied); + if current_size != 0 { + block_sizes.push(current_size); + } let mut buffers: Vec> = block_sizes .iter() @@ -380,7 +379,7 @@ fn interleave_views_compact( // SAFETY: inline views are unchanged, null views are zero, and every other // view addresses its copied source range with the original length and prefix. Ok(Arc::new(unsafe { - GenericByteViewArray::::new_unchecked(views.into(), buffers.into(), nulls.finish()) + GenericByteViewArray::::new_unchecked(views.into(), buffers.into(), interleaved.nulls) })) } @@ -1524,6 +1523,58 @@ mod tests { } } + #[test] + fn test_compact_values_at_and_above_buffer_limit() { + let exact = "x".repeat(32); + let oversized = "y".repeat(33); + let input = StringViewArray::from(vec![ + Some(exact.as_str()), + Some(oversized.as_str()), + Some("inline"), + None, + ]); + let indices = [ + (0, 2), + (0, 3), + (0, 1), + (0, 2), + (0, 0), + (0, 3), + (0, 1), + (0, 0), + ]; + for preserve in [false, true] { + let compact = if preserve { + interleave_views_compact::(&[&input], &indices, 32) + } else { + interleave_views_compact::(&[&input], &indices, 32) + } + .unwrap(); + assert_eq!( + compact.as_ref(), + interleave(&[&input], &indices).unwrap().as_ref() + ); + let compact = compact.as_string_view(); + let expected_sizes = if preserve { + vec![33, 32] + } else { + vec![33, 32, 33, 32] + }; + assert_eq!( + compact + .data_buffers() + .iter() + .map(|buffer| buffer.len()) + .collect::>(), + expected_sizes + ); + for buffer in compact.data_buffers().iter() { + assert_eq!(buffer.len(), buffer.capacity()); + } + compact.to_data().validate_full().unwrap(); + } + } + #[test] fn test_compact_empty_null_and_inline_selections() { let input = From d8726603c966e1bad0f781de443ca067e553d0dc Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 2 Oct 2026 23:49:39 +0000 Subject: [PATCH 5/7] Preserve the original free interleave dispatch path --- arrow-select/src/interleave.rs | 112 ++++++++++++++++++--------------- 1 file changed, 60 insertions(+), 52 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 12815a2dde0e..92dc8e962198 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -81,7 +81,58 @@ pub fn interleave( values: &[&dyn Array], indices: &[(usize, usize)], ) -> Result { - Interleaver::default().interleave(values, indices) + let data_type = validate_interleave_inputs(values)?; + if indices.is_empty() { + return Ok(new_empty_array(data_type)); + } + + downcast_primitive! { + data_type => (primitive_helper, values, indices, data_type), + DataType::Utf8 => interleave_bytes::(values, indices), + DataType::LargeUtf8 => interleave_bytes::(values, indices), + DataType::Binary => interleave_bytes::(values, indices), + DataType::LargeBinary => interleave_bytes::(values, indices), + DataType::BinaryView => interleave_views::(values, indices), + DataType::Utf8View => interleave_views::(values, indices), + DataType::Dictionary(k, _) => downcast_integer! { + k.as_ref() => (dict_helper, values, indices), + _ => unreachable!("illegal dictionary key type {k}") + }, + DataType::Struct(fields) => interleave_struct(fields, values, indices), + DataType::List(field) => interleave_list::(values, indices, field), + DataType::LargeList(field) => interleave_list::(values, indices, field), + DataType::FixedSizeList(field, size) => interleave_fixed_size_list(values, indices, field, *size), + DataType::Map(field, ordered) => interleave_map(values, indices, field, *ordered), + DataType::RunEndEncoded(r, _) => match r.data_type() { + DataType::Int16 => interleave_run_end::(values, indices), + DataType::Int32 => interleave_run_end::(values, indices), + DataType::Int64 => interleave_run_end::(values, indices), + t => unreachable!("illegal run-end type {t}"), + }, + DataType::ListView(field) => interleave_list_view::(values, indices, field), + DataType::LargeListView(field) => interleave_list_view::(values, indices, field), + _ => interleave_fallback(values, indices) + } +} + +#[inline] +fn validate_interleave_inputs<'a>(values: &[&'a dyn Array]) -> Result<&'a DataType, ArrowError> { + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + let data_type = values[0].data_type(); + for array in values.iter().skip(1) { + if array.data_type() != data_type { + return Err(ArrowError::InvalidArgumentError(format!( + "It is not possible to interleave arrays of different data types ({} and {})", + data_type, + array.data_type() + ))); + } + } + Ok(data_type) } /// Configurable interleaving of elements from multiple arrays. @@ -166,37 +217,16 @@ impl Interleaver { values: &[&dyn Array], indices: &[(usize, usize)], ) -> Result { - if values.is_empty() { - return Err(ArrowError::InvalidArgumentError( - "interleave requires input of at least one array".to_string(), - )); - } - let data_type = values[0].data_type(); - - for array in values.iter().skip(1) { - if array.data_type() != data_type { - return Err(ArrowError::InvalidArgumentError(format!( - "It is not possible to interleave arrays of different data types ({} and {})", - data_type, - array.data_type() - ))); - } - } - - if indices.is_empty() { - return Ok(new_empty_array(data_type)); - } - if self.compact_byte_views { - match data_type { - DataType::Utf8View => { + match values.first().map(|array| array.data_type()) { + Some(DataType::Utf8View) => { return self.interleave_views_compact::( values, indices, i32::MAX as usize, ); } - DataType::BinaryView => { + Some(DataType::BinaryView) => { return self.interleave_views_compact::( values, indices, @@ -207,33 +237,7 @@ impl Interleaver { } } - downcast_primitive! { - data_type => (primitive_helper, values, indices, data_type), - DataType::Utf8 => interleave_bytes::(values, indices), - DataType::LargeUtf8 => interleave_bytes::(values, indices), - DataType::Binary => interleave_bytes::(values, indices), - DataType::LargeBinary => interleave_bytes::(values, indices), - DataType::BinaryView => interleave_views::(values, indices), - DataType::Utf8View => interleave_views::(values, indices), - DataType::Dictionary(k, _) => downcast_integer! { - k.as_ref() => (dict_helper, values, indices), - _ => unreachable!("illegal dictionary key type {k}") - }, - DataType::Struct(fields) => interleave_struct(fields, values, indices), - DataType::List(field) => interleave_list::(values, indices, field), - DataType::LargeList(field) => interleave_list::(values, indices, field), - DataType::FixedSizeList(field, size) => interleave_fixed_size_list(values, indices, field, *size), - DataType::Map(field, ordered) => interleave_map(values, indices, field, *ordered), - DataType::RunEndEncoded(r, _) => match r.data_type() { - DataType::Int16 => interleave_run_end::(values, indices), - DataType::Int32 => interleave_run_end::(values, indices), - DataType::Int64 => interleave_run_end::(values, indices), - t => unreachable!("illegal run-end type {t}"), - }, - DataType::ListView(field) => interleave_list_view::(values, indices, field), - DataType::LargeListView(field) => interleave_list_view::(values, indices, field), - _ => interleave_fallback(values, indices) - } + interleave(values, indices) } fn interleave_views_compact( @@ -242,6 +246,10 @@ impl Interleaver { indices: &[(usize, usize)], max_buffer_size: usize, ) -> Result { + let data_type = validate_interleave_inputs(values)?; + if indices.is_empty() { + return Ok(new_empty_array(data_type)); + } if self.preserve_byte_view_sharing { interleave_views_compact::(values, indices, max_buffer_size) } else { From f7155b069f2d8cd63c74f51b5abe8cb3e9b846cd Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 5 Oct 2026 21:47:11 +0000 Subject: [PATCH 6/7] Address compact interleave benchmark and batch API feedback --- arrow-select/src/interleave.rs | 323 ++++++++++++++++++++++++++-- arrow/benches/interleave_kernels.rs | 13 +- 2 files changed, 318 insertions(+), 18 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 92dc8e962198..56190e03d4c5 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -211,7 +211,7 @@ impl Interleaver { /// /// # Panics /// - /// Panics if an array or row index is out of bounds. + /// May panic if an array or row index is out of bounds. pub fn interleave( &self, values: &[&dyn Array], @@ -240,6 +240,85 @@ impl Interleaver { interleave(values, indices) } + /// Interleaves rows from multiple [`RecordBatch`] instances using these options. + /// + /// Calls [`Self::interleave`] for each column and preserves the first batch's + /// schema, including metadata. All input schemas must be equal. Compaction + /// applies only to top-level byte-view columns, not to views nested in other + /// columns. Zero-column batches retain the number of selected rows. + /// + /// # Errors + /// + /// Returns an error if no batches are supplied, their schemas differ + /// (including metadata), or a column cannot be interleaved. + /// + /// # Panics + /// + /// May panic if a batch or row index is out of bounds. + /// + /// # Example + /// + /// ``` + /// use std::sync::Arc; + /// use arrow_array::{ArrayRef, RecordBatch, StringViewArray}; + /// use arrow_select::interleave::Interleaver; + /// + /// let batch = RecordBatch::try_from_iter([ + /// ("name", Arc::new(StringViewArray::from(vec!["a long selected value"])) as ArrayRef) + /// ])?; + /// let result = Interleaver::new() + /// .with_compact_byte_views(true) + /// .with_preserve_byte_view_sharing(true) + /// .interleave_record_batch(&[&batch], &[(0, 0), (0, 0)])?; + /// assert_eq!(result.num_rows(), 2); + /// assert_eq!(result.schema(), batch.schema()); + /// # Ok::<(), arrow_schema::ArrowError>(()) + /// ``` + pub fn interleave_record_batch( + &self, + record_batches: &[&RecordBatch], + indices: &[(usize, usize)], + ) -> Result { + let first = record_batches.first().ok_or_else(|| { + ArrowError::InvalidArgumentError( + "interleave_record_batch requires at least one batch".to_string(), + ) + })?; + let schema = first.schema(); + if record_batches + .iter() + .skip(1) + .any(|batch| batch.schema() != schema) + { + return Err(ArrowError::InvalidArgumentError( + "interleave_record_batch requires identical schemas".to_string(), + )); + } + if schema.fields().is_empty() { + // No column kernel will validate indices for zero-column batches. + for &(batch, row) in indices { + assert!( + row < record_batches[batch].num_rows(), + "row index out of bounds" + ); + } + } + let columns = (0..schema.fields().len()) + .map(|i| { + let values: Vec<&dyn Array> = record_batches + .iter() + .map(|batch| batch.column(i).as_ref()) + .collect(); + self.interleave(&values, indices) + }) + .collect::, _>>()?; + RecordBatch::try_new_with_options( + schema, + columns, + &RecordBatchOptions::new().with_row_count(Some(indices.len())), + ) + } + fn interleave_views_compact( &self, values: &[&dyn Array], @@ -1193,6 +1272,8 @@ fn interleave_fallback_dictionary( /// /// This function will call [`interleave`] on each array of the [`RecordBatch`] instances and assemble a new [`RecordBatch`]. /// +/// To configure byte-view compaction, use [`Interleaver::interleave_record_batch`]. +/// /// # Example /// ``` /// # use std::sync::Arc; @@ -1278,11 +1359,25 @@ mod tests { second.append_value("short"); second.append_value("fourth selected long string"); let second = second.finish().slice(1, 3); - assert!(first.data_buffers().len() > 1); - assert!(second.data_buffers().len() > 1); + // Slicing retains every source buffer, including the discarded rows. + assert_eq!(first.data_buffers().len(), 4); + assert_eq!(second.data_buffers().len(), 3); + let source_allocations: Vec<_> = first + .data_buffers() + .iter() + .chain(second.data_buffers().iter()) + .map(Buffer::data_ptr) + .collect(); - let indices = [(1, 2), (0, 1), (0, 2), (1, 1), (0, 0), (1, 2)]; + let indices = [(1, 2), (0, 1), (0, 2), (1, 1), (0, 0), (1, 2), (1, 0)]; let ordinary = interleave(&[&first, &second], &indices).unwrap(); + assert!( + ordinary + .as_string_view() + .data_buffers() + .iter() + .any(|buffer| source_allocations.contains(&buffer.data_ptr())) + ); for preserve in [false, true] { let compact = Interleaver::new() .with_compact_byte_views(true) @@ -1300,8 +1395,18 @@ mod tests { Some("short"), Some("first selected long string"), Some("fourth selected long string"), + Some("third selected long string"), ] ); + let copies = if preserve { 1 } else { 2 }; + let selected_payload = first.value(0).len() + + first.value(2).len() + + second.value(0).len() + + copies * second.value(2).len(); + assert_eq!(compact.data_buffers().len(), 1); + assert_eq!(compact.data_buffers()[0].len(), selected_payload); + assert_eq!(compact.data_buffers()[0].capacity(), selected_payload); + assert!(!source_allocations.contains(&compact.data_buffers()[0].data_ptr())); assert_eq!(compact.views()[0] == compact.views()[5], preserve); } } @@ -1346,13 +1451,15 @@ mod tests { let unselected = "u".repeat(32 * 1024); let hidden = "n".repeat(32 * 1024); let selected = "selected long string"; - let input = StringViewArray::from(vec![unselected.as_str(), selected, hidden.as_str()]); + let source = + StringViewArray::from(vec![unselected.as_str(), selected, hidden.as_str()]); let input = StringViewArray::try_new( - input.views().clone(), - input.data_buffers().clone(), + source.views().clone(), + source.data_buffers().clone(), Some(NullBuffer::from(vec![true, true, false])), ) .unwrap(); + drop(source); let compact = Interleaver::new() .with_compact_byte_views(true) .with_preserve_byte_view_sharing(preserve) @@ -1367,12 +1474,14 @@ mod tests { copies * selected.len() ); assert_eq!(compact.views()[1], 0); + assert_eq!( + compact.iter().collect::>(), + vec![Some(selected), None, Some(selected)] + ); + // Shared buffers also survive dropping input; compare live allocations instead. for output in compact.data_buffers().iter() { for source in input.data_buffers().iter() { - assert_ne!( - output.as_ptr() as usize - output.ptr_offset(), - source.as_ptr() as usize - source.ptr_offset() - ); + assert_ne!(output.data_ptr(), source.data_ptr()); } } assert_ne!(compact.views().as_ptr(), input.views().as_ptr()); @@ -1380,11 +1489,6 @@ mod tests { compact.nulls().unwrap().buffer().as_ptr(), input.nulls().unwrap().buffer().as_ptr() ); - drop(input); - assert_eq!( - compact.iter().collect::>(), - vec![Some(selected), None, Some(selected)] - ); } } @@ -1702,6 +1806,193 @@ mod tests { .as_ptr(), pointer ); + let batch = RecordBatch::try_from_iter([("nested", Arc::new(input) as ArrayRef)]).unwrap(); + for preserve in [false, true] { + let output = Interleaver::new() + .with_compact_byte_views(true) + .with_preserve_byte_view_sharing(preserve) + .interleave_record_batch(&[&batch], &[(0, 1)]) + .unwrap(); + assert_eq!(output.column(0).as_ref(), ordinary.as_ref()); + assert_eq!( + output + .column(0) + .as_list::() + .values() + .as_string_view() + .data_buffers()[0] + .as_ptr(), + pointer + ); + } + } + + #[test] + fn test_interleaver_record_batch_options() { + use arrow_schema::Schema; + + let schema = Arc::new(Schema::new_with_metadata( + vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8View, true) + .with_metadata([("field_key".to_string(), "field_value".to_string())]), + Field::new("bytes", DataType::BinaryView, true), + ], + [("schema_key".to_string(), "schema_value".to_string())], + )); + let first_value = "first selected long string"; + let second_value = "second selected long string"; + let first = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![10, 20])), + Arc::new(StringViewArray::from(vec![Some(first_value), None])), + Arc::new(BinaryViewArray::from(vec![ + Some(first_value.as_bytes()), + None, + ])), + ], + ) + .unwrap(); + let second = RecordBatch::try_new( + Arc::new(schema.as_ref().clone()), + vec![ + Arc::new(Int32Array::from(vec![30, 40])), + Arc::new(StringViewArray::from(vec![Some(second_value), None])), + Arc::new(BinaryViewArray::from(vec![ + Some(second_value.as_bytes()), + None, + ])), + ], + ) + .unwrap(); + let inputs = [&first, &second]; + let indices = [(1, 0), (0, 1), (0, 0), (1, 0)]; + let expected = interleave_record_batch(&inputs, &indices).unwrap(); + for compact in [false, true] { + for preserve in [false, true] { + let interleaver = Interleaver::new() + .with_compact_byte_views(compact) + .with_preserve_byte_view_sharing(preserve); + let output = interleaver + .interleave_record_batch(&inputs, &indices) + .unwrap(); + assert_eq!(output, expected); + assert!(Arc::ptr_eq(&output.schema(), &schema)); + for column in [1, 2] { + let output_data = output.column(column).to_data(); + let payload = &output_data.buffers()[1..]; + if compact { + let copies = if preserve { 1 } else { 2 }; + let size = first_value.len() + copies * second_value.len(); + assert_eq!(payload.len(), 1); + assert_eq!(payload[0].capacity(), size); + for input in inputs { + for source in &input.column(column).to_data().buffers()[1..] { + assert_ne!(payload[0].data_ptr(), source.data_ptr()); + } + } + } else { + let expected_data = expected.column(column).to_data(); + assert_eq!( + payload.iter().map(Buffer::data_ptr).collect::>(), + expected_data.buffers()[1..] + .iter() + .map(Buffer::data_ptr) + .collect::>() + ); + } + } + let empty = interleaver.interleave_record_batch(&inputs, &[]).unwrap(); + assert_eq!(empty.num_rows(), 0); + assert_eq!(empty.schema(), schema); + } + } + } + + #[test] + fn test_interleaver_record_batch_schema_errors() { + use arrow_schema::Schema; + + let field = Field::new("value", DataType::Utf8View, true); + let first = RecordBatch::try_new( + Arc::new(Schema::new(vec![field.clone()])), + vec![Arc::new(StringViewArray::from(vec!["selected long value"]))], + ) + .unwrap(); + let mismatches = [ + Schema::empty(), + Schema::new(vec![Field::new("other", DataType::Utf8View, true)]), + Schema::new(vec![Field::new("value", DataType::BinaryView, true)]), + Schema::new(vec![Field::new("value", DataType::Utf8View, false)]), + Schema::new(vec![ + field + .clone() + .with_metadata([("key".to_string(), "value".to_string())]), + ]), + Schema::new_with_metadata(vec![field], [("key".to_string(), "value".to_string())]), + ]; + for compact in [false, true] { + let interleaver = Interleaver::new().with_compact_byte_views(compact); + for indices in [&[][..], &[(0, 0)][..]] { + assert!(matches!( + interleaver.interleave_record_batch(&[], indices), + Err(ArrowError::InvalidArgumentError(_)) + )); + } + for schema in &mismatches { + let other = RecordBatch::new_empty(Arc::new(schema.clone())); + for indices in [&[][..], &[(0, 0)][..]] { + assert!(matches!( + interleaver.interleave_record_batch(&[&first, &other], indices), + Err(ArrowError::InvalidArgumentError(_)) + )); + } + } + } + } + + #[test] + fn test_interleaver_record_batch_zero_columns() { + use arrow_schema::Schema; + + let schema = Arc::new(Schema::empty()); + let first = RecordBatch::try_new_with_options( + schema.clone(), + vec![], + &RecordBatchOptions::new().with_row_count(Some(3)), + ) + .unwrap(); + let second = RecordBatch::try_new_with_options( + schema.clone(), + vec![], + &RecordBatchOptions::new().with_row_count(Some(2)), + ) + .unwrap(); + for compact in [false, true] { + let interleaver = Interleaver::new().with_compact_byte_views(compact); + let output = interleaver + .interleave_record_batch(&[&first, &second], &[(1, 1), (0, 2), (0, 0), (1, 1)]) + .unwrap(); + assert_eq!(output.num_rows(), 4); + assert_eq!(output.num_columns(), 0); + assert_eq!(output.schema(), schema); + assert_eq!( + interleaver + .interleave_record_batch(&[&first], &[]) + .unwrap() + .num_rows(), + 0 + ); + for invalid in [(2, 0), (0, 3), (1, 2), (usize::MAX, 0), (0, usize::MAX)] { + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + interleaver.interleave_record_batch(&[&first, &second], &[invalid]) + })) + .is_err() + ); + } + } } #[test] diff --git a/arrow/benches/interleave_kernels.rs b/arrow/benches/interleave_kernels.rs index 3c8f1481b0a6..61e4a0cc32d3 100644 --- a/arrow/benches/interleave_kernels.rs +++ b/arrow/benches/interleave_kernels.rs @@ -32,6 +32,7 @@ use arrow_buffer::ScalarBuffer; use arrow_select::interleave::{Interleaver, interleave}; use std::hint; use std::sync::Arc; +use std::time::Duration; fn do_bench( c: &mut Criterion, @@ -274,5 +275,13 @@ fn add_benchmark(c: &mut Criterion) { } } -criterion_group!(benches, add_benchmark, bench_byte_view_compaction); -criterion_main!(benches); +criterion_group!(benches, add_benchmark); +criterion_group! { + name = byte_view_benches; + config = Criterion::default() + .sample_size(20) + .warm_up_time(Duration::from_millis(100)) + .measurement_time(Duration::from_millis(300)); + targets = bench_byte_view_compaction +} +criterion_main!(benches, byte_view_benches); From c8ee97880e46af9257a41e8df355bd9bcf23f238 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 5 Oct 2026 22:10:58 +0000 Subject: [PATCH 7/7] Preserve inlined setup on the existing interleave path --- arrow-select/src/interleave.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 56190e03d4c5..c4ac12e1b0c5 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -81,7 +81,23 @@ pub fn interleave( values: &[&dyn Array], indices: &[(usize, usize)], ) -> Result { - let data_type = validate_interleave_inputs(values)?; + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + let data_type = values[0].data_type(); + + for array in values.iter().skip(1) { + if array.data_type() != data_type { + return Err(ArrowError::InvalidArgumentError(format!( + "It is not possible to interleave arrays of different data types ({} and {})", + data_type, + array.data_type() + ))); + } + } + if indices.is_empty() { return Ok(new_empty_array(data_type)); } @@ -504,6 +520,7 @@ struct Interleave<'a, T> { } impl<'a, T: Array + 'static> Interleave<'a, T> { + #[inline] fn new(values: &[&'a dyn Array], indices: &'a [(usize, usize)]) -> Self { let mut has_nulls = false; let arrays: Vec<&T> = values