PyArrow: Support view types when writing to a partitioned table

pyiceberg/io/pyarrow.py

@@ -3090,6 +3090,14 @@ 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 = pa.schema(
+        [field.with_type(_replace_view_types(field.type)) for field in arrow_table.schema],
+        metadata=arrow_table.schema.metadata,
+    )
+    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]
@@ -3141,6 +3149,29 @@ def _partition_value_filter(partition_field_name: str, value: Any) -> pc.Express
     return pc.field(partition_field_name) == value
 
 
+def _replace_view_types(arrow_type: pa.DataType) -> pa.DataType:
+    """Replace string_view and binary_view, including nested occurrences, with large_string and large_binary."""
+    if pa.types.is_string_view(arrow_type):
+        return pa.large_string()
+    elif pa.types.is_binary_view(arrow_type):
+        return pa.large_binary()
+    elif isinstance(arrow_type, pa.StructType):
+        return pa.struct([field.with_type(_replace_view_types(field.type)) for field in arrow_type])
+    elif isinstance(arrow_type, pa.MapType):
+        return pa.map_(
+            arrow_type.key_field.with_type(_replace_view_types(arrow_type.key_type)),
+            arrow_type.item_field.with_type(_replace_view_types(arrow_type.item_type)),
+            keys_sorted=arrow_type.keys_sorted,
+        )
+    elif isinstance(arrow_type, pa.LargeListType):
+        return pa.large_list(arrow_type.value_field.with_type(_replace_view_types(arrow_type.value_type)))
+    elif isinstance(arrow_type, pa.FixedSizeListType):
+        return pa.list_(arrow_type.value_field.with_type(_replace_view_types(arrow_type.value_type)), arrow_type.list_size)
+    elif isinstance(arrow_type, pa.ListType):
+        return pa.list_(arrow_type.value_field.with_type(_replace_view_types(arrow_type.value_type)))
+    return arrow_type
+
+
 def _get_field_from_arrow_table(arrow_table: pa.Table, field_path: str) -> pa.Array:
     """Get a field from an Arrow table, supporting both literal field names and nested field paths.
 

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),