diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 74e29defe835..c4ac12e1b0c5 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -29,10 +29,10 @@ use arrow_buffer::bit_util; use arrow_buffer::{ ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, 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 { @@ -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 +/// [`Interleaver::with_compact_byte_views`]. pub fn interleave( values: &[&dyn Array], indices: &[(usize, usize)], @@ -128,6 +131,384 @@ pub fn interleave( } } +#[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. +/// +/// 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 + /// + /// May panic if an array or row index is out of bounds. + pub fn interleave( + &self, + values: &[&dyn Array], + indices: &[(usize, usize)], + ) -> Result { + if self.compact_byte_views { + match values.first().map(|array| array.data_type()) { + Some(DataType::Utf8View) => { + return self.interleave_views_compact::( + values, + indices, + i32::MAX as usize, + ); + } + Some(DataType::BinaryView) => { + return self.interleave_views_compact::( + values, + indices, + i32::MAX as usize, + ); + } + _ => {} + } + } + + 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], + 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 { + interleave_views_compact::(values, indices, max_buffer_size) + } + } +} + +fn interleave_views_compact( + values: &[&dyn Array], + indices: &[(usize, usize)], + max_buffer_size: usize, +) -> Result { + let interleaved = Interleave::>::new(values, indices); + let arrays = &interleaved.arrays; + let mut views = Vec::with_capacity(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 interleaved + .nulls + .as_ref() + .is_some_and(|nulls| nulls.is_null(views.len())) + { + views.push(0); + continue; + } + let mut view = ByteView::from(raw); + if view.length <= MAX_INLINE_VIEW_LEN { + views.push(raw); + continue; + } + let len = view.length as usize; + if len > i32::MAX as usize { + return Err(ArrowError::OffsetOverflowError(len)); + } + 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()); + } + } + } + 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(current_size); + current_size = 0; + } + if PRESERVE_SHARING { + view.buffer_index = block_sizes.len() as u32; + view.offset = current_size as u32; + } + 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() + .map(|&size| Vec::with_capacity(size)) + .collect(); + 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; + } + *raw = view + .with_buffer_index(0) + .with_offset(buffer.len() as u32) + .as_u128(); + // 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 + // view addresses its copied source range with the original length and prefix. + Ok(Arc::new(unsafe { + GenericByteViewArray::::new_unchecked(views.into(), buffers.into(), interleaved.nulls) + })) +} + +/// 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 @@ -139,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 @@ -907,6 +1289,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; @@ -969,12 +1353,665 @@ 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); + // 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), (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) + .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"), + 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); + } + } + + #[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 source = + StringViewArray::from(vec![unselected.as_str(), selected, hidden.as_str()]); + let input = StringViewArray::try_new( + 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) + .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); + 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.data_ptr(), source.data_ptr()); + } + } + assert_ne!(compact.views().as_ptr(), input.views().as_ptr()); + assert_ne!( + compact.nulls().unwrap().buffer().as_ptr(), + input.nulls().unwrap().buffer().as_ptr() + ); + } + } + + #[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_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 = + 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_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![ + "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 + ); + 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] fn test_primitive() { let a = Int32Array::from_iter_values([1, 2, 3, 4]); 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..61e4a0cc32d3 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,9 +29,10 @@ 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; +use std::time::Duration; fn do_bench( c: &mut Criterion, @@ -66,6 +69,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(Arc::new(selected.gc()) as ArrayRef) + }) + }, + ); + 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); @@ -219,4 +276,12 @@ fn add_benchmark(c: &mut Criterion) { } criterion_group!(benches, add_benchmark); -criterion_main!(benches); +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);