PR Comments

Author: nssalian

pyiceberg/io/fileformat.py

@@ -28,6 +28,7 @@
 from pyiceberg.partitioning import PartitionField, PartitionSpec, partition_record_value
 from pyiceberg.schema import Schema
 from pyiceberg.typedef import Properties, Record
+from pyiceberg.types import NestedField
 
 if TYPE_CHECKING:
     import pyarrow as pa
@@ -161,6 +162,10 @@ def create_writer(
         properties: Properties,
     ) -> FileFormatWriter: ...
 
+    @abstractmethod
+    def add_field_metadata(self, field: NestedField, metadata: dict[bytes, bytes], include_field_ids: bool) -> None:
+        """Add format-specific Arrow field metadata."""
+
 
 class FileFormatFactory:
     """Registry of FileFormatModel implementations."""

pyiceberg/io/pyarrow.py

@@ -1886,7 +1886,7 @@ def _to_requested_schema(
     include_field_ids: bool = False,
     projected_missing_fields: dict[int, Any] = EMPTY_DICT,
     allow_timestamp_tz_mismatch: bool = False,
-    file_format: FileFormat = FileFormat.PARQUET,
+    format_model: FileFormatModel | None = None,
 ) -> pa.RecordBatch:
     # We could reuse some of these visitors
     struct_array = visit_with_partner(
@@ -1898,7 +1898,7 @@ def _to_requested_schema(
             include_field_ids,
             projected_missing_fields=projected_missing_fields,
             allow_timestamp_tz_mismatch=allow_timestamp_tz_mismatch,
-            file_format=file_format,
+            format_model=format_model,
         ),
         ArrowAccessor(file_schema),
     )
@@ -1911,7 +1911,7 @@ class ArrowProjectionVisitor(SchemaWithPartnerVisitor[pa.Array, pa.Array | None]
     _downcast_ns_timestamp_to_us: bool
     _projected_missing_fields: dict[int, Any]
     _allow_timestamp_tz_mismatch: bool
-    _file_format: FileFormat
+    _format_model: FileFormatModel | None
 
     def __init__(
         self,
@@ -1920,16 +1920,18 @@ def __init__(
         include_field_ids: bool = False,
         projected_missing_fields: dict[int, Any] = EMPTY_DICT,
         allow_timestamp_tz_mismatch: bool = False,
-        file_format: FileFormat = FileFormat.PARQUET,
+        format_model: FileFormatModel | None = None,
     ) -> None:
+        if include_field_ids and format_model is None:
+            raise ValueError("format_model is required when include_field_ids=True")
         self._file_schema = file_schema
         self._include_field_ids = include_field_ids
         self._downcast_ns_timestamp_to_us = downcast_ns_timestamp_to_us
         self._projected_missing_fields = projected_missing_fields
         # When True, allows projecting timestamptz (UTC) to timestamp (no tz).
         # Allowed for reading (aligns with Spark); disallowed for writing to enforce Iceberg spec's strict typing.
         self._allow_timestamp_tz_mismatch = allow_timestamp_tz_mismatch
-        self._file_format = file_format
+        self._format_model = format_model
 
     def _cast_if_needed(self, field: NestedField, values: pa.Array) -> pa.Array:
         file_field = self._file_schema.find_field(field.field_id)
@@ -1984,16 +1986,11 @@ def _cast_if_needed(self, field: NestedField, values: pa.Array) -> pa.Array:
         return values
 
     def _construct_field(self, field: NestedField, arrow_type: pa.DataType) -> pa.Field:
-        metadata = {}
+        metadata: dict[bytes, bytes] = {}
         if field.doc:
-            metadata[PYARROW_FIELD_DOC_KEY] = field.doc
-        if self._include_field_ids:
-            if self._file_format == FileFormat.ORC:
-                metadata[ORC_FIELD_ID_KEY] = str(field.field_id)
-            else:
-                metadata[PYARROW_PARQUET_FIELD_ID_KEY] = str(field.field_id)
-        if self._file_format == FileFormat.ORC:
-            metadata[ORC_FIELD_REQUIRED_KEY] = str(field.required).lower()
+            metadata[PYARROW_FIELD_DOC_KEY] = field.doc.encode()
+        if self._format_model is not None:
+            self._format_model.add_field_metadata(field, metadata, self._include_field_ids)
 
         return pa.field(
             name=field.name,
@@ -2675,6 +2672,10 @@ def create_writer(
     ) -> ParquetFormatWriter:
         return ParquetFormatWriter(output_file, file_schema, properties)
 
+    def add_field_metadata(self, field: NestedField, metadata: dict[bytes, bytes], include_field_ids: bool) -> None:
+        if include_field_ids:
+            metadata[PYARROW_PARQUET_FIELD_ID_KEY] = str(field.field_id).encode()
+
 
 FileFormatFactory.register(ParquetFormatModel())
 
@@ -2706,7 +2707,7 @@ def write_data_file(task: WriteTask) -> DataFile:
                 batch=batch,
                 downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us,
                 include_field_ids=True,
-                file_format=file_format,
+                format_model=format_model,
             )
             for batch in task.record_batches
         ]

tests/conftest.py

@@ -172,6 +172,20 @@ def table_schema_simple() -> Schema:
     )
 
 
