Decode dictionary-encoded arrays before bucket transform

Author: ebyhr

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)}")

tests/test_transforms.py

@@ -1682,6 +1682,13 @@ 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()))
+    expected = pa.array([6, 7, 6], type=pa.int32())
+    transform: Transform[Any, Any] = BucketTransform(num_buckets=10)
+    assert transform.pyarrow_transform(StringType())(input_arr) == expected
+
+
 @pytest.mark.parametrize(
     "source_type, input_arr, expected, width",
     [