diff --git a/pyiceberg/transforms.py b/pyiceberg/transforms.py index cd0d7cebcb..7fb0a66930 100644 --- a/pyiceberg/transforms.py +++ b/pyiceberg/transforms.py @@ -145,12 +145,17 @@ def _cast_if_needed(arr: "ArrayLike") -> "ArrayLike": else: return arr + def _decode_if_dict(arr: "pa.Array") -> "pa.Array": + if isinstance(arr.type, pa.DictionaryType): + return arr.dictionary_decode() + return arr + if isinstance(array, pa.Array): - return _cast_if_needed(transform_func(array, *args)) + return _cast_if_needed(transform_func(_decode_if_dict(array), *args)) elif isinstance(array, pa.ChunkedArray): result_chunks = [] for arr in array.iterchunks(): - result_chunks.append(_cast_if_needed(transform_func(arr, *args))) + result_chunks.append(_cast_if_needed(transform_func(_decode_if_dict(arr), *args))) return pa.chunked_array(result_chunks) else: raise ValueError(f"PyArrow array can only be of type pa.Array or pa.ChunkedArray, but found {type(array)}") diff --git a/tests/test_transforms.py b/tests/test_transforms.py index d296fcdb21..5d50e3ca63 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1682,6 +1682,23 @@ def test_bucket_pyarrow_void_transform() -> None: assert output_arr == VoidTransform().pyarrow_transform(IntegerType())(input_arr) +def test_bucket_pyarrow_transform_dictionary_encoded() -> None: + input_arr = pa.array(["foo", "bar", "foo"], type=pa.dictionary(pa.int32(), pa.utf8())) + assert BucketTransform(num_buckets=10).pyarrow_transform(StringType())(input_arr) == pa.array([6, 7, 6], type=pa.int32()) + + +def test_bucket_pyarrow_transform_dictionary_encoded_chunked() -> None: + input_chunked = pa.chunked_array( + [ + pa.array(["foo", "bar"], type=pa.dictionary(pa.int32(), pa.utf8())), + pa.array(["baz"], type=pa.dictionary(pa.int32(), pa.utf8())), + ] + ) + assert BucketTransform(num_buckets=10).pyarrow_transform(StringType())(input_chunked) == pa.chunked_array( + [pa.array([6, 7], type=pa.int32()), pa.array([4], type=pa.int32())] + ) + + @pytest.mark.parametrize( "source_type, input_arr, expected, width", [