Skip to content

Commit 0992d66

Browse files
authored
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
1 parent 267118b commit 0992d66

5 files changed

Lines changed: 380 additions & 25 deletions

File tree

‎pyiceberg/expressions/visitors.py‎

Lines changed: 9 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1822,31 +1822,32 @@ def _can_contain_nans(self, field_id: int) -> bool:
18221822

18231823

18241824
class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC):
1825-
"""Finds the residuals for an Expression the partitions in the given PartitionSpec.
1825+
"""Find residuals for an expression using partition values.
18261826
18271827
A residual expression is made by partially evaluating an expression using partition values.
18281828
For example, if a table is partitioned by day(utc_timestamp) and is read with a filter expression
1829-
utc_timestamp > a and utc_timestamp < b, then there are 4 possible residuals expressions
1829+
utc_timestamp > a and utc_timestamp < b, then there are 4 possible residual expressions
18301830
for the partition data, d:
18311831
1832-
1833-
1. If d > day(a) and d &lt; day(b), the residual is always true
1832+
1. If d > day(a) and d < day(b), the residual is always true
18341833
2. If d == day(a) and d != day(b), the residual is utc_timestamp > a
1835-
3. if d == day(b) and d != day(a), the residual is utc_timestamp < b
1834+
3. If d == day(b) and d != day(a), the residual is utc_timestamp < b
18361835
4. If d == day(a) == day(b), the residual is utc_timestamp > a and utc_timestamp < b
1837-
Partition data is passed using StructLike. Residuals are returned by residualFor(StructLike).
18381836
"""
18391837

18401838
schema: Schema
18411839
spec: PartitionSpec
18421840
case_sensitive: bool
18431841
expr: BooleanExpression
1842+
partition_schema: Schema
1843+
struct: Record
18441844

18451845
def __init__(self, schema: Schema, spec: PartitionSpec, case_sensitive: bool, expr: BooleanExpression) -> None:
18461846
self.schema = schema
18471847
self.spec = spec
18481848
self.case_sensitive = case_sensitive
18491849
self.expr = expr
1850+
self.partition_schema = Schema(*spec.partition_type(schema).fields)
18501851

18511852
def eval(self, partition_data: Record) -> BooleanExpression:
18521853
self.struct = partition_data
@@ -1967,17 +1968,12 @@ def visit_bound_predicate(self, predicate: BoundPredicate) -> BooleanExpression:
19671968
if parts == []:
19681969
return predicate
19691970

1970-
def struct_to_schema(struct: StructType) -> Schema:
1971-
return Schema(*struct.fields)
1972-
19731971
for part in parts:
19741972
strict_projection = part.transform.strict_project(part.name, predicate)
19751973
strict_result = None
19761974

19771975
if strict_projection is not None:
1978-
bound = strict_projection.bind(
1979-
struct_to_schema(self.spec.partition_type(self.schema)), case_sensitive=self.case_sensitive
1980-
)
1976+
bound = strict_projection.bind(self.partition_schema, case_sensitive=self.case_sensitive)
19811977
if isinstance(bound, BoundPredicate):
19821978
strict_result = super().visit_bound_predicate(bound)
19831979
else:
@@ -1990,9 +1986,7 @@ def struct_to_schema(struct: StructType) -> Schema:
19901986
inclusive_projection = part.transform.project(part.name, predicate)
19911987
inclusive_result = None
19921988
if inclusive_projection is not None:
1993-
bound_inclusive = inclusive_projection.bind(
1994-
struct_to_schema(self.spec.partition_type(self.schema)), case_sensitive=self.case_sensitive
1995-
)
1989+
bound_inclusive = inclusive_projection.bind(self.partition_schema, case_sensitive=self.case_sensitive)
19961990
if isinstance(bound_inclusive, BoundPredicate):
19971991
# using predicate method specific to inclusive
19981992
inclusive_result = super().visit_bound_predicate(bound_inclusive)

‎pyiceberg/table/__init__.py‎

Lines changed: 41 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from types import TracebackType
3131
from typing import TYPE_CHECKING, Any, TypeVar
3232

33+
from cachetools import LRUCache
3334
from pydantic import Field
3435

3536
from pyiceberg.exceptions import CommitFailedException, ValidationException
@@ -39,6 +40,7 @@
3940
_InclusiveMetricsEvaluator,
4041
bind,
4142
expression_evaluator,
43+
extract_field_ids,
4244
inclusive_projection,
4345
manifest_evaluator,
4446
)
@@ -141,6 +143,9 @@ class UpsertResult:
141143

142144

143145
class TableProperties:
146+
RESIDUAL_CACHE_MAX_SIZE = "read.residual-cache.max-size"
147+
RESIDUAL_CACHE_MAX_SIZE_DEFAULT = 128
148+
144149
PARQUET_ROW_GROUP_SIZE_BYTES = "write.parquet.row-group-size-bytes"
145150
PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT = 128 * 1024 * 1024 # 128 MB
146151

@@ -2841,7 +2846,39 @@ def plan_files(
28412846
data_entries: list[ManifestEntry] = []
28422847
delete_index = DeleteFileIndex()
28432848

2844-
residual_evaluators: dict[int, Callable[[DataFile], ResidualEvaluator]] = KeyDefaultDict(self._build_residual_evaluator)
2849+
residual_evaluators: dict[int, ResidualEvaluator] = KeyDefaultDict(self._build_residual_evaluator)
2850+
referenced_field_ids = extract_field_ids(
2851+
bind(self.table_metadata.schema(), self.row_filter, case_sensitive=self.case_sensitive)
2852+
)
2853+
partition_specs = self.table_metadata.specs()
2854+
residual_cache_key_positions: dict[int, tuple[int, ...]] = KeyDefaultDict(
2855+
lambda spec_id: tuple(
2856+
pos
2857+
for pos, partition_field in enumerate(partition_specs[spec_id].fields)
2858+
if partition_field.source_id in referenced_field_ids
2859+
)
2860+
)
2861+
# A residual can only depend on partition fields derived from source columns
2862+
# referenced by the scan filter. Keep the cache local and bounded.
2863+
residual_cache_max_size = property_as_int(
2864+
self.options,
2865+
TableProperties.RESIDUAL_CACHE_MAX_SIZE,
2866+
TableProperties.RESIDUAL_CACHE_MAX_SIZE_DEFAULT,
2867+
)
2868+
if residual_cache_max_size is None or residual_cache_max_size <= 0:
2869+
raise ValueError(f"{TableProperties.RESIDUAL_CACHE_MAX_SIZE} must be a positive integer")
2870+
residual_cache: LRUCache[tuple[int, tuple[Any, ...]], BooleanExpression] = LRUCache(maxsize=residual_cache_max_size)
2871+
2872+
def residual_for(data_file: DataFile) -> BooleanExpression:
2873+
partition = data_file.partition
2874+
partition_values = tuple(partition[pos] for pos in residual_cache_key_positions[data_file.spec_id])
2875+
cache_key = data_file.spec_id, partition_values
2876+
try:
2877+
return residual_cache[cache_key]
2878+
except KeyError:
2879+
residual = residual_evaluators[data_file.spec_id].residual_for(partition)
2880+
residual_cache[cache_key] = residual
2881+
return residual
28452882

28462883
for manifest_entry in chain.from_iterable(self.plan_manifest_entries(manifests)):
28472884
if not manifest_entry_filter(manifest_entry):
@@ -2865,9 +2902,7 @@ def plan_files(
28652902
data_entry.data_file,
28662903
partition_key=data_entry.data_file.partition,
28672904
),
2868-
residual=residual_evaluators[data_entry.data_file.spec_id](data_entry.data_file).residual_for(
2869-
data_entry.data_file.partition
2870-
),
2905+
residual=residual_for(data_entry.data_file),
28712906
)
28722907
for data_entry in data_entries
28732908
]
@@ -2904,15 +2939,12 @@ def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]:
29042939
include_empty_files,
29052940
).eval
29062941

2907-
def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], ResidualEvaluator]:
2942+
def _build_residual_evaluator(self, spec_id: int) -> ResidualEvaluator:
29082943
spec = self.table_metadata.specs()[spec_id]
29092944

29102945
from pyiceberg.expressions.visitors import residual_evaluator_of
29112946

2912-
# The lambda created here is run in multiple threads.
2913-
# So we avoid creating _EvaluatorExpression methods bound to a single
2914-
# shared instance across multiple threads.
2915-
return lambda datafile: residual_evaluator_of(
2947+
return residual_evaluator_of(
29162948
spec=spec,
29172949
expr=self.row_filter,
29182950
case_sensitive=self.case_sensitive,
Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""Benchmark residual planning with a realistic 15-leaf predicate.
18+
19+
Every file has a unique unreferenced partition-hash value. The repeated case
20+
measures cache reuse by relevant partition values, while the unique case forces
21+
cache misses.
22+
23+
Run with:
24+
uv run pytest tests/benchmark/test_residual_evaluator_benchmark.py -v -s -m benchmark
25+
"""
26+
27+
from __future__ import annotations
28+
29+
import statistics
30+
import timeit
31+
32+
import pytest
33+
34+
from pyiceberg.expressions import And, BooleanExpression, EqualTo, GreaterThanOrEqual, LessThanOrEqual, Or
35+
from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus
36+
from pyiceberg.partitioning import PartitionField, PartitionSpec
37+
from pyiceberg.schema import Schema
38+
from pyiceberg.table import ManifestGroupPlanner, Table
39+
from pyiceberg.table.metadata import TableMetadataV2
40+
from pyiceberg.transforms import IdentityTransform
41+
from pyiceberg.typedef import Record
42+
from pyiceberg.types import LongType, NestedField
43+
44+
45+
def _row_filter() -> BooleanExpression:
46+
"""Select five day ranges, each scoped to a region."""
47+
windows = ((0, 1, 1), (2, 3, 4), (4, 5, 7), (6, 7, 10), (8, 10, 13))
48+
branches = [
49+
And(
50+
And(GreaterThanOrEqual("event_day", start_day), LessThanOrEqual("event_day", end_day)),
51+
EqualTo("region_id", region_id),
52+
)
53+
for start_day, end_day, region_id in windows
54+
]
55+
56+
combined = branches[0]
57+
for branch in branches[1:]:
58+
combined = Or(combined, branch)
59+
return combined
60+
61+
62+
def _manifest_entry(file_number: int, relevant_partition: int) -> ManifestEntry:
63+
data_file = DataFile.from_args(
64+
content=DataFileContent.DATA,
65+
file_path=f"s3://bucket/data-{file_number}.parquet",
66+
file_format=FileFormat.PARQUET,
67+
partition=Record(relevant_partition, file_number),
68+
record_count=1,
69+
file_size_in_bytes=1,
70+
)
71+
data_file.spec_id = 0
72+
return ManifestEntry.from_args(
73+
status=ManifestEntryStatus.ADDED,
74+
snapshot_id=1,
75+
sequence_number=1,
76+
file_sequence_number=1,
77+
data_file=data_file,
78+
)
79+
80+
81+
@pytest.mark.benchmark
82+
@pytest.mark.parametrize(
83+
"num_relevant_partitions",
84+
[7, 2_000],
85+
ids=["repeated-relevant-partitions", "unique-relevant-partitions"],
86+
)
87+
def test_residual_planning(table_v2: Table, monkeypatch: pytest.MonkeyPatch, num_relevant_partitions: int) -> None:
88+
num_files = 2_000
89+
entries = [_manifest_entry(file_number, file_number % num_relevant_partitions) for file_number in range(num_files)]
90+
schema = Schema(
91+
NestedField(1, "event_day", LongType(), required=True),
92+
NestedField(2, "region_id", LongType(), required=True),
93+
NestedField(3, "partition_hash", LongType(), required=True),
94+
)
95+
spec = PartitionSpec(
96+
PartitionField(1, 1000, IdentityTransform(), "event_day"),
97+
PartitionField(3, 1001, IdentityTransform(), "partition_hash"),
98+
spec_id=0,
99+
)
100+
metadata = TableMetadataV2(
101+
location="s3://bucket/table",
102+
last_column_id=3,
103+
schemas=[schema],
104+
current_schema_id=schema.schema_id,
105+
partition_specs=[spec],
106+
default_spec_id=spec.spec_id,
107+
)
108+
planner = ManifestGroupPlanner(table_metadata=metadata, io=table_v2.io, row_filter=_row_filter())
109+
110+
monkeypatch.setattr(planner, "plan_manifest_entries", lambda _: iter([entries]))
111+
112+
timings = timeit.repeat(lambda: list(planner.plan_files([])), number=1, repeat=3)
113+
114+
assert len(list(planner.plan_files([]))) == num_files
115+
print(
116+
f"Planned {num_files} files across {num_relevant_partitions} relevant partitions "
117+
f"with a 15-leaf predicate in {statistics.mean(timings):.3f}s (best: {min(timings):.3f}s)"
118+
)

