PyArrow: Collect NaN value counts when writing Parquet files

pyiceberg/io/pyarrow.py

@@ -2574,6 +2574,7 @@ def data_file_statistics_from_parquet_metadata(
     parquet_metadata: pq.FileMetaData,
     stats_columns: dict[int, StatisticsCollector],
     parquet_column_mapping: dict[str, int],
+    nan_value_counts: dict[int, int] | None = None,
 ) -> DataFileStatistics:
     """
     Compute and return DataFileStatistics that includes the following.
@@ -2591,13 +2592,15 @@ def data_file_statistics_from_parquet_metadata(
         stats_columns (Dict[int, StatisticsCollector]): The statistics gathering plan. It is required to
             set the mode for column metrics collection
         parquet_column_mapping (Dict[str, int]): The mapping of the parquet file name to the field ID
+        nan_value_counts (Optional[Dict[int, int]]): The NaN counts collected from the written data,
+            as Parquet statistics do not track NaN values
     """
     column_sizes: dict[int, int] = {}
     value_counts: dict[int, int] = {}
     split_offsets: list[int] = []
 
     null_value_counts: dict[int, int] = {}
-    nan_value_counts: dict[int, int] = {}
+    nan_value_counts = dict(nan_value_counts) if nan_value_counts else {}
 
     col_aggs = {}
 
@@ -2689,6 +2692,40 @@ def data_file_statistics_from_parquet_metadata(
     )
 
 
+def _nan_value_counts(table: pa.Table, schema: Schema, stats_columns: dict[int, StatisticsCollector]) -> dict[int, int]:
+    """Count the NaN values of the float and double columns that collect metrics."""
+    nan_field_ids = {
+        field_id
+        for field_id, stats_col in stats_columns.items()
+        if isinstance(stats_col.iceberg_type, (FloatType, DoubleType)) and stats_col.mode != MetricsMode(MetricModeTypes.NONE)
+    }
+    nan_value_counts: dict[int, int] = {}
+    if not nan_field_ids:
+        return nan_value_counts
+
+    def count(field_id: int, field_type: IcebergType, array: pa.Array) -> None:
+        if isinstance(field_type, StructType):
+            for nested_field in field_type.fields:
+                if array.type.get_field_index(nested_field.name) >= 0:
+                    # struct_field also applies the validity of the parent struct
+                    count(nested_field.field_id, nested_field.field_type, pc.struct_field(array, nested_field.name))
+        elif isinstance(field_type, ListType):
+            count(field_type.element_id, field_type.element_type, array.flatten())
+        elif isinstance(field_type, MapType):
+            entries = array.cast(pa.list_(pa.struct([array.type.key_field, array.type.item_field]))).flatten()
+            count(field_type.key_id, field_type.key_type, entries.field(0))
+            count(field_type.value_id, field_type.value_type, entries.field(1))
+        elif field_id in nan_field_ids and pa.types.is_floating(array.type):
+            nan_value_counts[field_id] = nan_value_counts.get(field_id, 0) + (pc.sum(pc.is_nan(array)).as_py() or 0)
+
+    for field in schema.fields:
+        if table.schema.get_field_index(field.name) >= 0:
+            for chunk in table.column(field.name).chunks:
+                count(field.field_id, field.field_type, chunk)
+
+    return nan_value_counts
+
+
 class ParquetFormatWriter(FileFormatWriter):
     """Writes Arrow tables to a Parquet file."""
 
