perf: reuse partition evaluator within manifests

pyiceberg/table/__init__.py

@@ -95,6 +95,7 @@
     KeyDefaultDict,
     Properties,
     Record,
+    StructProtocol,
     TableVersion,
 )
 from pyiceberg.types import strtobool
@@ -2582,10 +2583,11 @@ def plan_manifest_entries(self, manifests: Iterable[ManifestFile]) -> Iterator[l
             manifest_file for manifest_file in manifests if manifest_evaluators[manifest_file.partition_spec_id](manifest_file)
         ]
 
-        # step 2: filter the data files in each manifest
-        # this filter depends on the partition spec used to write the manifest file
+        # step 2: filter the data files in each manifest using a task-local partition evaluator
 
-        partition_evaluators: dict[int, Callable[[DataFile], bool]] = KeyDefaultDict(self._build_partition_evaluator)
+        partition_evaluator_factories: dict[int, Callable[[], Callable[[DataFile], bool]]] = KeyDefaultDict(
+            self._build_partition_evaluator_factory
+        )
 
         min_sequence_number = _min_sequence_number(manifests)
 
@@ -2596,7 +2598,7 @@ def plan_manifest_entries(self, manifests: Iterable[ManifestFile]) -> Iterator[l
                 (
                     self.io,
                     manifest,
-                    partition_evaluators[manifest.partition_spec_id],
+                    partition_evaluator_factories[manifest.partition_spec_id](),
                     self._build_metrics_evaluator(),
                 )
                 for manifest in manifests
@@ -2659,16 +2661,26 @@ def _build_manifest_evaluator(self, spec_id: int) -> Callable[[ManifestFile], bo
         spec = self.table_metadata.specs()[spec_id]
         return manifest_evaluator(spec, self.table_metadata.schema(), self.partition_filters[spec_id], self.case_sensitive)
 
-    def _build_partition_evaluator(self, spec_id: int) -> Callable[[DataFile], bool]:
+    def _build_partition_evaluator_factory(self, spec_id: int) -> Callable[[], Callable[[DataFile], bool]]:
         spec = self.table_metadata.specs()[spec_id]
         partition_type = spec.partition_type(self.table_metadata.schema())
         partition_schema = Schema(*partition_type.fields)
         partition_expr = self.partition_filters[spec_id]
 
-        # 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)
+        # The schema and expression are immutable and can be cached per spec, while
+        # each manifest task gets its own mutable evaluator.
+        def partition_evaluator_factory() -> Callable[[DataFile], bool]:
+            evaluator: Callable[[StructProtocol], bool] | None = None
+
+            def partition_evaluator(data_file: DataFile) -> bool:
+                nonlocal evaluator
+                if evaluator is None:
+                    evaluator = expression_evaluator(partition_schema, partition_expr, self.case_sensitive)
+                return evaluator(data_file.partition)
+
+            return partition_evaluator
+
+        return partition_evaluator_factory
 
     def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]:
         schema = self.table_metadata.schema()

tests/benchmark/test_partition_evaluator_benchmark.py

@@ -0,0 +1,115 @@
+# 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 partition evaluation across the sequential entries in one manifest.
+
+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 _combined_filter() -> BooleanExpression:
+    branches: list[BooleanExpression] = []
+    for value in range(11):
+        branch: BooleanExpression = GreaterThanOrEqual("x", 0)
+        for predicate in (
+            LessThanOrEqual("x", 10),
+            EqualTo("y", value),
+            EqualTo("y", value + 1),
+            EqualTo("y", value + 2),
+            EqualTo("y", value + 3),
+        ):
+            branch = And(branch, predicate)
+        branches.append(branch)
+
+    combined = branches[0]
+    for branch in branches[1:]:
+        combined = Or(combined, branch)
+    return combined
+
+
+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,
+    )
+
+
+@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, "x", LongType(), required=True),
+        NestedField(2, "y", LongType(), required=True),
+    )
+    spec = PartitionSpec(
+        PartitionField(1, 1000, IdentityTransform(), "x"),
+        PartitionField(2, 1001, IdentityTransform(), "y"),
+        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=_combined_filter())
+    data_files = [_data_file(file_number) for file_number in range(num_files)]
+    partition_evaluator_factory = planner._build_partition_evaluator_factory(spec.spec_id)
+
+    def evaluate_files() -> int:
+        matches = 0
+        for start in range(0, num_files, files_per_manifest):
+            partition_evaluator = partition_evaluator_factory()
+            matches += sum(partition_evaluator(data_file) for data_file in data_files[start : start + files_per_manifest])
+        return matches
+
+    assert evaluate_files() == 0
+    timings = timeit.repeat(evaluate_files, number=1, 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 66-leaf predicate in "
+        f"{statistics.mean(timings):.3f}s (best: {min(timings):.3f}s)"
+    )

tests/table/test_partition_evaluator_planning.py

@@ -0,0 +1,134 @@
+# 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_reuses_instance_per_manifest_callable(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:
+            calls.append(struct[0])
+            return True
+
+        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))
+    first_file = _data_file(1, 1)
+    second_file = _data_file(2, 10)
+
+    partition_evaluator_factory = planner._build_partition_evaluator_factory(0)
+    first_callable = partition_evaluator_factory()
+    assert not evaluator_calls
+    assert first_callable(first_file)
+    assert first_callable(second_file)
+
+    second_callable = partition_evaluator_factory()
+    assert len(evaluator_calls) == 1
+    assert second_callable(first_file)
+
+    assert evaluator_calls == [[1, 10], [1]]
+
+
+def test_manifest_group_planner_creates_partition_evaluator_per_manifest(
+    table_v2: Table, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=GreaterThan("x", 5))
+    built_factories: list[int] = []
+    built_evaluators: list[Callable[[DataFile], bool]] = []
+    opened_evaluators: list[Callable[[DataFile], bool]] = []
+
+    def build_partition_evaluator_factory(spec_id: int) -> Callable[[], Callable[[DataFile], bool]]:
+        built_factories.append(spec_id)
+
+        def partition_evaluator_factory() -> Callable[[DataFile], bool]:
+            def partition_evaluator(data_file: DataFile) -> bool:
+                return True
+
+            built_evaluators.append(partition_evaluator)
+            return partition_evaluator
+
+        return partition_evaluator_factory
+
+    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_factory", build_partition_evaluator_factory)
+    monkeypatch.setattr(table_module, "_open_manifest", open_manifest)
+
+    list(planner.plan_manifest_entries([_manifest_file(1), _manifest_file(2)]))
+
+    assert built_factories == [0]
+    assert len(built_evaluators) == 2
+    assert {id(evaluator) for evaluator in opened_evaluators} == {id(evaluator) for evaluator in built_evaluators}
+
+
+def test_reused_partition_evaluator_replaces_file_state(table_v2: Table) -> None:
+    planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=GreaterThan("x", 5))
+    partition_evaluator = planner._build_partition_evaluator_factory(0)()
+
+    assert not partition_evaluator(_data_file(1, 1))
+    assert partition_evaluator(_data_file(2, 10))
+    assert not partition_evaluator(_data_file(3, 2))