diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index df7aaf824f..73d4c032e9 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -18,6 +18,7 @@ import itertools import logging +import math import random import time import uuid @@ -35,7 +36,7 @@ from pyiceberg.encryption.kms import KeyManagementClient 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, @@ -422,12 +423,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] diff --git a/tests/catalog/test_catalog_behaviors.py b/tests/catalog/test_catalog_behaviors.py index 1ed9ea4096..6d694a9448 100644 --- a/tests/catalog/test_catalog_behaviors.py +++ b/tests/catalog/test_catalog_behaviors.py @@ -21,7 +21,7 @@ import json import os -from collections.abc import Generator +from collections.abc import Callable, Generator from pathlib import Path from typing import Any @@ -43,7 +43,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 @@ -1376,6 +1376,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}" diff --git a/tests/table/test_commit_retry.py b/tests/table/test_commit_retry.py index ab95457e23..96e85b8c2b 100644 --- a/tests/table/test_commit_retry.py +++ b/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") diff --git a/tests/table/test_init.py b/tests/table/test_init.py index db6dc00e39..a866e6ffad 100644 --- a/tests/table/test_init.py +++ b/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 @@ -2267,6 +2269,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