diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index e859734089..1b933c6e88 100644 --- a/pyiceberg/io/pyarrow.py +++ b/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. diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index 6e66ff43e9..d43f55f3ba 100644 --- a/tests/io/test_pyarrow.py +++ b/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),