Transforms: Support dictionary-encoded PyArrow arrays in partition transforms (#3633) - #3841
Open
hedger9487 wants to merge 1 commit into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Updates PyIceberg’s PyArrow partition transform wrapper to transparently handle dictionary-encoded (pa.DictionaryArray) inputs by decoding them before invoking pyiceberg_core.transform, preventing “Unsupported data type … Dictionary(…, …)” errors and enabling transforms to work with read_dictionary-produced columns.
Changes:
- Normalize dictionary-encoded PyArrow arrays via
dictionary_decode()inside_pyiceberg_transform_wrapper(including forChunkedArraychunks). - Add a regression test exercising
BucketTransformandTruncateTransformon dictionary-encoded arrays and chunked dictionary arrays.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
pyiceberg/transforms.py |
Decode dictionary-encoded PyArrow arrays before passing them to pyiceberg_core transform functions. |
tests/test_transforms.py |
Add coverage for dictionary-encoded and chunked dictionary-encoded inputs for key partition transforms. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+1716
to
+1727
| dict_arr = pa.DictionaryArray.from_arrays(pa.array([0, 1, 0, None]), pa.array(["foo", "bar"])) | ||
| raw_arr = pa.array(["foo", "bar", "foo", None]) | ||
| bucket_transform = BucketTransform(num_buckets=10) | ||
| expected_bucket = bucket_transform.pyarrow_transform(StringType())(raw_arr) | ||
| assert bucket_transform.pyarrow_transform(StringType())(dict_arr) == expected_bucket | ||
|
|
||
| chunked_dict = pa.chunked_array([dict_arr, dict_arr]) | ||
| expected_chunked = pa.chunked_array([expected_bucket, expected_bucket]) | ||
| assert bucket_transform.pyarrow_transform(StringType())(chunked_dict) == expected_chunked | ||
|
|
||
| truncate_transform = TruncateTransform(width=3) | ||
| dict_truncate_arr = pa.DictionaryArray.from_arrays(pa.array([0, 1, 0]), pa.array(["developer", "iceberg"])) |
Comment on lines
+1715
to
+1720
| def test_pyarrow_transforms_dictionary_encoded() -> None: | ||
| dict_arr = pa.DictionaryArray.from_arrays(pa.array([0, 1, 0, None]), pa.array(["foo", "bar"])) | ||
| raw_arr = pa.array(["foo", "bar", "foo", None]) | ||
| bucket_transform = BucketTransform(num_buckets=10) | ||
| expected_bucket = bucket_transform.pyarrow_transform(StringType())(raw_arr) | ||
| assert bucket_transform.pyarrow_transform(StringType())(dict_arr) == expected_bucket |
This file contains hidden or 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
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.
Description
Fixes #3633.
When PyArrow tables containing dictionary-encoded columns (
pa.DictionaryArray) are passed to Iceberg partition transforms (such asBucketTransform,TruncateTransform, or time transforms),_pyiceberg_transform_wrapperforwards theDictionaryArraydirectly topyiceberg_core.transform, raisingValueError: Feature Unsupported => Unsupported data type for bucket transform: Dictionary(Int64, Utf8).This PR updates
_pyiceberg_transform_wrapperto normalize dictionary-encoded arrays viaarr.dictionary_decode()before invokingtransform_func, allowing all partition transforms to transparently handle dictionary-encoded PyArrow arrays and chunked arrays.Testing
test_pyarrow_transforms_dictionary_encodedintests/test_transforms.pycoveringDictionaryArrayandChunkedArrayonBucketTransformandTruncateTransform.