This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11696-f971a78518039c1108c8e8ffdc17e4bebc85dec5 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit d9d2d1edf46b1265d9fea9770c5e86933dff4ce8 Author: RIchard Baah <[email protected]> AuthorDate: Wed Oct 28 15:05:07 2026 +0000 perf(coalesce): specialize sparse byte array filters (#11696) # Which issue does this PR close? Part of #7764. prefix for #11563 # Why Very sparse filters currently materialize variable-width byte arrays before coalescing them, which copies the selected data twice. # What changed Adds a specialized coalescer for `Utf8`, `LargeUtf8`, `Binary`, and `LargeBinary`. It uses a single `GenericByteBuilder` for the in-progress output. Sparse index filters selecting at most `1 / 1024` of rows append selected values directly; other filters retain the materialized filtering path. # user facing? There are no public API changes. --- arrow-select/src/coalesce.rs | 80 ++++++++++++++- arrow-select/src/coalesce/byte.rs | 201 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 278 insertions(+), 3 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 0cd7d9c727..4efa8ca851 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -23,7 +23,7 @@ use crate::filter::{FilterBuilder, FilterPredicate, FilterSelection}; use crate::interleave::interleave_record_batch; use crate::take::take_record_batch; -use arrow_array::types::{BinaryViewType, StringViewType}; +use arrow_array::types::{BinaryViewType, GenericBinaryType, GenericStringType, StringViewType}; use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, downcast_primitive}; use arrow_schema::{ArrowError, DataType, SchemaRef}; use std::collections::VecDeque; @@ -32,12 +32,14 @@ use std::sync::Arc; // https://github.com/apache/datafusion/blob/9d2f04996604e709ee440b65f41e7b882f50b788/datafusion/physical-plan/src/coalesce/mod.rs#L26-L25 mod boolean; +mod byte; mod byte_view; mod fixed_size_binary; mod generic; mod primitive; use boolean::InProgressBooleanArray; +use byte::InProgressByteArray; use byte_view::InProgressByteViewArray; use fixed_size_binary::InProgressFixedSizeBinaryArray; use generic::GenericInProgressArray; @@ -47,7 +49,14 @@ fn has_sparse_filter_copy(data_type: &DataType) -> bool { data_type.is_primitive() || matches!( data_type, - DataType::Boolean | DataType::Utf8View | DataType::BinaryView + DataType::Boolean + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + // Unlike offset-based byte arrays, views can reuse their backing buffers. + | DataType::Utf8View + | DataType::BinaryView ) } @@ -62,6 +71,14 @@ fn should_use_sparse_filter_copy(filter_len: usize, selected_count: usize) -> bo selected_count <= filter_len / SPARSE_FILTER_COPY_MAX_SELECTIVITY_DENOMINATOR } +// Variable-width byte arrays only outperform the materialized filter kernel at +// very low selection densities. +const BYTE_ARRAY_DIRECT_COPY_MAX_SELECTIVITY_DENOMINATOR: usize = 1024; + +fn should_use_byte_array_direct_copy(filter_len: usize, selected_count: usize) -> bool { + selected_count <= filter_len / BYTE_ARRAY_DIRECT_COPY_MAX_SELECTIVITY_DENOMINATOR +} + /// Concatenate multiple [`RecordBatch`]es /// /// Implements the common pattern of incrementally creating output @@ -165,6 +182,8 @@ pub struct BatchCoalescer { in_progress_arrays: Vec<Box<dyn InProgressArray>>, /// True if some column still needs the materialized filter path. has_non_specialized_filter_columns: bool, + /// True if the schema contains a variable-width byte-array column. + has_byte_array_columns: bool, /// Buffered row count. Always less than `batch_size` buffered_rows: usize, /// Completed batches @@ -186,6 +205,12 @@ impl BatchCoalescer { .fields() .iter() .any(|field| !has_sparse_filter_copy(field.data_type())); + let has_byte_array_columns = schema.fields().iter().any(|field| { + matches!( + field.data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary + ) + }); let in_progress_arrays = schema .fields() .iter() @@ -197,6 +222,7 @@ impl BatchCoalescer { target_batch_size, in_progress_arrays, has_non_specialized_filter_columns, + has_byte_array_columns, // We will for sure store at least one completed batch completed: VecDeque::with_capacity(1), buffered_rows: 0, @@ -756,6 +782,8 @@ impl BatchCoalescer { let should_materialize_filter = exceeds_coalesce_limit || self.has_non_specialized_filter_columns || does_not_fit_buffer + || (self.has_byte_array_columns + && !should_use_byte_array_direct_copy(filter_len, selected_count)) || !should_use_sparse_filter_copy(filter_len, selected_count); if should_materialize_filter { @@ -798,6 +826,14 @@ fn create_in_progress_array(data_type: &DataType, batch_size: usize) -> Box<dyn Box::new(InProgressByteViewArray::<BinaryViewType>::new(batch_size)) } DataType::Boolean => Box::new(InProgressBooleanArray::new(batch_size)), + DataType::Utf8 => Box::new(InProgressByteArray::<GenericStringType<i32>>::new(batch_size)), + DataType::LargeUtf8 => { + Box::new(InProgressByteArray::<GenericStringType<i64>>::new(batch_size)) + } + DataType::Binary => Box::new(InProgressByteArray::<GenericBinaryType<i32>>::new(batch_size)), + DataType::LargeBinary => { + Box::new(InProgressByteArray::<GenericBinaryType<i64>>::new(batch_size)) + } DataType::FixedSizeBinary(size) => { Box::new(InProgressFixedSizeBinaryArray::new(*size, batch_size)) } @@ -1822,10 +1858,12 @@ mod tests { ])); let coalescer = BatchCoalescer::new(supported, 100); assert!(!coalescer.has_non_specialized_filter_columns); + assert!(!coalescer.has_byte_array_columns); let utf8 = Arc::new(Schema::new(vec![Field::new("utf8", DataType::Utf8, true)])); let coalescer = BatchCoalescer::new(utf8, 100); - assert!(coalescer.has_non_specialized_filter_columns); + assert!(!coalescer.has_non_specialized_filter_columns); + assert!(coalescer.has_byte_array_columns); let boolean = Arc::new(Schema::new(vec![Field::new( "boolean", @@ -1836,6 +1874,42 @@ mod tests { assert!(!coalescer.has_non_specialized_filter_columns); } + fn sparse_filter_byte_array_input() -> (RecordBatch, BooleanArray) { + let mut strings = vec![Some("padding"); 4096]; + strings[..4].copy_from_slice(&[Some("zero"), None, Some("two"), Some("three")]); + let batch = RecordBatch::try_from_iter([( + "bytes", + Arc::new(StringArray::from(strings)) as ArrayRef, + )]) + .unwrap(); + let filter = BooleanArray::from_iter((0..batch.num_rows()).map(|idx| Some(idx < 4))); + (batch, filter) + } + + #[test] + fn test_sparse_filter_byte_array_direct_copy() { + let (batch, filter) = sparse_filter_byte_array_input(); + assert!(should_use_byte_array_direct_copy(batch.num_rows(), 4)); + + let fallback_filter = + BooleanArray::from_iter((0..batch.num_rows()).map(|idx| Some(idx < 5))); + assert!(!should_use_byte_array_direct_copy(batch.num_rows(), 5)); + let mut coalescer = BatchCoalescer::new(batch.schema(), 6); + coalescer + .push_batch_with_filter(batch.clone(), &fallback_filter) + .unwrap(); + coalescer.finish_buffered_batch().unwrap(); + let expected = filter_record_batch(&batch, &fallback_filter).unwrap(); + assert_eq!(coalescer.next_completed_batch().unwrap(), expected); + + Test::new("coalesce_sparse_filter_utf8_direct_copy") + .with_batch(batch) + .with_filter(filter) + .with_batch_size(4) + .with_expected_output_sizes([4]) + .run(); + } + #[derive(Debug, Clone, PartialEq)] struct ExpectedLayout { len: usize, diff --git a/arrow-select/src/coalesce/byte.rs b/arrow-select/src/coalesce/byte.rs new file mode 100644 index 0000000000..2c623d0d11 --- /dev/null +++ b/arrow-select/src/coalesce/byte.rs @@ -0,0 +1,201 @@ +// 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 super::InProgressArray; +use crate::filter::{FilterIndices, FilterPredicate, FilterSelection}; +use arrow_array::builder::GenericByteBuilder; +use arrow_array::cast::AsArray; +use arrow_array::types::ByteArrayType; +use arrow_array::{Array, ArrayRef, GenericByteArray, OffsetSizeTrait}; +use arrow_buffer::ArrowNativeType; +use arrow_schema::ArrowError; +use std::sync::Arc; + +pub(crate) struct InProgressByteArray<T: ByteArrayType> { + source: Option<ArrayRef>, + batch_size: usize, + direct: GenericByteBuilder<T>, +} + +impl<T: ByteArrayType> std::fmt::Debug for InProgressByteArray<T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("InProgressByteArray") + .field( + "direct_rows", + &self.direct.offsets_slice().len().saturating_sub(1), + ) + .finish() + } +} + +impl<T: ByteArrayType> InProgressByteArray<T> +where + T::Offset: OffsetSizeTrait, +{ + pub(crate) fn new(batch_size: usize) -> Self { + Self { + source: None, + batch_size, + direct: GenericByteBuilder::with_capacity(0, 0), + } + } + + fn checked_values_len(current: usize, additional: usize) -> Result<usize, ArrowError> { + let new_len = current + .checked_add(additional) + .ok_or(ArrowError::OffsetOverflowError(usize::MAX))?; + if new_len > T::Offset::MAX_OFFSET { + return Err(ArrowError::OffsetOverflowError(new_len)); + } + Ok(new_len) + } + + fn ensure_capacity(&mut self, source: &GenericByteArray<T>, rows: usize) { + if self.direct.offsets_slice().len() == 1 && !source.is_empty() { + let offsets = source.value_offsets(); + let bytes = offsets[source.len()].as_usize() - offsets[0].as_usize(); + let rows = rows.min(self.batch_size); + let estimated = bytes.saturating_mul(rows) / source.len(); + // Later slices may contain slightly larger values than this first sample. + let capacity = estimated.saturating_add(estimated / 8); + self.direct = GenericByteBuilder::with_capacity(rows, capacity); + } + } + + fn append_filter_indices( + &mut self, + source: &GenericByteArray<T>, + indices: FilterIndices<'_>, + count: usize, + ) -> Result<(), ArrowError> { + self.ensure_capacity(source, count); + indices.try_for_each(|idx| { + if source.is_null(idx) { + self.direct.append_null(); + } else { + let value = source.value(idx); + Self::checked_values_len( + self.direct.values_slice().len(), + source.value_length(idx).as_usize(), + )?; + self.direct.append_value(value); + } + Ok::<(), ArrowError>(()) + })?; + Ok(()) + } +} + +impl<T: ByteArrayType + 'static> InProgressArray for InProgressByteArray<T> +where + T::Offset: OffsetSizeTrait, +{ + fn set_source(&mut self, source: Option<ArrayRef>) { + self.source = source; + } + + fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> { + let source = self.source.as_ref().ok_or_else(|| { + ArrowError::InvalidArgumentError( + "Internal Error: InProgressByteArray: source not set".to_string(), + ) + })?; + let source = source.as_bytes::<T>().slice(offset, len); + self.ensure_capacity(&source, len); + self.direct.append_array(&source) + } + + fn copy_rows_by_filter_from( + &mut self, + source: ArrayRef, + filter: &FilterPredicate, + ) -> Result<(), ArrowError> { + match filter.selection() { + FilterSelection::None => Ok(()), + // Index selections are sparse enough that copying values directly avoids the + // extra allocation of a filtered byte array. Dense contiguous selections use + // the filter kernel below, which copies their ranges in bulk. + FilterSelection::Indices(indices) => { + self.append_filter_indices(source.as_bytes::<T>(), indices, filter.count()) + } + FilterSelection::Slices(_) | FilterSelection::All { .. } => { + let filtered = filter.filter(source.as_ref())?; + self.direct.append_array(filtered.as_bytes::<T>()) + } + } + } + + fn finish(&mut self) -> Result<ArrayRef, ArrowError> { + Ok(Arc::new(self.direct.finish())) + } + + fn size(&self) -> usize { + self.source + .as_ref() + .map_or(0, |source| source.get_array_memory_size()) + + self.direct.offsets_capacity() * size_of::<T::Offset>() + + self.direct.values_capacity() + + self.direct.validity_capacity() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::filter::FilterBuilder; + use arrow_array::types::GenericStringType; + use arrow_array::{BooleanArray, StringArray}; + + #[test] + fn copy_rows_and_filter_preserve_order() { + let mut values = vec![Some("padding"); 4098]; + values[1..4].copy_from_slice(&[Some("zero"), None, Some("two")]); + values[2048] = Some("far"); + let source: ArrayRef = Arc::new(StringArray::from(values).slice(1, 4096)); + let mask = + BooleanArray::from_iter((0..source.len()).map(|idx| Some(matches!(idx, 0..=2 | 2047)))); + let predicate = FilterBuilder::new(&mask).optimize().build(); + let mut in_progress = InProgressByteArray::<GenericStringType<i32>>::new(source.len()); + in_progress.set_source(Some(Arc::clone(&source))); + in_progress.copy_rows(0, 1).unwrap(); + in_progress + .copy_rows_by_filter_from(Arc::clone(&source), &predicate) + .unwrap(); + in_progress.copy_rows(2, 1).unwrap(); + + let actual = in_progress.finish().unwrap(); + let expected = StringArray::from(vec![ + Some("zero"), + Some("zero"), + None, + Some("two"), + Some("far"), + Some("two"), + ]); + assert_eq!(actual.as_string::<i32>(), &expected); + } + + #[test] + fn checked_values_len_boundaries() { + type Bytes = InProgressByteArray<GenericStringType<i32>>; + let max = i32::MAX as usize; + + assert_eq!(Bytes::checked_values_len(max - 1, 1).unwrap(), max); + assert!(Bytes::checked_values_len(max, 1).is_err()); + assert!(Bytes::checked_values_len(usize::MAX, 1).is_err()); + } +}
