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.

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

Conversation

onursatici
Copy link
Contributor

@onursatici onursatici commented Nov 23, 2024

Which issue does this PR close?

Rationale for this change

Currently interleaving ByteViewArrays are done with the fallback implementation, which uses a MutableArrayBuilder. The extend 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 array

Are there any user-facing changes?

}

let array =
GenericByteViewArray::<T>::try_new(views_builder.into(), buffers, interleaved.nulls)?;
Copy link
Contributor Author

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?

Copy link
Contributor

Choose a reason for hiding this comment

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

Yes probably

@onursatici onursatici marked this pull request as ready for review November 24, 2024 10:22
@alamb
Copy link
Contributor

alamb commented Nov 24, 2024

Thanks @onursatici -- this looks very cool. I hope to review it in the next few days

Copy link
Contributor

@tustvold tustvold left a 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

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

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?

}

let array =
GenericByteViewArray::<T>::try_new(views_builder.into(), buffers, interleaved.nulls)?;
Copy link
Contributor

Choose a reason for hiding this comment

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

Yes probably

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

Choose a reason for hiding this comment

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

Suggested change
// 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();
Copy link
Contributor

@tustvold tustvold Nov 26, 2024

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

Copy link
Contributor

@alamb alamb left a 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
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

(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)

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

let mut views_builder = BufferBuilder::new(indices.len());
let mut buffers = Vec::with_capacity(values[0].len());

let mut buffer_lookup = HashMap::new();
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

);
}

#[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

@alamb alamb changed the title add byteview specific interleave Add support StringView / BinaryView in interleave kernel Dec 18, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
arrow Changes to the arrow crate
Projects
None yet
3 participants