Make partition expression evaluation stateless (#3664)

*I used CODEX to analyze this problem and create this PR. I've reviewed
the code and tests and stand by them. This summary is written completely
by a human (me) other than very light copy editing by an LLM.*

Currently, because `_ExpressionEvaluator` can't be safely shared one is
created per file. That's expensive relative to actually evaluating the
expression. This change makes `_ExpressionEvaluator` shareable by moving
the mutable parts into a new `_ExpressionEvaluationVisitor` which stays
local to the per-file work of evaluation.

These benchmarks are showing improvements in sub-second times, but in
the production workloads I tested this on it did translate to real wall
clock improvement. But irrespective of performance gains the new
`_ExpressionEvaluator` is safer since nothing was preventing its
accidental concurrent reuse in the future.

CODEX-generated summary follows:
-------
Partition pruning currently binds the same projected expression for
every data file because `_ExpressionEvaluator` stores the current record
as mutable instance state. That repeated preparation is expensive, but
sharing the existing evaluator across workers would allow concurrent
calls to mix records.

This separates preparation from per-record evaluation.
`_ExpressionEvaluator` binds the expression once, while every call
creates a private `_ExpressionEvaluationVisitor` containing that call's
record. Manifest planning can therefore prepare one evaluator per
partition spec and safely share it across manifests and workers. This
provides the construction-reuse benefit targeted by #3656 without
depending on files within a manifest remaining sequential, and it also
benefits one-file manifests.

## Summary

- Split the prepared expression evaluator from the call-local mutable
visitor.
- Reuse one prepared partition evaluator per partition spec during
manifest planning.
- Add deterministic concurrent-use, prepared-state, and planner-sharing
coverage.
- Add a benchmark using a realistic 15-leaf predicate across dense and
one-file manifests.

## Performance

I compared this branch with `main` using the partition-evaluator
workload in this PR. It evaluates 1,000 files with two identity
partition fields and a 15-leaf predicate modeling five `(event_day range
AND region_id)` branches. Timings are medians from seven samples of five
iterations.

| Manifest layout | `main` | This PR | Speedup |
|---|---:|---:|---:|
| 1 manifest × 1,000 files | 187.928 ms | 14.413 ms | 13.04× |
| 1,000 manifests × 1 file | 188.000 ms | 14.584 ms | 12.89× |

The improvement comes from binding the projected partition expression
once per partition spec instead of once per file. This is an isolated
partition-pruning benchmark rather than an end-to-end scan-planning
measurement.

## Testing

- `pytest tests/expressions/test_visitors.py
tests/table/test_partition_evaluator_planning.py
tests/table/test_init.py` (200 passed)
- `pytest tests/benchmark/test_partition_evaluator_benchmark.py -m
benchmark` (2 passed)

Author: dossett

pyiceberg/expressions/visitors.py

@@ -454,16 +454,25 @@ def expression_evaluator(schema: Schema, unbound: BooleanExpression, case_sensit
     return _ExpressionEvaluator(schema, unbound, case_sensitive).eval
 
 
-class _ExpressionEvaluator(BoundBooleanExpressionVisitor[bool]):
+class _ExpressionEvaluator:
+    """An evaluator that binds an expression once and keeps evaluation state local to each call."""
+
     bound: BooleanExpression
-    struct: StructProtocol
 
     def __init__(self, schema: Schema, unbound: BooleanExpression, case_sensitive: bool):
         self.bound = bind(schema, unbound, case_sensitive)
 
     def eval(self, struct: StructProtocol) -> bool:
+        return visit(self.bound, _ExpressionEvaluationVisitor(struct))
+
+
+class _ExpressionEvaluationVisitor(BoundBooleanExpressionVisitor[bool]):
+    """Evaluate a bound expression against one struct."""
+
+    struct: StructProtocol
+
+    def __init__(self, struct: StructProtocol):
         self.struct = struct
-        return visit(self.bound, self)
 
     def visit_in(self, term: BoundTerm, literals: set[L]) -> bool:
         return term.eval(self.struct) in literals

pyiceberg/table/__init__.py

@@ -2679,11 +2679,11 @@ def _build_partition_evaluator(self, spec_id: int) -> Callable[[DataFile], bool]
         partition_type = spec.partition_type(self.table_metadata.schema())
         partition_schema = Schema(*partition_type.fields)
         partition_expr = self.partition_filters[spec_id]
+        evaluator = expression_evaluator(partition_schema, partition_expr, self.case_sensitive)
 
-        # 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 data_file: expression_evaluator(partition_schema, partition_expr, self.case_sensitive)(data_file.partition)
+        # Expression evaluators keep input-specific state local to each call, so the
+        # prepared evaluator can be shared by every manifest using this spec.
+        return lambda data_file: evaluator(data_file.partition)
 
     def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]:
         schema = self.table_metadata.schema()

tests/benchmark/test_partition_evaluator_benchmark.py

@@ -0,0 +1,112 @@
+# 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 a realistic 15-leaf partition predicate when a prepared evaluator is shared across manifests.
+
+Run with:
+    uv run pytest tests/benchmark/test_partition_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, FileFormat
+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 _data_file(file_number: int) -> DataFile:
+    return DataFile.from_args(
+        file_path=f"s3://bucket/data-{file_number}.parquet",
+        file_format=FileFormat.PARQUET,
+        partition=Record(file_number % 11, file_number % 15),
+        record_count=100,
+        file_size_in_bytes=1,
+    )
+
+
+def _partition_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
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    "files_per_manifest",
+    [1_000, 1],
+    ids=["many-files-per-manifest", "one-file-per-manifest"],
+)
+def test_partition_evaluator_reuse(table_v2: Table, files_per_manifest: int) -> None:
+    num_files = 1_000
+    schema = Schema(
+        NestedField(1, "event_day", LongType(), required=True),
+        NestedField(2, "region_id", LongType(), required=True),
+    )
+    spec = PartitionSpec(
+        PartitionField(1, 1000, IdentityTransform(), "event_day"),
+        PartitionField(2, 1001, IdentityTransform(), "region_id"),
+        spec_id=0,
+    )
+    metadata = TableMetadataV2(
+        location="s3://bucket/table",
+        last_column_id=2,
+        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=_partition_filter())
+    data_files = [_data_file(file_number) for file_number in range(num_files)]
+
+    def evaluate_files() -> int:
+        partition_evaluator = planner._build_partition_evaluator(spec.spec_id)
+        matches = 0
+        for start in range(0, num_files, files_per_manifest):
+            matches += sum(partition_evaluator(data_file) for data_file in data_files[start : start + files_per_manifest])
+        return matches
+
+    assert evaluate_files() == 67
+    iterations = 100
+    timings_ms = [timing * 1_000 / iterations for timing in timeit.repeat(evaluate_files, number=iterations, repeat=3)]
+    file_label = "file" if files_per_manifest == 1 else "files"
+
+    print(
+        f"Evaluated partitions for {num_files} files with {files_per_manifest} {file_label} per manifest "
+        f"and a 15-leaf predicate in "
+        f"{statistics.mean(timings_ms):.3f}ms (best: {min(timings_ms):.3f}ms)"
+    )

