This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
The following commit(s) were added to refs/heads/main by this push:
new 0992d661a Cache residuals with stateless evaluators (#3666)
0992d661a is described below
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([]))
