diff --git a/paimon-python/pypaimon/common/predicate.py b/paimon-python/pypaimon/common/predicate.py index 80f1d0b91edc..5e69bdef30e4 100644 --- a/paimon-python/pypaimon/common/predicate.py +++ b/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 diff --git a/paimon-python/pypaimon/tests/predicates_test.py b/paimon-python/pypaimon/tests/predicates_test.py index 684ea038d54e..ca40118a4061 100644 --- a/paimon-python/pypaimon/tests/predicates_test.py +++ b/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):