Transforms: Support dictionary-encoded PyArrow arrays in partition transforms (#3633) - #3841
hedger9487 wants to merge 1 commit into
Conversation
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.
| 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"])) |
| 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 pull request has been marked as stale due to 30 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that's incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@iceberg.apache.org list. Thank you for your contributions. |
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.