Fix scans dropping NaN rows for IsNaN, NotEqualTo and NotIn

pyiceberg/io/pyarrow.py

@@ -71,7 +71,18 @@
 from pyiceberg.conversions import to_bytes
 from pyiceberg.encryption.key_metadata import StandardKeyMetadata
 from pyiceberg.exceptions import ResolveError
-from pyiceberg.expressions import AlwaysTrue, BooleanExpression, BoundIsNaN, BoundIsNull, BoundTerm, Not, Or
+from pyiceberg.expressions import (
+    AlwaysTrue,
+    And,
+    BooleanExpression,
+    BoundIsNaN,
+    BoundIsNull,
+    BoundNotEqualTo,
+    BoundNotIn,
+    BoundTerm,
+    Not,
+    Or,
+)
 from pyiceberg.expressions.literals import Literal
 from pyiceberg.expressions.visitors import (
     BoundBooleanExpressionVisitor,
@@ -1734,6 +1745,27 @@ def _get_column_projection_values(
     return projected_missing_fields
 
 
+def _has_nan_sensitive_predicate(expr: BooleanExpression) -> bool:
+    """Return whether a predicate in the expression can match NaN values that pyarrow may prune away.
+
+    Parquet statistics exclude NaN, so a row group whose non-NaN values are all equal reports min == max.
+    pyarrow then prunes it for predicates that the NaN values satisfy: IsNaN, NotEqualTo and NotIn on a
+    float or double column.
+    """
+    if isinstance(expr, (And, Or)):
+        return _has_nan_sensitive_predicate(expr.left) or _has_nan_sensitive_predicate(expr.right)
+    return isinstance(expr, (BoundIsNaN, BoundNotEqualTo, BoundNotIn)) and isinstance(
+        expr.term.ref().field.field_type, (FloatType, DoubleType)
+    )
+
+
+def _nan_safe_conjuncts(expr: BooleanExpression) -> BooleanExpression:
+    """Return the top-level conjuncts that pyarrow can prune on without dropping NaN rows, joined with AND."""
+    if isinstance(expr, And):
+        return And(_nan_safe_conjuncts(expr.left), _nan_safe_conjuncts(expr.right))
+    return AlwaysTrue() if _has_nan_sensitive_predicate(expr) else expr
+
+
 def _task_to_record_batches(
     io: FileIO,
     task: FileScanTask,
@@ -1774,21 +1806,36 @@ def _task_to_record_batches(
         )
 
         pyarrow_filter = None
+        scanner_filter = None
+        # Set when the scanner does not apply the whole filter, so each batch is filtered after the read
+        filter_batches = False
         if bound_row_filter is not AlwaysTrue():
             translated_row_filter = translate_column_names(
                 bound_row_filter, file_schema, case_sensitive=case_sensitive, projected_field_values=projected_missing_fields
             )
             bound_file_filter = bind(file_schema, translated_row_filter, case_sensitive=case_sensitive)
             pyarrow_filter = expression_to_pyarrow(bound_file_filter, file_schema)
+            if positional_deletes:
+                # Positional deletes must be applied before any row is filtered
+                filter_batches = True
+            elif _has_nan_sensitive_predicate(bound_file_filter):
+                # Parquet statistics exclude NaN, so pyarrow can prune a row group that holds NaN rows matching
+                # a NaN-sensitive conjunct (https://github.com/apache/arrow/issues/51491). Only the other
+                # conjuncts are pushed down.
+                filter_batches = True
+                nan_safe_filter = _nan_safe_conjuncts(bound_file_filter)
+                if nan_safe_filter is not AlwaysTrue():
+                    scanner_filter = expression_to_pyarrow(nan_safe_filter, file_schema)
+            else:
+                scanner_filter = pyarrow_filter
 
         file_project_schema = prune_columns(file_schema, projected_field_ids, select_full_types=False)
 
         fragment_scanner = ds.Scanner.from_fragment(
             fragment=fragment,
             schema=physical_schema,
             # This will push down the query to Arrow.
-            # But in case there are positional deletes, we have to apply them first
-            filter=pyarrow_filter if not positional_deletes else None,
+            filter=scanner_filter,
             columns=[col.name for col in file_project_schema.columns],
         )
 
@@ -1803,16 +1850,17 @@ def _task_to_record_batches(
                 # Create the mask of indices that we're interested in
                 indices = _combine_positional_deletes(positional_deletes, current_index, current_index + len(batch))
                 current_batch = current_batch.take(indices)
-                if pyarrow_filter is not None:
-                    # Temporary fix until PyArrow 21 is the minimum supported version
-                    # (https://github.com/apache/arrow/pull/46057): RecordBatch.filter raises
-                    # IndexError on PyArrow <21 when the result is empty; Table.filter does not.
-                    table = pa.Table.from_batches([current_batch])
-                    table = table.filter(pyarrow_filter)
-                    if table.num_rows == 0:
-                        current_batch = current_batch.slice(0, 0)
-                    else:
-                        current_batch = table.combine_chunks().to_batches()[0]
+
+            if filter_batches and pyarrow_filter is not None:
+                # Temporary fix until PyArrow 21 is the minimum supported version
+                # (https://github.com/apache/arrow/pull/46057): RecordBatch.filter raises
+                # IndexError on PyArrow <21 when the result is empty; Table.filter does not.
+                table = pa.Table.from_batches([current_batch])
+                table = table.filter(pyarrow_filter)
+                if table.num_rows == 0:
+                    current_batch = current_batch.slice(0, 0)
+                else:
+                    current_batch = table.combine_chunks().to_batches()[0]
 
             # skip empty batches
             if current_batch.num_rows == 0:

tests/catalog/test_catalog_behaviors.py

@@ -1415,6 +1415,28 @@ def test_replace_files_in_nan_identity_partition(catalog: Catalog) -> None:
         assert sorted(tbl.scan(row_filter="value is nan").to_arrow()["id"].to_pylist()) == expected_nan_ids, name
 
 
+def test_scan_nan_row_with_equal_finite_statistics(catalog: Catalog) -> None:
+    catalog.create_namespace("default")
+    identifier = f"default.scan_nan_equal_statistics_{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)
+    # Parquet statistics ignore NaN, so the single row group reports min == max == 5.0 for this file.
+    tbl.append(
+        pa.Table.from_pydict(
+            {"id": [1, 2, 3], "value": [float("nan"), 5.0, 5.0]},
+            schema=schema_to_pyarrow(iceberg_schema),
+        )
+    )
+
+    assert tbl.scan(row_filter="value is nan").to_arrow()["id"].to_pylist() == [1]
+    assert tbl.scan(row_filter="value != 5.0").to_arrow()["id"].to_pylist() == [1]
+    assert tbl.scan(row_filter="value not in (5.0, 6.0)").to_arrow()["id"].to_pylist() == [1]
+    assert tbl.scan(row_filter="value is not nan").to_arrow()["id"].to_pylist() == [2, 3]
+
+
 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/io/test_pyarrow.py

@@ -69,8 +69,13 @@
     BoundStartsWith,
     EqualTo,
     GreaterThan,
+    IsNaN,
     IsNull,
+    LessThan,
     Not,
+    NotEqualTo,
+    NotIn,
+    NotNaN,
     Or,
 )
 from pyiceberg.expressions.literals import literal
@@ -4303,6 +4308,62 @@ def test_task_to_record_batches_filter_after_positional_deletes_empty_result(tmp
     assert result_batches == []
 
 
+def test_task_to_record_batches_keeps_nan_rows_for_nan_sensitive_filters(tmpdir: str) -> None:
+    """Regression test: pyarrow must not prune a row group that holds NaN rows matching the filter.
+
+    Parquet statistics exclude NaN, so the row group [NaN, 5.0, 5.0] reports min == max == 5.0.
+    pyarrow prunes such a row group for IsNaN, NotEqualTo and NotIn, which drops the NaN row.
+    NotNaN is a control.
+    """
+    from pyiceberg.expressions.visitors import bind
+
+    for arrow_type, iceberg_type in [(pa.float32(), FloatType()), (pa.float64(), DoubleType())]:
+        arrow_schema = pa.schema(
+            (
+                pa.field("id", pa.int32(), nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"}),
+                pa.field("x", arrow_type, nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "2"}),
+            )
+        )
+        arrow_table = pa.table(
+            {"id": pa.array([0, 1, 2], type=pa.int32()), "x": pa.array([float("nan"), 5.0, 5.0], type=arrow_type)},
+            schema=arrow_schema,
+        )
+        data_file = _write_table_to_data_file(
+            f"{tmpdir}/nan_row_group_equal_bounds_{iceberg_type}.parquet", arrow_schema, arrow_table
+        )
+
+        table_schema = Schema(
+            NestedField(1, "id", IntegerType(), required=False),
+            NestedField(2, "x", iceberg_type, required=False),
+        )
+        projected_schema = Schema(NestedField(1, "id", IntegerType(), required=False))
+
+        cases: list[tuple[BooleanExpression, list[int]]] = [
+            (IsNaN("x"), [0]),
+            (NotEqualTo("x", 5.0), [0]),
+            (NotIn("x", {5.0, 6.0}), [0]),
+            (And(IsNaN("x"), EqualTo("id", 0)), [0]),
+            (Or(IsNaN("x"), EqualTo("id", 2)), [0, 2]),
+            (NotNaN("x"), [1, 2]),
+        ]
+        for row_filter, expected_ids in cases:
+            # The filter column is not selected, so the filter must still see it after the read.
+            result_batches = list(
+                _task_to_record_batches(
+                    PyArrowFileIO(),
+                    FileScanTask(data_file),
+                    bound_row_filter=bind(table_schema, row_filter, case_sensitive=True),
+                    projected_schema=projected_schema,
+                    table_schema=table_schema,
+                    projected_field_ids={1, 2},
+                    positional_deletes=None,
+                    case_sensitive=True,
+                )
+            )
+            actual_ids = [value for batch in result_batches for value in batch.column(0).to_pylist()]
+            assert actual_ids == expected_ids, f"{iceberg_type}: {row_filter}"
+
+
 def test_task_to_record_batches_encrypted_by_iceberg_rust() -> None:
     """The data file was encrypted by iceberg-rust, see tests/table/encryption/README.md."""
     data_file = DataFile.from_args(
@@ -4559,7 +4620,7 @@ def test_orc_row_filtering_predicate_pushdown(tmp_path: Path) -> None:
     import pyarrow as pa
     import pyarrow.orc as orc
 
-    from pyiceberg.expressions import And, EqualTo, GreaterThan, In, LessThan, Or
+    from pyiceberg.expressions import And, EqualTo, GreaterThan, In, Or
     from pyiceberg.io.pyarrow import ArrowScan, PyArrowFileIO
     from pyiceberg.manifest import DataFileContent, FileFormat
     from pyiceberg.partitioning import PartitionSpec