@@ -2698,6 +2735,8 @@ def __init__(self, output_file: OutputFile, file_schema: Schema, properties: Pro
         self._properties = properties
         self._writer: pq.ParquetWriter | None = None
         self._fos: OutputStream | None = None
+        self._stats_columns = compute_statistics_plan(file_schema, properties)
+        self._nan_value_counts: dict[int, int] = {}
         self._parquet_writer_kwargs = _get_parquet_writer_kwargs(properties)
         self._row_group_size = property_as_int(
             properties=properties,
@@ -2720,6 +2759,9 @@ def write(self, table: pa.Table) -> None:
                 raise
             self._fos = fos
         self._writer.write(table, row_group_size=self._row_group_size)
+        # Parquet statistics do not track NaN values, so count them while writing
+        for field_id, nan_count in _nan_value_counts(table, self._file_schema, self._stats_columns).items():
+            self._nan_value_counts[field_id] = self._nan_value_counts.get(field_id, 0) + nan_count
 
     def close(self) -> DataFileStatistics:
         if self._result is not None:
@@ -2730,8 +2772,9 @@ def close(self) -> DataFileStatistics:
             self._writer.close()
             self._result = data_file_statistics_from_parquet_metadata(
                 parquet_metadata=self._writer.writer.metadata,
-                stats_columns=compute_statistics_plan(self._file_schema, self._properties),
+                stats_columns=self._stats_columns,
                 parquet_column_mapping=parquet_path_to_id_mapping(self._file_schema),
+                nan_value_counts=self._nan_value_counts,
             )
         return self._result
 

tests/catalog/test_catalog_behaviors.py

@@ -1345,6 +1345,27 @@ def test_append_nan_to_identity_partitioned_table(catalog: Catalog) -> None:
     assert sorted(tbl.scan(row_filter="value is nan").to_arrow()["id"].to_pylist()) == [2, 4]
 
 
+def test_delete_keeps_nan_rows_that_do_not_match_the_filter(catalog: Catalog) -> None:
+    catalog.create_namespace("default")
+    identifier = f"default.delete_keeps_nan_rows_{catalog.name}"
+    iceberg_schema = Schema(
+        NestedField(1, "id", IntegerType(), required=False),
+        NestedField(2, "value", DoubleType(), required=False),
+    )
+    tbl = catalog.create_table(identifier=identifier, schema=iceberg_schema)
+    tbl.append(
+        pa.Table.from_pydict(
+            {"id": [1, 2], "value": [1.0, float("nan")]},
+            schema=schema_to_pyarrow(iceberg_schema),
+        )
+    )
+
+    # NaN < 5.0 is false, so the NaN row must survive and the file cannot be dropped as a whole
+    tbl.delete("value < 5.0")
+
+    assert tbl.scan().to_arrow()["id"].to_pylist() == [2]
+
+
 def test_record_batch_reader_consumed_exactly_once(catalog: Catalog) -> None:
     """The streaming path must consume the underlying generator exactly once.
     A regression that drained the reader twice (e.g. an extra .schema access

tests/integration/test_writes/test_writes.py

@@ -1687,7 +1687,7 @@ def test_merge_manifests_file_content(session_catalog: Catalog, arrow_table_with
             (11, b"\x01"),
             (12, b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"),
         ]
-        assert tbl_a_data_file["nan_value_counts"] == []
+        assert tbl_a_data_file["nan_value_counts"] == [(6, 0), (7, 0)]
         assert tbl_a_data_file["null_value_counts"] == [
             (1, 1),
             (2, 1),

tests/io/test_format_writers.py

@@ -24,10 +24,10 @@
 import pytest
 
 from pyiceberg.io.fileformat import FileFormatFactory, FileFormatModel
-from pyiceberg.io.pyarrow import PyArrowFileIO
+from pyiceberg.io.pyarrow import PyArrowFileIO, schema_to_pyarrow
 from pyiceberg.manifest import FileFormat
 from pyiceberg.schema import Schema
-from pyiceberg.types import LongType, NestedField
+from pyiceberg.types import DoubleType, FloatType, ListType, LongType, MapType, NestedField, StringType, StructType
 
 
 @pytest.fixture(params=FileFormatFactory.available_formats(), ids=lambda f: f.name.lower())
@@ -73,6 +73,57 @@ def test_null_handling(format_model: FileFormatModel, table_schema_simple: Schem
     assert stats.null_value_counts.get(1) == 1
 
 
+def test_nan_value_counts(format_model: FileFormatModel, tmp_path: Path) -> None:
+    """Float and double columns produce nan_value_counts, including nested and empty counts."""
+    nan = float("nan")
+    schema = Schema(
+        NestedField(1, "id", LongType(), required=False),
+        NestedField(2, "f", FloatType(), required=False),
+        NestedField(3, "d", DoubleType(), required=False),
+        NestedField(4, "no_nans", DoubleType(), required=False),
+        NestedField(5, "s", StructType(NestedField(6, "x", DoubleType(), required=False)), required=False),
+        NestedField(7, "l", ListType(element_id=8, element_type=DoubleType(), element_required=False), required=False),
+        NestedField(
+            9,
+            "m",
+            MapType(key_id=10, key_type=StringType(), value_id=11, value_type=DoubleType(), value_required=False),
+            required=False,
+        ),
+    )
+    table = pa.Table.from_pylist(
+        [
+            {"id": 1, "f": nan, "d": nan, "no_nans": 1.0, "s": {"x": nan}, "l": [nan, 1.0], "m": {"a": nan}},
+            {"id": 2, "f": 1.0, "d": nan, "no_nans": None, "s": None, "l": None, "m": None},
+            {"id": 3, "f": None, "d": 2.0, "no_nans": 2.0, "s": {"x": 1.0}, "l": [nan, nan], "m": {"b": 1.0, "c": nan}},
+        ],
+        schema=schema_to_pyarrow(schema),
+    )
+    file_path = str(tmp_path / f"test.{format_model.file_extension()}")
+    writer = format_model.create_writer(PyArrowFileIO().new_output(file_path), schema, {})
+    writer.write(table)
+    writer.write(table.slice(0, 1))
+    stats = writer.close()
+
+    assert stats.nan_value_counts == {2: 2, 3: 3, 4: 0, 6: 2, 8: 4, 11: 3}
+
+
+def test_nan_value_counts_respect_metrics_mode(format_model: FileFormatModel, tmp_path: Path) -> None:
+    """Columns with metrics mode none do not get a NaN count."""
+    schema = Schema(
+        NestedField(1, "a", DoubleType(), required=False),
+        NestedField(2, "b", DoubleType(), required=False),
+    )
+    table = pa.Table.from_pydict({"a": [float("nan")], "b": [float("nan")]}, schema=schema_to_pyarrow(schema))
+    file_path = str(tmp_path / f"test.{format_model.file_extension()}")
+    writer = format_model.create_writer(
+        PyArrowFileIO().new_output(file_path), schema, {"write.metadata.metrics.column.b": "none"}
+    )
+    writer.write(table)
+    stats = writer.close()
+
+    assert stats.nan_value_counts == {1: 1}
+
+
 def test_context_manager_caches_result(
     format_model: FileFormatModel, table_schema_simple: Schema, arrow_table_simple: pa.Table, tmp_path: Path
 ) -> None: