This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-3666-267118b21d04c5bb26063b2417b3fe28726a253e in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
commit 0992d661a9bb540265fde4e8f397fecdfb55a510 Author: Aaron Niskode-Dossett <[email protected]> AuthorDate: Wed Oct 28 22:23:40 2026 +0000 Cache residuals with stateless evaluators (#3666) * Make residual evaluation stateless * Cache residuals during scan planning * Benchmark residual planning * Address residual evaluator review feedback * Preserve residual visitor compatibility --- pyiceberg/expressions/visitors.py | 24 +-- pyiceberg/table/__init__.py | 50 +++++- .../benchmark/test_residual_evaluator_benchmark.py | 118 +++++++++++++ tests/expressions/test_residual_evaluator.py | 24 ++- tests/table/test_residual_evaluator_planning.py | 189 +++++++++++++++++++++ 5 files changed, 380 insertions(+), 25 deletions(-) diff --git a/pyiceberg/expressions/visitors.py b/pyiceberg/expressions/visitors.py index c8ef8d4c4..a8de08866 100644 --- a/pyiceberg/expressions/visitors.py +++ b/pyiceberg/expressions/visitors.py @@ -1822,31 +1822,32 @@ class _StrictMetricsEvaluationVisitor(_MetricsEvaluationVisitor): class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC): - """Finds the residuals for an Expression the partitions in the given PartitionSpec. + """Find residuals for an expression using partition values. A residual expression is made by partially evaluating an expression using partition values. For example, if a table is partitioned by day(utc_timestamp) and is read with a filter expression - utc_timestamp > a and utc_timestamp < b, then there are 4 possible residuals expressions + utc_timestamp > a and utc_timestamp < b, then there are 4 possible residual expressions for the partition data, d: - - 1. If d > day(a) and d < day(b), the residual is always true + 1. If d > day(a) and d < day(b), the residual is always true 2. If d == day(a) and d != day(b), the residual is utc_timestamp > a - 3. if d == day(b) and d != day(a), the residual is utc_timestamp < b + 3. If d == day(b) and d != day(a), the residual is utc_timestamp < b 4. If d == day(a) == day(b), the residual is utc_timestamp > a and utc_timestamp < b - Partition data is passed using StructLike. Residuals are returned by residualFor(StructLike). """ 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: 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 @@ -1967,17 +1968,12 @@ class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC): 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: @@ -1990,9 +1986,7 @@ class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC): 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 7da0e59d2..8d79472af 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -30,6 +30,7 @@ from itertools import chain from types import TracebackType from typing import TYPE_CHECKING, Any, TypeVar +from cachetools import LRUCache from pydantic import Field from pyiceberg.exceptions import CommitFailedException, ValidationException @@ -39,6 +40,7 @@ from pyiceberg.expressions.visitors import ( _InclusiveMetricsEvaluator, bind, expression_evaluator, + extract_field_ids, inclusive_projection, manifest_evaluator, ) @@ -141,6 +143,9 @@ class UpsertResult: class TableProperties: + RESIDUAL_CACHE_MAX_SIZE = "read.residual-cache.max-size" + RESIDUAL_CACHE_MAX_SIZE_DEFAULT = 128 + PARQUET_ROW_GROUP_SIZE_BYTES = "write.parquet.row-group-size-bytes" PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT = 128 * 1024 * 1024 # 128 MB @@ -2841,7 +2846,39 @@ class ManifestGroupPlanner: 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_max_size = property_as_int( + self.options, + TableProperties.RESIDUAL_CACHE_MAX_SIZE, + TableProperties.RESIDUAL_CACHE_MAX_SIZE_DEFAULT, + ) + if residual_cache_max_size is None or residual_cache_max_size <= 0: + raise ValueError(f"{TableProperties.RESIDUAL_CACHE_MAX_SIZE} must be a positive integer") + 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): @@ -2865,9 +2902,7 @@ class ManifestGroupPlanner: 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 ] @@ -2904,15 +2939,12 @@ class ManifestGroupPlanner: include_empty_files, ).eval - 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_residual_evaluator_benchmark.py b/tests/benchmark/test_residual_evaluator_benchmark.py new file mode 100644 index 000000000..7c8330db5 --- /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, + ) + + [email protected] [email protected]( + "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)" + ) diff --git a/tests/expressions/test_residual_evaluator.py b/tests/expressions/test_residual_evaluator.py index fbd6a993b..4c6295e9e 100644 --- a/tests/expressions/test_residual_evaluator.py +++ b/tests/expressions/test_residual_evaluator.py @@ -21,6 +21,7 @@ from pyiceberg.expressions import ( AlwaysFalse, AlwaysTrue, And, + BooleanExpression, EqualTo, GreaterThan, GreaterThanOrEqual, @@ -36,7 +37,7 @@ from pyiceberg.expressions import ( StartsWith, ) from pyiceberg.expressions.literals import literal -from pyiceberg.expressions.visitors import residual_evaluator_of +from pyiceberg.expressions.visitors import ResidualVisitor, residual_evaluator_of from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema from pyiceberg.transforms import DayTransform, IdentityTransform @@ -88,6 +89,27 @@ def test_identity_transform_residual() -> None: assert residual == AlwaysFalse() +def test_residual_visitor_preserves_public_eval_api() -> None: + schema = Schema(NestedField(1, "a", IntegerType())) + spec = PartitionSpec(PartitionField(1, 1001, IdentityTransform(), "a_part")) + visitor = ResidualVisitor(schema=schema, spec=spec, case_sensitive=True, expr=EqualTo("a", 1)) + + assert visitor.eval(Record(1)) == AlwaysTrue() + assert visitor.eval(Record(0)) == AlwaysFalse() + + +def test_residual_visitor_subclass_can_customize_evaluation() -> None: + class FalseForTrueResidualVisitor(ResidualVisitor): + def visit_true(self) -> BooleanExpression: + return AlwaysFalse() + + schema = Schema(NestedField(1, "a", IntegerType())) + spec = PartitionSpec(PartitionField(1, 1001, IdentityTransform(), "a_part")) + visitor = FalseForTrueResidualVisitor(schema=schema, spec=spec, case_sensitive=True, expr=AlwaysTrue()) + + assert visitor.eval(Record(1)) == AlwaysFalse() + + def test_case_insensitive_identity_transform_residuals() -> None: schema = Schema(NestedField(50, "dateint", IntegerType()), NestedField(51, "hour", IntegerType())) diff --git a/tests/table/test_residual_evaluator_planning.py b/tests/table/test_residual_evaluator_planning.py new file mode 100644 index 000000000..7fe2ffda9 --- /dev/null +++ b/tests/table/test_residual_evaluator_planning.py @@ -0,0 +1,189 @@ +# 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 Iterable, Iterator +from typing import Any + +import pytest + +from pyiceberg.expressions import AlwaysFalse, AlwaysTrue, And, BooleanExpression, EqualTo +from pyiceberg.io import FileIO +from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus, ManifestFile +from pyiceberg.partitioning import PartitionField, PartitionSpec +from pyiceberg.table import ManifestGroupPlanner, Table, TableProperties +from pyiceberg.table.metadata import TableMetadata +from pyiceberg.transforms import BucketTransform, IdentityTransform +from pyiceberg.typedef import EMPTY_DICT, Properties, Record +from pyiceberg.types import LongType + + +class _ManifestEntriesPlanner(ManifestGroupPlanner): + def __init__( + self, + table_metadata: TableMetadata, + io: FileIO, + row_filter: BooleanExpression, + entries: list[ManifestEntry], + options: Properties = EMPTY_DICT, + ) -> None: + super().__init__(table_metadata=table_metadata, io=io, row_filter=row_filter, options=options) + self.entries = entries + + def plan_manifest_entries(self, _manifests: Iterable[ManifestFile]) -> Iterator[list[ManifestEntry]]: + return iter([self.entries]) + + +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, + entries: list[ManifestEntry], + *partition_specs: PartitionSpec, + options: Properties = EMPTY_DICT, +) -> _ManifestEntriesPlanner: + metadata = table_v2.metadata.model_copy(update={"partition_specs": list(partition_specs)}) + return _ManifestEntriesPlanner( + table_metadata=metadata, + io=table_v2.io, + row_filter=row_filter, + entries=entries, + options=options, + ) + + +def test_plan_files_returns_correct_residuals_for_repeated_relevant_partitions(table_v2: Table) -> 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)), + ] + planner = _planner(table_v2, EqualTo("x", 1), entries, _identity_spec(0, 1, 2)) + + tasks = list(planner.plan_files([])) + + assert [task.residual for task in tasks] == [AlwaysTrue(), AlwaysTrue(), AlwaysFalse()] + + +def test_plan_files_distinguishes_each_referenced_partition_field(table_v2: Table) -> None: + entries = [ + _manifest_entry(0, spec_id=0, partition=(1, 10)), + _manifest_entry(1, spec_id=0, partition=(1, 20)), + ] + planner = _planner( + table_v2, + And(EqualTo("x", 1), EqualTo("y", 10)), + entries, + _identity_spec(0, 1, 2), + ) + + tasks = list(planner.plan_files([])) + + assert [task.residual for task in tasks] == [AlwaysTrue(), AlwaysFalse()] + + +def test_plan_files_isolates_residuals_by_partition_spec(table_v2: Table) -> None: + predicate = EqualTo("x", 1) + entries = [ + _manifest_entry(0, spec_id=0, partition=(1,)), + _manifest_entry(1, spec_id=1, partition=(1,)), + ] + planner = _planner( + table_v2, + predicate, + entries, + _identity_spec(0, 1), + _identity_spec(1, 2), + ) + + tasks = list(planner.plan_files([])) + + assert [task.residual for task in tasks] == [AlwaysTrue(), predicate] + + +def test_plan_files_distinguishes_each_transform_for_a_referenced_field(table_v2: Table) -> None: + bucket_7: BucketTransform[int] = BucketTransform(7) + bucket_5: BucketTransform[int] = BucketTransform(5) + x_bucket_7 = bucket_7.transform(LongType())(1) + x_bucket_5 = bucket_5.transform(LongType())(1) + assert x_bucket_7 is not None + assert x_bucket_5 is not None + + spec = PartitionSpec( + PartitionField(1, 1000, bucket_7, "x_bucket_7"), + PartitionField(1, 1001, bucket_5, "x_bucket_5"), + PartitionField(2, 1002, IdentityTransform(), "partition_hash"), + spec_id=0, + ) + predicate = EqualTo("x", 1) + entries = [ + _manifest_entry(0, spec_id=0, partition=(x_bucket_7, x_bucket_5, 10)), + _manifest_entry(1, spec_id=0, partition=(x_bucket_7, (x_bucket_5 + 1) % 5, 20)), + ] + planner = _planner(table_v2, predicate, entries, spec) + + tasks = list(planner.plan_files([])) + + assert [task.residual for task in tasks] == [predicate, AlwaysFalse()] + + [email protected]("cache_size", ["0", "-1"]) +def test_plan_files_rejects_non_positive_residual_cache_size(table_v2: Table, cache_size: str) -> None: + planner = _planner( + table_v2, + EqualTo("x", 1), + [_manifest_entry(0, spec_id=0, partition=(1,))], + _identity_spec(0, 1), + options={TableProperties.RESIDUAL_CACHE_MAX_SIZE: cache_size}, + ) + + with pytest.raises(ValueError, match="read.residual-cache.max-size must be a positive integer"): + list(planner.plan_files([]))
