(iceberg-python) 03/10: default-value bug fixes and adding more tests

[ke...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kevinjqliu pushed a commit to branch specPartitionIdentity
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git

commit e2f42ef4c13f8eb4953b74f04dad6be5747ea6bd
Author: Gabriel Igliozzi <gi...@netflix.com>
AuthorDate: Thu Dec 19 00:18:46 2024 -0500

    default-value bug fixes and adding more tests
---
 pyiceberg/io/pyarrow.py  | 19 +++++++----------
 tests/io/test_pyarrow.py | 54 +++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 12 deletions(-)

diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py
index 93c5f868..eb570592 100644
--- a/pyiceberg/io/pyarrow.py
+++ b/pyiceberg/io/pyarrow.py
@@ -1331,20 +1331,17 @@ def _task_to_record_batches(
         if file_schema is None:
             raise ValueError(f"Missing Iceberg schema in Metadata for file: {path}")
 
-        # Apply column projection rules for missing partitions and default values (ref: https://iceberg.apache.org/spec/#column-projection)
-        projected_missing_fields = None
+        # Apply column projection rules for missing partitions and default values
+        # https://iceberg.apache.org/spec/#column-projection
+        projected_missing_fields = {}
         for field_id in projected_field_ids.difference(file_project_schema.field_ids):
             for partition_field in partition_spec.fields_by_source_id(field_id):
                 if isinstance(partition_field.transform, IdentityTransform) and task.file.partition is not None:
-                    projected_missing_fields = (partition_field.name, task.file.partition[0])
+                    projected_missing_fields[partition_field.name] = task.file.partition[0]
 
-            if (
-                nested_field := projected_schema.find_field(field_id)
-                and projected_missing_fields is None
-                and task.file.partition is None
-            ):
+            if nested_field := projected_schema.find_field(field_id):
                 if nested_field.initial_default is not None:
-                    projected_missing_fields(nested_field.name, nested_field)
+                    projected_missing_fields[nested_field.name] = nested_field.initial_default
 
         fragment_scanner = ds.Scanner.from_fragment(
             fragment=fragment,
@@ -1392,8 +1389,8 @@ def _task_to_record_batches(
                     use_large_types=use_large_types,
                 )
 
-                if projected_missing_fields is not None:
-                    name, value = projected_missing_fields
+                # Inject projected column values if available
+                for name, value in projected_missing_fields.items():
                     result_batch = result_batch.set_column(result_batch.schema.get_field_index(name), name, [value])
 
                 yield result_batch
diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py
index 0507d9cf..b294dd19 100644
--- a/tests/io/test_pyarrow.py
+++ b/tests/io/test_pyarrow.py
@@ -1128,7 +1128,7 @@ def test_projection_concat_files(schema_int: Schema, file_int: str) -> None:
     assert repr(result_table.schema) == "id: int32"
 
 
-def test_projection_partition_inference(tmp_path: str, example_task: FileScanTask):
+def test_projection_partition_inference(tmp_path: str):
     schema = Schema(
         NestedField(1, "partition_field", IntegerType(), required=False),
         NestedField(2, "other_field", StringType(), required=False),
@@ -1180,6 +1180,58 @@ other_field: [["x"]]"""
     )
 
 
+def test_projection_initial_default_inference(tmp_path: str) -> None:
+    schema = Schema(
+        NestedField(1, "other_field", StringType(), required=False),
+        NestedField(2, "other_field_1", StringType(), required=False, initial_default="foo"),
+    )
+
+    table = TableMetadataV2(
+        location="file://a/b/c.json",
+        last_column_id=2,
+        format_version=2,
+        current_schema_id=0,
+        schemas=[schema],
+        partition_specs=[PartitionSpec()],
+        properties={TableProperties.DEFAULT_NAME_MAPPING: create_mapping_from_schema(schema).model_dump_json()},
+    )
+
+    pa_schema = pa.schema([pa.field("other_field", pa.string())])
+    pa_table = pa.table({"other_field": ["x"]}, schema=pa_schema)
+    pq.write_table(pa_table, f"{tmp_path}/datafile.parquet")
+
+    data_file = DataFile(
+        content=DataFileContent.DATA,
+        file_path=f"{tmp_path}/datafile.parquet",
+        file_format=FileFormat.PARQUET,
+        partition=Record(),
+        file_size_in_bytes=os.path.getsize(f"{tmp_path}/datafile.parquet"),
+        sort_order_id=None,
+        spec_id=0,
+        equality_ids=None,
+        key_metadata=None,
+    )
+
+    table_result_scan = ArrowScan(
+        table_metadata=table,
+        io=load_file_io(),
+        projected_schema=schema,
+        row_filter=AlwaysTrue(),
+    ).to_table(tasks=[FileScanTask(data_file=data_file)])
+
+    print(str(table_result_scan))
+
+    assert (
+        str(table_result_scan)
+        == """pyarrow.Table
+other_field: large_string
+other_field_1: string
+----
+other_field: [["x"]]
+other_field_1: [["foo"]]"""
+    )
+
+
 def test_projection_filter(schema_int: Schema, file_int: str) -> None:
     result_table = project(schema_int, [file_int], GreaterThan("id", 4))
     assert len(result_table.columns[0]) == 0