From 3cceb73495736fce5952bee7ca875ac53f6927e7 Mon Sep 17 00:00:00 2001 From: Sylwester Lachiewicz Date: Mon, 2 Nov 2026 21:36:46 +0000 Subject: [PATCH 1/2] Project identity partition values into required columns missing from the data file Column projection filled a missing identity partition column only when the field was optional, so a required one raised ResolveError even though the manifest carries its value. Closes #4267. --- pyiceberg/io/pyarrow.py | 9 +++++++-- tests/io/test_pyarrow.py | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 1dd816d5d3..5b53963360 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -2131,8 +2131,13 @@ def struct( array = self._cast_if_needed(field, field_array) field_arrays.append(array) fields.append(self._construct_field(field, array.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 + elif ( + field.optional + or field.initial_default is not None + or self._projected_missing_fields.get(field.field_id) is not None + ): + # When an optional field is added, when a required field with a non-null initial default is added, + # or when a required field is an identity partition column that is not stored in the data file 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] diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index abb7d384d5..81e3537d52 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -1550,6 +1550,45 @@ 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.""" + table_schema = Schema( + NestedField(1, "other_field", StringType(), required=False), + NestedField(2, "partition_col", IntegerType(), required=True), + ) + partition_spec = PartitionSpec(PartitionField(2, 1000, IdentityTransform(), "partition_col")) + table_metadata = TableMetadataV2( + location=str(tmpdir), last_column_id=2, format_version=2, schemas=[table_schema], partition_specs=[partition_spec] + ) + + 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) + ) + + def scan(partition_value: Any, row_filter: BooleanExpression) -> pa.Table: + 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)]) + + result = scan(1, AlwaysTrue()) + assert result.to_pydict() == {"other_field": ["foo", "bar"], "partition_col": [1, 1]} + assert not result.schema.field("partition_col").nullable + assert scan(1, EqualTo("partition_col", 1)).num_rows == 2 + assert scan(1, EqualTo("partition_col", 2)).num_rows == 0 + + # A null partition value cannot fill a required column + with pytest.raises(ResolveError, match="Field is required, and could not be found in the file"): + scan(None, AlwaysTrue()) + + @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 From a648233500d8e7cad550d8ccb4f591928205fdb9 Mon Sep 17 00:00:00 2001 From: Sylwester Lachiewicz Date: Mon, 2 Nov 2026 22:00:58 +0000 Subject: [PATCH 2/2] Reject a null identity partition value for a required field unless its parent is null A null projected value no longer reaches a required field through its initial-default, and it fills a required child of an optional struct when that struct is null in every row. --- pyiceberg/io/pyarrow.py | 24 ++++++----- tests/io/test_pyarrow.py | 90 ++++++++++++++++++++++++++++++++-------- 2 files changed, 85 insertions(+), 29 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 5b53963360..a430983525 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -2131,18 +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.optional - or field.initial_default is not None - or self._projected_missing_fields.get(field.field_id) is not None - ): - # When an optional field is added, when a required field with a non-null initial default is added, - # or when a required field is an identity partition column that is not stored in the data file + 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))) diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index 81e3537d52..c5a3568801 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -1552,21 +1552,16 @@ def scan(projected_schema: Schema, row_filter: BooleanExpression) -> pa.Table: 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.""" - table_schema = Schema( - NestedField(1, "other_field", StringType(), required=False), - NestedField(2, "partition_col", IntegerType(), required=True), - ) - partition_spec = PartitionSpec(PartitionField(2, 1000, IdentityTransform(), "partition_col")) - table_metadata = TableMetadataV2( - location=str(tmpdir), last_column_id=2, format_version=2, schemas=[table_schema], partition_specs=[partition_spec] - ) - 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(partition_value: Any, row_filter: BooleanExpression) -> pa.Table: + 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, @@ -1578,15 +1573,74 @@ def scan(partition_value: Any, row_filter: BooleanExpression) -> pa.Table: ) return ArrowScan(table_metadata, PyArrowFileIO(), table_schema, row_filter).to_table([FileScanTask(data_file)]) - result = scan(1, AlwaysTrue()) - assert result.to_pydict() == {"other_field": ["foo", "bar"], "partition_col": [1, 1]} - assert not result.schema.field("partition_col").nullable - assert scan(1, EqualTo("partition_col", 1)).num_rows == 2 - assert scan(1, EqualTo("partition_col", 2)).num_rows == 0 + 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 - with pytest.raises(ResolveError, match="Field is required, and could not be found in the file"): - scan(None, AlwaysTrue()) + # 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"])