diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 1dd816d5d3..a430983525 100644 --- a/pyiceberg/io/pyarrow.py +++ b/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))) diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index abb7d384d5..c5a3568801 100644 --- a/tests/io/test_pyarrow.py +++ b/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