Support inferring partition from Hive-style path in add_files

Author: ebyhr

pyiceberg/io/pyarrow.py

@@ -2921,11 +2921,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,

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

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}"

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