perf: cache residuals during scan planning

pyiceberg/expressions/visitors.py

@@ -1805,12 +1805,14 @@ class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC):
     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:
         self.struct = partition_data
@@ -1931,17 +1933,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 +1951,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)

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,35 @@ def plan_files(
         data_entries: list[ManifestEntry] = []
         delete_index = DeleteFileIndex()
 
-        residual_evaluators: dict[int, Callable[[DataFile], ResidualEvaluator]] = KeyDefaultDict(self._build_residual_evaluator)
+        residual_evaluator_factories: dict[int, Callable[[DataFile], 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: tuple[int, tuple[Any, ...]] = data_file.spec_id, partition_values
+            try:
+                return residual_cache[cache_key]
+            except KeyError:
+                residual_evaluator = residual_evaluator_factories[data_file.spec_id](data_file)
+                residual = residual_evaluator.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 +2677,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
         ]
@@ -2689,10 +2720,9 @@ def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], Residu
 
         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(
+        # ResidualEvaluator stores the current partition while evaluating. Return a
+        # factory so every cache miss uses a fresh, unshared evaluator instance.
+        return lambda _: residual_evaluator_of(
             spec=spec,
             expr=self.row_filter,
             case_sensitive=self.case_sensitive,

tests/benchmark/test_scan_planning_benchmark.py

@@ -0,0 +1,109 @@
+# 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 evaluation with a high-cardinality irrelevant partition field.
+
+Every file has a unique unreferenced partition-hash value. The repeated case
+measures reuse by relevant partition values, while the unique case forces misses.
+
+Run with:
+    uv run pytest tests/benchmark/test_scan_planning_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.table import ManifestGroupPlanner, Table
+from pyiceberg.transforms import IdentityTransform
+from pyiceberg.typedef import Record
+
+
+def _combined_filter() -> BooleanExpression:
+    branches: list[BooleanExpression] = []
+    for source in range(11):
+        branch: BooleanExpression = GreaterThanOrEqual("x", 0)
+        for predicate in (
+            LessThanOrEqual("x", 6),
+            EqualTo("y", source),
+            EqualTo("y", source + 1),
+            EqualTo("y", source + 2),
+            EqualTo("y", source + 3),
+        ):
+            branch = And(branch, predicate)
+        branches.append(branch)
+
+    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)]
+    spec = PartitionSpec(
+        PartitionField(1, 1000, IdentityTransform(), "x"),
+        PartitionField(3, 1001, IdentityTransform(), "partition_hash"),
+        spec_id=0,
+    )
+    metadata = table_v2.metadata.model_copy(update={"partition_specs": [spec]})
+    planner = ManifestGroupPlanner(
+        table_metadata=metadata,
+        io=table_v2.io,
+        row_filter=_combined_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 in "
+        f"{statistics.mean(timings):.3f}s (best: {min(timings):.3f}s)"
+    )

tests/expressions/test_residual_evaluator.py

@@ -41,7 +41,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, TimestampType
+from pyiceberg.types import DoubleType, FloatType, IntegerType, NestedField, StringType, StructType, TimestampType
 
 
 def test_identity_transform_residual() -> None:
@@ -88,6 +88,26 @@ def test_identity_transform_residual() -> None:
     assert residual == 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)
+
+    res_eval = residual_evaluator_of(spec=spec, expr=EqualTo("dateint", 20170815), case_sensitive=True, schema=schema)
+
+    assert res_eval.residual_for(Record(20170815)) == AlwaysTrue()
+    assert res_eval.residual_for(Record(20170816)) == AlwaysFalse()
+    assert partition_type_calls == 1
+
+
 def test_case_insensitive_identity_transform_residuals() -> None:
     schema = Schema(NestedField(50, "dateint", IntegerType()), NestedField(51, "hour", IntegerType()))