[python] Keep files with NaN float stats in stats-based pruning

paimon-python/pypaimon/common/predicate.py

@@ -21,6 +21,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from decimal import Decimal
+from numbers import Real
 from typing import Any, ClassVar, Dict, List, Optional, Sequence, cast
 
 import pyarrow
@@ -309,6 +310,15 @@ def _null_literal(literals, arity=1) -> bool:
         literal is None for literal in literals[:arity])
 
 
+def _is_nan(value) -> bool:
+    return isinstance(value, Real) and value != value
+
+
+def _has_nan(min_v, max_v, literals, arity) -> bool:
+    # NaN stats cannot safely prune with IEEE row comparisons.
+    return _is_nan(min_v) or _is_nan(max_v) or any(_is_nan(v) for v in (literals or [])[:arity])
+
+
 class Equal(Tester):
     name = 'equal'
     literal_arity = 1
@@ -319,6 +329,8 @@ def test_by_value(self, val, literals) -> bool:
         return val == literals[0]
 
     def test_by_stats(self, min_v, max_v, literals) -> bool:
+        if _has_nan(min_v, max_v, literals, 1):
+            return True
         if _null_literal(literals):
             return False
         return min_v <= literals[0] <= max_v
@@ -355,6 +367,8 @@ def test_by_value(self, val, literals) -> bool:
         return val < literals[0]
 
     def test_by_stats(self, min_v, max_v, literals) -> bool:
+        if _has_nan(min_v, max_v, literals, 1):
+            return True
         if _null_literal(literals):
             return False
         return literals[0] > min_v
@@ -373,6 +387,8 @@ def test_by_value(self, val, literals) -> bool:
         return val <= literals[0]
 
     def test_by_stats(self, min_v, max_v, literals) -> bool:
+        if _has_nan(min_v, max_v, literals, 1):
+            return True
         if _null_literal(literals):
             return False
         return literals[0] >= min_v
@@ -391,6 +407,8 @@ def test_by_value(self, val, literals) -> bool:
         return val > literals[0]
 
     def test_by_stats(self, min_v, max_v, literals) -> bool:
+        if _has_nan(min_v, max_v, literals, 1):
+            return True
         if _null_literal(literals):
             return False
         return literals[0] < max_v
@@ -409,6 +427,8 @@ def test_by_value(self, val, literals) -> bool:
         return val >= literals[0]
 
     def test_by_stats(self, min_v, max_v, literals) -> bool:
+        if _has_nan(min_v, max_v, literals, 1):
+            return True
         if _null_literal(literals):
             return False
         return literals[0] <= max_v
@@ -426,8 +446,10 @@ def test_by_value(self, val, literals) -> bool:
         return val in literals
 
     def test_by_stats(self, min_v, max_v, literals) -> bool:
+        if _is_nan(min_v) or _is_nan(max_v):
+            return True
         return any(
-            min_v <= literal <= max_v
+            _is_nan(literal) or min_v <= literal <= max_v
             for literal in literals
             if literal is not None
         )
@@ -471,6 +493,8 @@ def test_by_value(self, val, literals) -> bool:
         return literals[0] <= val <= literals[1]
 
     def test_by_stats(self, min_v, max_v, literals) -> bool:
+        if _has_nan(min_v, max_v, literals, 2):
+            return True
         if _null_literal(literals, 2):
             return False
         return literals[0] <= max_v and literals[1] >= min_v
@@ -567,6 +591,8 @@ def test_by_value(self, val, literals) -> bool:
         return val < literals[0] or val > literals[1]
 
     def test_by_stats(self, min_v, max_v, literals) -> bool:
+        if _has_nan(min_v, max_v, literals, 2):
+            return True
         if _null_literal(literals, 2):
             return False
         return literals[0] > min_v or literals[1] < max_v

paimon-python/pypaimon/tests/predicates_test.py

@@ -29,8 +29,9 @@
 from pypaimon import CatalogFactory, Schema
 from pypaimon.common.predicate import Predicate
 from pypaimon.manifest.schema.simple_stats import SimpleStats
-from pypaimon.schema.data_types import DataField
-from pypaimon.table.row.generic_row import GenericRow, GenericRowDeserializer
+from pypaimon.schema.data_types import AtomicType, DataField
+from pypaimon.table.row.binary_row import BinaryRow
+from pypaimon.table.row.generic_row import GenericRow, GenericRowDeserializer, GenericRowSerializer
 from pypaimon.table.row.offset_row import OffsetRow
 from pypaimon.table.row.projected_row import ProjectedRow
 
@@ -47,6 +48,58 @@ def _random_format():
     return random.choice(['parquet', 'avro', 'orc'])
 
 
+@pytest.mark.parametrize('data_type', ['FLOAT', 'DOUBLE'])
+@pytest.mark.parametrize('method,arity,finite_result', [
+    ('equal', 1, True), ('notEqual', 1, False),
+    ('lessThan', 1, False), ('lessOrEqual', 1, True),
+    ('greaterThan', 1, False), ('greaterOrEqual', 1, True),
+    ('in', 2, True), ('notIn', 2, False),
+    ('between', 2, True), ('notBetween', 2, False),
+])
+def test_by_simple_stats_with_nan(data_type, method, arity, finite_result):
+    fields = [DataField(0, 'd', AtomicType(data_type))]
+    cases = [([-0.0, -0.0] + [0.0] * arity, finite_result),
+             ([float('nan')] * (arity + 2), True)]
+    for index in range(arity + 2):
+        values = [0.0, 1.0] + [0.0] * arity
+        values[index] = float('nan')
+        cases.append((values, True))
+    for values, expected in cases:
+        min_row, max_row = (
+            BinaryRow(GenericRowSerializer.to_bytes(GenericRow([v], fields)), fields) for v in values[:2])
+        stats = SimpleStats(min_row, max_row, [0])
+        predicate = Predicate(method, 0, 'd', values[2:])
+        assert predicate.test_by_simple_stats(stats, 1) == expected, values
+
+
+@pytest.mark.python_plan
+def test_append_read_with_nan_stats(tmp_path):
+    catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+    catalog.create_database('default', False)
+    pa_schema = pa.schema([('d', pa.float64())])
+    catalog.create_table('default.t', Schema.from_pyarrow_schema(
+        pa_schema, options={'file.format': 'avro', 'metadata.stats-mode': 'full'}), False)
+    table = catalog.get_table('default.t')
+    builder = table.new_batch_write_builder()
+    write, commit = builder.new_write(), builder.new_commit()
+    try:
+        write.write_arrow(pa.Table.from_pydict({'d': [1.0, float('nan'), 3.0]}, schema=pa_schema))
+        messages = write.prepare_commit()
+        # Java's full stats collector includes NaN as the maximum.
+        messages[0].new_files[0].value_stats = SimpleStats(
+            GenericRow([1.0], table.fields), GenericRow([float('nan')], table.fields), [0])
+        commit.commit(messages)
+    finally:
+        write.close()
+        commit.close()
+
+    read = table.new_read_builder()
+    read.with_filter(read.new_predicate_builder().greater_than('d', 2.0))
+    splits = read.new_scan().plan().splits()
+    assert len(splits) == 1
+    assert read.new_read().to_arrow(splits).column('d').to_pylist() == [3.0]
+
+
 class PredicateTest(unittest.TestCase):
 
     def test_large_boolean_arrow_expression_is_balanced(self):