Reuse parquet_path_to_id_mapping

Author: renaudb

pyiceberg/io/pyarrow.py

@@ -2479,19 +2479,6 @@ def parquet_path_to_id_mapping(
     return result
 
 
-def id_to_parquet_path_mapping(schema: Schema) -> dict[int, str]:
-    """
-    Compute the mapping of Iceberg column ID to parquet column path.
-
-    Args:
-        schema (pyiceberg.schema.Schema): The current table schema.
-    """
-    result: dict[int, str] = {}
-    for pair in pre_order_visit(schema, ID2ParquetPathVisitor()):
-        result[pair.field_id] = pair.parquet_path
-    return result
-
-
 @dataclass(frozen=True)
 class BloomFilterOptions:
     parquet_path: str
@@ -2573,17 +2560,17 @@ def primitive(self, primitive: PrimitiveType) -> builtins.list[BloomFilterOption
 def get_bloom_filter_options(
     schema: Schema,
     table_properties: dict[str, str],
+    id_to_parquet_path: dict[int, str],
 ) -> dict[str, dict[str, Any]]:
     """
     Get the bloom filter options from the table properties.
 
     Args:
         schema (pyiceberg.schema.Schema): The current table schema.
         table_properties (dict[str, str]): The table properties.
+        id_to_parquet_path (Dict[int, str]): The mapping of the field ID to the parquet file name.
     """
-    bloom_filter_options = pre_order_visit(
-        schema, BloomFilterOptionsCollector(schema, table_properties, id_to_parquet_path_mapping(schema))
-    )
+    bloom_filter_options = pre_order_visit(schema, BloomFilterOptionsCollector(schema, table_properties, id_to_parquet_path))
     result: dict[str, dict[str, Any]] = {}
     for bf_opts in bloom_filter_options:
         result[bf_opts.parquet_path] = {
@@ -2731,7 +2718,8 @@ def write_parquet(task: WriteTask) -> DataFile:
         else:
             file_schema = table_schema
 
-        parquet_writer_kwargs = _get_parquet_writer_kwargs(table_metadata.properties, file_schema)
+        parquet_column_mapping = parquet_path_to_id_mapping(file_schema)
+        parquet_writer_kwargs = _get_parquet_writer_kwargs(table_metadata.properties, file_schema, parquet_column_mapping)
 
         downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False
         batches = [
@@ -2758,7 +2746,7 @@ def write_parquet(task: WriteTask) -> DataFile:
         statistics = data_file_statistics_from_parquet_metadata(
             parquet_metadata=writer.writer.metadata,
             stats_columns=compute_statistics_plan(file_schema, table_metadata.properties),
-            parquet_column_mapping=parquet_path_to_id_mapping(file_schema),
+            parquet_column_mapping=parquet_column_mapping,
         )
         data_file = DataFile.from_args(
             content=DataFileContent.DATA,
@@ -2877,7 +2865,9 @@ def parquet_file_to_data_file(io: FileIO, table_metadata: TableMetadata, file_pa
 PYARROW_UNCOMPRESSED_CODEC = "none"
 
 
-def _get_parquet_writer_kwargs(table_properties: Properties, file_schema: Schema) -> dict[str, Any]:
+def _get_parquet_writer_kwargs(
+    table_properties: Properties, file_schema: Schema, parquet_column_mapping: dict[str, int]
+) -> dict[str, Any]:
     from pyiceberg.table import TableProperties
 
     unsupported_key_patterns = [
@@ -2908,7 +2898,8 @@ def _get_parquet_writer_kwargs(table_properties: Properties, file_schema: Schema
     if compression_codec == ICEBERG_UNCOMPRESSED_CODEC:
         compression_codec = PYARROW_UNCOMPRESSED_CODEC
 
-    bloom_filter_options = get_bloom_filter_options(file_schema, table_properties)
+    id_to_parquet_path = {field_id: parquet_path for parquet_path, field_id in parquet_column_mapping.items()}
+    bloom_filter_options = get_bloom_filter_options(file_schema, table_properties, id_to_parquet_path)
 
     return {
         "compression": compression_codec,

tests/table/test_metadata.py

@@ -26,7 +26,7 @@
 import pytest
 
 from pyiceberg.exceptions import ValidationError
-from pyiceberg.io.pyarrow import get_bloom_filter_options
+from pyiceberg.io.pyarrow import get_bloom_filter_options, parquet_path_to_id_mapping
 from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.serializers import FromByteStream
@@ -907,7 +907,8 @@ def test_get_bloom_filter_options() -> None:
         "write.parquet.bloom-filter-ndv.column.qux.quux": "3000",
     }
 
-    bloom_filter_options = get_bloom_filter_options(schema, table_properties)
+    id_to_parquet_path = {field_id: parquet_path for parquet_path, field_id in parquet_path_to_id_mapping(schema).items()}
+    bloom_filter_options = get_bloom_filter_options(schema, table_properties, id_to_parquet_path)
     assert bloom_filter_options == {
         "foo": {"fpp": 0.01, "ndv": 1000},
         "qux.quux": {"fpp": 0.03, "ndv": 3000},