Add support for partitioning by nested columns

Author: geruh

pyiceberg/io/pyarrow.py

@@ -2728,9 +2728,11 @@ def _determine_partitions(spec: PartitionSpec, schema: Schema, arrow_table: pa.T
 
     for partition, name in zip(spec.fields, partition_fields):
         source_field = schema.find_field(partition.source_id)
-        arrow_table = arrow_table.append_column(
-            name, partition.transform.pyarrow_transform(source_field.field_type)(arrow_table[source_field.name])
-        )
+        full_field_name = schema.find_column_name(partition.source_id)
+        if full_field_name is None:
+            raise ValueError(f"Could not find column name for field ID: {partition.source_id}")
+        field_array = _get_field_from_arrow_table(arrow_table, full_field_name)
+        arrow_table = arrow_table.append_column(name, partition.transform.pyarrow_transform(source_field.field_type)(field_array))
 
     unique_partition_fields = arrow_table.select(partition_fields).group_by(partition_fields).aggregate([])
 
@@ -2765,3 +2767,22 @@ def _determine_partitions(spec: PartitionSpec, schema: Schema, arrow_table: pa.T
         )
 
     return table_partitions
+
+
+def _get_field_from_arrow_table(arrow_table: pa.Table, field_path: str) -> pa.Array:
+    """Get a nested field from an Arrow table struct type field using dot notation.
+
+    Args:
+        arrow_table: The Arrow table containing the field
+        field_path: Dot-separated field path (e.g., "name" or "bar.baz.timestamp")
+
+    Returns:
+        The unnested field as a PyArrow Array
+    """
+    if "." not in field_path:
+        return arrow_table[field_path]
+
+    path_parts = field_path.split(".")
+    field_array = arrow_table[path_parts[0]]
+    field_array = pc.struct_field(field_array, path_parts[1:])  # type: ignore
+    return field_array

tests/io/test_pyarrow.py

@@ -84,7 +84,7 @@
 from pyiceberg.table import FileScanTask, TableProperties
 from pyiceberg.table.metadata import TableMetadataV2
 from pyiceberg.table.name_mapping import create_mapping_from_schema
-from pyiceberg.transforms import IdentityTransform
+from pyiceberg.transforms import HourTransform, IdentityTransform
 from pyiceberg.typedef import UTF8, Properties, Record
 from pyiceberg.types import (
     BinaryType,
@@ -2350,6 +2350,72 @@ def test_partition_for_demo() -> None:
     )
 
 
+def test_partition_for_nested_field() -> None:
+    schema = Schema(
+        NestedField(id=1, name="foo", field_type=StringType(), required=True),
+        NestedField(
+            id=2,
+            name="bar",
+            field_type=StructType(
+                NestedField(id=3, name="baz", field_type=TimestampType(), required=False),
+                NestedField(id=4, name="qux", field_type=IntegerType(), required=False),
+            ),
+            required=True,
+        ),
+    )
+
+    spec = PartitionSpec(PartitionField(source_id=3, field_id=1000, transform=HourTransform(), name="ts"))
+
+    from datetime import datetime
+
+    t1 = datetime(2025, 7, 11, 9, 30, 0)
+    t2 = datetime(2025, 7, 11, 10, 30, 0)
+
+    test_data = [
+        {"foo": "a", "bar": {"baz": t1, "qux": 1}},
+        {"foo": "b", "bar": {"baz": t2, "qux": 2}},
+    ]
+
+    arrow_table = pa.Table.from_pylist(test_data, schema=schema.as_arrow())
+    partitions = _determine_partitions(spec, schema, arrow_table)
+    partition_values = {p.partition_key.partition[0] for p in partitions}
+
+    assert partition_values == {486729, 486730}
+
+
+def test_partition_for_deep_nested_field() -> None:
+    schema = Schema(
+        NestedField(
+            id=1,
+            name="foo",
+            field_type=StructType(
+                NestedField(
+                    id=2,
+                    name="bar",
+                    field_type=StructType(NestedField(id=3, name="baz", field_type=StringType(), required=False)),
+                    required=True,
+                )
+            ),
+            required=True,
+        )
+    )
+
+    spec = PartitionSpec(PartitionField(source_id=3, field_id=1000, transform=IdentityTransform(), name="qux"))
+
+    test_data = [
+        {"foo": {"bar": {"baz": "data-1"}}},
+        {"foo": {"bar": {"baz": "data-2"}}},
+        {"foo": {"bar": {"baz": "data-1"}}},
+    ]
+
+    arrow_table = pa.Table.from_pylist(test_data, schema=schema.as_arrow())
+    partitions = _determine_partitions(spec, schema, arrow_table)
+
+    assert len(partitions) == 2  # 2 unique partitions
+    partition_values = {p.partition_key.partition[0] for p in partitions}
+    assert partition_values == {"data-1", "data-2"}
+
+
 def test_identity_partition_on_multi_columns() -> None:
     test_pa_schema = pa.schema([("born_year", pa.int64()), ("n_legs", pa.int64()), ("animal", pa.string())])
     test_schema = Schema(