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,
@@ -420,12 +421,16 @@ def _build_partition_predicate(self, partition_records: set[Record], partition_f
 
         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)
-            ]
+            predicates: list[BooleanExpression] = []
+            for pos, partition_field in enumerate(partition_fields):
+                ref = Reference(partition_field)
+                value = partition_record[pos]
+                if value is None:
+                    predicates.append(IsNull(ref))
+                elif isinstance(value, float) and math.isnan(value):
+                    predicates.append(IsNaN(ref))
+                else:
+                    predicates.append(EqualTo(ref, value))
             per_record_exprs.append(And(*predicates) if len(predicates) > 1 else predicates[0])
 
         return Or(*per_record_exprs) if len(per_record_exprs) > 1 else per_record_exprs[0]

pyiceberg/typedef.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import math
 from abc import abstractmethod
 from collections.abc import Callable
 from datetime import date, datetime, time
@@ -173,6 +174,10 @@ class IcebergRootModel(RootModel[T], Generic[T]):
     model_config = ConfigDict(frozen=True)
 
 
+def _is_nan(value: Any) -> bool:
+    return isinstance(value, float) and math.isnan(value)
+
+
 class Record(StructProtocol):
     __slots__ = ("_data",)
     _data: list[Any]
@@ -196,8 +201,19 @@ def __getitem__(self, pos: int) -> Any:
         return self._data[pos]
 
     def __eq__(self, other: Any) -> bool:
-        """Return the equality of two instances of the Record class."""
-        return self._data == other._data if isinstance(other, Record) else False
+        """Compare records, treating NaN floats at the same position as equal."""
+        if not isinstance(other, Record):
+            return False
+        if self._data == other._data:
+            return True
+        if len(self._data) != len(other._data):
+            return False
+        for a, b in zip(self._data, other._data, strict=True):
+            if a is b or a == b:
+                continue
+            if not (_is_nan(a) and _is_nan(b)):
+                return False
+        return True
 
     def __repr__(self) -> str:
         """Return the string representation of the Record class."""

tests/catalog/test_catalog_behaviors.py

