Cache residuals with stateless evaluators

pyiceberg/expressions/visitors.py

@@ -1785,36 +1785,28 @@ 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):
-    """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.
-    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 residuals expressions
-    for the partition data, d:
-
-
-    1. If d > day(a) and d &lt; 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
-    Partition data is passed using StructLike. Residuals are returned by residualFor(StructLike).
-    """
+class _ResidualEvaluationVisitor(BoundBooleanExpressionVisitor[BooleanExpression]):
+    """Evaluate a residual expression for one partition."""
 
     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 +1923,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 +1941,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,6 +1970,46 @@ 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)

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,
 )
@@ -100,7 +102,7 @@
 from pyiceberg.types import strtobool
 from pyiceberg.utils.concurrent import ExecutorFactory
 from pyiceberg.utils.config import Config
-from pyiceberg.utils.properties import property_as_bool
+from pyiceberg.utils.properties import property_as_bool, property_as_int
 
 if TYPE_CHECKING:
     import bodo.pandas as bd
@@ -128,6 +130,9 @@ class UpsertResult:
 
 
 class TableProperties:
+    RESIDUAL_CACHE_MAX_SIZE = "read.residual-cache.max-size"
+    RESIDUAL_CACHE_MAX_SIZE_DEFAULT = 128
+
     PARQUET_ROW_GROUP_SIZE_BYTES = "write.parquet.row-group-size-bytes"
     PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT = 128 * 1024 * 1024  # 128 MB
 
@@ -2620,7 +2625,39 @@ 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_max_size = property_as_int(
+            self.options,
+            TableProperties.RESIDUAL_CACHE_MAX_SIZE,
+            TableProperties.RESIDUAL_CACHE_MAX_SIZE_DEFAULT,
+        )
+        if residual_cache_max_size is None or residual_cache_max_size <= 0:
+            raise ValueError(f"{TableProperties.RESIDUAL_CACHE_MAX_SIZE} must be a positive integer")
+        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 +2681,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 +2719,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)"
+    )