-
Notifications
You must be signed in to change notification settings - Fork 821
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?
Conversation
arrow-select/src/interleave.rs
Outdated
} | ||
|
||
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 comment
The reason will be displayed to describe this comment to others. Learn more.
should we do new_unchecked
here instead similar to other specific interleave methods?
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.
Yes probably
Thanks @onursatici -- this looks very cool. I hope to review it in the next few days |
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.
I've not been following the StringView work very closely but this appears to be a similar approach to #6154 and have the same broad issue.
I think we probably need a more holistic think about how we should be handling deduplicating array views, that can then be applied across the selection kernels, much like we have for dictionaries.
There are lots of ways this could be done, and it is certainly possible that no general purpose solution is possible and we need a more sophisticated API for this #6692
arrow-select/src/interleave.rs
Outdated
) -> 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()); |
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?
arrow-select/src/interleave.rs
Outdated
} | ||
|
||
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 comment
The reason will be displayed to describe this comment to others. Learn more.
Yes probably
arrow-select/src/interleave.rs
Outdated
let fallback_result = fallback.as_string_view(); | ||
// as of commit 97055631, assertion below, commented out to not block future improvements, passes: | ||
// 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); |
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.
// assert_eq!(fallback_result.data_buffers().len(), 2); | |
assert_eq!(fallback_result.data_buffers().len(), 2); |
Ultimately if this starts failing it indicates improvements have been made that might make the specialized impl redundant
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 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.
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.
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 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
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.
Thank you @onursatici -- I apologize for the delay in reviewing.
I have reviewed the code and I think it looks quite good. I left some comments on this PR about some additional testing I think is needed.
If you are willing to implement that additional testing, I think this PR will be good to go (I agree it is potentially complimentary with #6808)
let result = values.as_string_view(); | ||
assert_eq!(result.data_buffers().len(), 1); | ||
|
||
// Test fallback implementation |
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.
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
(1, 0), // "test" | ||
(0, 4), // "baz" | ||
(1, 3), // "views" | ||
(0, 1), // "world_long_string_not_inlined" |
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.
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)
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); |
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.
other code that manipulates vies checks the length before casting
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 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
); | ||
} | ||
|
||
#[test] |
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.
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
StringView
/ BinaryView
in interleave
kernel
Which issue does this PR close?
Rationale for this change
Currently interleaving
ByteViewArray
s are done with the fallback implementation, which uses aMutableArrayBuilder
. Theextend
method on this builder is copying over all variadic buffers because it doesn't know if there are buffers not referenced by any views in the array. Especially on datafusion's TopK implementation, which uses an heap that interleaves arrow arrays to produce the top k rows, current interleave implementation results in an explosion of variadic buffer count for byte view arrays, adding the same set of buffers over and over again. Where this becomes really problematic is when sending such arrays over flight, current encoder materialises all variadic buffers.What changes are included in this PR?
Add a
ByteViewArray
specific interleave implementation that does not add a previously referenced variadic buffer when building the interleaved arrayAre there any user-facing changes?