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,39 @@ def test_bucket_pyarrow_void_transform() -> None:
     assert output_arr == VoidTransform().pyarrow_transform(IntegerType())(input_arr)
 
 
+# Regression test for https://github.com/apache/iceberg-python/issues/3633
+@pytest.mark.parametrize(
+    "source_type, input_arr, expected, num_buckets",
+    [
+        (
+            StringType(),
+            pa.array(["foo", "bar", "foo"], type=pa.dictionary(pa.int32(), pa.utf8())),
+            pa.array([6, 7, 6], type=pa.int32()),
+            10,
+        ),
+        (
+            StringType(),
+            pa.chunked_array(
+                [
+                    pa.array(["foo", "bar"], type=pa.dictionary(pa.int32(), pa.utf8())),
+                    pa.array(["baz"], type=pa.dictionary(pa.int32(), pa.utf8())),
+                ]
+            ),
+            pa.chunked_array([pa.array([6, 7], type=pa.int32()), pa.array([4], type=pa.int32())]),
+            10,
+        ),
+    ],
+)
+def test_bucket_pyarrow_transform_dictionary_encoded(
+    source_type: PrimitiveType,
+    input_arr: pa.Array | pa.ChunkedArray,
+    expected: pa.Array | pa.ChunkedArray,
+    num_buckets: int,
+) -> None:
+    transform: Transform[Any, Any] = BucketTransform(num_buckets=num_buckets)
+    assert expected == transform.pyarrow_transform(source_type)(input_arr)
+
+
 @pytest.mark.parametrize(
     "source_type, input_arr, expected, width",
     [