@@ -20,7 +20,7 @@
 """
 
 import os
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from pathlib import Path
 from typing import Any
 
@@ -42,7 +42,7 @@
 from pyiceberg.io.pyarrow import _dataframe_to_data_files, schema_to_pyarrow
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
-from pyiceberg.table import TableProperties
+from pyiceberg.table import Table, TableProperties
 from pyiceberg.table.snapshots import Operation
 from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder
 from pyiceberg.table.update import AddSchemaUpdate, SetCurrentSchemaUpdate
@@ -1345,6 +1345,45 @@ 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 _nan_partitioned_table(catalog: Catalog, name: str) -> Table:
+    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(f"default.{name}_{catalog.name}", schema=iceberg_schema, partition_spec=partition_spec)
+    tbl.append(_id_value_rows(tbl, [1, 2, 3, 4], [1.0, 1.0, float("nan"), float("nan")]))
+    return tbl
+
+
+def _id_value_rows(tbl: Table, ids: list[int], values: list[float]) -> pa.Table:
+    return pa.Table.from_pydict({"id": ids, "value": values}, schema=schema_to_pyarrow(tbl.schema()))
+
+
+def test_replace_files_in_nan_identity_partition(catalog: Catalog) -> None:
+    catalog.create_namespace("default")
+    cases: list[tuple[str, Callable[[Table], None], list[int], list[int]]] = [
+        ("delete_nan_file", lambda t: t.delete("value is nan"), [1, 2], []),
+        ("rewrite_nan_file", lambda t: t.delete("id = 3"), [1, 2, 4], [4]),
+        ("overwrite", lambda t: t.overwrite(_id_value_rows(t, [9], [2.0])), [9], []),
+        (
+            "dynamic_overwrite",
+            lambda t: t.dynamic_partition_overwrite(_id_value_rows(t, [9], [float("nan")])),
+            [1, 2, 9],
+            [9],
+        ),
+    ]
+    for name, operation, expected_ids, expected_nan_ids in cases:
+        tbl = _nan_partitioned_table(catalog, name)
+
+        operation(tbl)
+
+        assert sorted(tbl.scan().to_arrow()["id"].to_pylist()) == expected_ids, name
+        assert sorted(tbl.scan(row_filter="value is nan").to_arrow()["id"].to_pylist()) == expected_nan_ids, name
+
+
 def test_scan_integer_column_with_decimal_literal(catalog: Catalog) -> None:
     catalog.create_namespace("default")
     identifier = f"default.scan_integer_decimal_literal_{catalog.name}"

tests/table/test_commit_retry.py

@@ -23,9 +23,10 @@
 from pyiceberg.catalog import Catalog
 from pyiceberg.exceptions import CommitFailedException, CommitStateUnknownException, ValidationException
 from pyiceberg.schema import Schema
-from pyiceberg.table import TableProperties, Transaction
+from pyiceberg.table import Table, TableProperties, Transaction
 from pyiceberg.table.snapshots import IsolationLevel, Operation
-from pyiceberg.types import LongType, NestedField, StringType
+from pyiceberg.typedef import Record
+from pyiceberg.types import DoubleType, LongType, NestedField, StringType
 
 
 def test_isolation_level_enum() -> None:
@@ -395,6 +396,85 @@ def test_file_overwrite_validates_concurrent_file_delete(
     assert sorted(result["value"].to_pylist()) == expected_values
 
 
+def _commit_partition_position_delete(table: Table, partition: Record, file_path: str, pos: int) -> None:
+    """PyIceberg cannot write delete files, so commit one without path bounds, as other engines write them."""
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestContent, ManifestWriter, ManifestWriterV2
+    from pyiceberg.partitioning import PartitionSpec
+    from pyiceberg.table.update.snapshot import _OverwriteFiles
+
+    class DeleteManifestWriter(ManifestWriterV2):
+        def content(self) -> ManifestContent:
+            return ManifestContent.DELETES
+
+        @property
+        def _meta(self) -> dict[str, str]:
+            return {**super()._meta, "content": "deletes"}
+
+    def new_delete_manifest_writer(producer: _OverwriteFiles, spec: PartitionSpec) -> ManifestWriter:
+        return DeleteManifestWriter(
+            spec, producer.schema(), producer.new_manifest_output(), producer._snapshot_id, producer._compression
+        )
+
+    delete_path = table.location_provider().new_data_location(f"pos-delete-{uuid.uuid4()}.parquet")
+    with table.io.new_output(delete_path).create() as output:
+        pq.write_table(pa.table({"file_path": [file_path], "pos": [pos]}), output)
+    delete_file = DataFile.from_args(
+        content=DataFileContent.POSITION_DELETES,
+        file_path=delete_path,
+        file_format=FileFormat.PARQUET,
+        partition=partition,
+        record_count=1,
+        file_size_in_bytes=len(table.io.new_input(delete_path)),
+        spec_id=table.spec().spec_id,
+    )
+
+    with patch.object(_OverwriteFiles, "new_manifest_writer", new_delete_manifest_writer):
+        with table.transaction() as transaction, transaction.update_snapshot().overwrite() as overwrite:
+            overwrite.append_data_file(delete_file)
+
+
+def test_file_overwrite_validates_concurrent_position_delete_in_nan_partition(catalog: Catalog) -> None:
+    """A file replacement in a NaN partition must fail when a delete for that file was committed concurrently."""
+    import pyarrow as pa
+
+    from pyiceberg.io.pyarrow import _dataframe_to_data_files
+    from pyiceberg.partitioning import PartitionField, PartitionSpec
+    from pyiceberg.transforms import IdentityTransform
+
+    catalog.create_namespace("default")
+    schema = Schema(
+        NestedField(1, "id", LongType(), required=False),
+        NestedField(2, "value", DoubleType(), required=False),
+    )
+    spec = PartitionSpec(PartitionField(source_id=2, field_id=1000, transform=IdentityTransform(), name="value"))
+    identifier = "default.concurrent_nan_position_delete"
+    table = catalog.create_table(identifier, schema=schema, partition_spec=spec, properties={"format-version": "2"})
+    rows = pa.table({"id": pa.array([1, 2], pa.int64()), "value": [float("nan"), float("nan")]})
+    table.append(rows)
+    file_to_replace = next(task.file for task in table.scan().plan_files())
+
+    replacing_table = catalog.load_table(identifier)
+    deleting_table = catalog.load_table(identifier)
+
+    replacement_file = list(
+        _dataframe_to_data_files(table_metadata=replacing_table.metadata, df=rows, io=replacing_table.io, write_uuid=uuid.uuid4())
+    )[0]
+    replacing_transaction = replacing_table.transaction()
+    with replacing_transaction.update_snapshot().overwrite() as overwrite:
+        overwrite.delete_data_file(file_to_replace)
+        overwrite.append_data_file(replacement_file)
+
+    _commit_partition_position_delete(deleting_table, Record(float("nan")), file_to_replace.file_path, pos=0)
+
+    with pytest.raises(ValidationException, match="found new delete for replace data file"):
+        replacing_transaction.commit_transaction()
+
+    assert catalog.load_table(identifier).scan().to_arrow()["id"].to_pylist() == [2]
+
+
 def test_concurrent_overwrite_append_retries_successfully(catalog: Catalog) -> None:
     """Append after a concurrent overwrite should succeed via retry."""
     catalog.create_namespace("default")

tests/table/test_delete_file_index.py

@@ -233,3 +233,10 @@ def test_record_equality_for_partition_lookup() -> None:
 
     assert len(index.for_data_file(1, data_file, partition_b)) == 1
     assert len(index.for_data_file(1, data_file, partition_c)) == 0
+
+
+def test_record_equality_for_nan_partition_lookup() -> None:
+    index = DeleteFileIndex()
+    index.add_delete_file(_create_partition_delete(sequence_number=2, partition=Record(float("nan"))), Record(float("nan")))
+
+    assert len(index.for_data_file(1, _create_data_file(), Record(float("nan")))) == 1

tests/table/test_init.py

@@ -33,6 +33,8 @@
     And,
     EqualTo,
     In,
+    IsNaN,
+    IsNull,
 )
 from pyiceberg.expressions.visitors import bind
 from pyiceberg.io import PY_IO_IMPL, FileIO, load_file_io
@@ -2177,6 +2179,16 @@ def test_build_large_partition_predicate(table_v2: Table) -> None:
     bind(table_v2.metadata.schema(), expr, case_sensitive=True)
 
 
+def test_build_partition_predicate_with_nan_and_null(table_v2: Table) -> None:
+    with table_v2.transaction() as tx:
+        expr = tx._build_partition_predicate(
+            partition_records={Record(float("nan"), None, 1.0)},
+            partition_fields=["a", "b", "c"],
+        )
+
+    assert expr == And(IsNaN("a"), IsNull("b"), EqualTo("c", 1.0))
+
+
 def test_overwrite_delete_data_file_on_bucket_partition(catalog: Catalog) -> None:
     """Delete a data file without projecting its already-transformed partition value again."""
     import pyarrow as pa

tests/table/test_validate.py

@@ -30,11 +30,13 @@
     _added_delete_files,
     _deleted_data_files,
     _validate_added_data_files,
+    _validate_data_files_exist,
     _validate_deleted_data_files,
     _validate_no_new_delete_files,
     _validate_no_new_deletes_for_data_files,
     _validation_history,
 )
+from pyiceberg.typedef import Record
 
 
 @pytest.fixture
@@ -235,6 +237,37 @@ class DummyEntry:
             )
 
 
+def test_validate_data_files_exist_raises_on_conflict_in_nan_partition(
+    table_v2_with_extensive_snapshots_and_manifests: tuple[Table, dict[int, list[ManifestFile]]],
+) -> None:
+    table, mock_manifests = table_v2_with_extensive_snapshots_and_manifests
+    oldest_snapshot = table.snapshots()[0]
+    newest_snapshot = cast(Snapshot, table.current_snapshot())
+
+    snapshots = table.snapshots()
+    snapshots[-2] = snapshots[-2].model_copy(update={"summary": Summary(operation=Operation.DELETE)})
+    table.metadata = table.metadata.model_copy(update={"snapshots": snapshots})
+
+    def data_file() -> DataFile:
+        return DataFile.from_args(file_path="s3://bucket/data.parquet", partition=Record(float("nan")), spec_id=0)
+
+    deleted_entry = ManifestEntry.from_args(
+        status=ManifestEntryStatus.DELETED, snapshot_id=snapshots[-2].snapshot_id, data_file=data_file()
+    )
+
+    with (
+        patch("pyiceberg.table.snapshots.Snapshot.manifests", new=lambda self, io: mock_manifests.get(self.snapshot_id, [])),
+        patch("pyiceberg.manifest.ManifestFile.fetch_manifest_entry", return_value=[deleted_entry]),
+    ):
+        with pytest.raises(ValidationException, match="Data files were concurrently deleted"):
+            _validate_data_files_exist(
+                table=table,
+                starting_snapshot=newest_snapshot,
+                data_files={data_file()},
+                parent_snapshot=oldest_snapshot,
+            )
+
+
 @pytest.mark.parametrize("operation", [Operation.APPEND, Operation.OVERWRITE])
 def test_validate_added_data_files_conflicting_count(
     table_v2_with_extensive_snapshots_and_manifests: tuple[Table, dict[int, list[ManifestFile]]],

tests/test_typedef.py

@@ -65,3 +65,12 @@ def computed(self) -> int:
 
     with pytest.raises(TypeError, match="Unexpected RecordWithProperty fields: computed"):
         RecordWithProperty._bind(StructType(), computed=1)
+
+
+def test_record_equality_with_nan() -> None:
+    assert Record(1, float("nan")) == Record(1, float("nan"))
+    assert Record(Record(float("nan"))) == Record(Record(float("nan")))
+    assert len({Record(float("nan")), Record(float("nan"))}) == 1
+
+    assert Record(float("nan")) != Record(1.0)
+    assert Record(1, float("nan")) != Record(2, float("nan"))