From a0e7f4c1116eb01498a7e5fcef985113462cbacf Mon Sep 17 00:00:00 2001 From: Aaron Niskode-Dossett Date: Tue, 14 Jul 2026 14:14:44 -0500 Subject: [PATCH] perf: reuse partition evaluator within manifests --- pyiceberg/table/__init__.py | 30 ++-- .../test_partition_evaluator_benchmark.py | 115 +++++++++++++++ .../test_partition_evaluator_planning.py | 134 ++++++++++++++++++ 3 files changed, 270 insertions(+), 9 deletions(-) create mode 100644 tests/benchmark/test_partition_evaluator_benchmark.py create mode 100644 tests/table/test_partition_evaluator_planning.py diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 63b87d290e..c59a75f13b 100644 --- a/pyiceberg/table/__init__.py +++ b/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() diff --git a/tests/benchmark/test_partition_evaluator_benchmark.py b/tests/benchmark/test_partition_evaluator_benchmark.py new file mode 100644 index 0000000000..337e57b250 --- /dev/null +++ b/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)" + ) diff --git a/tests/table/test_partition_evaluator_planning.py b/tests/table/test_partition_evaluator_planning.py new file mode 100644 index 0000000000..4bc7896bff --- /dev/null +++ b/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))