PyArrow: Support view types when writing to a partitioned table

pyiceberg/io/pyarrow.py

@@ -1611,6 +1611,51 @@ def primitive(self, primitive: pa.DataType) -> pa.DataType:
         return primitive
 
 
+class _ConvertViewTypes(PyArrowSchemaVisitor[pa.Schema | pa.DataType]):
+    """Replace string_view and binary_view, including nested occurrences, with large_string and large_binary.
+
+    Everything else is kept as is: list flavour, field names, nullability, metadata and dictionary encoding.
+    """
+
+    def schema(self, schema: pa.Schema, struct_result: pa.StructType) -> pa.Schema:
+        return pa.schema(struct_result, metadata=schema.metadata)
+
+    def struct(self, struct: pa.StructType, field_results: builtins.list[pa.Field]) -> pa.StructType:
+        return pa.struct(field_results)
+
+    def field(self, field: pa.Field, field_result: pa.DataType) -> pa.Field:
+        return field.with_type(_keep_dictionary(field.type, field_result))
+
+    def list(self, list_type: pa.ListType | pa.LargeListType | pa.FixedSizeListType, element_result: pa.DataType) -> pa.DataType:
+        value_field = list_type.value_field.with_type(_keep_dictionary(list_type.value_type, element_result))
+        if isinstance(list_type, pa.FixedSizeListType):
+            return pa.list_(value_field, list_type.list_size)
+        elif isinstance(list_type, pa.LargeListType):
+            return pa.large_list(value_field)
+        return pa.list_(value_field)
+
+    def map(self, map_type: pa.MapType, key_result: pa.DataType, value_result: pa.DataType) -> pa.DataType:
+        return pa.map_(
+            map_type.key_field.with_type(_keep_dictionary(map_type.key_type, key_result)),
+            map_type.item_field.with_type(_keep_dictionary(map_type.item_type, value_result)),
+            keys_sorted=map_type.keys_sorted,
+        )
+
+    def primitive(self, primitive: pa.DataType) -> pa.DataType:
+        if pa.types.is_string_view(primitive):
+            return pa.large_string()
+        elif pa.types.is_binary_view(primitive):
+            return pa.large_binary()
+        return primitive
+
+
+def _keep_dictionary(original: pa.DataType, result: pa.DataType) -> pa.DataType:
+    """Re-apply dictionary encoding that visit_pyarrow strips before visiting the value type."""
+    if isinstance(original, pa.DictionaryType):
+        return pa.dictionary(original.index_type, result, ordered=original.ordered)
+    return result
+
+
 class _ConvertToIcebergWithoutIDs(_ConvertToIceberg):
     """
     Converts PyArrowSchema to Iceberg Schema with all -1 ids.
@@ -3090,6 +3135,11 @@ def _determine_partitions(spec: PartitionSpec, schema: Schema, arrow_table: pa.T
     - Then we produce a set of partitions by filtering on each of the combinations
     - We combine the chunks to create a copy to avoid GIL congestion on the original table
     """
+    # PyArrow has no group by, filter or take kernels for view types, so replace them with their large counterparts
+    arrow_schema = visit_pyarrow(arrow_table.schema, _ConvertViewTypes())
+    if arrow_schema != arrow_table.schema:
+        arrow_table = arrow_table.cast(arrow_schema)
+
     # Assign unique names to columns where the partition transform has been applied
     # to avoid conflicts
     partition_fields = [f"_partition_{field.name}" for field in spec.fields]

tests/io/test_pyarrow.py

@@ -93,7 +93,7 @@
 from pyiceberg.table import FileScanTask, TableProperties, WriteTask
 from pyiceberg.table.metadata import TableMetadataV2
 from pyiceberg.table.name_mapping import create_mapping_from_schema
-from pyiceberg.transforms import HourTransform, IdentityTransform
+from pyiceberg.transforms import BucketTransform, HourTransform, IdentityTransform, Transform, TruncateTransform
 from pyiceberg.typedef import UTF8, Properties, Record, TableVersion
 from pyiceberg.types import (
     BinaryType,
@@ -2868,6 +2868,77 @@ def test_determine_partitions_identity_nan() -> None:
     assert rows_by_partition == {"1.0": [1], "nan": [2, 4], "None": [3]}
 
 
+@pytest.mark.parametrize("transform", [IdentityTransform(), BucketTransform(4), TruncateTransform(1)])
+def test_determine_partitions_with_view_types(transform: Transform[Any, Any]) -> None:
+    schema = Schema(
+        NestedField(id=1, name="region", field_type=StringType(), required=False),
+        NestedField(id=2, name="payload", field_type=BinaryType(), required=False),
+        NestedField(
+            id=3,
+            name="info",
+            field_type=StructType(
+                NestedField(id=4, name="tags", field_type=ListType(5, StringType(), element_required=False), required=False),
+                NestedField(id=6, name="attrs", field_type=MapType(7, StringType(), 8, BinaryType()), required=False),
+            ),
+            required=False,
+        ),
+    )
+    spec = PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=transform, name="region_part"))
+    arrow_schema = pa.schema(
+        [
+            pa.field("region", pa.string_view()),
+            pa.field("payload", pa.binary_view()),
+            pa.field(
+                "info",
+                pa.struct(
+                    [
+                        pa.field("tags", pa.list_(pa.string_view())),
+                        pa.field("attrs", pa.map_(pa.string_view(), pa.binary_view())),
+                    ]
+                ),
+            ),
+        ]
+    )
+    rows = [
+        {"region": "us", "payload": b"a", "info": {"tags": ["x"], "attrs": [("k", b"v")]}},
+        {"region": "eu", "payload": None, "info": None},
+        {"region": "us", "payload": b"b", "info": {"tags": None, "attrs": None}},
+        {"region": None, "payload": b"c", "info": {"tags": ["y", None], "attrs": []}},
+    ]
+    arrow_table = pa.Table.from_pylist(rows, schema=arrow_schema)
+
+    partitions = list(_determine_partitions(spec, schema, arrow_table))
+
+    assert {p.partition_key.partition[0] for p in partitions} == {
+        transform.transform(StringType())(region) for region in ("us", "eu", None)
+    }
+    written = pa.concat_tables([p.arrow_table_partition for p in partitions])
+    assert sorted(written.to_pylist(), key=str) == sorted(rows, key=str)
+
+
+def test_append_view_types_to_partitioned_table(catalog: InMemoryCatalog) -> None:
+    schema = Schema(
+        NestedField(id=1, name="id", field_type=LongType(), required=False),
+        NestedField(id=2, name="region", field_type=StringType(), required=False),
+        NestedField(id=3, name="payload", field_type=BinaryType(), required=False),
+    )
+    spec = PartitionSpec(PartitionField(source_id=2, field_id=1000, transform=IdentityTransform(), name="region"))
+    catalog.create_namespace("default")
+    table = catalog.create_table("default.test_append_view_types", schema=schema, partition_spec=spec)
+
+    arrow_table = pa.table(
+        {
+            "id": pa.array([1, 2, 3], type=pa.int64()),
+            "region": pa.array(["us", "eu", "us"], type=pa.string_view()),
+            "payload": pa.array([b"a", None, b"c"], type=pa.binary_view()),
+        }
+    )
+    table.append(arrow_table)
+
+    assert len(table.inspect.partitions()) == 2
+    assert sorted(table.scan().to_arrow().to_pylist(), key=lambda row: row["id"]) == arrow_table.to_pylist()
+
+
 def test_inspect_partition_for_nested_field(catalog: InMemoryCatalog) -> None:
     schema = Schema(
         NestedField(id=1, name="foo", field_type=StringType(), required=True),