+@pytest.fixture(scope="session")
+def arrow_table_simple() -> "pa.Table":
+    """Pyarrow table that pairs with `table_schema_simple` (3 rows, no nulls)."""
+    import pyarrow as pa
+
+    return pa.table(
+        {
+            "foo": ["a", "b", "c"],
+            "bar": pa.array([1, 2, 3], type=pa.int32()),
+            "baz": [True, False, True],
+        }
+    )
+
+
 @pytest.fixture(scope="session")
 def table_schema_with_full_nested_fields() -> Schema:
     return Schema(

tests/io/test_fileformat.py

@@ -51,6 +51,9 @@ def file_extension(self) -> str:
         def create_writer(self, output_file: Any, file_schema: Any, properties: Any) -> Any:
             raise NotImplementedError
 
+        def add_field_metadata(self, field: Any, metadata: Any, include_field_ids: bool) -> None:
+            pass
+
     original = dict(FileFormatFactory._registry)
     try:
         model = _DummyModel()

tests/io/test_format_writers.py

@@ -35,51 +35,27 @@ def format_model(request: pytest.FixtureRequest) -> FileFormatModel:
     return FileFormatFactory.get(request.param)
 
 
-@pytest.fixture
-def simple_table() -> pa.Table:
-    return pa.table(
-        {
-            "foo": ["a", "b", "c"],
-            "bar": pa.array([1, 2, 3], type=pa.int32()),
-            "baz": [True, False, True],
-        }
-    )
-
-
 def test_parquet_registered() -> None:
     """ParquetFormatModel is registered in the factory."""
     model = FileFormatFactory.get(FileFormat.PARQUET)
     assert model.format == FileFormat.PARQUET
     assert model.file_extension() == "parquet"
 
 
-def test_round_trip(format_model: FileFormatModel, table_schema_simple: Schema, simple_table: pa.Table, tmp_path: Path) -> None:
+def test_round_trip(
+    format_model: FileFormatModel, table_schema_simple: Schema, arrow_table_simple: pa.Table, tmp_path: Path
+) -> None:
     """Write a table and read it back, to verify equality and record count."""
     file_path = str(tmp_path / f"test.{format_model.file_extension()}")
     writer = format_model.create_writer(PyArrowFileIO().new_output(file_path), table_schema_simple, {})
-    writer.write(simple_table)
+    writer.write(arrow_table_simple)
     statistics = writer.close()
 
     result = ds.dataset(file_path).to_table()
-    assert result.equals(simple_table)
+    assert result.equals(arrow_table_simple)
     assert statistics.record_count == 3
 
 
-def test_statistics_record_count(format_model: FileFormatModel, table_schema_simple: Schema, tmp_path: Path) -> None:
-    """close() returns DataFileStatistics with correct record count."""
-    table = pa.table(
-        {
-            "foo": ["a", "b", "c", "d", "e"],
-            "bar": pa.array([10, 20, 30, 40, 50], type=pa.int32()),
-            "baz": [True] * 5,
-        }
-    )
-    file_path = str(tmp_path / f"test.{format_model.file_extension()}")
-    writer = format_model.create_writer(PyArrowFileIO().new_output(file_path), table_schema_simple, {})
-    writer.write(table)
-    assert writer.close().record_count == 5
-
-
 def test_null_handling(format_model: FileFormatModel, table_schema_simple: Schema, tmp_path: Path) -> None:
     """Nullable columns produce correct null_value_counts in statistics."""
     table = pa.table(
@@ -98,23 +74,23 @@ def test_null_handling(format_model: FileFormatModel, table_schema_simple: Schem
 
 
 def test_context_manager_caches_result(
-    format_model: FileFormatModel, table_schema_simple: Schema, simple_table: pa.Table, tmp_path: Path
+    format_model: FileFormatModel, table_schema_simple: Schema, arrow_table_simple: pa.Table, tmp_path: Path
 ) -> None:
     """writer.result() returns cached statistics after context manager exit."""
     file_path = str(tmp_path / f"test.{format_model.file_extension()}")
     writer = format_model.create_writer(PyArrowFileIO().new_output(file_path), table_schema_simple, {})
     with writer:
-        writer.write(simple_table)
+        writer.write(arrow_table_simple)
     assert writer.result().record_count == 3
 
 
 def test_close_is_idempotent(
-    format_model: FileFormatModel, table_schema_simple: Schema, simple_table: pa.Table, tmp_path: Path
+    format_model: FileFormatModel, table_schema_simple: Schema, arrow_table_simple: pa.Table, tmp_path: Path
 ) -> None:
     """Calling close() twice returns the same cached statistics object."""
     file_path = str(tmp_path / f"test.{format_model.file_extension()}")
     writer = format_model.create_writer(PyArrowFileIO().new_output(file_path), table_schema_simple, {})
-    writer.write(simple_table)
+    writer.write(arrow_table_simple)
     stats1 = writer.close()
     stats2 = writer.close()
     assert stats1 is stats2
@@ -128,28 +104,16 @@ def test_close_without_write_raises(format_model: FileFormatModel, table_schema_
         writer.close()
 
 
-def test_construct_field_uses_orc_field_id_key() -> None:
-    """ArrowProjectionVisitor uses ORC field ID and required keys when file_format is ORC."""
-    from pyiceberg.io.pyarrow import (
-        ORC_FIELD_ID_KEY,
-        ORC_FIELD_REQUIRED_KEY,
-        PYARROW_PARQUET_FIELD_ID_KEY,
-        ArrowProjectionVisitor,
-    )
+def test_parquet_format_model_adds_field_id_metadata() -> None:
+    """ParquetFormatModel.add_field_metadata writes the Parquet field-id key when requested."""
+    from pyiceberg.io.pyarrow import PYARROW_PARQUET_FIELD_ID_KEY, ParquetFormatModel
+
+    field = NestedField(field_id=1, name="x", field_type=LongType(), required=True)
+
+    metadata: dict[bytes, bytes] = {}
+    ParquetFormatModel().add_field_metadata(field, metadata, include_field_ids=True)
+    assert metadata == {PYARROW_PARQUET_FIELD_ID_KEY: b"1"}
 
-    schema = Schema(NestedField(field_id=1, name="x", field_type=LongType(), required=True))
-
-    visitor = ArrowProjectionVisitor(schema, include_field_ids=True, file_format=FileFormat.ORC)
-    field = visitor._construct_field(schema.find_field(1), pa.int64())
-    assert field.metadata is not None
-    assert ORC_FIELD_ID_KEY in field.metadata
-    assert ORC_FIELD_REQUIRED_KEY in field.metadata
-    assert field.metadata[ORC_FIELD_REQUIRED_KEY] == b"true"
-    assert PYARROW_PARQUET_FIELD_ID_KEY not in field.metadata
-
-    visitor_pq = ArrowProjectionVisitor(schema, include_field_ids=True, file_format=FileFormat.PARQUET)
-    field_pq = visitor_pq._construct_field(schema.find_field(1), pa.int64())
-    assert field_pq.metadata is not None
-    assert PYARROW_PARQUET_FIELD_ID_KEY in field_pq.metadata
-    assert ORC_FIELD_ID_KEY not in field_pq.metadata
-    assert ORC_FIELD_REQUIRED_KEY not in field_pq.metadata
+    metadata_no_ids: dict[bytes, bytes] = {}
+    ParquetFormatModel().add_field_metadata(field, metadata_no_ids, include_field_ids=False)
+    assert metadata_no_ids == {}