(iceberg-python) branch main updated: Make partition expression evaluation stateless (#3664)

[fo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

Fokko pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git


The following commit(s) were added to refs/heads/main by this push:
     new 23aa9a93c Make partition expression evaluation stateless (#3664)
23aa9a93c is described below

commit 23aa9a93c0996cbcbc93dc4354d0d86512283f61
Author: Aaron Niskode-Dossett <do...@gmail.com>
AuthorDate: Mon Jul 27 13:38:37 2026 -0500

    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)
---
 pyiceberg/expressions/visitors.py                  |  15 ++-
 pyiceberg/table/__init__.py                        |   8 +-
 .../test_partition_evaluator_benchmark.py          | 112 +++++++++++++++++++++
 tests/expressions/test_visitors.py                 |  54 ++++++++++
 tests/table/test_evaluator_planning.py             |  65 ++++++++++++
 5 files changed, 247 insertions(+), 7 deletions(-)

diff --git a/pyiceberg/expressions/visitors.py b/pyiceberg/expressions/visitors.py
index 320cd3110..fa663179b 100644
--- a/pyiceberg/expressions/visitors.py
+++ b/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
diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py
index 5270b9a9e..c0adce84d 100644
--- a/pyiceberg/table/__init__.py
+++ b/pyiceberg/table/__init__.py
@@ -2679,11 +2679,11 @@ class ManifestGroupPlanner:
         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()
diff --git a/tests/benchmark/test_partition_evaluator_benchmark.py b/tests/benchmark/test_partition_evaluator_benchmark.py
new file mode 100644
index 000000000..428ad7cff
--- /dev/null
+++ b/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)"
+    )
diff --git a/tests/expressions/test_visitors.py b/tests/expressions/test_visitors.py
index a0634d683..7cee3836e 100644
--- a/tests/expressions/test_visitors.py
+++ b/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 @@ from pyiceberg.expressions.visitors import (
     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")
diff --git a/tests/table/test_evaluator_planning.py b/tests/table/test_evaluator_planning.py
new file mode 100644
index 000000000..a4ad7bddd
--- /dev/null
+++ b/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]]