Arrow: Avoid buffer-overflow by avoid doing a sort #1539
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Fokko
force-pushed
the
fd-fix-overflowing-buffer
branch
2 times, most recently
from
e548117 to
4658c3c
Compare
3 tasks
kevinjqliu
approved these changes
Jan 20, 2025
| y = ["fixed_string"] * 30_000 | ||
| tb = pa.chunked_array([y] * 10_000) | ||
| # Create pa.table | ||
| arrow_table = pa.table({"a": ta, "b": tb}) |
There was a problem hiding this comment.
it wasnt obv to me that this test offset is beyond 32 bits, but i ran it and 4800280000 is >2^32/4294967296
>>> len(arrow_table)
300000000
>>> arrow_table.get_total_buffer_size()
4800280000
04a8218 to
3841fe7
Compare
kevinjqliu
reviewed
Jan 20, 2025
pyiceberg/partitioning.py
Outdated
Comment on lines
418
to
419
| # When adding files, it can be that we still need to convert from logical types to physical types | ||
| iceberg_typed_value = _to_partition_representation(iceberg_type, value) |
There was a problem hiding this comment.
is this due to the fact that we already transform the partition key value
partition.transform.pyarrow_transform(source_field.field_type)(arrow_table[source_field.name])
and this expects the untransformed value?
if thats the case, can we just omit the transformation before the group_by?
There was a problem hiding this comment.
Ah, of course. We want to know the output tuples after the transform, so omitting the transformation is not possible. I think we could do a follow-up PR where we split out the logic for the write path, and the add-files path. Since after this PR, this is not needed when doing partitioned writes, we just need it to preprocess when importing partitions.
Fokko
force-pushed
the
fd-fix-overflowing-buffer
branch
from
3841fe7 to
c84dd8d
Compare
|
Ugh, accidentally pushed |
|
:'( |
Fokko
added a commit
that referenced
this pull request
Jan 23, 2025
Second attempt of #1539 This was already being discussed back here: #208 (comment) This PR changes from doing a sort, and then a single pass over the table to the approach where we determine the unique partition tuples filter on them individually. Fixes #1491 Because the sort caused buffers to be joined where it would overflow in Arrow. I think this is an issue on the Arrow side, and it should automatically break up into smaller buffers. The `combine_chunks` method does this correctly. Now: ``` 0.42877754200890195 Run 1 took: 0.2507691659993725 Run 2 took: 0.24833179199777078 Run 3 took: 0.24401691700040828 Run 4 took: 0.2419595829996979 Average runtime of 0.28 seconds ``` Before: ``` Run 0 took: 1.0768639159941813 Run 1 took: 0.8784021250030492 Run 2 took: 0.8486490420036716 Run 3 took: 0.8614017910003895 Run 4 took: 0.8497851670108503 Average runtime of 0.9 seconds ``` So it comes with a nice speedup as well :) --------- Co-authored-by: Kevin Liu <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This was already being discussed back here: #208 (comment)
This PR changes from doing a sort, and then a single pass over the table to the approach where we determine the unique partition tuples filter on them individually.
Fixes #1491
Because the sort caused buffers to be joined where it would overflow in Arrow. I think this is an issue on the Arrow side, and it should automatically break up into smaller buffers. The
combine_chunksmethod does this correctly.Now:
Before:
So it comes with a nice speedup as well :)