This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 007894ef0d Reuse buffers when interleaving byte views (#11281)
007894ef0d is described below
commit 007894ef0db0eb3b4dc3226d7b490b1104352231
Author: Peter L <[email protected]>
AuthorDate: Thu Oct 1 12:40:39 2026 +0930
Reuse buffers when interleaving byte views (#11281)
# Which issue does this PR close?
Fixes https://github.com/apache/arrow-rs/issues/11280
And also helps with: https://github.com/apache/datafusion/pull/23565
# Rationale for this change
See the original issue, but, basically the interleave kernel can make
sub-optimal string views especially when buffers are shared.
# What changes are included in this PR?
We use a btreemap on the seen buffers (not individual values) and use
that to prevent us from cloning the same buffer multiple times.
# Are these changes tested?
Yes tested, and locally benchmarks have shown a moderate speed increase
in some of the interleave benchmarks, on my dev machine which has an
**AMD Ryzen 9 9950X3D** processor:
| benchmark | main | this PR | change |
|-----------------------------------------------------------------------|----------|----------|--------|
| interleave str_view(0.0) 100 [0..100, 100..230, 450..1000] | 284.5 ns
| 268.6 ns | -5.6% |
| interleave str_view(0.0) 400 [0..100, 100..230, 450..1000] | 706.9 ns
| 563.9 ns | -20.2% |
| interleave str_view(0.0) 1024 [0..100, 100..230, 450..1000] | 1.717 µs
| 1.250 µs | -27.2% |
| interleave str_view(0.0) 1024 [0..100, 100..230, 450..1000, 0..1000] |
1.783 µs | 1.274 µs | -28.6% |
# Are there any user-facing changes?
No user facing changes, just an optimization.
---
arrow-select/src/interleave.rs | 84 ++++++++++++++++++++++++++++++++++++++----
1 file changed, 76 insertions(+), 8 deletions(-)
diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs
index 54f4412c5b..74e29defe8 100644
--- a/arrow-select/src/interleave.rs
+++ b/arrow-select/src/interleave.rs
@@ -26,10 +26,13 @@ use arrow_array::types::*;
use arrow_array::*;
use arrow_buffer::bit_mask::set_bits;
use arrow_buffer::bit_util;
-use arrow_buffer::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer,
OffsetBuffer};
+use arrow_buffer::{
+ ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer,
OffsetBuffer,
+};
use arrow_data::ByteView;
use arrow_data::transform::MutableArrayData;
use arrow_schema::{ArrowError, DataType, FieldRef, Fields};
+use std::collections::BTreeMap;
use std::sync::Arc;
macro_rules! primitive_helper {
@@ -314,8 +317,17 @@ fn interleave_views<T: ByteViewType>(
offsets.push(total_buffers);
}
+ // Marks a buffer in `buffer_to_new_index` that has not yet been
referenced.
+ //
+ // The view's buffer index is a signed 32-bit integer in the Arrow
specification,
+ // so `u32::MAX` can never be a valid buffer index
+ const UNASSIGNED: u32 = u32::MAX;
+
// contains the mapping from old buffer index to new buffer index
- let mut buffer_to_new_index = vec![None; total_buffers];
+ let mut buffer_to_new_index = vec![UNASSIGNED; total_buffers];
+
+ // Contains the index in `buffers` of each buffer already emitted, keyed
by identity
+ let mut seen_buffers = BTreeMap::new();
let views: Vec<u128> = indices
.iter()
@@ -329,12 +341,15 @@ fn interleave_views<T: ByteViewType>(
// value is big enough to be in a variadic buffer
let view = ByteView::from(*view);
let buffer_to_new_idx = offsets[*array_idx] + view.buffer_index as
usize;
- let new_buffer_idx: u32 =
- *buffer_to_new_index[buffer_to_new_idx].get_or_insert_with(|| {
- buffers.push(array.data_buffers()[view.buffer_index as
usize].clone());
- (buffers.len() - 1) as u32
- });
- view.with_buffer_index(new_buffer_idx).as_u128()
+ let new_buffer_idx = &mut buffer_to_new_index[buffer_to_new_idx];
+ if *new_buffer_idx == UNASSIGNED {
+ *new_buffer_idx = push_buffer(
+ &mut seen_buffers,
+ &mut buffers,
+ &array.data_buffers()[view.buffer_index as usize],
+ );
+ }
+ view.with_buffer_index(*new_buffer_idx).as_u128()
})
.collect();
@@ -344,6 +359,24 @@ fn interleave_views<T: ByteViewType>(
Ok(Arc::new(array))
}
+/// Returns the index of `buffer` in `buffers`, adding it if not already
present.
+///
+/// Multiple input arrays, e.g. slices of the same array, may reference the
same
+/// buffer, so this checks by identity to ensure it is only emitted once
+#[inline(never)]
+fn push_buffer(
+ seen: &mut BTreeMap<(*const u8, usize), u32>,
+ buffers: &mut Vec<Buffer>,
+ buffer: &Buffer,
+) -> u32 {
+ *seen
+ .entry((buffer.as_ptr(), buffer.len()))
+ .or_insert_with(|| {
+ buffers.push(buffer.clone());
+ (buffers.len() - 1) as u32
+ })
+}
+
fn interleave_struct(
fields: &Fields,
values: &[&dyn Array],
@@ -1674,6 +1707,41 @@ mod tests {
);
}
+ #[test]
+ fn test_interleave_views_shared_buffers() {
+ let long = |i: usize| format!("long_string_not_inlined_{i}");
+ let base = StringViewArray::from_iter_values((0..10).map(long));
+ assert_eq!(base.data_buffers().len(), 1);
+
+ // Slices share the same buffers
+ let a = base.slice(0, 5);
+ let b = base.slice(5, 5);
+ // A distinct buffer list containing the same underlying buffer
+ let c = StringViewArray::try_new(base.views().clone(),
base.data_buffers().to_vec(), None)
+ .unwrap();
+
+ let indices = &[(0, 1), (1, 2), (2, 3), (0, 4), (1, 0), (2, 9)];
+ let values = interleave(&[&a, &b, &c], indices).unwrap();
+ let result = values.as_string_view();
+ assert_eq!(result.data_buffers().len(), 1);
+
+ let expected: Vec<_> = [1, 7, 3, 4, 5,
9].into_iter().map(long).collect();
+ let actual: Vec<_> = result.iter().map(|x|
x.unwrap().to_string()).collect();
+ assert_eq!(actual, expected);
+
+ // Many slices, exceeding the linear scan threshold
+ let slices: Vec<_> = (0..40).map(|i| base.slice(i % 10, 1)).collect();
+ let arrays: Vec<&dyn Array> = slices.iter().map(|x| x as &dyn
Array).collect();
+ let indices: Vec<_> = (0..40).rev().map(|i| (i, 0)).collect();
+ let values = interleave(&arrays, &indices).unwrap();
+ let result = values.as_string_view();
+ assert_eq!(result.data_buffers().len(), 1);
+
+ let expected: Vec<_> = (0..40).rev().map(|i| long(i % 10)).collect();
+ let actual: Vec<_> = result.iter().map(|x|
x.unwrap().to_string()).collect();
+ assert_eq!(actual, expected);
+ }
+
#[test]
fn test_interleave_views_multiple_buffers() {
let str1 = "very_long_string_from_first_buffer".as_bytes();
