From bca46fc3c81865380229362a0cb10d6eaea78a7c Mon Sep 17 00:00:00 2001 From: Aaron Niskode-Dossett Date: Wed, 15 Jul 2026 11:42:30 -0500 Subject: [PATCH 1/3] Make residual evaluation stateless --- pyiceberg/expressions/visitors.py | 58 ++++++++---- tests/expressions/test_residual_evaluator.py | 92 +++++++++++++++++++- 2 files changed, 131 insertions(+), 19 deletions(-) diff --git a/pyiceberg/expressions/visitors.py b/pyiceberg/expressions/visitors.py index 320cd3110e..04814add27 100644 --- a/pyiceberg/expressions/visitors.py +++ b/pyiceberg/expressions/visitors.py @@ -1785,7 +1785,7 @@ def _can_contain_nans(self, field_id: int) -> bool: return (nan_count := self.nan_counts.get(field_id)) is not None and nan_count > 0 -class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC): +class _ResidualEvaluationVisitor(BoundBooleanExpressionVisitor[BooleanExpression]): """Finds the residuals for an Expression the partitions in the given PartitionSpec. A residual expression is made by partially evaluating an expression using partition values. @@ -1804,17 +1804,22 @@ class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC): schema: Schema spec: PartitionSpec case_sensitive: bool - expr: BooleanExpression + partition_schema: Schema + struct: Record - def __init__(self, schema: Schema, spec: PartitionSpec, case_sensitive: bool, expr: BooleanExpression) -> None: + def __init__( + self, + schema: Schema, + spec: PartitionSpec, + case_sensitive: bool, + partition_schema: Schema, + partition_data: Record, + ) -> None: self.schema = schema self.spec = spec self.case_sensitive = case_sensitive - self.expr = expr - - def eval(self, partition_data: Record) -> BooleanExpression: + self.partition_schema = partition_schema self.struct = partition_data - return visit(self.expr, visitor=self) def visit_true(self) -> BooleanExpression: return AlwaysTrue() @@ -1931,17 +1936,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 +1954,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) @@ -1985,9 +1983,33 @@ def visit_unbound_predicate(self, predicate: UnboundPredicate) -> BooleanExpress return bound -class ResidualEvaluator(ResidualVisitor): +class ResidualEvaluator: + """Prepare residual evaluation once while keeping partition state local to each call.""" + + schema: Schema + 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 residual_for(self, partition_data: Record) -> BooleanExpression: - return self.eval(partition_data) + return visit( + self.expr, + visitor=_ResidualEvaluationVisitor( + schema=self.schema, + spec=self.spec, + case_sensitive=self.case_sensitive, + partition_schema=self.partition_schema, + partition_data=partition_data, + ), + ) class UnpartitionedResidualEvaluator(ResidualEvaluator): diff --git a/tests/expressions/test_residual_evaluator.py b/tests/expressions/test_residual_evaluator.py index 375639ee7b..b5027d94dd 100644 --- a/tests/expressions/test_residual_evaluator.py +++ b/tests/expressions/test_residual_evaluator.py @@ -15,6 +15,10 @@ # specific language governing permissions and limitations # under the License. # pylint:disable=redefined-outer-name +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from typing import Any + import pytest from pyiceberg.expressions import ( @@ -41,7 +45,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 +92,92 @@ def test_identity_transform_residual() -> None: assert residual == AlwaysFalse() +def test_residual_evaluator_does_not_mutate_prepared_state() -> None: + schema = Schema(NestedField(1, "a", IntegerType()), NestedField(2, "b", IntegerType())) + spec = PartitionSpec( + PartitionField(1, 1001, IdentityTransform(), "a_part"), + PartitionField(2, 1002, IdentityTransform(), "b_part"), + ) + evaluator = residual_evaluator_of( + spec=spec, + expr=And(EqualTo("a", 1), EqualTo("b", 1)), + case_sensitive=True, + schema=schema, + ) + initial_state = vars(evaluator).copy() + + assert evaluator.residual_for(Record(1, 1)) == AlwaysTrue() + assert evaluator.residual_for(Record(0, 0)) == AlwaysFalse() + assert evaluator.residual_for(Record(1, 1)) == AlwaysTrue() + + assert vars(evaluator) == initial_state + + +def test_residual_evaluator_concurrent_calls_do_not_share_partitions() -> 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 residual evaluations") + return value + + schema = Schema(NestedField(1, "a", IntegerType()), NestedField(2, "b", IntegerType())) + spec = PartitionSpec( + PartitionField(1, 1001, IdentityTransform(), "a_part"), + PartitionField(2, 1002, IdentityTransform(), "b_part"), + ) + evaluator = residual_evaluator_of( + spec=spec, + expr=And(EqualTo("a", 1), EqualTo("b", 1)), + case_sensitive=True, + schema=schema, + ) + first_read = Event() + release_first_read = Event() + + with ThreadPoolExecutor(max_workers=2) as executor: + matching_result = executor.submit( + evaluator.residual_for, + BlockingRecord(first_read, release_first_read, 1, 1), + ) + assert first_read.wait(timeout=5) + + try: + non_matching_result = executor.submit(evaluator.residual_for, Record(0, 0)).result(timeout=5) + finally: + release_first_read.set() + + assert matching_result.result(timeout=5) == AlwaysTrue() + assert non_matching_result == 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) + + evaluator = residual_evaluator_of(spec=spec, expr=EqualTo("dateint", 20170815), case_sensitive=True, schema=schema) + + assert evaluator.residual_for(Record(20170815)) == AlwaysTrue() + assert evaluator.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())) From 5676096db56afd27394f48056d4d44e8d8c0886d Mon Sep 17 00:00:00 2001 From: Aaron Niskode-Dossett Date: Wed, 15 Jul 2026 11:50:27 -0500 Subject: [PATCH 2/3] Cache residuals during scan planning --- pyiceberg/table/__init__.py | 43 +++- .../table/test_residual_evaluator_planning.py | 197 ++++++++++++++++++ 2 files changed, 231 insertions(+), 9 deletions(-) create mode 100644 tests/table/test_residual_evaluator_planning.py diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 63b87d290e..b7ee50923f 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 @@ -37,6 +38,7 @@ _InclusiveMetricsEvaluator, bind, expression_evaluator, + extract_field_ids, inclusive_projection, manifest_evaluator, ) @@ -117,6 +119,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 relevant partition values without adding +# unbounded key storage when scans contain a distinct value for every data file. +_RESIDUAL_CACHE_MAX_SIZE = 128 @dataclass() @@ -2620,7 +2625,32 @@ 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) + 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 residual_cache_key_positions[data_file.spec_id]) + cache_key = data_file.spec_id, partition_values + try: + return residual_cache[cache_key] + except KeyError: + residual = residual_evaluators[data_file.spec_id].residual_for(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 +2674,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 +2712,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/table/test_residual_evaluator_planning.py b/tests/table/test_residual_evaluator_planning.py new file mode 100644 index 0000000000..f8eba4fda0 --- /dev/null +++ b/tests/table/test_residual_evaluator_planning.py @@ -0,0 +1,197 @@ +# 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 typing import Any + +import pytest + +import pyiceberg.table as table_module +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[tuple[Any, ...]] = [] + + def residual_for(self, partition: Record) -> BooleanExpression: + 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: 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), + 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 _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,)), + ] + evaluators = {0: _CountingResidualEvaluator(0), 1: _CountingResidualEvaluator(1)} + evaluator_builds: list[int] = [] + planner = _planner(table_v2, EqualTo("x", 1), _identity_spec(0, 1), _identity_spec(1, 1)) + + def build_evaluator(spec_id: int) -> _CountingResidualEvaluator: + evaluator_builds.append(spec_id) + return evaluators[spec_id] + + monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries])) + monkeypatch.setattr(planner, "_build_residual_evaluator", build_evaluator) + + tasks = list(planner.plan_files([])) + + assert evaluator_builds == [0, 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), + EqualTo("x", 11), + EqualTo("x", 11), + EqualTo("x", 2), + ] + + +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,)), + ] + evaluator = _CountingResidualEvaluator(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) + + 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 e9fc91070d573bf2c13c1ffbb7f9ba73d83969b4 Mon Sep 17 00:00:00 2001 From: Aaron Niskode-Dossett Date: Wed, 15 Jul 2026 11:51:37 -0500 Subject: [PATCH 3/3] Benchmark residual planning --- .../test_residual_evaluator_benchmark.py | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 tests/benchmark/test_residual_evaluator_benchmark.py diff --git a/tests/benchmark/test_residual_evaluator_benchmark.py b/tests/benchmark/test_residual_evaluator_benchmark.py new file mode 100644 index 0000000000..7c8330db51 --- /dev/null +++ b/tests/benchmark/test_residual_evaluator_benchmark.py @@ -0,0 +1,118 @@ +# 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 planning with a realistic 15-leaf predicate. + +Every file has a unique unreferenced partition-hash value. The repeated case +measures cache reuse by relevant partition values, while the unique case forces +cache misses. + +Run with: + uv run pytest tests/benchmark/test_residual_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, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus +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 _row_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 + + +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(relevant_partition, file_number), + 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_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_relevant_partitions) for file_number in range(num_files)] + schema = Schema( + NestedField(1, "event_day", LongType(), required=True), + NestedField(2, "region_id", LongType(), required=True), + NestedField(3, "partition_hash", LongType(), required=True), + ) + spec = PartitionSpec( + PartitionField(1, 1000, IdentityTransform(), "event_day"), + PartitionField(3, 1001, IdentityTransform(), "partition_hash"), + spec_id=0, + ) + metadata = TableMetadataV2( + location="s3://bucket/table", + last_column_id=3, + 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=_row_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_relevant_partitions} relevant partitions " + f"with a 15-leaf predicate in {statistics.mean(timings):.3f}s (best: {min(timings):.3f}s)" + )