Match NaN partition values with IsNaN in partition predicates

pyiceberg/table/__init__.py

@@ -18,6 +18,7 @@
 
 import itertools
 import logging
+import math
 import random
 import time
 import uuid
@@ -34,7 +35,7 @@
 from pydantic import Field
 
 from pyiceberg.exceptions import CommitFailedException, ValidationException
-from pyiceberg.expressions import AlwaysFalse, AlwaysTrue, And, BooleanExpression, EqualTo, IsNull, Or, Reference
+from pyiceberg.expressions import AlwaysFalse, AlwaysTrue, And, BooleanExpression, EqualTo, IsNaN, IsNull, Or, Reference
 from pyiceberg.expressions.visitors import (
     ResidualEvaluator,
     _InclusiveMetricsEvaluator,
@@ -418,13 +419,18 @@ def _build_partition_predicate(self, partition_records: set[Record], partition_f
         if not partition_records or not partition_fields:
             return AlwaysFalse()
 
+        def _matches(partition_field: str, value: Any) -> BooleanExpression:
+            if value is None:
+                return IsNull(Reference(partition_field))
+            # Expression literals reject NaN, so match a NaN partition value with IsNaN
+            if isinstance(value, float) and math.isnan(value):
+                return IsNaN(Reference(partition_field))
+            return EqualTo(Reference(partition_field), value)
+
         per_record_exprs: list[BooleanExpression] = []
         for partition_record in partition_records:
             predicates: list[BooleanExpression] = [
-                EqualTo(Reference(partition_field), partition_record[pos])
-                if partition_record[pos] is not None
-                else IsNull(Reference(partition_field))
-                for pos, partition_field in enumerate(partition_fields)
+                _matches(partition_field, partition_record[pos]) for pos, partition_field in enumerate(partition_fields)
             ]
             per_record_exprs.append(And(*predicates) if len(predicates) > 1 else predicates[0])
 

tests/catalog/test_catalog_behaviors.py

@@ -1345,6 +1345,36 @@ 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_and_overwrite_nan_identity_partition(catalog: Catalog) -> None:
+    catalog.create_namespace("default")
+    identifier = f"default.delete_nan_identity_partition_{catalog.name}"
+    iceberg_schema = Schema(
+        NestedField(1, "id", IntegerType(), required=False),
+        NestedField(2, "value", DoubleType(), required=False),
+    )
+    partition_spec = PartitionSpec(
+        PartitionField(source_id=2, field_id=1000, transform=IdentityTransform(), name="value"),
+    )
+    tbl = catalog.create_table(identifier=identifier, schema=iceberg_schema, partition_spec=partition_spec)
+    arrow_schema = schema_to_pyarrow(iceberg_schema)
+    tbl.append(pa.Table.from_pydict({"id": [1, 2, 3, 4], "value": [1.0, 1.0, float("nan"), float("nan")]}, schema=arrow_schema))
+
+    # Rewrites the NaN partition's data file
+    tbl.delete("id = 3")
+    assert sorted(tbl.scan().to_arrow()["id"].to_pylist()) == [1, 2, 4]
+
+    # Drops the NaN partition's data file
+    tbl.delete("value is nan")
+    assert sorted(tbl.scan().to_arrow()["id"].to_pylist()) == [1, 2]
+
+    tbl.append(pa.Table.from_pydict({"id": [5], "value": [float("nan")]}, schema=arrow_schema))
+    tbl.dynamic_partition_overwrite(pa.Table.from_pydict({"id": [6], "value": [float("nan")]}, schema=arrow_schema))
+    assert sorted(tbl.scan().to_arrow()["id"].to_pylist()) == [1, 2, 6]
+
+    tbl.overwrite(pa.Table.from_pydict({"id": [7], "value": [2.0]}, schema=arrow_schema))
+    assert tbl.scan().to_arrow()["id"].to_pylist() == [7]
+
+
 def test_scan_integer_column_with_decimal_literal(catalog: Catalog) -> None:
     catalog.create_namespace("default")
     identifier = f"default.scan_integer_decimal_literal_{catalog.name}"