(iceberg-python) 05/10: Fix lint issues, reorder partition spec to be of higher importance than initial-default

[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 d57af796f175e69428effae838fb33038b485d36
Author: Gabriel Igliozzi <gi...@netflix.com>
AuthorDate: Thu Dec 19 10:10:10 2024 -0500

    Fix lint issues, reorder partition spec to be of higher importance than initial-default
---
 pyiceberg/io/pyarrow.py  | 12 +++++-------
 tests/io/test_pyarrow.py |  2 +-
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py
index bf8398ab..63dc99da 100644
--- a/pyiceberg/io/pyarrow.py
+++ b/pyiceberg/io/pyarrow.py
@@ -1308,7 +1308,7 @@ def _task_to_record_batches(
     case_sensitive: bool,
     name_mapping: Optional[NameMapping] = None,
     use_large_types: bool = True,
-    partition_spec: PartitionSpec = None,
+    partition_spec: Optional[PartitionSpec] = None,
 ) -> Iterator[pa.RecordBatch]:
     _, _, path = _parse_location(task.file.file_path)
     arrow_format = ds.ParquetFileFormat(pre_buffer=True, buffer_size=(ONE_MEGABYTE * 8))
@@ -1332,19 +1332,17 @@ def _task_to_record_batches(
 
         # Apply column projection rules for missing partitions and default values
         # https://iceberg.apache.org/spec/#column-projection
-
         file_project_schema = prune_columns(file_schema, projected_field_ids, select_full_types=False)
         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]
-                    continue
-
             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.initial_default
+            if partition_spec is not None:
+                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]
 
         fragment_scanner = ds.Scanner.from_fragment(
             fragment=fragment,
diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py
index b294dd19..48646759 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):
+def test_projection_partition_inference(tmp_path: str) -> None:
     schema = Schema(
         NestedField(1, "partition_field", IntegerType(), required=False),
         NestedField(2, "other_field", StringType(), required=False),