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..819a8df431 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,35 @@ def plan_files( data_entries: list[ManifestEntry] = [] delete_index = DeleteFileIndex() - residual_evaluators: dict[int, Callable[[DataFile], 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) + ) + 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: tuple[int, tuple[Any, ...]] = data_file.spec_id, partition_values + try: + return residual_cache[cache_key] + except KeyError: + residual_evaluator = residual_evaluator_factories[data_file.spec_id](data_file) + residual = residual_evaluator.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 +2677,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 ] @@ -2689,10 +2720,9 @@ def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], Residu 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( + # 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/benchmark/test_scan_planning_benchmark.py b/tests/benchmark/test_scan_planning_benchmark.py new file mode 100644 index 0000000000..4d7d393571 --- /dev/null +++ b/tests/benchmark/test_scan_planning_benchmark.py @@ -0,0 +1,109 @@ +# 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 with a high-cardinality irrelevant partition field. + +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 +""" + +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.table import ManifestGroupPlanner, Table +from pyiceberg.transforms import IdentityTransform +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("y", source + 1), + EqualTo("y", source + 2), + EqualTo("y", source + 3), + ): + 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, 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)] + 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=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_relevant_partitions} relevant 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..e360e7d2fc --- /dev/null +++ b/tests/table/test_scan_planning.py @@ -0,0 +1,202 @@ +# 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]) + + +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, + 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,)), + ] + 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: evaluator_factories[spec_id]) + + tasks = list(planner.plan_files([])) + + 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), + 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_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_factory) + + tasks = list(planner.plan_files([])) + + 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)] + + +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_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_factory) + + list(planner.plan_files([])) + + 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( + 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_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_factory) + + list(planner.plan_files([])) + + 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: + 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_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_factory) + + tasks = list(planner.plan_files([])) + + 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), + EqualTo("x", 3), + EqualTo("x", 1), + ]