Project identity partition values into required columns missing from the data file

pyiceberg/io/pyarrow.py

@@ -2131,13 +2131,20 @@ def struct(
                 array = self._cast_if_needed(field, field_array)
                 field_arrays.append(array)
                 fields.append(self._construct_field(field, array.type))
+            elif field.field_id in self._projected_missing_fields:
+                # An identity partition column that is not stored in the data file takes its value from the partition
+                projected_value = self._projected_missing_fields[field.field_id]
+                # A null can fill a required field only where its parent struct is null as well
+                parent_is_null = not isinstance(struct_array, pa.RecordBatch) and struct_array.null_count == len(struct_array)
+                if projected_value is None and field.required and not parent_is_null:
+                    raise ResolveError(f"Field is required, but its identity partition value is null: {field}")
+                arrow_type = schema_to_pyarrow(field.field_type, include_field_ids=self._include_field_ids)
+                field_arrays.append(pa.repeat(pa.scalar(projected_value, type=arrow_type), len(struct_array)))
+                fields.append(self._construct_field(field, arrow_type))
             elif field.optional or field.initial_default is not None:
                 # When an optional field is added, or when a required field with a non-null initial default is added
                 arrow_type = schema_to_pyarrow(field.field_type, include_field_ids=self._include_field_ids)
-                if field.field_id in self._projected_missing_fields:
-                    projected_value = self._projected_missing_fields[field.field_id]
-                    field_arrays.append(pa.repeat(pa.scalar(projected_value, type=arrow_type), len(struct_array)))
-                elif field.initial_default is None:
+                if field.initial_default is None:
                     field_arrays.append(pa.nulls(len(struct_array), type=arrow_type))
                 else:
                     field_arrays.append(pa.repeat(pa.scalar(field.initial_default, type=arrow_type), len(struct_array)))

tests/io/test_pyarrow.py

@@ -1550,6 +1550,99 @@ def scan(projected_schema: Schema, row_filter: BooleanExpression) -> pa.Table:
         assert scan(projected_schema, EqualTo("partition_col", 7)).num_rows == 0
 
 
+def test_identity_transform_column_projection_required_field(tmpdir: str) -> None:
+    """A required identity partition column missing from the data file must project from the partition value."""
+    arrow_schema = pa.schema([pa.field("other_field", pa.string(), metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"})])
+    file_path = _write_table_to_file(
+        f"{tmpdir}/required_partition.parquet", arrow_schema, pa.table({"other_field": ["foo", "bar"]}, schema=arrow_schema)
+    )
+    partition_spec = PartitionSpec(PartitionField(2, 1000, IdentityTransform(), "partition_col"))
+
+    def scan(table_schema: Schema, partition_value: Any, row_filter: BooleanExpression) -> pa.Table:
+        table_metadata = TableMetadataV3(
+            location=str(tmpdir), last_column_id=2, schemas=[table_schema], partition_specs=[partition_spec], next_row_id=0
+        )
+        data_file = DataFile.from_args(
+            content=DataFileContent.DATA,
+            file_path=file_path,
+            file_format=FileFormat.PARQUET,
+            partition=Record(partition_value),
+            record_count=2,
+            file_size_in_bytes=os.path.getsize(file_path),
+            spec_id=partition_spec.spec_id,
+        )
+        return ArrowScan(table_metadata, PyArrowFileIO(), table_schema, row_filter).to_table([FileScanTask(data_file)])
+
+    for initial_default in (None, 7):
+        table_schema = Schema(
+            NestedField(1, "other_field", StringType(), required=False),
+            NestedField(2, "partition_col", IntegerType(), required=True, initial_default=initial_default),
+        )
+
+        result = scan(table_schema, 1, AlwaysTrue())
+        assert result.to_pydict() == {"other_field": ["foo", "bar"], "partition_col": [1, 1]}
+        assert not result.schema.field("partition_col").nullable
+        assert scan(table_schema, 1, EqualTo("partition_col", 1)).num_rows == 2
+        assert scan(table_schema, 1, EqualTo("partition_col", 2)).num_rows == 0
+
+        # A null partition value cannot fill a required column, whether or not the field has an initial-default
+        with pytest.raises(ResolveError, match="Field is required, but its identity partition value is null"):
+            scan(table_schema, None, AlwaysTrue())
+
+
+def test_identity_transform_column_projection_required_field_under_null_parent(tmpdir: str) -> None:
+    """A null partition value fills a required nested column only where its optional parent struct is null."""
+    table_schema = Schema(
+        NestedField(1, "id", IntegerType(), required=False),
+        NestedField(
+            2,
+            "location",
+            StructType(
+                NestedField(3, "city", StringType(), required=False),
+                NestedField(4, "region", StringType(), required=True),
+            ),
+            required=False,
+        ),
+    )
+    partition_spec = PartitionSpec(PartitionField(4, 1000, IdentityTransform(), "region"))
+    table_metadata = TableMetadataV2(
+        location=str(tmpdir), last_column_id=4, format_version=2, schemas=[table_schema], partition_specs=[partition_spec]
+    )
+
+    location_type = pa.struct([pa.field("city", pa.string(), metadata={PYARROW_PARQUET_FIELD_ID_KEY: "3"})])
+    arrow_schema = pa.schema(
+        [
+            pa.field("id", pa.int32(), metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"}),
+            pa.field("location", location_type, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "2"}),
+        ]
+    )
+
+    def scan(locations: list[dict[str, str] | None], partition_value: Any) -> pa.Table:
+        file_path = _write_table_to_file(
+            f"{tmpdir}/nested_{sum(location is not None for location in locations)}_{partition_value}.parquet",
+            arrow_schema,
+            pa.table({"id": [1, 2], "location": pa.array(locations, type=location_type)}, schema=arrow_schema),
+        )
+        data_file = DataFile.from_args(
+            content=DataFileContent.DATA,
+            file_path=file_path,
+            file_format=FileFormat.PARQUET,
+            partition=Record(partition_value),
+            record_count=2,
+            file_size_in_bytes=os.path.getsize(file_path),
+            spec_id=partition_spec.spec_id,
+        )
+        return ArrowScan(table_metadata, PyArrowFileIO(), table_schema, AlwaysTrue()).to_table([FileScanTask(data_file)])
+
+    assert scan([None, None], None).to_pydict() == {"id": [1, 2], "location": [None, None]}
+    assert scan([{"city": "Bern"}, None], "north").to_pydict() == {
+        "id": [1, 2],
+        "location": [{"city": "Bern", "region": "north"}, None],
+    }
+    with pytest.raises(ResolveError, match="Field is required, but its identity partition value is null"):
+        scan([{"city": "Bern"}, None], None)
+
+
 @pytest.mark.parametrize("with_field_ids", [False, True], ids=["name-mapping", "field-ids"])
 def test_identity_transform_columns_projection(tmp_path: str, catalog: InMemoryCatalog, with_field_ids: bool) -> None:
     # Test by adding a non-partitioned data file to a multi-partitioned table, verifying partition value