From f9c00ff3864ebed3904842e4c6350462546e7612 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 17 Jun 2025 22:31:56 +0200 Subject: [PATCH 1/5] interleave_views faster --- arrow-select/src/interleave.rs | 58 +++++++++++++++++------------ arrow/benches/interleave_kernels.rs | 1 + 2 files changed, 36 insertions(+), 23 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index b09de13fee66..081b66df1cf8 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -18,7 +18,7 @@ //! Interleave elements from multiple arrays use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values}; -use arrow_array::builder::{BooleanBufferBuilder, BufferBuilder, PrimitiveBuilder}; +use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder}; use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; @@ -26,7 +26,6 @@ use arrow_buffer::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, Of use arrow_data::transform::MutableArrayData; use arrow_data::ByteView; use arrow_schema::{ArrowError, DataType}; -use std::collections::HashMap; use std::sync::Arc; macro_rules! primitive_helper { @@ -238,32 +237,45 @@ fn interleave_views( indices: &[(usize, usize)], ) -> Result { let interleaved = Interleave::<'_, GenericByteViewArray>::new(values, indices); - let mut views_builder = BufferBuilder::new(indices.len()); let mut buffers = Vec::new(); // (input array_index, input buffer_index) -> output buffer_index - let mut buffer_lookup: HashMap<(usize, u32), u32> = HashMap::new(); - for (array_idx, value_idx) in indices { - let array = interleaved.arrays[*array_idx]; - let raw_view = array.views().get(*value_idx).unwrap(); - let view_len = *raw_view as u32; - if view_len <= 12 { - views_builder.append(*raw_view); - continue; - } - // value is big enough to be in a variadic buffer - let view = ByteView::from(*raw_view); - let new_buffer_idx: &mut u32 = buffer_lookup - .entry((*array_idx, view.buffer_index)) - .or_insert_with(|| { - buffers.push(array.data_buffers()[view.buffer_index as usize].clone()); - (buffers.len() - 1) as u32 - }); - views_builder.append(view.with_buffer_index(*new_buffer_idx).into()); - } + // A mapping from (input array_index, input buffer_index) -> output buffer_index + // The outer vec corresponds to the input array index. + // The inner vec corresponds to the buffer index within that input array. + // The value is the index of the buffer in the output array. + let mut buffer_remap: Vec>> = interleaved + .arrays + .iter() + .map(|a| vec![None; a.data_buffers().len()]) + .collect(); + + let views: Vec = indices + .iter() + .map(|(array_idx, value_idx)| { + let array = interleaved.arrays[*array_idx]; + let raw_view = array.views().get(*value_idx).unwrap(); + let view_len = *raw_view as u32; + if view_len <= 12 { + return *raw_view; + } + // value is big enough to be in a variadic buffer + let view = ByteView::from(*raw_view); + let new_buffer_idx = match &mut buffer_remap[*array_idx][view.buffer_index as usize] { + Some(idx) => *idx, + opt => { + buffers.push(array.data_buffers()[view.buffer_index as usize].clone()); + let new_idx = (buffers.len() - 1) as u32; + *opt = Some(new_idx); + new_idx + } + }; + view.with_buffer_index(new_buffer_idx).as_u128() + }) + .collect(); let array = unsafe { - GenericByteViewArray::::new_unchecked(views_builder.into(), buffers, interleaved.nulls) + GenericByteViewArray::::new_unchecked(views.into(), buffers, interleaved.nulls) }; Ok(Arc::new(array)) } diff --git a/arrow/benches/interleave_kernels.rs b/arrow/benches/interleave_kernels.rs index 77dc9500ea06..fa71747e729e 100644 --- a/arrow/benches/interleave_kernels.rs +++ b/arrow/benches/interleave_kernels.rs @@ -84,6 +84,7 @@ fn add_benchmark(c: &mut Criterion) { ("str(20, 0.5)", &string_opt), ("dict(20, 0.0)", &dict), ("dict_sparse(20, 0.0)", &sparse_dict), + ("string view", &create_string_view_array(1024, 0.0)), ]; for (prefix, base) in cases { From 19bfcbbb1d0c104673eb431ca8c2dc43e2a253b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 17 Jun 2025 22:48:19 +0200 Subject: [PATCH 2/5] interleave_views bench --- arrow/benches/interleave_kernels.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/arrow/benches/interleave_kernels.rs b/arrow/benches/interleave_kernels.rs index fa71747e729e..60125a4ee364 100644 --- a/arrow/benches/interleave_kernels.rs +++ b/arrow/benches/interleave_kernels.rs @@ -77,6 +77,8 @@ fn add_benchmark(c: &mut Criterion) { let values = create_string_array_with_len::(1024, 0.0, 20); let sparse_dict = create_sparse_dict_from_values::(1024, 0.0, &values, 10..20); + let string_view = create_string_view_array(1024, 0.0); + let cases: &[(&str, &dyn Array)] = &[ ("i32(0.0)", &i32), ("i32(0.5)", &i32_opt), @@ -84,7 +86,7 @@ fn add_benchmark(c: &mut Criterion) { ("str(20, 0.5)", &string_opt), ("dict(20, 0.0)", &dict), ("dict_sparse(20, 0.0)", &sparse_dict), - ("string view", &create_string_view_array(1024, 0.0)), + ("str_view(0.0)", &string_view), ]; for (prefix, base) in cases { From 4b7d77152e5a2881c8d39265f344eef6fac86c36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 17 Jun 2025 23:09:45 +0200 Subject: [PATCH 3/5] Use more cache-friendly datastructure --- arrow-select/src/interleave.rs | 44 ++++++++++++++++------------------ 1 file changed, 21 insertions(+), 23 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 081b66df1cf8..46b00783cba0 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -239,37 +239,35 @@ fn interleave_views( let interleaved = Interleave::<'_, GenericByteViewArray>::new(values, indices); let mut buffers = Vec::new(); - // (input array_index, input buffer_index) -> output buffer_index - // A mapping from (input array_index, input buffer_index) -> output buffer_index - // The outer vec corresponds to the input array index. - // The inner vec corresponds to the buffer index within that input array. - // The value is the index of the buffer in the output array. - let mut buffer_remap: Vec>> = interleaved - .arrays - .iter() - .map(|a| vec![None; a.data_buffers().len()]) - .collect(); + // Contains the offsets of start buffer in `buffer_to_new_index` + let mut offsets = Vec::with_capacity(interleaved.arrays.len() + 1); + offsets.push(0); + let mut total_buffers = 0; + for a in interleaved.arrays.iter() { + total_buffers += a.data_buffers().len(); + offsets.push(total_buffers); + } + + // contains the mapping from old buffer index to new buffer index + let mut buffer_to_new_index = vec![None; total_buffers]; let views: Vec = indices .iter() .map(|(array_idx, value_idx)| { let array = interleaved.arrays[*array_idx]; - let raw_view = array.views().get(*value_idx).unwrap(); - let view_len = *raw_view as u32; + let view = array.views().get(*value_idx).unwrap(); + let view_len = *view as u32; if view_len <= 12 { - return *raw_view; + return *view; } // value is big enough to be in a variadic buffer - let view = ByteView::from(*raw_view); - let new_buffer_idx = match &mut buffer_remap[*array_idx][view.buffer_index as usize] { - Some(idx) => *idx, - opt => { - buffers.push(array.data_buffers()[view.buffer_index as usize].clone()); - let new_idx = (buffers.len() - 1) as u32; - *opt = Some(new_idx); - new_idx - } - }; + let view = ByteView::from(*view); + let remap_idx = offsets[*array_idx] + view.buffer_index as usize; + let new_buffer_idx: u32 = *buffer_to_new_index[remap_idx].get_or_insert_with(|| { + buffers.push(array.data_buffers()[view.buffer_index as usize].clone()); + let new_idx = (buffers.len() - 1) as u32; + new_idx + }); view.with_buffer_index(new_buffer_idx).as_u128() }) .collect(); From a4ac7058d0bbc69f451e0159c3b27a47d84e5eec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 17 Jun 2025 23:14:01 +0200 Subject: [PATCH 4/5] Clippy --- arrow-select/src/interleave.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 46b00783cba0..67e69c979731 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -265,8 +265,7 @@ fn interleave_views( let remap_idx = offsets[*array_idx] + view.buffer_index as usize; let new_buffer_idx: u32 = *buffer_to_new_index[remap_idx].get_or_insert_with(|| { buffers.push(array.data_buffers()[view.buffer_index as usize].clone()); - let new_idx = (buffers.len() - 1) as u32; - new_idx + (buffers.len() - 1) as u32 }); view.with_buffer_index(new_buffer_idx).as_u128() }) From 7c113e48decb1bed6c27bc7be435dc0515fb91f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 17 Jun 2025 23:15:13 +0200 Subject: [PATCH 5/5] Style --- arrow-select/src/interleave.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 67e69c979731..0e77d6610fc4 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -262,11 +262,12 @@ fn interleave_views( } // value is big enough to be in a variadic buffer let view = ByteView::from(*view); - let remap_idx = offsets[*array_idx] + view.buffer_index as usize; - let new_buffer_idx: u32 = *buffer_to_new_index[remap_idx].get_or_insert_with(|| { - buffers.push(array.data_buffers()[view.buffer_index as usize].clone()); - (buffers.len() - 1) as u32 - }); + 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() }) .collect();