From 089dcea39ca792d8e0e570016759a0d64c19e041 Mon Sep 17 00:00:00 2001 From: Aaron Niskode-Dossett Date: Tue, 14 Jul 2026 11:42:19 -0500 Subject: [PATCH 1/3] perf: cache residuals during scan planning --- pyiceberg/expressions/visitors.py | 13 +-- pyiceberg/table/__init__.py | 31 +++-- .../benchmark/test_scan_planning_benchmark.py | 97 ++++++++++++++++ tests/expressions/test_residual_evaluator.py | 22 +++- tests/table/test_scan_planning.py | 108 ++++++++++++++++++ 5 files changed, 252 insertions(+), 19 deletions(-) create mode 100644 tests/benchmark/test_scan_planning_benchmark.py create mode 100644 tests/table/test_scan_planning.py diff --git a/pyiceberg/expressions/visitors.py b/pyiceberg/expressions/visitors.py index 320cd3110e..ef7fd826d1 100644 --- a/pyiceberg/expressions/visitors.py +++ b/pyiceberg/expressions/visitors.py @@ -1805,12 +1805,14 @@ class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC): spec: PartitionSpec case_sensitive: bool expr: BooleanExpression + partition_schema: Schema def __init__(self, schema: Schema, spec: PartitionSpec, case_sensitive: bool, expr: BooleanExpression) -> None: self.schema = schema self.spec = spec self.case_sensitive = case_sensitive self.expr = expr + self.partition_schema = Schema(*spec.partition_type(schema).fields) def eval(self, partition_data: Record) -> BooleanExpression: self.struct = partition_data @@ -1931,17 +1933,12 @@ def visit_bound_predicate(self, predicate: BoundPredicate) -> BooleanExpression: if parts == []: return predicate - def struct_to_schema(struct: StructType) -> Schema: - return Schema(*struct.fields) - for part in parts: strict_projection = part.transform.strict_project(part.name, predicate) strict_result = None if strict_projection is not None: - bound = strict_projection.bind( - struct_to_schema(self.spec.partition_type(self.schema)), case_sensitive=self.case_sensitive - ) + bound = strict_projection.bind(self.partition_schema, case_sensitive=self.case_sensitive) if isinstance(bound, BoundPredicate): strict_result = super().visit_bound_predicate(bound) else: @@ -1954,9 +1951,7 @@ def struct_to_schema(struct: StructType) -> Schema: inclusive_projection = part.transform.project(part.name, predicate) inclusive_result = None if inclusive_projection is not None: - bound_inclusive = inclusive_projection.bind( - struct_to_schema(self.spec.partition_type(self.schema)), case_sensitive=self.case_sensitive - ) + bound_inclusive = inclusive_projection.bind(self.partition_schema, case_sensitive=self.case_sensitive) if isinstance(bound_inclusive, BoundPredicate): # using predicate method specific to inclusive inclusive_result = super().visit_bound_predicate(bound_inclusive) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 63b87d290e..29c8316fe9 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -28,6 +28,7 @@ from types import TracebackType from typing import TYPE_CHECKING, Any, TypeVar +from cachetools import LRUCache from pydantic import Field import pyiceberg.expressions.parser as parser @@ -117,6 +118,9 @@ ALWAYS_TRUE = AlwaysTrue() DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE = "downcast-ns-timestamp-to-us-on-write" +# Retain a small working set for repeated partitions without adding unbounded key +# storage when scans contain a distinct partition value for every data file. +_RESIDUAL_CACHE_MAX_SIZE = 128 @dataclass() @@ -2620,7 +2624,21 @@ def plan_files( data_entries: list[ManifestEntry] = [] delete_index = DeleteFileIndex() - residual_evaluators: dict[int, Callable[[DataFile], ResidualEvaluator]] = KeyDefaultDict(self._build_residual_evaluator) + residual_evaluators: dict[int, ResidualEvaluator] = KeyDefaultDict(self._build_residual_evaluator) + # Residuals depend only on the scan configuration, partition spec, and partition value. + # Keep the cache local to this planning call and bounded for high-cardinality partition specs. + residual_cache: LRUCache[tuple[int, tuple[Any, ...]], BooleanExpression] = LRUCache(maxsize=_RESIDUAL_CACHE_MAX_SIZE) + + def residual_for(data_file: DataFile) -> BooleanExpression: + partition = data_file.partition + partition_values = tuple(partition[pos] for pos in range(len(partition))) + cache_key: tuple[int, tuple[Any, ...]] = data_file.spec_id, partition_values + try: + return residual_cache[cache_key] + except KeyError: + residual = residual_evaluators[data_file.spec_id].residual_for(data_file.partition) + residual_cache[cache_key] = residual + return residual for manifest_entry in chain.from_iterable(self.plan_manifest_entries(manifests)): if not manifest_entry_filter(manifest_entry): @@ -2644,9 +2662,7 @@ def plan_files( data_entry.data_file, partition_key=data_entry.data_file.partition, ), - residual=residual_evaluators[data_entry.data_file.spec_id](data_entry.data_file).residual_for( - data_entry.data_file.partition - ), + residual=residual_for(data_entry.data_file), ) for data_entry in data_entries ] @@ -2684,15 +2700,12 @@ def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]: include_empty_files, ).eval(data_file) - def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], ResidualEvaluator]: + def _build_residual_evaluator(self, spec_id: int) -> ResidualEvaluator: spec = self.table_metadata.specs()[spec_id] from pyiceberg.expressions.visitors import residual_evaluator_of - # 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 datafile: residual_evaluator_of( + return residual_evaluator_of( spec=spec, expr=self.row_filter, case_sensitive=self.case_sensitive, diff --git a/tests/benchmark/test_scan_planning_benchmark.py b/tests/benchmark/test_scan_planning_benchmark.py new file mode 100644 index 0000000000..c0a2724f34 --- /dev/null +++ b/tests/benchmark/test_scan_planning_benchmark.py @@ -0,0 +1,97 @@ +# 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 residual evaluation for repeated and high-cardinality partitions. + +The repeated-partition case measures cache reuse. The unique-partition case +exercises cache eviction and exposes overhead when every lookup misses. + +Run with: + uv run pytest tests/benchmark/test_scan_planning_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, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus +from pyiceberg.table import ManifestGroupPlanner, Table +from pyiceberg.typedef import Record + + +def _combined_filter() -> BooleanExpression: + branches: list[BooleanExpression] = [] + for source in range(11): + branch: BooleanExpression = GreaterThanOrEqual("x", 0) + for predicate in ( + LessThanOrEqual("x", 6), + EqualTo("y", source), + EqualTo("z", source), + EqualTo("y", source + 1), + EqualTo("z", source + 1), + ): + branch = And(branch, predicate) + branches.append(branch) + + combined = branches[0] + for branch in branches[1:]: + combined = Or(combined, branch) + return combined + + +def _manifest_entry(file_number: int, partition: int) -> ManifestEntry: + data_file = DataFile.from_args( + content=DataFileContent.DATA, + file_path=f"s3://bucket/data-{file_number}.parquet", + file_format=FileFormat.PARQUET, + partition=Record(partition), + record_count=1, + file_size_in_bytes=1, + ) + data_file.spec_id = 0 + return ManifestEntry.from_args( + status=ManifestEntryStatus.ADDED, + snapshot_id=1, + sequence_number=1, + file_sequence_number=1, + data_file=data_file, + ) + + +@pytest.mark.benchmark +@pytest.mark.parametrize("num_partitions", [7, 2_000], ids=["repeated-partitions", "unique-partitions"]) +def test_residual_planning(table_v2: Table, monkeypatch: pytest.MonkeyPatch, num_partitions: int) -> None: + num_files = 2_000 + entries = [_manifest_entry(file_number, file_number % num_partitions) for file_number in range(num_files)] + planner = ManifestGroupPlanner( + table_metadata=table_v2.metadata, + io=table_v2.io, + row_filter=_combined_filter(), + ) + + monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) + + timings = timeit.repeat(lambda: list(planner.plan_files([])), number=1, repeat=3) + + assert len(list(planner.plan_files([]))) == num_files + print( + f"Planned {num_files} files across {num_partitions} partitions in " + f"{statistics.mean(timings):.3f}s (best: {min(timings):.3f}s)" + ) diff --git a/tests/expressions/test_residual_evaluator.py b/tests/expressions/test_residual_evaluator.py index 375639ee7b..5ab21a1865 100644 --- a/tests/expressions/test_residual_evaluator.py +++ b/tests/expressions/test_residual_evaluator.py @@ -41,7 +41,7 @@ from pyiceberg.schema import Schema from pyiceberg.transforms import DayTransform, IdentityTransform from pyiceberg.typedef import Record -from pyiceberg.types import DoubleType, FloatType, IntegerType, NestedField, StringType, TimestampType +from pyiceberg.types import DoubleType, FloatType, IntegerType, NestedField, StringType, StructType, TimestampType def test_identity_transform_residual() -> None: @@ -88,6 +88,26 @@ def test_identity_transform_residual() -> None: assert residual == AlwaysFalse() +def test_partition_schema_reused_across_residuals(monkeypatch: pytest.MonkeyPatch) -> None: + schema = Schema(NestedField(50, "dateint", IntegerType())) + spec = PartitionSpec(PartitionField(50, 1050, IdentityTransform(), "dateint_part")) + partition_type_calls = 0 + original_partition_type = PartitionSpec.partition_type + + def counting_partition_type(self: PartitionSpec, schema: Schema) -> StructType: + nonlocal partition_type_calls + partition_type_calls += 1 + return original_partition_type(self, schema) + + monkeypatch.setattr(PartitionSpec, "partition_type", counting_partition_type) + + res_eval = residual_evaluator_of(spec=spec, expr=EqualTo("dateint", 20170815), case_sensitive=True, schema=schema) + + assert res_eval.residual_for(Record(20170815)) == AlwaysTrue() + assert res_eval.residual_for(Record(20170816)) == AlwaysFalse() + assert partition_type_calls == 1 + + def test_case_insensitive_identity_transform_residuals() -> None: schema = Schema(NestedField(50, "dateint", IntegerType()), NestedField(51, "hour", IntegerType())) diff --git a/tests/table/test_scan_planning.py b/tests/table/test_scan_planning.py new file mode 100644 index 0000000000..a9f8d43c1e --- /dev/null +++ b/tests/table/test_scan_planning.py @@ -0,0 +1,108 @@ +# 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 + +import pytest + +import pyiceberg.table as table_module +from pyiceberg.expressions import AlwaysTrue, BooleanExpression, EqualTo +from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus +from pyiceberg.table import ManifestGroupPlanner, Table +from pyiceberg.typedef import Record + + +class _CountingResidualEvaluator: + def __init__(self, marker: int) -> None: + self.marker = marker + self.calls: list[int] = [] + + def residual_for(self, partition: Record) -> BooleanExpression: + partition_value = partition[0] + self.calls.append(partition_value) + return EqualTo("x", self.marker * 10 + partition_value) + + +def _manifest_entry(file_number: int, spec_id: int, partition: int) -> ManifestEntry: + data_file = DataFile.from_args( + content=DataFileContent.DATA, + file_path=f"s3://bucket/data-{file_number}.parquet", + file_format=FileFormat.PARQUET, + partition=Record(partition), + record_count=1, + file_size_in_bytes=1, + ) + data_file.spec_id = spec_id + return ManifestEntry.from_args( + status=ManifestEntryStatus.ADDED, + snapshot_id=1, + sequence_number=1, + file_sequence_number=1, + data_file=data_file, + ) + + +def test_manifest_group_planner_reuses_residuals_by_spec_and_partition(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None: + entries = [ + _manifest_entry(0, spec_id=0, partition=1), + _manifest_entry(1, spec_id=0, partition=1), + _manifest_entry(2, spec_id=1, partition=1), + _manifest_entry(3, spec_id=1, partition=1), + _manifest_entry(4, spec_id=0, partition=2), + ] + evaluators = {0: _CountingResidualEvaluator(0), 1: _CountingResidualEvaluator(1)} + planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=AlwaysTrue()) + + monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda spec_id: evaluators[spec_id]) + + tasks = list(planner.plan_files([])) + + assert evaluators[0].calls == [1, 2] + assert evaluators[1].calls == [1] + assert [task.residual for task in tasks] == [ + EqualTo("x", 1), + EqualTo("x", 1), + EqualTo("x", 11), + EqualTo("x", 11), + EqualTo("x", 2), + ] + + +def test_manifest_group_planner_bounds_residual_cache(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None: + entries = [ + _manifest_entry(0, spec_id=0, partition=1), + _manifest_entry(1, spec_id=0, partition=2), + _manifest_entry(2, spec_id=0, partition=3), + _manifest_entry(3, spec_id=0, partition=1), + ] + evaluator = _CountingResidualEvaluator(0) + planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=AlwaysTrue()) + + monkeypatch.setattr(table_module, "_RESIDUAL_CACHE_MAX_SIZE", 2) + monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator) + + tasks = list(planner.plan_files([])) + + assert evaluator.calls == [1, 2, 3, 1] + assert [task.residual for task in tasks] == [ + EqualTo("x", 1), + EqualTo("x", 2), + EqualTo("x", 3), + EqualTo("x", 1), + ] From 0429ffe818a07c04984633691166fe26b8c7c1f0 Mon Sep 17 00:00:00 2001 From: Aaron Niskode-Dossett Date: Tue, 14 Jul 2026 12:29:35 -0500 Subject: [PATCH 2/3] perf: key residual cache by referenced partitions --- pyiceberg/table/__init__.py | 22 ++- .../benchmark/test_scan_planning_benchmark.py | 36 +++-- tests/table/test_scan_planning.py | 125 +++++++++++++++--- 3 files changed, 145 insertions(+), 38 deletions(-) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 29c8316fe9..37879c6c0f 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -38,6 +38,7 @@ _InclusiveMetricsEvaluator, bind, expression_evaluator, + extract_field_ids, inclusive_projection, manifest_evaluator, ) @@ -118,8 +119,8 @@ ALWAYS_TRUE = AlwaysTrue() DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE = "downcast-ns-timestamp-to-us-on-write" -# Retain a small working set for repeated partitions without adding unbounded key -# storage when scans contain a distinct partition value for every data file. +# Retain a small working set for repeated relevant partition values without adding +# unbounded key storage when scans contain a distinct value for every data file. _RESIDUAL_CACHE_MAX_SIZE = 128 @@ -2625,13 +2626,24 @@ def plan_files( delete_index = DeleteFileIndex() residual_evaluators: dict[int, ResidualEvaluator] = KeyDefaultDict(self._build_residual_evaluator) - # Residuals depend only on the scan configuration, partition spec, and partition value. - # Keep the cache local to this planning call and bounded for high-cardinality partition specs. + referenced_field_ids = extract_field_ids( + bind(self.table_metadata.schema(), self.row_filter, case_sensitive=self.case_sensitive) + ) + partition_specs = self.table_metadata.specs() + residual_cache_key_positions: dict[int, tuple[int, ...]] = KeyDefaultDict( + lambda spec_id: tuple( + pos + for pos, partition_field in enumerate(partition_specs[spec_id].fields) + if partition_field.source_id in referenced_field_ids + ) + ) + # A residual can only depend on partition fields derived from source columns + # referenced by the scan filter. Keep the cache local and bounded. residual_cache: LRUCache[tuple[int, tuple[Any, ...]], BooleanExpression] = LRUCache(maxsize=_RESIDUAL_CACHE_MAX_SIZE) def residual_for(data_file: DataFile) -> BooleanExpression: partition = data_file.partition - partition_values = tuple(partition[pos] for pos in range(len(partition))) + partition_values = tuple(partition[pos] for pos in residual_cache_key_positions[data_file.spec_id]) cache_key: tuple[int, tuple[Any, ...]] = data_file.spec_id, partition_values try: return residual_cache[cache_key] diff --git a/tests/benchmark/test_scan_planning_benchmark.py b/tests/benchmark/test_scan_planning_benchmark.py index c0a2724f34..4d7d393571 100644 --- a/tests/benchmark/test_scan_planning_benchmark.py +++ b/tests/benchmark/test_scan_planning_benchmark.py @@ -14,10 +14,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Benchmark residual evaluation for repeated and high-cardinality partitions. +"""Benchmark residual evaluation with a high-cardinality irrelevant partition field. -The repeated-partition case measures cache reuse. The unique-partition case -exercises cache eviction and exposes overhead when every lookup misses. +Every file has a unique unreferenced partition-hash value. The repeated case +measures reuse by relevant partition values, while the unique case forces misses. Run with: uv run pytest tests/benchmark/test_scan_planning_benchmark.py -v -s -m benchmark @@ -32,7 +32,9 @@ from pyiceberg.expressions import And, BooleanExpression, EqualTo, GreaterThanOrEqual, LessThanOrEqual, Or from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus +from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.table import ManifestGroupPlanner, Table +from pyiceberg.transforms import IdentityTransform from pyiceberg.typedef import Record @@ -43,9 +45,9 @@ def _combined_filter() -> BooleanExpression: for predicate in ( LessThanOrEqual("x", 6), EqualTo("y", source), - EqualTo("z", source), EqualTo("y", source + 1), - EqualTo("z", source + 1), + EqualTo("y", source + 2), + EqualTo("y", source + 3), ): branch = And(branch, predicate) branches.append(branch) @@ -56,12 +58,12 @@ def _combined_filter() -> BooleanExpression: return combined -def _manifest_entry(file_number: int, partition: int) -> ManifestEntry: +def _manifest_entry(file_number: int, relevant_partition: int) -> ManifestEntry: data_file = DataFile.from_args( content=DataFileContent.DATA, file_path=f"s3://bucket/data-{file_number}.parquet", file_format=FileFormat.PARQUET, - partition=Record(partition), + partition=Record(relevant_partition, file_number), record_count=1, file_size_in_bytes=1, ) @@ -76,12 +78,22 @@ def _manifest_entry(file_number: int, partition: int) -> ManifestEntry: @pytest.mark.benchmark -@pytest.mark.parametrize("num_partitions", [7, 2_000], ids=["repeated-partitions", "unique-partitions"]) -def test_residual_planning(table_v2: Table, monkeypatch: pytest.MonkeyPatch, num_partitions: int) -> None: +@pytest.mark.parametrize( + "num_relevant_partitions", + [7, 2_000], + ids=["repeated-relevant-partitions", "unique-relevant-partitions"], +) +def test_residual_planning(table_v2: Table, monkeypatch: pytest.MonkeyPatch, num_relevant_partitions: int) -> None: num_files = 2_000 - entries = [_manifest_entry(file_number, file_number % num_partitions) for file_number in range(num_files)] + entries = [_manifest_entry(file_number, file_number % num_relevant_partitions) for file_number in range(num_files)] + spec = PartitionSpec( + PartitionField(1, 1000, IdentityTransform(), "x"), + PartitionField(3, 1001, IdentityTransform(), "partition_hash"), + spec_id=0, + ) + metadata = table_v2.metadata.model_copy(update={"partition_specs": [spec]}) planner = ManifestGroupPlanner( - table_metadata=table_v2.metadata, + table_metadata=metadata, io=table_v2.io, row_filter=_combined_filter(), ) @@ -92,6 +104,6 @@ def test_residual_planning(table_v2: Table, monkeypatch: pytest.MonkeyPatch, num assert len(list(planner.plan_files([]))) == num_files print( - f"Planned {num_files} files across {num_partitions} partitions in " + f"Planned {num_files} files across {num_relevant_partitions} relevant partitions in " f"{statistics.mean(timings):.3f}s (best: {min(timings):.3f}s)" ) diff --git a/tests/table/test_scan_planning.py b/tests/table/test_scan_planning.py index a9f8d43c1e..951b03e758 100644 --- a/tests/table/test_scan_planning.py +++ b/tests/table/test_scan_planning.py @@ -17,32 +17,36 @@ from __future__ import annotations +from typing import Any + import pytest import pyiceberg.table as table_module -from pyiceberg.expressions import AlwaysTrue, BooleanExpression, EqualTo +from pyiceberg.expressions import And, BooleanExpression, EqualTo from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus +from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.table import ManifestGroupPlanner, Table +from pyiceberg.transforms import BucketTransform, IdentityTransform from pyiceberg.typedef import Record class _CountingResidualEvaluator: def __init__(self, marker: int) -> None: self.marker = marker - self.calls: list[int] = [] + self.calls: list[tuple[Any, ...]] = [] def residual_for(self, partition: Record) -> BooleanExpression: - partition_value = partition[0] - self.calls.append(partition_value) - return EqualTo("x", self.marker * 10 + partition_value) + partition_values = tuple(partition[pos] for pos in range(len(partition))) + self.calls.append(partition_values) + return EqualTo("x", self.marker * 10 + partition[0]) -def _manifest_entry(file_number: int, spec_id: int, partition: int) -> ManifestEntry: +def _manifest_entry(file_number: int, spec_id: int, partition: tuple[Any, ...]) -> ManifestEntry: data_file = DataFile.from_args( content=DataFileContent.DATA, file_path=f"s3://bucket/data-{file_number}.parquet", file_format=FileFormat.PARQUET, - partition=Record(partition), + partition=Record(*partition), record_count=1, file_size_in_bytes=1, ) @@ -56,24 +60,44 @@ def _manifest_entry(file_number: int, spec_id: int, partition: int) -> ManifestE ) +def _identity_spec(spec_id: int, *source_ids: int) -> PartitionSpec: + return PartitionSpec( + *( + PartitionField( + source_id, + 1000 + spec_id * 10 + pos, + IdentityTransform(), + f"field_{source_id}_{pos}", + ) + for pos, source_id in enumerate(source_ids) + ), + spec_id=spec_id, + ) + + +def _planner(table_v2: Table, row_filter: BooleanExpression, *partition_specs: PartitionSpec) -> ManifestGroupPlanner: + metadata = table_v2.metadata.model_copy(update={"partition_specs": list(partition_specs)}) + return ManifestGroupPlanner(table_metadata=metadata, io=table_v2.io, row_filter=row_filter) + + def test_manifest_group_planner_reuses_residuals_by_spec_and_partition(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None: entries = [ - _manifest_entry(0, spec_id=0, partition=1), - _manifest_entry(1, spec_id=0, partition=1), - _manifest_entry(2, spec_id=1, partition=1), - _manifest_entry(3, spec_id=1, partition=1), - _manifest_entry(4, spec_id=0, partition=2), + _manifest_entry(0, spec_id=0, partition=(1,)), + _manifest_entry(1, spec_id=0, partition=(1,)), + _manifest_entry(2, spec_id=1, partition=(1,)), + _manifest_entry(3, spec_id=1, partition=(1,)), + _manifest_entry(4, spec_id=0, partition=(2,)), ] evaluators = {0: _CountingResidualEvaluator(0), 1: _CountingResidualEvaluator(1)} - planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=AlwaysTrue()) + planner = _planner(table_v2, EqualTo("x", 1), _identity_spec(0, 1), _identity_spec(1, 1)) monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) monkeypatch.setattr(planner, "_build_residual_evaluator", lambda spec_id: evaluators[spec_id]) tasks = list(planner.plan_files([])) - assert evaluators[0].calls == [1, 2] - assert evaluators[1].calls == [1] + assert evaluators[0].calls == [(1,), (2,)] + assert evaluators[1].calls == [(1,)] assert [task.residual for task in tasks] == [ EqualTo("x", 1), EqualTo("x", 1), @@ -83,15 +107,74 @@ def test_manifest_group_planner_reuses_residuals_by_spec_and_partition(table_v2: ] +def test_manifest_group_planner_ignores_unreferenced_partition_fields(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None: + entries = [ + _manifest_entry(0, spec_id=0, partition=(1, 10)), + _manifest_entry(1, spec_id=0, partition=(1, 20)), + _manifest_entry(2, spec_id=0, partition=(2, 30)), + ] + evaluator = _CountingResidualEvaluator(0) + planner = _planner(table_v2, EqualTo("x", 1), _identity_spec(0, 1, 2)) + + monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator) + + tasks = list(planner.plan_files([])) + + assert evaluator.calls == [(1, 10), (2, 30)] + assert [task.residual for task in tasks] == [EqualTo("x", 1), EqualTo("x", 1), EqualTo("x", 2)] + + +def test_manifest_group_planner_includes_referenced_partition_fields(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None: + entries = [ + _manifest_entry(0, spec_id=0, partition=(1, 10)), + _manifest_entry(1, spec_id=0, partition=(1, 20)), + ] + evaluator = _CountingResidualEvaluator(0) + planner = _planner(table_v2, And(EqualTo("x", 1), EqualTo("y", 10)), _identity_spec(0, 1, 2)) + + monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator) + + list(planner.plan_files([])) + + assert evaluator.calls == [(1, 10), (1, 20)] + + +def test_manifest_group_planner_includes_all_partition_transforms_for_referenced_source( + table_v2: Table, monkeypatch: pytest.MonkeyPatch +) -> None: + spec = PartitionSpec( + PartitionField(1, 1000, BucketTransform(7), "x_bucket_7"), + PartitionField(1, 1001, BucketTransform(5), "x_bucket_5"), + PartitionField(2, 1002, IdentityTransform(), "partition_hash"), + spec_id=0, + ) + entries = [ + _manifest_entry(0, spec_id=0, partition=(5, 0, 10)), + _manifest_entry(1, spec_id=0, partition=(5, 1, 20)), + _manifest_entry(2, spec_id=0, partition=(5, 0, 30)), + ] + evaluator = _CountingResidualEvaluator(0) + planner = _planner(table_v2, EqualTo("x", 1), spec) + + monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator) + + list(planner.plan_files([])) + + assert evaluator.calls == [(5, 0, 10), (5, 1, 20)] + + def test_manifest_group_planner_bounds_residual_cache(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None: entries = [ - _manifest_entry(0, spec_id=0, partition=1), - _manifest_entry(1, spec_id=0, partition=2), - _manifest_entry(2, spec_id=0, partition=3), - _manifest_entry(3, spec_id=0, partition=1), + _manifest_entry(0, spec_id=0, partition=(1,)), + _manifest_entry(1, spec_id=0, partition=(2,)), + _manifest_entry(2, spec_id=0, partition=(3,)), + _manifest_entry(3, spec_id=0, partition=(1,)), ] evaluator = _CountingResidualEvaluator(0) - planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=AlwaysTrue()) + planner = _planner(table_v2, EqualTo("x", 1), _identity_spec(0, 1)) monkeypatch.setattr(table_module, "_RESIDUAL_CACHE_MAX_SIZE", 2) monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) @@ -99,7 +182,7 @@ def test_manifest_group_planner_bounds_residual_cache(table_v2: Table, monkeypat tasks = list(planner.plan_files([])) - assert evaluator.calls == [1, 2, 3, 1] + assert evaluator.calls == [(1,), (2,), (3,), (1,)] assert [task.residual for task in tasks] == [ EqualTo("x", 1), EqualTo("x", 2), From 2d93f07b9a9c97ce1dffd30311e69cbb9cbae257 Mon Sep 17 00:00:00 2001 From: Aaron Niskode-Dossett Date: Tue, 14 Jul 2026 15:24:37 -0500 Subject: [PATCH 3/3] perf: create residual evaluators on cache misses --- pyiceberg/table/__init__.py | 13 +++++++--- tests/table/test_scan_planning.py | 43 +++++++++++++++++++------------ 2 files changed, 36 insertions(+), 20 deletions(-) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 37879c6c0f..819a8df431 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -2625,7 +2625,9 @@ def plan_files( data_entries: list[ManifestEntry] = [] delete_index = DeleteFileIndex() - residual_evaluators: dict[int, ResidualEvaluator] = KeyDefaultDict(self._build_residual_evaluator) + residual_evaluator_factories: dict[int, Callable[[DataFile], ResidualEvaluator]] = KeyDefaultDict( + self._build_residual_evaluator + ) referenced_field_ids = extract_field_ids( bind(self.table_metadata.schema(), self.row_filter, case_sensitive=self.case_sensitive) ) @@ -2648,7 +2650,8 @@ def residual_for(data_file: DataFile) -> BooleanExpression: try: return residual_cache[cache_key] except KeyError: - residual = residual_evaluators[data_file.spec_id].residual_for(data_file.partition) + residual_evaluator = residual_evaluator_factories[data_file.spec_id](data_file) + residual = residual_evaluator.residual_for(partition) residual_cache[cache_key] = residual return residual @@ -2712,12 +2715,14 @@ def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]: include_empty_files, ).eval(data_file) - def _build_residual_evaluator(self, spec_id: int) -> ResidualEvaluator: + def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], ResidualEvaluator]: spec = self.table_metadata.specs()[spec_id] from pyiceberg.expressions.visitors import residual_evaluator_of - return residual_evaluator_of( + # ResidualEvaluator stores the current partition while evaluating. Return a + # factory so every cache miss uses a fresh, unshared evaluator instance. + return lambda _: residual_evaluator_of( spec=spec, expr=self.row_filter, case_sensitive=self.case_sensitive, diff --git a/tests/table/test_scan_planning.py b/tests/table/test_scan_planning.py index 951b03e758..e360e7d2fc 100644 --- a/tests/table/test_scan_planning.py +++ b/tests/table/test_scan_planning.py @@ -41,6 +41,17 @@ def residual_for(self, partition: Record) -> BooleanExpression: return EqualTo("x", self.marker * 10 + partition[0]) +class _CountingResidualEvaluatorFactory: + def __init__(self, marker: int) -> None: + self.marker = marker + self.evaluators: list[_CountingResidualEvaluator] = [] + + def __call__(self, _: DataFile) -> _CountingResidualEvaluator: + evaluator = _CountingResidualEvaluator(self.marker) + self.evaluators.append(evaluator) + return evaluator + + def _manifest_entry(file_number: int, spec_id: int, partition: tuple[Any, ...]) -> ManifestEntry: data_file = DataFile.from_args( content=DataFileContent.DATA, @@ -88,16 +99,16 @@ def test_manifest_group_planner_reuses_residuals_by_spec_and_partition(table_v2: _manifest_entry(3, spec_id=1, partition=(1,)), _manifest_entry(4, spec_id=0, partition=(2,)), ] - evaluators = {0: _CountingResidualEvaluator(0), 1: _CountingResidualEvaluator(1)} + evaluator_factories = {0: _CountingResidualEvaluatorFactory(0), 1: _CountingResidualEvaluatorFactory(1)} planner = _planner(table_v2, EqualTo("x", 1), _identity_spec(0, 1), _identity_spec(1, 1)) monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) - monkeypatch.setattr(planner, "_build_residual_evaluator", lambda spec_id: evaluators[spec_id]) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda spec_id: evaluator_factories[spec_id]) tasks = list(planner.plan_files([])) - assert evaluators[0].calls == [(1,), (2,)] - assert evaluators[1].calls == [(1,)] + assert [evaluator.calls for evaluator in evaluator_factories[0].evaluators] == [[(1,)], [(2,)]] + assert [evaluator.calls for evaluator in evaluator_factories[1].evaluators] == [[(1,)]] assert [task.residual for task in tasks] == [ EqualTo("x", 1), EqualTo("x", 1), @@ -113,15 +124,15 @@ def test_manifest_group_planner_ignores_unreferenced_partition_fields(table_v2: _manifest_entry(1, spec_id=0, partition=(1, 20)), _manifest_entry(2, spec_id=0, partition=(2, 30)), ] - evaluator = _CountingResidualEvaluator(0) + evaluator_factory = _CountingResidualEvaluatorFactory(0) planner = _planner(table_v2, EqualTo("x", 1), _identity_spec(0, 1, 2)) monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) - monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator_factory) tasks = list(planner.plan_files([])) - assert evaluator.calls == [(1, 10), (2, 30)] + assert [evaluator.calls for evaluator in evaluator_factory.evaluators] == [[(1, 10)], [(2, 30)]] assert [task.residual for task in tasks] == [EqualTo("x", 1), EqualTo("x", 1), EqualTo("x", 2)] @@ -130,15 +141,15 @@ def test_manifest_group_planner_includes_referenced_partition_fields(table_v2: T _manifest_entry(0, spec_id=0, partition=(1, 10)), _manifest_entry(1, spec_id=0, partition=(1, 20)), ] - evaluator = _CountingResidualEvaluator(0) + evaluator_factory = _CountingResidualEvaluatorFactory(0) planner = _planner(table_v2, And(EqualTo("x", 1), EqualTo("y", 10)), _identity_spec(0, 1, 2)) monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) - monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator_factory) list(planner.plan_files([])) - assert evaluator.calls == [(1, 10), (1, 20)] + assert [evaluator.calls for evaluator in evaluator_factory.evaluators] == [[(1, 10)], [(1, 20)]] def test_manifest_group_planner_includes_all_partition_transforms_for_referenced_source( @@ -155,15 +166,15 @@ def test_manifest_group_planner_includes_all_partition_transforms_for_referenced _manifest_entry(1, spec_id=0, partition=(5, 1, 20)), _manifest_entry(2, spec_id=0, partition=(5, 0, 30)), ] - evaluator = _CountingResidualEvaluator(0) + evaluator_factory = _CountingResidualEvaluatorFactory(0) planner = _planner(table_v2, EqualTo("x", 1), spec) monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) - monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator_factory) list(planner.plan_files([])) - assert evaluator.calls == [(5, 0, 10), (5, 1, 20)] + assert [evaluator.calls for evaluator in evaluator_factory.evaluators] == [[(5, 0, 10)], [(5, 1, 20)]] def test_manifest_group_planner_bounds_residual_cache(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None: @@ -173,16 +184,16 @@ def test_manifest_group_planner_bounds_residual_cache(table_v2: Table, monkeypat _manifest_entry(2, spec_id=0, partition=(3,)), _manifest_entry(3, spec_id=0, partition=(1,)), ] - evaluator = _CountingResidualEvaluator(0) + evaluator_factory = _CountingResidualEvaluatorFactory(0) planner = _planner(table_v2, EqualTo("x", 1), _identity_spec(0, 1)) monkeypatch.setattr(table_module, "_RESIDUAL_CACHE_MAX_SIZE", 2) monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) - monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator) + monkeypatch.setattr(planner, "_build_residual_evaluator", lambda _: evaluator_factory) tasks = list(planner.plan_files([])) - assert evaluator.calls == [(1,), (2,), (3,), (1,)] + assert [evaluator.calls for evaluator in evaluator_factory.evaluators] == [[(1,)], [(2,)], [(3,)], [(1,)]] assert [task.residual for task in tasks] == [ EqualTo("x", 1), EqualTo("x", 2),