‎tests/expressions/test_residual_evaluator.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
AlwaysFalse,
2222
AlwaysTrue,
2323
And,
24+
BooleanExpression,
2425
EqualTo,
2526
GreaterThan,
2627
GreaterThanOrEqual,
@@ -36,7 +37,7 @@
3637
StartsWith,
3738
)
3839
from pyiceberg.expressions.literals import literal
39-
from pyiceberg.expressions.visitors import residual_evaluator_of
40+
from pyiceberg.expressions.visitors import ResidualVisitor, residual_evaluator_of
4041
from pyiceberg.partitioning import PartitionField, PartitionSpec
4142
from pyiceberg.schema import Schema
4243
from pyiceberg.transforms import DayTransform, IdentityTransform
@@ -88,6 +89,27 @@ def test_identity_transform_residual() -> None:
8889
assert residual == AlwaysFalse()
8990

9091

92+
def test_residual_visitor_preserves_public_eval_api() -> None:
93+
schema = Schema(NestedField(1, "a", IntegerType()))
94+
spec = PartitionSpec(PartitionField(1, 1001, IdentityTransform(), "a_part"))
95+
visitor = ResidualVisitor(schema=schema, spec=spec, case_sensitive=True, expr=EqualTo("a", 1))
96+
97+
assert visitor.eval(Record(1)) == AlwaysTrue()
98+
assert visitor.eval(Record(0)) == AlwaysFalse()
99+
100+
101+
def test_residual_visitor_subclass_can_customize_evaluation() -> None:
102+
class FalseForTrueResidualVisitor(ResidualVisitor):
103+
def visit_true(self) -> BooleanExpression:
104+
return AlwaysFalse()
105+
106+
schema = Schema(NestedField(1, "a", IntegerType()))
107+
spec = PartitionSpec(PartitionField(1, 1001, IdentityTransform(), "a_part"))
108+
visitor = FalseForTrueResidualVisitor(schema=schema, spec=spec, case_sensitive=True, expr=AlwaysTrue())
109+
110+
assert visitor.eval(Record(1)) == AlwaysFalse()
111+
112+
91113
def test_case_insensitive_identity_transform_residuals() -> None:
92114
schema = Schema(NestedField(50, "dateint", IntegerType()), NestedField(51, "hour", IntegerType()))
93115

0 commit comments

Comments
 (0)