This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-4082-bd98ff1de88febaca645d54004a83de37a2d5ce9 in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
commit c1c49a228c27c777f9e63f7a46d801d85fcfd285 Author: Anupam Srivastava <[email protected]> AuthorDate: Sat Oct 24 21:25:56 2026 +0000 Fix filtering on identity partition columns outside the scan projection (#4082) * Fix filtering on identity partitions outside the scan projection Resolve manifest partition values for filter fields as well as output fields so scans with narrow projections keep matching rows. Cover name-mapped and field-ID Parquet files and falsy partition values. * Rename missing projection field IDs for clarity * Clarify projection docstring and cover count() with filter-only partition column Co-authored-by: Copilot App <[email protected]> --------- Co-authored-by: Kevin Liu <[email protected]> Co-authored-by: Copilot App <[email protected]> --- pyiceberg/io/pyarrow.py | 12 ++++----- tests/io/test_pyarrow.py | 70 +++++++++++++++++++++++++++++++++++++----------- 2 files changed, 60 insertions(+), 22 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 843dc0449..70dd37146 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1669,21 +1669,21 @@ class _ConvertToIcebergWithoutIDs(_ConvertToIceberg): def _get_column_projection_values( file: DataFile, - projected_schema: Schema, + projected_field_ids: set[int], table_schema: Schema, partition_spec: PartitionSpec | None, file_project_field_ids: set[int], ) -> dict[int, Any]: - """Apply Column Projection rules to File Schema.""" - project_schema_diff = projected_schema.field_ids.difference(file_project_field_ids) - if len(project_schema_diff) == 0 or partition_spec is None: + """Return identity partition values for projected field IDs missing from the data file.""" + missing_field_ids = projected_field_ids.difference(file_project_field_ids) + if len(missing_field_ids) == 0 or partition_spec is None: return EMPTY_DICT partition_schema = partition_spec.partition_type(table_schema) accessors = build_position_accessors(partition_schema) projected_missing_fields = {} - for field_id in project_schema_diff: + 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) @@ -1728,7 +1728,7 @@ def _task_to_record_batches( # Apply column projection rules: https://iceberg.apache.org/spec/#column-projection projected_missing_fields = _get_column_projection_values( - task.file, projected_schema, table_schema, partition_spec, file_schema.field_ids + task.file, projected_field_ids, table_schema, partition_spec, file_schema.field_ids ) pyarrow_filter = None diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index e6a085b8d..d665df9a7 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -1334,11 +1334,10 @@ def test_projection_concat_files(schema_int: Schema, file_int: str) -> None: assert repr(result_table.schema) == "id: int32" -def test_identity_transform_column_projection(tmp_path: str, catalog: InMemoryCatalog) -> None: [email protected]("with_field_ids", [False, True], ids=["name-mapping", "field-ids"]) +def test_identity_transform_column_projection(tmp_path: str, catalog: InMemoryCatalog, with_field_ids: bool) -> None: # Test by adding a non-partitioned data file to a partitioned table, verifying partition value # projection from manifest metadata. - # TODO: Update to use a data file created by writing data to an unpartitioned table once add_files supports field IDs. - # (context: https://github.com/apache/iceberg-python/pull/1443#discussion_r1901374875) schema = Schema( NestedField(1, "other_field", StringType(), required=False), NestedField(2, "partition_id", IntegerType(), required=False) @@ -1353,12 +1352,23 @@ def test_identity_transform_column_projection(tmp_path: str, catalog: InMemoryCa "default.test_projection_partition", schema=schema, partition_spec=partition_spec, - properties={TableProperties.DEFAULT_NAME_MAPPING: create_mapping_from_schema(schema).model_dump_json()}, + properties={} + if with_field_ids + else {TableProperties.DEFAULT_NAME_MAPPING: create_mapping_from_schema(schema).model_dump_json()}, ) file_data = pa.array(["foo", "bar", "baz"], type=pa.string()) - file_loc = f"{tmp_path}/test.parquet" - pq.write_table(pa.table([file_data], names=["other_field"]), file_loc) + file_table = pa.table([file_data], names=["other_field"]) + if with_field_ids: + source_table = catalog.create_table( + "default.unpartitioned_source", schema=Schema(schema.fields[0]), location=f"{tmp_path}/source" + ) + source_table.append(file_table) + file_loc = next(iter(source_table.scan().plan_files())).file.file_path + assert pq.read_schema(file_loc).field("other_field").metadata[PYARROW_PARQUET_FIELD_ID_KEY] == b"1" + else: + file_loc = f"{tmp_path}/test.parquet" + pq.write_table(file_table, file_loc) statistics = data_file_statistics_from_parquet_metadata( parquet_metadata=pq.read_metadata(file_loc), @@ -1403,13 +1413,21 @@ def test_identity_transform_column_projection(tmp_path: str, catalog: InMemoryCa # Test that row filter does not return any rows for a non-existing partition value assert len(table.scan(row_filter="partition_id = -1").to_arrow()) == 0 + assert table.scan(row_filter="partition_id = 1", selected_fields=("other_field",)).to_arrow() == pa.table( + {"other_field": ["foo", "bar", "baz"]} + ) + assert table.scan( + row_filter="partition_id = 1 AND other_field = 'bar'", selected_fields=("other_field",) + ).to_arrow() == pa.table({"other_field": ["bar"]}) + assert table.scan(row_filter="partition_id = 1 AND other_field = 'bar'", selected_fields=("other_field",)).count() == 1 + @pytest.mark.parametrize( - "partition_field_type, arrow_partition_type, partition_value", + "partition_field_type, arrow_partition_type, partition_value, row_filter", [ - (IntegerType(), pa.int32(), 0), - (StringType(), pa.large_string(), ""), - (IntegerType(), pa.int32(), None), + (IntegerType(), pa.int32(), 0, "partition_col = 0"), + (StringType(), pa.large_string(), "", "partition_col = ''"), + (IntegerType(), pa.int32(), None, "partition_col IS NULL"), ], ) def test_identity_transform_column_projection_with_falsy_value( @@ -1418,6 +1436,7 @@ def test_identity_transform_column_projection_with_falsy_value( partition_field_type: PrimitiveType, arrow_partition_type: pa.DataType, partition_value: Any, + row_filter: str, ) -> None: """Partition value projection must preserve falsy values (0, "") and still render None as null.""" schema = Schema( @@ -1471,13 +1490,15 @@ def test_identity_transform_column_projection_with_falsy_value( }, schema=expected_schema, ) + assert table.scan(row_filter=row_filter, selected_fields=("other_field",)).to_arrow() == pa.table( + {"other_field": ["foo", "bar"]} + ) -def test_identity_transform_columns_projection(tmp_path: str, catalog: InMemoryCatalog) -> None: [email protected]("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 # projection from manifest metadata. - # TODO: Update to use a data file created by writing data to an unpartitioned table once add_files supports field IDs. - # (context: https://github.com/apache/iceberg-python/pull/1443#discussion_r1901374875) schema = Schema( NestedField(1, "field_1", StringType(), required=False), NestedField(2, "field_2", IntegerType(), required=False), @@ -1494,12 +1515,23 @@ def test_identity_transform_columns_projection(tmp_path: str, catalog: InMemoryC "default.test_projection_partitions", schema=schema, partition_spec=partition_spec, - properties={TableProperties.DEFAULT_NAME_MAPPING: create_mapping_from_schema(schema).model_dump_json()}, + properties={} + if with_field_ids + else {TableProperties.DEFAULT_NAME_MAPPING: create_mapping_from_schema(schema).model_dump_json()}, ) file_data = pa.array(["foo"], type=pa.string()) - file_loc = f"{tmp_path}/test.parquet" - pq.write_table(pa.table([file_data], names=["field_1"]), file_loc) + file_table = pa.table([file_data], names=["field_1"]) + if with_field_ids: + source_table = catalog.create_table( + "default.unpartitioned_source", schema=Schema(schema.fields[0]), location=f"{tmp_path}/source" + ) + source_table.append(file_table) + file_loc = next(iter(source_table.scan().plan_files())).file.file_path + assert pq.read_schema(file_loc).field("field_1").metadata[PYARROW_PARQUET_FIELD_ID_KEY] == b"1" + else: + file_loc = f"{tmp_path}/test.parquet" + pq.write_table(file_table, file_loc) statistics = data_file_statistics_from_parquet_metadata( parquet_metadata=pq.read_metadata(file_loc), @@ -1536,6 +1568,12 @@ field_1: [["foo"]] field_2: [[2]] field_3: [[3]]""" ) + assert table.scan(row_filter="field_2 = 2 AND field_3 = 3", selected_fields=("field_1",)).to_arrow() == pa.table( + {"field_1": ["foo"]} + ) + assert table.scan(row_filter="field_2 = 2 AND field_3 = 3", selected_fields=("field_1", "field_2")).to_arrow() == pa.table( + {"field_1": ["foo"], "field_2": pa.array([2], type=pa.int32())} + ) @pytest.fixture
