Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions pyiceberg/io/pyarrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -1686,9 +1686,7 @@ def _get_column_projection_values(
for field_id in missing_field_ids:
for partition_field in partition_spec.fields_by_source_id(field_id):
if isinstance(partition_field.transform, IdentityTransform):
partition_value = accessors[partition_field.field_id].get(file.partition)
if partition_value is not None:
projected_missing_fields[field_id] = partition_value
projected_missing_fields[field_id] = accessors[partition_field.field_id].get(file.partition)
Comment thread
kevinjqliu marked this conversation as resolved.

return projected_missing_fields

Expand Down Expand Up @@ -2130,8 +2128,8 @@ def struct(
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)
projected_value = self._projected_missing_fields.get(field.field_id)
if projected_value is not None:
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:
field_arrays.append(pa.nulls(len(struct_array), type=arrow_type))
Expand Down
38 changes: 38 additions & 0 deletions tests/io/test_pyarrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,9 @@
BoundNotStartsWith,
BoundReference,
BoundStartsWith,
EqualTo,
GreaterThan,
IsNull,
Not,
Or,
)
Expand Down Expand Up @@ -1495,6 +1497,42 @@ def test_identity_transform_column_projection_with_falsy_value(
)


def test_identity_transform_null_partition_value_takes_precedence_over_initial_default(tmpdir: str) -> None:
"""A null identity partition value must project as null, not as the field's initial-default."""
table_schema = Schema(
NestedField(1, "other_field", StringType(), required=False),
NestedField(2, "partition_col", IntegerType(), required=False, initial_default=7),
)
partition_spec = PartitionSpec(PartitionField(2, 1000, IdentityTransform(), "partition_col"))
table_metadata = TableMetadataV3(
location=str(tmpdir), last_column_id=2, schemas=[table_schema], partition_specs=[partition_spec], next_row_id=0
)

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}/null_partition.parquet", arrow_schema, pa.table({"other_field": ["foo", "bar"]}, schema=arrow_schema)
)
data_file = DataFile.from_args(
content=DataFileContent.DATA,
file_path=file_path,
file_format=FileFormat.PARQUET,
partition=Record(None),
record_count=2,
file_size_in_bytes=os.path.getsize(file_path),
spec_id=partition_spec.spec_id,
)

def scan(projected_schema: Schema, row_filter: BooleanExpression) -> pa.Table:
return ArrowScan(table_metadata, PyArrowFileIO(), projected_schema, row_filter).to_table([FileScanTask(data_file)])

assert scan(table_schema, AlwaysTrue()).to_pydict() == {"other_field": ["foo", "bar"], "partition_col": [None, None]}

# Filters on the partition column see null whether or not the column is projected
for projected_schema in (table_schema, Schema(table_schema.fields[0])):
assert scan(projected_schema, IsNull("partition_col"))["other_field"].to_pylist() == ["foo", "bar"]
assert scan(projected_schema, EqualTo("partition_col", 7)).num_rows == 0


@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
Expand Down
Loading