Preserve residual visitor compatibility

Author: dossett

pyiceberg/expressions/visitors.py

@@ -1785,28 +1785,37 @@ def _can_contain_nans(self, field_id: int) -> bool:
         return (nan_count := self.nan_counts.get(field_id)) is not None and nan_count > 0
 
 
-class _ResidualEvaluationVisitor(BoundBooleanExpressionVisitor[BooleanExpression]):
-    """Evaluate a residual expression for one partition."""
+class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC):
+    """Find residuals for an expression using partition values.
+
+    A residual expression is made by partially evaluating an expression using partition values.
+    For example, if a table is partitioned by day(utc_timestamp) and is read with a filter expression
+    utc_timestamp > a and utc_timestamp < b, then there are 4 possible residual expressions
+    for the partition data, d:
+
+    1. If d > day(a) and d < day(b), the residual is always true
+    2. If d == day(a) and d != day(b), the residual is utc_timestamp > a
+    3. If d == day(b) and d != day(a), the residual is utc_timestamp < b
+    4. If d == day(a) == day(b), the residual is utc_timestamp > a and utc_timestamp < b
+    """
 
     schema: Schema
     spec: PartitionSpec
     case_sensitive: bool
+    expr: BooleanExpression
     partition_schema: Schema
     struct: Record
 
-    def __init__(
-        self,
-        schema: Schema,
-        spec: PartitionSpec,
-        case_sensitive: bool,
-        partition_schema: Schema,
-        partition_data: Record,
-    ) -> None:
+    def __init__(self, schema: Schema, spec: PartitionSpec, case_sensitive: bool, expr: BooleanExpression) -> None:
         self.schema = schema
         self.spec = spec
         self.case_sensitive = case_sensitive
-        self.partition_schema = partition_schema
+        self.expr = expr
+        self.partition_schema = Schema(*spec.partition_type(schema).fields)
+
+    def eval(self, partition_data: Record) -> BooleanExpression:
         self.struct = partition_data
+        return visit(self.expr, visitor=self)
 
     def visit_true(self) -> BooleanExpression:
         return AlwaysTrue()
@@ -1844,10 +1853,10 @@ def visit_is_nan(self, term: BoundTerm) -> BooleanExpression:
 
     def visit_not_nan(self, term: BoundTerm) -> BooleanExpression:
         val = term.eval(self.struct)
-        if isinstance(val, SupportsFloat) and not math.isnan(val):
-            return self.visit_true()
-        else:
+        if isinstance(val, SupportsFloat) and math.isnan(val):
             return self.visit_false()
+        else:
+            return self.visit_true()
 
     def visit_less_than(self, term: BoundTerm, literal: LiteralValue) -> BooleanExpression:
         if term.eval(self.struct) < literal.value:
@@ -1970,46 +1979,6 @@ def visit_unbound_predicate(self, predicate: UnboundPredicate) -> BooleanExpress
         return bound
 
 
-class ResidualVisitor:
-    """Find residuals for an expression using partition values.
-
-    A residual expression is made by partially evaluating an expression using partition values.
-    For example, if a table is partitioned by day(utc_timestamp) and is read with a filter expression
-    utc_timestamp > a and utc_timestamp < b, then there are 4 possible residual expressions
-    for the partition data, d:
-
-    1. If d > day(a) and d < day(b), the residual is always true
-    2. If d == day(a) and d != day(b), the residual is utc_timestamp > a
-    3. If d == day(b) and d != day(a), the residual is utc_timestamp < b
-    4. If d == day(a) == day(b), the residual is utc_timestamp > a and utc_timestamp < b
-    """
-
-    schema: Schema
-    spec: PartitionSpec
-    case_sensitive: bool
-    expr: BooleanExpression
-    partition_schema: Schema
-
-    def __init__(self, schema: Schema, spec: PartitionSpec, case_sensitive: bool, expr: BooleanExpression) -> None:
-        self.schema = schema
-        self.spec = spec
-        self.case_sensitive = case_sensitive
-        self.expr = expr
-        self.partition_schema = Schema(*spec.partition_type(schema).fields)
-
-    def eval(self, partition_data: Record) -> BooleanExpression:
-        return visit(
-            self.expr,
-            visitor=_ResidualEvaluationVisitor(
-                schema=self.schema,
-                spec=self.spec,
-                case_sensitive=self.case_sensitive,
-                partition_schema=self.partition_schema,
-                partition_data=partition_data,
-            ),
-        )
-
-
 class ResidualEvaluator(ResidualVisitor):
     def residual_for(self, partition_data: Record) -> BooleanExpression:
         return self.eval(partition_data)

