-
Notifications
You must be signed in to change notification settings - Fork 824
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add support StringView
/ BinaryView
in interleave
kernel
#6779
base: main
Are you sure you want to change the base?
Changes from 1 commit
27ac654
2052cb8
8b11191
d5daedc
6ff5c0a
763a792
7aa95d2
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -22,9 +22,15 @@ use arrow_array::builder::{BooleanBufferBuilder, BufferBuilder, PrimitiveBuilder | |
use arrow_array::cast::AsArray; | ||
use arrow_array::types::*; | ||
use arrow_array::*; | ||
use arrow_buffer::{ArrowNativeType, MutableBuffer, NullBuffer, NullBufferBuilder, OffsetBuffer}; | ||
use arrow_buffer::{ | ||
ArrowNativeType, Buffer, MutableBuffer, NullBuffer, NullBufferBuilder, OffsetBuffer, | ||
ScalarBuffer, | ||
}; | ||
use arrow_data::transform::MutableArrayData; | ||
use arrow_data::ByteView; | ||
use arrow_schema::{ArrowError, DataType}; | ||
use builder::{ArrayBuilder, GenericByteViewBuilder}; | ||
use std::collections::HashMap; | ||
use std::sync::Arc; | ||
|
||
macro_rules! primitive_helper { | ||
|
@@ -97,6 +103,8 @@ pub fn interleave( | |
DataType::LargeUtf8 => interleave_bytes::<LargeUtf8Type>(values, indices), | ||
DataType::Binary => interleave_bytes::<BinaryType>(values, indices), | ||
DataType::LargeBinary => interleave_bytes::<LargeBinaryType>(values, indices), | ||
DataType::BinaryView => interleave_views::<BinaryViewType>(values, indices), | ||
DataType::Utf8View => interleave_views::<StringViewType>(values, indices), | ||
DataType::Dictionary(k, _) => downcast_integer! { | ||
k.as_ref() => (dict_helper, values, indices), | ||
_ => unreachable!("illegal dictionary key type {k}") | ||
|
@@ -231,6 +239,39 @@ fn interleave_dictionaries<K: ArrowDictionaryKeyType>( | |
Ok(Arc::new(array)) | ||
} | ||
|
||
fn interleave_views<T: ByteViewType>( | ||
values: &[&dyn Array], | ||
indices: &[(usize, usize)], | ||
) -> Result<ArrayRef, ArrowError> { | ||
let interleaved = Interleave::<'_, GenericByteViewArray<T>>::new(values, indices); | ||
let mut views_builder = BufferBuilder::new(indices.len()); | ||
let mut buffers = Vec::with_capacity(values[0].len()); | ||
|
||
let mut buffer_lookup = HashMap::new(); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This misunderstands #6780, the issue isn't not skipping buffers that aren't referenced, it is that the arrays being interleaved may contain the same actual buffers, e.g. sourced from the same parquet dictionary page. This needs to deduplicate based on the underlying Buffer pointers. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I see, I have filed #6808 to deduplicate the buffers while building the interleaved array in the fallback implementation. I am not sure if that would completely remove the need for this PR though, it should guarantee that no duplicate buffers should exist on the interleaved array, but the fallback implementation would still have 2 buffers in the test below, because those two buffers are unique. It feels like this PR and #6808 are complementary There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you also please add some comments about what the buffer lookup represents? I think it is // (input array_index, input buffer_index) --> output array buffer_index |
||
for (array_idx, value_idx) in indices { | ||
let array = interleaved.arrays[*array_idx]; | ||
let raw_view = array.views().get(*value_idx).unwrap(); | ||
let view = ByteView::from(*raw_view); | ||
onursatici marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
if view.length <= 12 { | ||
views_builder.append(*raw_view); | ||
continue; | ||
} | ||
// value is big enough to be in a variadic buffer | ||
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()); | ||
} | ||
|
||
let array = | ||
GenericByteViewArray::<T>::try_new(views_builder.into(), buffers, interleaved.nulls)?; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. should we do There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes probably |
||
Ok(Arc::new(array)) | ||
} | ||
|
||
/// Fallback implementation of interleave using [`MutableArrayData`] | ||
fn interleave_fallback( | ||
values: &[&dyn Array], | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why is this the capacity?