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

Implement dictionary support for reading ByteView from parquet #5973

Merged
merged 2 commits into from
Jul 3, 2024

Conversation

XiangpengHao
Copy link
Contributor

@XiangpengHao XiangpengHao commented Jun 28, 2024

Which issue does this PR close?

Part of #5904 , sequel to #5968 and #5970 and #5972

This PR is not ready to review until the above PRs are merged.

Rationale for this change

Implement dictionary encoding support for reading ByteView from parquet.

To show that reading ByteViewArray is faster than reading ByteArray:

cargo bench --bench arrow_reader --features="arrow test_common experimental" "arrow_array_reader/Binary.*Array/dictionary encoded"

We should get something like (~3x faster than ByteArray, many times faster than current implementation):

arrow_array_reader/BinaryArray/dictionary encoded, mandatory, no NULLs
                        time:   [270.32 µs 270.34 µs 270.36 µs]

arrow_array_reader/BinaryArray/dictionary encoded, optional, no NULLs
                        time:   [268.11 µs 268.14 µs 268.18 µs]

arrow_array_reader/BinaryArray/dictionary encoded, optional, half NULLs
                        time:   [372.87 µs 372.91 µs 372.95 µs]

arrow_array_reader/BinaryViewArray/dictionary encoded, mandatory, no NULLs
                        time:   [110.81 µs 110.82 µs 110.83 µs]

arrow_array_reader/BinaryViewArray/dictionary encoded, optional, no NULLs
                        time:   [111.89 µs 111.91 µs 111.92 µs]

arrow_array_reader/BinaryViewArray/dictionary encoded, optional, half NULLs
                        time:   [115.55 µs 115.56 µs 115.57 µs]

This is expected as the new implementation never copies data; it tries its best to reuse the buffer from the parquet decoder.

What changes are included in this PR?

This PR is (even) less straightforward than the previous ones, mostly because we need to map the dictionary view buffer indexes to the output view buffer indexes. Please let me know if any point of the code is too difficult to maintain!

Are there any user-facing changes?

@github-actions github-actions bot added parquet Changes to the parquet crate arrow Changes to the arrow crate and removed arrow Changes to the arrow crate labels Jun 28, 2024
@XiangpengHao XiangpengHao force-pushed the parquet-string-view-3 branch 2 times, most recently from cef542d to e1f53df Compare July 2, 2024 03:21
@XiangpengHao XiangpengHao marked this pull request as ready for review July 2, 2024 03:22
@alamb
Copy link
Contributor

alamb commented Jul 2, 2024

I rebased this branch so it only had the changes to support dictionary reading

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.

Thanks @XiangpengHao -- this looks very clever 👌

I wonder if there is some way this code should be tested? Or will be be tested as part of a larger end to end parquet reading exercise later?

}

// Check if the last few buffer of `output`` are the same as the `dict` buffer
// This is to avoid creating a new buffers if the same dictionary is used for multiple `read`
Copy link
Contributor

Choose a reason for hiding this comment

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

That is a clever optimization

parquet/src/arrow/array_reader/byte_view_array.rs Outdated Show resolved Hide resolved
};

self.decoder.read(len, |keys| {
for k in keys {
Copy link
Contributor

Choose a reason for hiding this comment

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

Do we know for sure k are valid dictionary indices (aka since this is from untrusted input, do we have to verify that k is within the bound of the dictionary)?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

The input is untrusted, but I think we can do very little to make it safer, e.g., if the view is maliciously crafted we can easily run into memory issues due to too large offset or invalid buffer idx.
But I agree that having this early check is useful, if it doesn't make the code significantly slower. I added the check here, will come back to this when I file the PR to optimize the performance.

@XiangpengHao
Copy link
Contributor Author

Or will be be tested as part of a larger end to end parquet reading exercise later?

The code to enable and test dictionary encoding is here:
https://github.com/apache/arrow-rs/pull/5973/files#diff-50e48ffbe585053165f5fdc7421dc30b1d1c9abf3ada6aff43dea9f7907489acR482-R486

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.

Looks good to me -- thanks @XiangpengHao

@alamb alamb merged commit e7a0008 into apache:master Jul 3, 2024
16 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
parquet Changes to the parquet crate
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants