Skip to content

Commit 41006bf

Browse files
sungwykevinjqliuCopilot
authored
Bug Fix: handle column projection edge case with null identity partition (#4083)
* fix column projection edge case with null identity partition * Cover filter-only null identity partitions and simplify test Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Kevin Liu <kevin.jq.liu@gmail.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 945c8db commit 41006bf

2 files changed

Lines changed: 41 additions & 5 deletions

File tree

‎pyiceberg/io/pyarrow.py‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1686,9 +1686,7 @@ def _get_column_projection_values(
16861686
for field_id in missing_field_ids:
16871687
for partition_field in partition_spec.fields_by_source_id(field_id):
16881688
if isinstance(partition_field.transform, IdentityTransform):
1689-
partition_value = accessors[partition_field.field_id].get(file.partition)
1690-
if partition_value is not None:
1691-
projected_missing_fields[field_id] = partition_value
1689+
projected_missing_fields[field_id] = accessors[partition_field.field_id].get(file.partition)
16921690

16931691
return projected_missing_fields
16941692

@@ -2130,8 +2128,8 @@ def struct(
21302128
elif field.optional or field.initial_default is not None:
21312129
# When an optional field is added, or when a required field with a non-null initial default is added
21322130
arrow_type = schema_to_pyarrow(field.field_type, include_field_ids=self._include_field_ids)
2133-
projected_value = self._projected_missing_fields.get(field.field_id)
2134-
if projected_value is not None:
2131+
if field.field_id in self._projected_missing_fields:
2132+
projected_value = self._projected_missing_fields[field.field_id]
21352133
field_arrays.append(pa.repeat(pa.scalar(projected_value, type=arrow_type), len(struct_array)))
21362134
elif field.initial_default is None:
21372135
field_arrays.append(pa.nulls(len(struct_array), type=arrow_type))

‎tests/io/test_pyarrow.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,9 @@
6363
BoundNotStartsWith,
6464
BoundReference,
6565
BoundStartsWith,
66+
EqualTo,
6667
GreaterThan,
68+
IsNull,
6769
Not,
6870
Or,
6971
)
@@ -1495,6 +1497,42 @@ def test_identity_transform_column_projection_with_falsy_value(
14951497
)
14961498

14971499

1500+
def test_identity_transform_null_partition_value_takes_precedence_over_initial_default(tmpdir: str) -> None:
1501+
"""A null identity partition value must project as null, not as the field's initial-default."""
1502+
table_schema = Schema(
1503+
NestedField(1, "other_field", StringType(), required=False),
1504+
NestedField(2, "partition_col", IntegerType(), required=False, initial_default=7),
1505+
)
1506+
partition_spec = PartitionSpec(PartitionField(2, 1000, IdentityTransform(), "partition_col"))
1507+
table_metadata = TableMetadataV3(
1508+
location=str(tmpdir), last_column_id=2, schemas=[table_schema], partition_specs=[partition_spec], next_row_id=0
1509+
)
1510+
1511+
arrow_schema = pa.schema([pa.field("other_field", pa.string(), metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"})])
1512+
file_path = _write_table_to_file(
1513+
f"{tmpdir}/null_partition.parquet", arrow_schema, pa.table({"other_field": ["foo", "bar"]}, schema=arrow_schema)
1514+
)
1515+
data_file = DataFile.from_args(
1516+
content=DataFileContent.DATA,
1517+
file_path=file_path,
1518+
file_format=FileFormat.PARQUET,
1519+
partition=Record(None),
1520+
record_count=2,
1521+
file_size_in_bytes=os.path.getsize(file_path),
1522+
spec_id=partition_spec.spec_id,
1523+
)
1524+
1525+
def scan(projected_schema: Schema, row_filter: BooleanExpression) -> pa.Table:
1526+
return ArrowScan(table_metadata, PyArrowFileIO(), projected_schema, row_filter).to_table([FileScanTask(data_file)])
1527+
1528+
assert scan(table_schema, AlwaysTrue()).to_pydict() == {"other_field": ["foo", "bar"], "partition_col": [None, None]}
1529+
1530+
# Filters on the partition column see null whether or not the column is projected
1531+
for projected_schema in (table_schema, Schema(table_schema.fields[0])):
1532+
assert scan(projected_schema, IsNull("partition_col"))["other_field"].to_pylist() == ["foo", "bar"]
1533+
assert scan(projected_schema, EqualTo("partition_col", 7)).num_rows == 0
1534+
1535+
14981536
@pytest.mark.parametrize("with_field_ids", [False, True], ids=["name-mapping", "field-ids"])
14991537
def test_identity_transform_columns_projection(tmp_path: str, catalog: InMemoryCatalog, with_field_ids: bool) -> None:
15001538
# Test by adding a non-partitioned data file to a multi-partitioned table, verifying partition value

0 commit comments

Comments
 (0)