Cache residuals with stateless evaluators

pyiceberg/expressions/visitors.py

@@ -1785,7 +1785,7 @@ 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 ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC):
+class _ResidualEvaluationVisitor(BoundBooleanExpressionVisitor[BooleanExpression]):
     """Finds the residuals for an Expression the partitions in the given PartitionSpec.
 
     A residual expression is made by partially evaluating an expression using partition values.
@@ -1804,17 +1804,22 @@ class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC):
     schema: Schema
     spec: PartitionSpec
     case_sensitive: bool
-    expr: BooleanExpression
+    partition_schema: Schema
+    struct: Record
 
-    def __init__(self, schema: Schema, spec: PartitionSpec, case_sensitive: bool, expr: BooleanExpression) -> None:
+    def __init__(
+        self,
+        schema: Schema,
+        spec: PartitionSpec,
+        case_sensitive: bool,
+        partition_schema: Schema,
+        partition_data: Record,
+    ) -> None:
         self.schema = schema
         self.spec = spec
         self.case_sensitive = case_sensitive
-        self.expr = expr
-
-    def eval(self, partition_data: Record) -> BooleanExpression:
+        self.partition_schema = partition_schema
         self.struct = partition_data
-        return visit(self.expr, visitor=self)
 
     def visit_true(self) -> BooleanExpression:
         return AlwaysTrue()
@@ -1931,17 +1936,12 @@ def visit_bound_predicate(self, predicate: BoundPredicate) -> BooleanExpression:
         if parts == []:
             return predicate
 
-        def struct_to_schema(struct: StructType) -> Schema:
-            return Schema(*struct.fields)
-
         for part in parts:
             strict_projection = part.transform.strict_project(part.name, predicate)
             strict_result = None
 
             if strict_projection is not None:
-                bound = strict_projection.bind(
-                    struct_to_schema(self.spec.partition_type(self.schema)), case_sensitive=self.case_sensitive
-                )
+                bound = strict_projection.bind(self.partition_schema, case_sensitive=self.case_sensitive)
                 if isinstance(bound, BoundPredicate):
                     strict_result = super().visit_bound_predicate(bound)
                 else:
@@ -1954,9 +1954,7 @@ def struct_to_schema(struct: StructType) -> Schema:
             inclusive_projection = part.transform.project(part.name, predicate)
             inclusive_result = None
             if inclusive_projection is not None:
-                bound_inclusive = inclusive_projection.bind(
-                    struct_to_schema(self.spec.partition_type(self.schema)), case_sensitive=self.case_sensitive
-                )
+                bound_inclusive = inclusive_projection.bind(self.partition_schema, case_sensitive=self.case_sensitive)
                 if isinstance(bound_inclusive, BoundPredicate):
                     # using predicate method specific to inclusive
                     inclusive_result = super().visit_bound_predicate(bound_inclusive)
@@ -1985,9 +1983,33 @@ def visit_unbound_predicate(self, predicate: UnboundPredicate) -> BooleanExpress
         return bound
 
 
-class ResidualEvaluator(ResidualVisitor):
+class ResidualEvaluator:
+    """Prepare residual evaluation once while keeping partition state local to each call."""
+
+    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 residual_for(self, partition_data: Record) -> BooleanExpression:
-        return self.eval(partition_data)
+        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 UnpartitionedResidualEvaluator(ResidualEvaluator):

pyiceberg/table/__init__.py

@@ -28,6 +28,7 @@
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, TypeVar
 
+from cachetools import LRUCache
 from pydantic import Field
 
 import pyiceberg.expressions.parser as parser
@@ -37,6 +38,7 @@
     _InclusiveMetricsEvaluator,
     bind,
     expression_evaluator,
+    extract_field_ids,
     inclusive_projection,
     manifest_evaluator,
 )
@@ -117,6 +119,9 @@
 
 ALWAYS_TRUE = AlwaysTrue()
 DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE = "downcast-ns-timestamp-to-us-on-write"
+# Retain a small working set for repeated relevant partition values without adding
+# unbounded key storage when scans contain a distinct value for every data file.
+_RESIDUAL_CACHE_MAX_SIZE = 128
 
 
 @dataclass()
@@ -2620,7 +2625,32 @@ def plan_files(
         data_entries: list[ManifestEntry] = []
         delete_index = DeleteFileIndex()
 
-        residual_evaluators: dict[int, Callable[[DataFile], ResidualEvaluator]] = KeyDefaultDict(self._build_residual_evaluator)
+        residual_evaluators: dict[int, ResidualEvaluator] = KeyDefaultDict(self._build_residual_evaluator)
+        referenced_field_ids = extract_field_ids(
+            bind(self.table_metadata.schema(), self.row_filter, case_sensitive=self.case_sensitive)
+        )
+        partition_specs = self.table_metadata.specs()
+        residual_cache_key_positions: dict[int, tuple[int, ...]] = KeyDefaultDict(
+            lambda spec_id: tuple(
+                pos
+                for pos, partition_field in enumerate(partition_specs[spec_id].fields)
+                if partition_field.source_id in referenced_field_ids
+            )
+        )
+        # A residual can only depend on partition fields derived from source columns
+        # referenced by the scan filter. Keep the cache local and bounded.
+        residual_cache: LRUCache[tuple[int, tuple[Any, ...]], BooleanExpression] = LRUCache(maxsize=_RESIDUAL_CACHE_MAX_SIZE)
+
+        def residual_for(data_file: DataFile) -> BooleanExpression:
+            partition = data_file.partition
+            partition_values = tuple(partition[pos] for pos in residual_cache_key_positions[data_file.spec_id])
+            cache_key = data_file.spec_id, partition_values
+            try:
+                return residual_cache[cache_key]
+            except KeyError:
+                residual = residual_evaluators[data_file.spec_id].residual_for(partition)
+                residual_cache[cache_key] = residual
+                return residual
 
         for manifest_entry in chain.from_iterable(self.plan_manifest_entries(manifests)):
             if not manifest_entry_filter(manifest_entry):
@@ -2644,9 +2674,7 @@ def plan_files(
                     data_entry.data_file,
                     partition_key=data_entry.data_file.partition,
                 ),
-                residual=residual_evaluators[data_entry.data_file.spec_id](data_entry.data_file).residual_for(
-                    data_entry.data_file.partition
-                ),
+                residual=residual_for(data_entry.data_file),
             )
             for data_entry in data_entries
         ]
@@ -2684,15 +2712,12 @@ def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]:
             include_empty_files,
         ).eval(data_file)
 
-    def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], ResidualEvaluator]:
+    def _build_residual_evaluator(self, spec_id: int) -> ResidualEvaluator:
         spec = self.table_metadata.specs()[spec_id]
 
         from pyiceberg.expressions.visitors import residual_evaluator_of
 
-        # The lambda created here is run in multiple threads.
-        # So we avoid creating _EvaluatorExpression methods bound to a single
-        # shared instance across multiple threads.
-        return lambda datafile: residual_evaluator_of(
+        return residual_evaluator_of(
             spec=spec,
             expr=self.row_filter,
             case_sensitive=self.case_sensitive,

tests/benchmark/test_residual_evaluator_benchmark.py

@@ -0,0 +1,118 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Benchmark residual planning with a realistic 15-leaf predicate.
+
+Every file has a unique unreferenced partition-hash value. The repeated case
+measures cache reuse by relevant partition values, while the unique case forces
+cache misses.
+
+Run with:
+    uv run pytest tests/benchmark/test_residual_evaluator_benchmark.py -v -s -m benchmark
+"""
+
+from __future__ import annotations
+
+import statistics
+import timeit
+
+import pytest
+
+from pyiceberg.expressions import And, BooleanExpression, EqualTo, GreaterThanOrEqual, LessThanOrEqual, Or
+from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus
+from pyiceberg.partitioning import PartitionField, PartitionSpec
+from pyiceberg.schema import Schema
+from pyiceberg.table import ManifestGroupPlanner, Table
+from pyiceberg.table.metadata import TableMetadataV2
+from pyiceberg.transforms import IdentityTransform
+from pyiceberg.typedef import Record
+from pyiceberg.types import LongType, NestedField
+
+
+def _row_filter() -> BooleanExpression:
+    """Select five day ranges, each scoped to a region."""
+    windows = ((0, 1, 1), (2, 3, 4), (4, 5, 7), (6, 7, 10), (8, 10, 13))
+    branches = [
+        And(
+            And(GreaterThanOrEqual("event_day", start_day), LessThanOrEqual("event_day", end_day)),
+            EqualTo("region_id", region_id),
+        )
+        for start_day, end_day, region_id in windows
+    ]
+
+    combined = branches[0]
+    for branch in branches[1:]:
+        combined = Or(combined, branch)
+    return combined
+
+
+def _manifest_entry(file_number: int, relevant_partition: int) -> ManifestEntry:
+    data_file = DataFile.from_args(
+        content=DataFileContent.DATA,
+        file_path=f"s3://bucket/data-{file_number}.parquet",
+        file_format=FileFormat.PARQUET,
+        partition=Record(relevant_partition, file_number),
+        record_count=1,
+        file_size_in_bytes=1,
+    )
+    data_file.spec_id = 0
+    return ManifestEntry.from_args(
+        status=ManifestEntryStatus.ADDED,
+        snapshot_id=1,
+        sequence_number=1,
+        file_sequence_number=1,
+        data_file=data_file,
+    )
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    "num_relevant_partitions",
+    [7, 2_000],
+    ids=["repeated-relevant-partitions", "unique-relevant-partitions"],
+)
+def test_residual_planning(table_v2: Table, monkeypatch: pytest.MonkeyPatch, num_relevant_partitions: int) -> None:
+    num_files = 2_000
+    entries = [_manifest_entry(file_number, file_number % num_relevant_partitions) for file_number in range(num_files)]
+    schema = Schema(
+        NestedField(1, "event_day", LongType(), required=True),
+        NestedField(2, "region_id", LongType(), required=True),
+        NestedField(3, "partition_hash", LongType(), required=True),
+    )
+    spec = PartitionSpec(
+        PartitionField(1, 1000, IdentityTransform(), "event_day"),
+        PartitionField(3, 1001, IdentityTransform(), "partition_hash"),
+        spec_id=0,
+    )
+    metadata = TableMetadataV2(
+        location="s3://bucket/table",
+        last_column_id=3,
+        schemas=[schema],
+        current_schema_id=schema.schema_id,
+        partition_specs=[spec],
+        default_spec_id=spec.spec_id,
+    )
+    planner = ManifestGroupPlanner(table_metadata=metadata, io=table_v2.io, row_filter=_row_filter())
+
+    monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries]))
+
+    timings = timeit.repeat(lambda: list(planner.plan_files([])), number=1, repeat=3)
+
+    assert len(list(planner.plan_files([]))) == num_files
+    print(
+        f"Planned {num_files} files across {num_relevant_partitions} relevant partitions "
+        f"with a 15-leaf predicate in {statistics.mean(timings):.3f}s (best: {min(timings):.3f}s)"
+    )