diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index e859734089..2c5a29fae4 100644 --- a/pyiceberg/io/pyarrow.py +++ b/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] 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),