Make partition expression evaluation stateless

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

@@ -2664,11 +2664,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_partition_evaluator_planning.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.
+
+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.io import FileIO
+from pyiceberg.manifest import DataFile, FileFormat, ManifestContent, ManifestEntry, ManifestFile
+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 _manifest_file(file_number: int) -> ManifestFile:
+    return ManifestFile.from_args(
+        manifest_path=f"s3://bucket/manifest-{file_number}.avro",
+        manifest_length=1,
+        partition_spec_id=0,
+        content=ManifestContent.DATA,
+        sequence_number=1,
+        min_sequence_number=1,
+        added_snapshot_id=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]]
+
+
+def test_manifest_group_planner_shares_partition_evaluator_across_manifests(
+    table_v2: Table, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=GreaterThan("x", 5))
+    built_specs: list[int] = []
+    opened_evaluators: list[Callable[[DataFile], bool]] = []
+
+    def build_partition_evaluator(spec_id: int) -> Callable[[DataFile], bool]:
+        built_specs.append(spec_id)
+        return lambda _: True
+
+    def open_manifest(
+        io: FileIO,
+        manifest: ManifestFile,
+        partition_evaluator: Callable[[DataFile], bool],
+        metrics_evaluator: Callable[[DataFile], bool],
+    ) -> list[ManifestEntry]:
+        opened_evaluators.append(partition_evaluator)
+        return []
+
+    monkeypatch.setattr(planner, "_build_manifest_evaluator", lambda _: lambda _: True)
+    monkeypatch.setattr(planner, "_build_partition_evaluator", build_partition_evaluator)
+    monkeypatch.setattr(table_module, "_open_manifest", open_manifest)
+
+    list(planner.plan_manifest_entries([_manifest_file(1), _manifest_file(2)]))
+
+    assert built_specs == [0]
+    assert len(opened_evaluators) == 2
+    assert opened_evaluators[0] is opened_evaluators[1]