diff --git a/pyiceberg/transforms.py b/pyiceberg/transforms.py index 5e0027a829..10825f8f90 100644 --- a/pyiceberg/transforms.py +++ b/pyiceberg/transforms.py @@ -145,12 +145,17 @@ def _cast_if_needed(arr: "ArrayLike") -> "ArrayLike": else: return arr + def _normalize_array(arr: "pa.Array") -> "pa.Array": + if pa.types.is_dictionary(arr.type): + 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(_normalize_array(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(_normalize_array(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 c977fcea14..ef43aeb967 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1710,3 +1710,21 @@ def test_calling_pyarrow_transform_without_pyiceberg_core_installed_correctly_ra with pytest.raises(NotInstalledError): transform.pyarrow_transform(StringType()) + + +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 + + 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"])) + raw_truncate_arr = pa.array(["developer", "iceberg", "developer"]) + expected_truncate = truncate_transform.pyarrow_transform(StringType())(raw_truncate_arr) + assert truncate_transform.pyarrow_transform(StringType())(dict_truncate_arr) == expected_truncate