diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..247729efa9 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -2899,11 +2899,15 @@ def parquet_file_to_data_file(io: FileIO, table_metadata: TableMetadata, file_pa stats_columns=compute_statistics_plan(schema, table_metadata.properties), parquet_column_mapping=parquet_path_to_id_mapping(schema), ) + partition_spec = table_metadata.spec() + partition = partition_spec.partition_from_path(file_path, schema) + if partition is None: + partition = statistics.partition(partition_spec, table_metadata.schema()) data_file = DataFile.from_args( content=DataFileContent.DATA, file_path=file_path, file_format=FileFormat.PARQUET, - partition=statistics.partition(table_metadata.spec(), table_metadata.schema()), + partition=partition, file_size_in_bytes=len(input_file), sort_order_id=None, spec_id=table_metadata.default_spec_id, diff --git a/pyiceberg/partitioning.py b/pyiceberg/partitioning.py index 3d06287b6c..49f6b2a2f8 100644 --- a/pyiceberg/partitioning.py +++ b/pyiceberg/partitioning.py @@ -22,7 +22,7 @@ from datetime import date, datetime, time from functools import cached_property, singledispatch from typing import Annotated, Any, Generic, TypeVar -from urllib.parse import quote_plus +from urllib.parse import quote_plus, unquote_plus from pydantic import ( BeforeValidator, @@ -40,6 +40,7 @@ HourTransform, IdentityTransform, MonthTransform, + TimeTransform, Transform, TruncateTransform, UnknownTransform, @@ -49,7 +50,9 @@ ) from pyiceberg.typedef import IcebergBaseModel, Record from pyiceberg.types import ( + BinaryType, DateType, + FixedType, IcebergType, NestedField, PrimitiveType, @@ -255,6 +258,56 @@ def partition_to_path(self, data: Record, schema: Schema) -> str: path = "/".join([field_str + "=" + value_str for field_str, value_str in zip(field_strs, value_strs, strict=True)]) return path + def partition_from_path(self, location: str, schema: Schema) -> Record | None: + """Infer a partition Record from a Hive-style path (trailing key=value dirs). + + Supports identity, bucket, and truncate (non-binary) transforms, since their + to_human_string output is parseable back into the transform's result value. + Returns None for unsupported transforms or a non-Hive-style path, so the + caller can fall back to another inference strategy. + """ + from pyiceberg.conversions import partition_to_py + + if self.is_unpartitioned(): + return None + + for field in self.fields: + if isinstance(field.transform, (TimeTransform, VoidTransform, UnknownTransform)): + return None + if isinstance(field.transform, TruncateTransform): + source_field = schema.find_field(field.source_id) + if isinstance(source_field.field_type, (FixedType, BinaryType)): + # base64-encoded in the path, not decodable back to bytes + return None + + segments = [segment for segment in location.split("/") if segment] + + partition_size = len(self.fields) + partition_end_index_exclusive = len(segments) - 1 # exclude the file name + partition_start_index = partition_end_index_exclusive - partition_size + if partition_start_index < 0: + return None + + partition_segments = segments[partition_start_index:partition_end_index_exclusive] + if not all("=" in segment for segment in partition_segments): + return None + + partition_type = self.partition_type(schema) + field_types = partition_type.fields + + values = [] + for partition_field, segment, field_type in zip(self.fields, partition_segments, field_types, strict=True): + key, _, raw_value = segment.partition("=") + key = unquote_plus(key) + if key != partition_field.name: + return None + + value_str = unquote_plus(raw_value) + value = partition_to_py(field_type.field_type, value_str) if value_str else None + values.append(value) + + return Record(*values) + def check_compatible(self, schema: Schema, allow_missing_fields: bool = False) -> None: # if the underlying field is dropped, we cannot check they are compatible -- continue schema_fields = schema._lazy_id_to_field diff --git a/tests/integration/test_add_files.py b/tests/integration/test_add_files.py index a1d45451d8..bc9640f92f 100644 --- a/tests/integration/test_add_files.py +++ b/tests/integration/test_add_files.py @@ -1009,6 +1009,119 @@ def test_add_files_hour_transform(session_catalog: Catalog) -> None: tbl.add_files(file_paths=[file_path]) +@pytest.mark.integration +def test_add_files_infers_partition_from_hive_style_path( + spark: SparkSession, session_catalog: Catalog, format_version: int +) -> None: + identifier = f"default.partitioned_table_hive_style_path_v{format_version}" + + partition_spec = PartitionSpec( + PartitionField(source_id=4, field_id=1000, transform=IdentityTransform(), name="baz"), + spec_id=0, + ) + + tbl = _create_table(session_catalog, identifier, format_version, partition_spec) + + # File's own "baz" values (900, 901) disagree with the path (baz=123): path must win. + file_paths = [ + f"s3://warehouse/default/partitioned_table_hive_style_path/v{format_version}/baz=123/test-{i}.parquet" for i in range(2) + ] + for i, file_path in enumerate(file_paths): + fo = tbl.io.new_output(file_path) + with fo.create(overwrite=True) as fos: + with pq.ParquetWriter(fos, schema=ARROW_SCHEMA) as writer: + writer.write_table( + pa.Table.from_pylist( + [ + { + "foo": True, + "bar": "bar_string", + "baz": 900 + i, + "qux": date(2024, 3, 7), + } + ], + schema=ARROW_SCHEMA, + ) + ) + + tbl.add_files(file_paths=file_paths) + + partition_rows = spark.sql(f"SELECT partition, record_count, file_count FROM {identifier}.partitions").collect() + assert [row.record_count for row in partition_rows] == [2] + assert [row.file_count for row in partition_rows] == [2] + assert [row.partition.baz for row in partition_rows] == [123] + + assert len(tbl.scan().to_arrow()) == 2, "Expected 2 rows" + + +@pytest.mark.integration +def test_add_files_infers_bucket_partition_from_hive_style_path( + spark: SparkSession, session_catalog: Catalog, format_version: int +) -> None: + identifier = f"default.partitioned_table_bucket_hive_style_path_v{format_version}" + + partition_spec = PartitionSpec( + PartitionField(source_id=4, field_id=1000, transform=BucketTransform(num_buckets=3), name="baz_bucket_3"), + spec_id=0, + ) + + tbl = _create_table(session_catalog, identifier, format_version, partition_spec) + + # Without a Hive-style path this spec fails, see test_add_files_to_bucket_partitioned_table_fails. + file_path = f"s3://warehouse/default/partitioned_table_bucket_hive_style_path/v{format_version}/baz_bucket_3=1/test.parquet" + fo = tbl.io.new_output(file_path) + with fo.create(overwrite=True) as fos: + with pq.ParquetWriter(fos, schema=ARROW_SCHEMA) as writer: + writer.write_table( + pa.Table.from_pylist( + [ + {"foo": True, "bar": "bar_string", "baz": 0, "qux": date(2024, 3, 7)}, + {"foo": True, "bar": "bar_string", "baz": 1, "qux": date(2024, 3, 7)}, + ], + schema=ARROW_SCHEMA, + ) + ) + + tbl.add_files(file_paths=[file_path]) + + partition_rows = spark.sql(f"SELECT partition, record_count, file_count FROM {identifier}.partitions").collect() + assert [row.record_count for row in partition_rows] == [2] + assert [row.partition.baz_bucket_3 for row in partition_rows] == [1] + + assert len(tbl.scan().to_arrow()) == 2, "Expected 2 rows" + + +@pytest.mark.integration +def test_add_files_falls_back_to_stats_when_path_is_not_hive_style( + spark: SparkSession, session_catalog: Catalog, format_version: int +) -> None: + identifier = f"default.partitioned_table_non_hive_style_path_v{format_version}" + + partition_spec = PartitionSpec( + PartitionField(source_id=4, field_id=1000, transform=IdentityTransform(), name="baz"), + spec_id=0, + ) + + tbl = _create_table(session_catalog, identifier, format_version, partition_spec) + + # No "baz=" dir, so falls back to stats-based inference. + file_path = f"s3://warehouse/default/partitioned_table_non_hive_style_path/v{format_version}/test.parquet" + fo = tbl.io.new_output(file_path) + with fo.create(overwrite=True) as fos: + with pq.ParquetWriter(fos, schema=ARROW_SCHEMA) as writer: + writer.write_table( + pa.Table.from_pylist( + [{"foo": True, "bar": "bar_string", "baz": 123, "qux": date(2024, 3, 7)}], + schema=ARROW_SCHEMA, + ) + ) + + tbl.add_files(file_paths=[file_path]) + + partition_rows = spark.sql(f"SELECT partition, record_count FROM {identifier}.partitions").collect() + assert [row.partition.baz for row in partition_rows] == [123] + + @pytest.mark.integration def test_add_files_to_branch(spark: SparkSession, session_catalog: Catalog, format_version: int) -> None: identifier = f"default.test_add_files_branch_v{format_version}" diff --git a/tests/table/test_partitioning.py b/tests/table/test_partitioning.py index 57d3bc1c26..9bf8a2d7c0 100644 --- a/tests/table/test_partitioning.py +++ b/tests/table/test_partitioning.py @@ -31,6 +31,7 @@ IdentityTransform, MonthTransform, TruncateTransform, + VoidTransform, YearTransform, ) from pyiceberg.typedef import Record @@ -194,6 +195,89 @@ def test_partition_spec_to_path_dropped_source_id() -> None: assert spec.partition_to_path(record, schema) == "my%23str%25bucket=my%2Bstr/other+str%2Bbucket=%28+%29/my%21int%3Abucket=10" +def test_partition_from_path_identity() -> None: + schema = Schema( + NestedField(field_id=1, name="foo", field_type=StringType(), required=False), + NestedField(field_id=2, name="baz", field_type=IntegerType(), required=True), + ) + spec = PartitionSpec( + PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="foo"), + PartitionField(source_id=2, field_id=1001, transform=IdentityTransform(), name="baz"), + spec_id=0, + ) + + assert spec.partition_from_path("s3://bucket/table/data/foo=hello/baz=123/00000-0.parquet", schema) == Record("hello", 123) + + +def test_partition_from_path_unpartitioned() -> None: + schema = Schema(NestedField(field_id=1, name="foo", field_type=StringType(), required=False)) + assert UNPARTITIONED_PARTITION_SPEC.partition_from_path("s3://bucket/table/data/00000-0.parquet", schema) is None + + +def test_partition_from_path_not_hive_style() -> None: + schema = Schema(NestedField(field_id=1, name="foo", field_type=StringType(), required=False)) + spec = PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="foo"), spec_id=0) + + assert spec.partition_from_path("s3://bucket/table/data/00000-0.parquet", schema) is None + + +def test_partition_from_path_field_name_mismatch() -> None: + schema = Schema(NestedField(field_id=1, name="foo", field_type=StringType(), required=False)) + spec = PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="foo"), spec_id=0) + + assert spec.partition_from_path("s3://bucket/table/data/wrong=hello/00000-0.parquet", schema) is None + + +def test_partition_from_path_bucket_transform() -> None: + schema = Schema(NestedField(field_id=1, name="int", field_type=IntegerType(), required=True)) + spec = PartitionSpec( + PartitionField(source_id=1, field_id=1000, transform=BucketTransform(num_buckets=3), name="int_bucket"), spec_id=0 + ) + + assert spec.partition_from_path("s3://bucket/table/data/int_bucket=1/00000-0.parquet", schema) == Record(1) + + +def test_partition_from_path_truncate_transform() -> None: + schema = Schema(NestedField(field_id=1, name="str", field_type=StringType(), required=False)) + spec = PartitionSpec( + PartitionField(source_id=1, field_id=1000, transform=TruncateTransform(width=3), name="str_trunc"), spec_id=0 + ) + + assert spec.partition_from_path("s3://bucket/table/data/str_trunc=abc/00000-0.parquet", schema) == Record("abc") + + +def test_partition_from_path_truncate_binary_transform_unsupported() -> None: + schema = Schema(NestedField(field_id=1, name="bin", field_type=BinaryType(), required=False)) + spec = PartitionSpec( + PartitionField(source_id=1, field_id=1000, transform=TruncateTransform(width=3), name="bin_trunc"), spec_id=0 + ) + + # base64-encoded path value, not decodable back to bytes + assert spec.partition_from_path("s3://bucket/table/data/bin_trunc=YWJj/00000-0.parquet", schema) is None + + +def test_partition_from_path_time_transform_unsupported() -> None: + schema = Schema(NestedField(field_id=1, name="date", field_type=DateType(), required=False)) + spec = PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=MonthTransform(), name="date_month"), spec_id=0) + + # calendar string in path, not the raw int + assert spec.partition_from_path("s3://bucket/table/data/date_month=2024-03/00000-0.parquet", schema) is None + + +def test_partition_from_path_void_transform_unsupported() -> None: + schema = Schema(NestedField(field_id=1, name="foo", field_type=StringType(), required=False)) + spec = PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=VoidTransform(), name="foo_null"), spec_id=0) + + assert spec.partition_from_path("s3://bucket/table/data/foo_null=null/00000-0.parquet", schema) is None + + +def test_partition_from_path_url_encoded_value() -> None: + schema = Schema(NestedField(field_id=1, name="foo", field_type=StringType(), required=False)) + spec = PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="foo"), spec_id=0) + + assert spec.partition_from_path("s3://bucket/table/data/foo=a%2Bb/00000-0.parquet", schema) == Record("a+b") + + def test_partition_type(table_schema_simple: Schema) -> None: spec = PartitionSpec( PartitionField(source_id=1, field_id=1000, transform=TruncateTransform(width=19), name="str_truncate"),