tests/expressions/test_visitors.py

@@ -16,6 +16,8 @@
 # under the License.
 # pylint:disable=redefined-outer-name
 
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
 from typing import Any
 
 import pytest
@@ -67,6 +69,7 @@
     BindVisitor,
     BooleanExpressionVisitor,
     BoundBooleanExpressionVisitor,
+    _ExpressionEvaluator,
     _ManifestEvalVisitor,
     expression_evaluator,
     expression_to_plain_format,
@@ -1630,6 +1633,57 @@ def test_expression_evaluator_null() -> None:
     assert expression_evaluator(schema, NotStartsWith("a", 1), case_sensitive=True)(struct) is True
 
 
+def test_expression_evaluator_does_not_mutate_prepared_state() -> None:
+    schema = Schema(
+        NestedField(1, "a", IntegerType(), required=True),
+        NestedField(2, "b", IntegerType(), required=True),
+    )
+    evaluator = _ExpressionEvaluator(schema, And(EqualTo("a", 1), EqualTo("b", 1)), case_sensitive=True)
+    initial_state = vars(evaluator).copy()
+
+    assert evaluator.eval(Record(1, 1)) is True
+    assert evaluator.eval(Record(0, 0)) is False
+    assert evaluator.eval(Record(1, 1)) is True
+
+    assert vars(evaluator) == initial_state
+
+
+def test_expression_evaluator_concurrent_calls_do_not_share_records() -> 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 expression evaluations")
+            return value
+
+    schema = Schema(
+        NestedField(1, "a", IntegerType(), required=True),
+        NestedField(2, "b", IntegerType(), required=True),
+    )
+    evaluator = expression_evaluator(schema, And(EqualTo("a", 1), EqualTo("b", 1)), case_sensitive=True)
+    first_read = Event()
+    release_first_read = Event()
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        matching_result = executor.submit(evaluator, BlockingRecord(first_read, release_first_read, 1, 1))
+        assert first_read.wait(timeout=5)
+
+        try:
+            non_matching_result = executor.submit(evaluator, Record(0, 0)).result(timeout=5)
+        finally:
+            release_first_read.set()
+
+        assert matching_result.result(timeout=5) is True
+        assert non_matching_result is False
+
+
 def test_expression_evaluator_binary_starts_with() -> None:
     schema = Schema(NestedField(1, "x", BinaryType(), required=False), schema_id=1)
     struct = Record(b"aa")

tests/table/test_evaluator_planning.py

@@ -0,0 +1,65 @@
+# 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.
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import pytest
+
+import pyiceberg.table as table_module
+from pyiceberg.expressions import BooleanExpression, GreaterThan
+from pyiceberg.manifest import DataFile, FileFormat
+from pyiceberg.schema import Schema
+from pyiceberg.table import ManifestGroupPlanner, Table
+from pyiceberg.typedef import Record, StructProtocol
+
+
+def _data_file(file_number: int, partition_value: int) -> DataFile:
+    return DataFile.from_args(
+        file_path=f"s3://bucket/data-{file_number}.parquet",
+        file_format=FileFormat.PARQUET,
+        partition=Record(partition_value),
+        record_count=100,
+        file_size_in_bytes=1,
+    )
+
+
+def test_partition_evaluator_prepares_once_per_spec(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None:
+    evaluator_calls: list[list[int]] = []
+
+    def counting_expression_evaluator(
+        schema: Schema, unbound: BooleanExpression, case_sensitive: bool
+    ) -> Callable[[StructProtocol], bool]:
+        calls: list[int] = []
+        evaluator_calls.append(calls)
+
+        def evaluate(struct: StructProtocol) -> bool:
+            value = struct[0]
+            calls.append(value)
+            return value > 5
+
+        return evaluate
+
+    monkeypatch.setattr(table_module, "expression_evaluator", counting_expression_evaluator)
+    planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=GreaterThan("x", 5))
+    partition_evaluator = planner._build_partition_evaluator(0)
+
+    assert len(evaluator_calls) == 1
+    assert not partition_evaluator(_data_file(1, 1))
+    assert partition_evaluator(_data_file(2, 10))
+    assert evaluator_calls == [[1, 10]]