Partition-scoped position deletes are not applied in NaN identity partitions #4279

import os
import tempfile

import pyarrow as pa
import pyarrow.parquet as pq

from pyiceberg.catalog.sql import SqlCatalog
from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestContent, ManifestWriterV2
from pyiceberg.partitioning import PartitionField, PartitionSpec
from pyiceberg.schema import Schema
from pyiceberg.transforms import IdentityTransform
from pyiceberg.typedef import Record
from pyiceberg.types import DoubleType, IntegerType, NestedField

a, b = Record(float("nan")), Record(float("nan"))
print(f"Record(nan) == Record(nan): {a == b}; hashes equal: {hash(a) == hash(b)}")

warehouse = tempfile.mkdtemp()
catalog = SqlCatalog("default", uri=f"sqlite:///{warehouse}/catalog.db", warehouse=f"file://{warehouse}")
catalog.create_namespace("default")
schema = Schema(NestedField(1, "id", IntegerType()), NestedField(2, "value", DoubleType()))
spec = PartitionSpec(PartitionField(source_id=2, field_id=1000, transform=IdentityTransform(), name="value"))
tbl = catalog.create_table("default.t", schema=schema, partition_spec=spec, properties={"format-version": "2"})

pos_delete_schema = pa.schema(
    [
        pa.field("file_path", pa.string(), nullable=False, metadata={"PARQUET:field_id": "2147483546"}),
        pa.field("pos", pa.int64(), nullable=False, metadata={"PARQUET:field_id": "2147483545"}),
    ]
)


def write_file(name: str, rows: pa.Table, value: float, content: DataFileContent) -> DataFile:
    path = f"{warehouse}/{name}.parquet"
    pq.write_table(rows, path)
    return DataFile.from_args(
        content=content,
        file_path=path,
        file_format=FileFormat.PARQUET,
        partition=Record(value),
        record_count=rows.num_rows,
        file_size_in_bytes=os.path.getsize(path),
    )


class DeleteManifestWriter(ManifestWriterV2):
    def content(self) -> ManifestContent:
        return ManifestContent.DELETES

    @property
    def _meta(self) -> dict[str, str]:  # private: the manifest header names its content
        return {**super()._meta, "content": "deletes"}


# Write files the way Spark or Java would: PyIceberg 0.12 cannot append to a NaN partition, and PyIceberg
# cannot write delete files. Each partition gets two data files and one position delete file that removes
# row 0 of both. A delete file covering two data files has no file_path bounds, so it is partition-scoped.
data_files, delete_files = [], []
for value, ids in [(1.0, [1, 2, 3, 4]), (float("nan"), [5, 6, 7, 8])]:
    for pair in (ids[:2], ids[2:]):
        rows = pa.Table.from_pydict({"id": pair, "value": [value] * 2}, schema=schema.as_arrow())
        data_files.append(write_file(f"data-{pair[0]}", rows, value, DataFileContent.DATA))
    deletes = pa.Table.from_pydict({"file_path": [f.file_path for f in data_files[-2:]], "pos": [0, 0]}, schema=pos_delete_schema)
    delete_files.append(write_file(f"deletes-{ids[0]}", deletes, value, DataFileContent.POSITION_DELETES))

with tbl.transaction() as tx, tx.update_snapshot().fast_append() as append:
    for data_file in data_files:
        append.append_data_file(data_file)

with tbl.transaction() as tx, tx.update_snapshot().overwrite() as overwrite:
    # private: PyIceberg has no delete-manifest writer, so swap one into the snapshot producer
    overwrite.new_manifest_writer = lambda spec: DeleteManifestWriter(
        spec, tbl.schema(), overwrite.new_manifest_output(), overwrite._snapshot_id, "deflate"
    )
    for delete_file in delete_files:
        overwrite.append_data_file(delete_file)

print("expected ids:", [2, 4, 6, 8])
print("actual ids:  ", sorted(tbl.scan().to_arrow()["id"].to_pylist()))