diff --git a/pyiceberg/expressions/visitors.py b/pyiceberg/expressions/visitors.py index 320cd3110e..fa663179b0 100644 --- a/pyiceberg/expressions/visitors.py +++ b/pyiceberg/expressions/visitors.py @@ -454,16 +454,25 @@ def expression_evaluator(schema: Schema, unbound: BooleanExpression, case_sensit return _ExpressionEvaluator(schema, unbound, case_sensitive).eval -class _ExpressionEvaluator(BoundBooleanExpressionVisitor[bool]): +class _ExpressionEvaluator: + """An evaluator that binds an expression once and keeps evaluation state local to each call.""" + bound: BooleanExpression - struct: StructProtocol def __init__(self, schema: Schema, unbound: BooleanExpression, case_sensitive: bool): self.bound = bind(schema, unbound, case_sensitive) def eval(self, struct: StructProtocol) -> bool: + return visit(self.bound, _ExpressionEvaluationVisitor(struct)) + + +class _ExpressionEvaluationVisitor(BoundBooleanExpressionVisitor[bool]): + """Evaluate a bound expression against one struct.""" + + struct: StructProtocol + + def __init__(self, struct: StructProtocol): self.struct = struct - return visit(self.bound, self) def visit_in(self, term: BoundTerm, literals: set[L]) -> bool: return term.eval(self.struct) in literals diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 63b87d290e..e4aa545031 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -2664,11 +2664,11 @@ def _build_partition_evaluator(self, spec_id: int) -> Callable[[DataFile], bool] partition_type = spec.partition_type(self.table_metadata.schema()) partition_schema = Schema(*partition_type.fields) partition_expr = self.partition_filters[spec_id] + evaluator = expression_evaluator(partition_schema, partition_expr, self.case_sensitive) - # 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 data_file: expression_evaluator(partition_schema, partition_expr, self.case_sensitive)(data_file.partition) + # Expression evaluators keep input-specific state local to each call, so the + # prepared evaluator can be shared by every manifest using this spec. + return lambda data_file: evaluator(data_file.partition) def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]: schema = self.table_metadata.schema() diff --git a/tests/benchmark/test_partition_evaluator_benchmark.py b/tests/benchmark/test_partition_evaluator_benchmark.py new file mode 100644 index 0000000000..428ad7cff8 --- /dev/null +++ b/tests/benchmark/test_partition_evaluator_benchmark.py @@ -0,0 +1,112 @@ +# 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 a realistic 15-leaf partition predicate when a prepared evaluator is shared across manifests. + +Run with: + uv run pytest tests/benchmark/test_partition_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, FileFormat +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 _data_file(file_number: int) -> DataFile: + return DataFile.from_args( + file_path=f"s3://bucket/data-{file_number}.parquet", + file_format=FileFormat.PARQUET, + partition=Record(file_number % 11, file_number % 15), + record_count=100, + file_size_in_bytes=1, + ) + + +def _partition_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 + + +@pytest.mark.benchmark +@pytest.mark.parametrize( + "files_per_manifest", + [1_000, 1], + ids=["many-files-per-manifest", "one-file-per-manifest"], +) +def test_partition_evaluator_reuse(table_v2: Table, files_per_manifest: int) -> None: + num_files = 1_000 + schema = Schema( + NestedField(1, "event_day", LongType(), required=True), + NestedField(2, "region_id", LongType(), required=True), + ) + spec = PartitionSpec( + PartitionField(1, 1000, IdentityTransform(), "event_day"), + PartitionField(2, 1001, IdentityTransform(), "region_id"), + spec_id=0, + ) + metadata = TableMetadataV2( + location="s3://bucket/table", + last_column_id=2, + 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=_partition_filter()) + data_files = [_data_file(file_number) for file_number in range(num_files)] + + def evaluate_files() -> int: + partition_evaluator = planner._build_partition_evaluator(spec.spec_id) + matches = 0 + for start in range(0, num_files, files_per_manifest): + matches += sum(partition_evaluator(data_file) for data_file in data_files[start : start + files_per_manifest]) + return matches + + assert evaluate_files() == 67 + iterations = 100 + timings_ms = [timing * 1_000 / iterations for timing in timeit.repeat(evaluate_files, number=iterations, repeat=3)] + file_label = "file" if files_per_manifest == 1 else "files" + + print( + f"Evaluated partitions for {num_files} files with {files_per_manifest} {file_label} per manifest " + f"and a 15-leaf predicate in " + f"{statistics.mean(timings_ms):.3f}ms (best: {min(timings_ms):.3f}ms)" + ) diff --git a/tests/expressions/test_visitors.py b/tests/expressions/test_visitors.py index a0634d683c..7cee3836eb 100644 --- a/tests/expressions/test_visitors.py +++ b/tests/expressions/test_visitors.py @@ -16,6 +16,8 @@ # under the License. # pylint:disable=redefined-outer-name +from concurrent.futures import ThreadPoolExecutor +from threading import Event from typing import Any import pytest @@ -67,6 +69,7 @@ BindVisitor, BooleanExpressionVisitor, BoundBooleanExpressionVisitor, + _ExpressionEvaluator, _ManifestEvalVisitor, expression_evaluator, expression_to_plain_format, @@ -1630,6 +1633,57 @@ def test_expression_evaluator_null() -> None: assert expression_evaluator(schema, NotStartsWith("a", 1), case_sensitive=True)(struct) is True +def test_expression_evaluator_does_not_mutate_prepared_state() -> None: + schema = Schema( + NestedField(1, "a", IntegerType(), required=True), + NestedField(2, "b", IntegerType(), required=True), + ) + evaluator = _ExpressionEvaluator(schema, And(EqualTo("a", 1), EqualTo("b", 1)), case_sensitive=True) + initial_state = vars(evaluator).copy() + + assert evaluator.eval(Record(1, 1)) is True + assert evaluator.eval(Record(0, 0)) is False + assert evaluator.eval(Record(1, 1)) is True + + assert vars(evaluator) == initial_state + + +def test_expression_evaluator_concurrent_calls_do_not_share_records() -> 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 expression evaluations") + return value + + schema = Schema( + NestedField(1, "a", IntegerType(), required=True), + NestedField(2, "b", IntegerType(), required=True), + ) + evaluator = expression_evaluator(schema, And(EqualTo("a", 1), EqualTo("b", 1)), case_sensitive=True) + first_read = Event() + release_first_read = Event() + + with ThreadPoolExecutor(max_workers=2) as executor: + matching_result = executor.submit(evaluator, BlockingRecord(first_read, release_first_read, 1, 1)) + assert first_read.wait(timeout=5) + + try: + non_matching_result = executor.submit(evaluator, Record(0, 0)).result(timeout=5) + finally: + release_first_read.set() + + assert matching_result.result(timeout=5) is True + assert non_matching_result is False + + def test_expression_evaluator_binary_starts_with() -> None: schema = Schema(NestedField(1, "x", BinaryType(), required=False), schema_id=1) struct = Record(b"aa") diff --git a/tests/table/test_partition_evaluator_planning.py b/tests/table/test_partition_evaluator_planning.py new file mode 100644 index 0000000000..b79762e5f6 --- /dev/null +++ b/tests/table/test_partition_evaluator_planning.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. + +from __future__ import annotations + +from collections.abc import Callable + +import pytest + +import pyiceberg.table as table_module +from pyiceberg.expressions import BooleanExpression, GreaterThan +from pyiceberg.io import FileIO +from pyiceberg.manifest import DataFile, FileFormat, ManifestContent, ManifestEntry, ManifestFile +from pyiceberg.schema import Schema +from pyiceberg.table import ManifestGroupPlanner, Table +from pyiceberg.typedef import Record, StructProtocol + + +def _data_file(file_number: int, partition_value: int) -> DataFile: + return DataFile.from_args( + file_path=f"s3://bucket/data-{file_number}.parquet", + file_format=FileFormat.PARQUET, + partition=Record(partition_value), + record_count=100, + file_size_in_bytes=1, + ) + + +def _manifest_file(file_number: int) -> ManifestFile: + return ManifestFile.from_args( + manifest_path=f"s3://bucket/manifest-{file_number}.avro", + manifest_length=1, + partition_spec_id=0, + content=ManifestContent.DATA, + sequence_number=1, + min_sequence_number=1, + added_snapshot_id=1, + ) + + +def test_partition_evaluator_prepares_once_per_spec(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None: + evaluator_calls: list[list[int]] = [] + + def counting_expression_evaluator( + schema: Schema, unbound: BooleanExpression, case_sensitive: bool + ) -> Callable[[StructProtocol], bool]: + calls: list[int] = [] + evaluator_calls.append(calls) + + def evaluate(struct: StructProtocol) -> bool: + value = struct[0] + calls.append(value) + return value > 5 + + return evaluate + + monkeypatch.setattr(table_module, "expression_evaluator", counting_expression_evaluator) + planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=GreaterThan("x", 5)) + partition_evaluator = planner._build_partition_evaluator(0) + + assert len(evaluator_calls) == 1 + assert not partition_evaluator(_data_file(1, 1)) + assert partition_evaluator(_data_file(2, 10)) + assert evaluator_calls == [[1, 10]] + + +def test_manifest_group_planner_shares_partition_evaluator_across_manifests( + table_v2: Table, monkeypatch: pytest.MonkeyPatch +) -> None: + planner = ManifestGroupPlanner(table_metadata=table_v2.metadata, io=table_v2.io, row_filter=GreaterThan("x", 5)) + built_specs: list[int] = [] + opened_evaluators: list[Callable[[DataFile], bool]] = [] + + def build_partition_evaluator(spec_id: int) -> Callable[[DataFile], bool]: + built_specs.append(spec_id) + return lambda _: True + + def open_manifest( + io: FileIO, + manifest: ManifestFile, + partition_evaluator: Callable[[DataFile], bool], + metrics_evaluator: Callable[[DataFile], bool], + ) -> list[ManifestEntry]: + opened_evaluators.append(partition_evaluator) + return [] + + monkeypatch.setattr(planner, "_build_manifest_evaluator", lambda _: lambda _: True) + monkeypatch.setattr(planner, "_build_partition_evaluator", build_partition_evaluator) + monkeypatch.setattr(table_module, "_open_manifest", open_manifest) + + list(planner.plan_manifest_entries([_manifest_file(1), _manifest_file(2)])) + + assert built_specs == [0] + assert len(opened_evaluators) == 2 + assert opened_evaluators[0] is opened_evaluators[1]