tests/expressions/test_residual_evaluator.py

@@ -15,16 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint:disable=redefined-outer-name
-from concurrent.futures import ThreadPoolExecutor
-from threading import Event
-from typing import Any
-
 import pytest
 
 from pyiceberg.expressions import (
     AlwaysFalse,
     AlwaysTrue,
     And,
+    BooleanExpression,
     EqualTo,
     GreaterThan,
     GreaterThanOrEqual,
@@ -45,7 +42,7 @@
 from pyiceberg.schema import Schema
 from pyiceberg.transforms import DayTransform, IdentityTransform
 from pyiceberg.typedef import Record
-from pyiceberg.types import DoubleType, FloatType, IntegerType, NestedField, StringType, StructType, TimestampType
+from pyiceberg.types import DoubleType, FloatType, IntegerType, NestedField, StringType, TimestampType
 
 
 def test_identity_transform_residual() -> None:
@@ -92,102 +89,25 @@ def test_identity_transform_residual() -> None:
     assert residual == AlwaysFalse()
 
 
-def test_residual_evaluator_does_not_mutate_prepared_state() -> None:
-    schema = Schema(NestedField(1, "a", IntegerType()), NestedField(2, "b", IntegerType()))
-    spec = PartitionSpec(
-        PartitionField(1, 1001, IdentityTransform(), "a_part"),
-        PartitionField(2, 1002, IdentityTransform(), "b_part"),
-    )
-    evaluator = residual_evaluator_of(
-        spec=spec,
-        expr=And(EqualTo("a", 1), EqualTo("b", 1)),
-        case_sensitive=True,
-        schema=schema,
-    )
-    initial_state = vars(evaluator).copy()
-
-    assert evaluator.residual_for(Record(1, 1)) == AlwaysTrue()
-    assert evaluator.residual_for(Record(0, 0)) == AlwaysFalse()
-    assert evaluator.residual_for(Record(1, 1)) == AlwaysTrue()
-
-    assert isinstance(evaluator, ResidualVisitor)
-    assert vars(evaluator) == initial_state
-
-
 def test_residual_visitor_preserves_public_eval_api() -> None:
     schema = Schema(NestedField(1, "a", IntegerType()))
     spec = PartitionSpec(PartitionField(1, 1001, IdentityTransform(), "a_part"))
     visitor = ResidualVisitor(schema=schema, spec=spec, case_sensitive=True, expr=EqualTo("a", 1))
-    initial_state = vars(visitor).copy()
 
     assert visitor.eval(Record(1)) == AlwaysTrue()
     assert visitor.eval(Record(0)) == AlwaysFalse()
-    assert vars(visitor) == initial_state
-
-
-def test_residual_evaluator_concurrent_calls_do_not_share_partitions() -> None:
-    class BlockingRecord(Record):
-        def __init__(self, first_read: Event, release_first_read: Event, *values: Any) -> None:
-            super().__init__(*values)
-            self.first_read = first_read
-            self.release_first_read = release_first_read
-
-        def __getitem__(self, pos: int) -> Any:
-            value = super().__getitem__(pos)
-            if pos == 0:
-                self.first_read.set()
-                if not self.release_first_read.wait(timeout=5):
-                    raise TimeoutError("Timed out waiting to interleave residual evaluations")
-            return value
-
-    schema = Schema(NestedField(1, "a", IntegerType()), NestedField(2, "b", IntegerType()))
-    spec = PartitionSpec(
-        PartitionField(1, 1001, IdentityTransform(), "a_part"),
-        PartitionField(2, 1002, IdentityTransform(), "b_part"),
-    )
-    evaluator = residual_evaluator_of(
-        spec=spec,
-        expr=And(EqualTo("a", 1), EqualTo("b", 1)),
-        case_sensitive=True,
-        schema=schema,
-    )
-    first_read = Event()
-    release_first_read = Event()
 
-    with ThreadPoolExecutor(max_workers=2) as executor:
-        matching_result = executor.submit(
-            evaluator.residual_for,
-            BlockingRecord(first_read, release_first_read, 1, 1),
-        )
-        assert first_read.wait(timeout=5)
 
-        try:
-            non_matching_result = executor.submit(evaluator.residual_for, Record(0, 0)).result(timeout=5)
-        finally:
-            release_first_read.set()
+def test_residual_visitor_subclass_can_customize_evaluation() -> None:
+    class FalseForTrueResidualVisitor(ResidualVisitor):
+        def visit_true(self) -> BooleanExpression:
+            return AlwaysFalse()
 
-        assert matching_result.result(timeout=5) == AlwaysTrue()
-        assert non_matching_result == AlwaysFalse()
-
-
-def test_partition_schema_reused_across_residuals(monkeypatch: pytest.MonkeyPatch) -> None:
-    schema = Schema(NestedField(50, "dateint", IntegerType()))
-    spec = PartitionSpec(PartitionField(50, 1050, IdentityTransform(), "dateint_part"))
-    partition_type_calls = 0
-    original_partition_type = PartitionSpec.partition_type
-
-    def counting_partition_type(self: PartitionSpec, schema: Schema) -> StructType:
-        nonlocal partition_type_calls
-        partition_type_calls += 1
-        return original_partition_type(self, schema)
-
-    monkeypatch.setattr(PartitionSpec, "partition_type", counting_partition_type)
-
-    evaluator = residual_evaluator_of(spec=spec, expr=EqualTo("dateint", 20170815), case_sensitive=True, schema=schema)
+    schema = Schema(NestedField(1, "a", IntegerType()))
+    spec = PartitionSpec(PartitionField(1, 1001, IdentityTransform(), "a_part"))
+    visitor = FalseForTrueResidualVisitor(schema=schema, spec=spec, case_sensitive=True, expr=AlwaysTrue())
 
-    assert evaluator.residual_for(Record(20170815)) == AlwaysTrue()
-    assert evaluator.residual_for(Record(20170816)) == AlwaysFalse()
-    assert partition_type_calls == 1
+    assert visitor.eval(Record(1)) == AlwaysFalse()
 
 
 def test_case_insensitive_identity_transform_residuals() -> None:
@@ -313,6 +233,9 @@ def test_is_not_nan() -> None:
     res_eval = residual_evaluator_of(spec=spec, expr=predicate, case_sensitive=True, schema=schema)
 
     residual = res_eval.residual_for(Record(None))
+    assert residual == AlwaysTrue()
+
+    residual = res_eval.residual_for(Record(float("nan")))
     assert residual == AlwaysFalse()
 
     residual = res_eval.residual_for(Record(2))
@@ -325,6 +248,9 @@ def test_is_not_nan() -> None:
     res_eval = residual_evaluator_of(spec=spec, expr=predicate, case_sensitive=True, schema=schema)
 
     residual = res_eval.residual_for(Record(None))
+    assert residual == AlwaysTrue()
+
+    residual = res_eval.residual_for(Record(float("nan")))
     assert residual == AlwaysFalse()
 
     residual = res_eval.residual_for(Record(2))