Skip to content
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.

Sign up for GitHub

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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Add support StringView / BinaryView in interleave kernel #6779

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
27ac654
add byteview specific interleave
onursatici Nov 23, 2024
2052cb8
clippy
onursatici Nov 24, 2024
8b11191
test
onursatici Nov 24, 2024
d5daedc
more clippy
onursatici Nov 24, 2024
6ff5c0a
more test coverage
onursatici Nov 25, 2024
763a792
enable assertion, remove explicit vector capacity
onursatici Nov 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
158 changes: 158 additions & 0 deletions arrow-select/src/interleave.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,9 @@ use arrow_array::types::*;
use arrow_array::*;
use arrow_buffer::{ArrowNativeType, MutableBuffer, NullBuffer, NullBufferBuilder, OffsetBuffer};
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 {
Expand Down Expand Up @@ -97,6 +99,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}")
Expand Down Expand Up @@ -231,6 +235,40 @@ 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::new();

let mut buffer_lookup = HashMap::new();
Copy link
Contributor

@tustvold tustvold Nov 26, 2024 •
edited
Loading

Choose a reason for hiding this comment

The 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.

Copy link
Contributor Author

Choose a reason for hiding this comment

The 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

Copy link
Contributor

Choose a reason for hiding this comment

The 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);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

other code that manipulates vies checks the length before casting


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 = unsafe {
GenericByteViewArray::<T>::new_unchecked(views_builder.into(), buffers, interleaved.nulls)
};
Ok(Arc::new(array))
}

/// Fallback implementation of interleave using [`MutableArrayData`]
fn interleave_fallback(
values: &[&dyn Array],
Expand Down Expand Up @@ -461,4 +499,124 @@ mod tests {
DictionaryArray::<Int32Type>::from_iter(vec![Some("0"), Some("1"), Some("2"), None]);
assert_eq!(array.as_ref(), &expected)
}

#[test]
fn test_interleave_views() {
let values = StringArray::from_iter_values([
"hello",
"world_long_string_not_inlined",
"foo",
"bar",
"baz",
]);
let view_a = StringViewArray::from(&values);

let values = StringArray::from_iter_values([
"test",
"data",
"more_long_string_not_inlined",
"views",
"here",
]);
let view_b = StringViewArray::from(&values);

let indices = &[
(0, 2), // "foo"
(1, 0), // "test"
(0, 4), // "baz"
(1, 3), // "views"
(0, 1), // "world_long_string_not_inlined"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since the test only picks one long string, only one buffer is going to be copied (the other strings are inlined)

To provoke the issue described in the ticket, I think you need to interleave that same long string from the two different buffers. Something like

        let indices = &[
            (0, 1), // "world_long_string_not_inlined" (view_a)
            (1, 2), // "world_long_string_not_inlined" (view_b)
            (0, 1), // "world_long_string_not_inlined" (view_a)
            (1, 2), // "world_long_string_not_inlined" (view_b)
            (0, 1), // "world_long_string_not_inlined" (view_a)
            (1, 2), // "world_long_string_not_inlined" (view_b)
            (0, 1), // "world_long_string_not_inlined" (view_a)
            (1, 2), // "world_long_string_not_inlined" (view_b)
            (0, 1), // "world_long_string_not_inlined" (view_a)
            (1, 2), // "world_long_string_not_inlined" (view_b)

And make sure the output view only has 2 buffers (the one from view_a and the one from view_b)

];

// Test specialized implementation
let values = interleave(&[&view_a, &view_b], indices).unwrap();
let result = values.as_string_view();
assert_eq!(result.data_buffers().len(), 1);

// Test fallback implementation
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since interleave_fallback isn't pub I don't think we should be testing it directly. Intead you should arrange the paramters to call it directly

let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
let fallback_result = fallback.as_string_view();
// note that fallback_result has 2 buffers, but only one long enough string to warrant a buffer
assert_eq!(fallback_result.data_buffers().len(), 2);

// Convert to strings for easier assertion
let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();

let fallback_collected: Vec<_> = fallback_result
.iter()
.map(|x| x.map(|s| s.to_string()))
.collect();

assert_eq!(&collected, &fallback_collected);

assert_eq!(
&collected,
&[
Some("foo".to_string()),
Some("test".to_string()),
Some("baz".to_string()),
Some("views".to_string()),
Some("world_long_string_not_inlined".to_string()),
]
);
}

#[test]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think another important test case is testing with a StringViewArray that has multiple buffers (otherwise the logic to handle mapping different buffers is not exercised)

So the input array should have multiple buffers as well

fn test_interleave_views_with_nulls() {
let values = StringArray::from_iter([
Some("hello"),
None,
Some("foo_long_string_not_inlined"),
Some("bar"),
None,
]);
let view_a = StringViewArray::from(&values);

let values = StringArray::from_iter([
Some("test"),
Some("data_long_string_not_inlined"),
None,
None,
Some("here"),
]);
let view_b = StringViewArray::from(&values);

let indices = &[
(0, 1), // null
(1, 2), // null
(0, 2), // "foo_long_string_not_inlined"
(1, 3), // null
(0, 4), // null
];

// Test specialized implementation
let values = interleave(&[&view_a, &view_b], indices).unwrap();
let result = values.as_string_view();
assert_eq!(result.data_buffers().len(), 1);

// Test fallback implementation
let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
let fallback_result = fallback.as_string_view();

// Convert to strings for easier assertion
let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();

let fallback_collected: Vec<_> = fallback_result
.iter()
.map(|x| x.map(|s| s.to_string()))
.collect();

assert_eq!(&collected, &fallback_collected);

assert_eq!(
&collected,
&[
None,
None,
Some("foo_long_string_not_inlined".to_string()),
None,
None,
]
);
}
}
Loading