Make metrics evaluation stateless

pyiceberg/expressions/visitors.py

@@ -16,7 +16,7 @@
 # under the License.
 import math
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from functools import singledispatch
 from typing import (
     Any,
@@ -1112,12 +1112,19 @@ def expression_to_plain_format(
     return [visit(expression, visitor) for expression in expressions]
 
 
-class _MetricsEvaluator(BoundBooleanExpressionVisitor[bool], ABC):
-    value_counts: dict[int, int]
-    null_counts: dict[int, int]
-    nan_counts: dict[int, int]
-    lower_bounds: dict[int, bytes]
-    upper_bounds: dict[int, bytes]
+class _MetricsEvaluationVisitor(BoundBooleanExpressionVisitor[bool], ABC):
+    value_counts: Mapping[int, int]
+    null_counts: Mapping[int, int]
+    nan_counts: Mapping[int, int]
+    lower_bounds: Mapping[int, bytes]
+    upper_bounds: Mapping[int, bytes]
+
+    def __init__(self, file: DataFile) -> None:
+        self.value_counts = file.value_counts or EMPTY_DICT
+        self.null_counts = file.null_value_counts or EMPTY_DICT
+        self.nan_counts = file.nan_value_counts or EMPTY_DICT
+        self.lower_bounds = file.lower_bounds or EMPTY_DICT
+        self.upper_bounds = file.upper_bounds or EMPTY_DICT
 
     def visit_true(self) -> bool:
         # all rows match
@@ -1154,14 +1161,15 @@ def _is_nan(self, val: Any) -> bool:
             return False
 
 
-class _InclusiveMetricsEvaluator(_MetricsEvaluator):
-    struct: StructType
+class _InclusiveMetricsEvaluator:
+    """Bind an inclusive metrics expression once and evaluate files without mutating prepared state."""
+
     expr: BooleanExpression
+    include_empty_files: bool
 
     def __init__(
         self, schema: Schema, expr: BooleanExpression, case_sensitive: bool = True, include_empty_files: bool = False
     ) -> None:
-        self.struct = schema.as_struct()
         self.include_empty_files = include_empty_files
         self.expr = bind(schema, rewrite_not(expr), case_sensitive)
 
@@ -1176,13 +1184,11 @@ def eval(self, file: DataFile) -> bool:
             # be updated once we implemented and set correct record count.
             return ROWS_MIGHT_MATCH
 
-        self.value_counts = file.value_counts or EMPTY_DICT
-        self.null_counts = file.null_value_counts or EMPTY_DICT
-        self.nan_counts = file.nan_value_counts or EMPTY_DICT
-        self.lower_bounds = file.lower_bounds or EMPTY_DICT
-        self.upper_bounds = file.upper_bounds or EMPTY_DICT
+        return visit(self.expr, _InclusiveMetricsEvaluationVisitor(file))
+
 
-        return visit(self.expr, self)
+class _InclusiveMetricsEvaluationVisitor(_MetricsEvaluationVisitor):
+    """Evaluate inclusive metrics for one data file."""
 
     def _may_contain_null(self, field_id: int) -> bool:
         return self.null_counts is None or (field_id in self.null_counts and self.null_counts.get(field_id) is not None)
@@ -1489,9 +1495,12 @@ def visit_bound_predicate(self, predicate: BoundPredicate) -> BooleanExpression:
         return result
 
 
-class _StrictMetricsEvaluator(_MetricsEvaluator):
+class _StrictMetricsEvaluator:
+    """Bind a strict metrics expression once and evaluate files without mutating prepared state."""
+
     struct: StructType
     expr: BooleanExpression
+    include_empty_files: bool
 
     def __init__(
         self, schema: Schema, expr: BooleanExpression, case_sensitive: bool = True, include_empty_files: bool = False
@@ -1515,13 +1524,17 @@ def eval(self, file: DataFile) -> bool:
             # be updated once we implemented and set correct record count.
             return ROWS_MUST_MATCH
 
-        self.value_counts = file.value_counts or EMPTY_DICT
-        self.null_counts = file.null_value_counts or EMPTY_DICT
-        self.nan_counts = file.nan_value_counts or EMPTY_DICT
-        self.lower_bounds = file.lower_bounds or EMPTY_DICT
-        self.upper_bounds = file.upper_bounds or EMPTY_DICT
+        return visit(self.expr, _StrictMetricsEvaluationVisitor(self.struct, file))
+
+
+class _StrictMetricsEvaluationVisitor(_MetricsEvaluationVisitor):
+    """Evaluate strict metrics for one data file."""
 
-        return visit(self.expr, self)
+    struct: StructType
+
+    def __init__(self, struct: StructType, file: DataFile) -> None:
+        super().__init__(file)
+        self.struct = struct
 
     def visit_is_null(self, term: BoundTerm) -> bool:
         # no need to check whether the field is required because binding evaluates that case

pyiceberg/table/__init__.py

@@ -2586,6 +2586,7 @@ def plan_manifest_entries(self, manifests: Iterable[ManifestFile]) -> Iterator[l
         # this filter depends on the partition spec used to write the manifest file
 
         partition_evaluators: dict[int, Callable[[DataFile], bool]] = KeyDefaultDict(self._build_partition_evaluator)
+        metrics_evaluator = self._build_metrics_evaluator()
 
         min_sequence_number = _min_sequence_number(manifests)
 
@@ -2597,7 +2598,7 @@ def plan_manifest_entries(self, manifests: Iterable[ManifestFile]) -> Iterator[l
                     self.io,
                     manifest,
                     partition_evaluators[manifest.partition_spec_id],
-                    self._build_metrics_evaluator(),
+                    metrics_evaluator,
                 )
                 for manifest in manifests
                 if self._check_sequence_number(min_sequence_number, manifest)
@@ -2674,15 +2675,14 @@ def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]:
         schema = self.table_metadata.schema()
         include_empty_files = strtobool(self.options.get("include_empty_files", "false"))
 
-        # The lambda created here is run in multiple threads.
-        # So we avoid creating _InclusiveMetricsEvaluator methods bound to a single
-        # shared instance across multiple threads.
-        return lambda data_file: _InclusiveMetricsEvaluator(
+        # Metrics evaluators keep file-specific state local to each call, so one
+        # prepared evaluator can be shared across all manifest tasks in this plan.
+        return _InclusiveMetricsEvaluator(
             schema,
             self.row_filter,
             self.case_sensitive,
             include_empty_files,
-        ).eval(data_file)
+        ).eval
 
     def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], ResidualEvaluator]:
         spec = self.table_metadata.specs()[spec_id]

tests/benchmark/test_metrics_evaluator_benchmark.py

@@ -0,0 +1,153 @@
+# 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 metrics predicate through manifest planning.
+
+Run with:
+    uv run pytest tests/benchmark/test_metrics_evaluator_benchmark.py -v -s -m benchmark
+"""
+
+from __future__ import annotations
+
+import statistics
+import timeit
+from collections.abc import Callable
+from typing import Any
+
+import pytest
+
+import pyiceberg.table as table_module
+from pyiceberg.conversions import to_bytes
+from pyiceberg.expressions import And, BooleanExpression, EqualTo, GreaterThanOrEqual, LessThanOrEqual, Or
+from pyiceberg.manifest import DataFile, FileFormat, ManifestContent, ManifestFile
+from pyiceberg.schema import Schema
+from pyiceberg.table import ManifestGroupPlanner, Table
+from pyiceberg.typedef import Record
+from pyiceberg.types import LongType, NestedField
+from pyiceberg.utils.concurrent import ExecutorFactory
+
+
+def _data_file(file_number: int) -> DataFile:
+    event_day = file_number % 11
+    region_id = file_number % 15
+    event_day_bytes = to_bytes(LongType(), event_day)
+    region_id_bytes = to_bytes(LongType(), region_id)
+    return DataFile.from_args(
+        file_path=f"s3://bucket/data-{file_number}.parquet",
+        file_format=FileFormat.PARQUET,
+        partition=Record(),
+        record_count=100,
+        file_size_in_bytes=1,
+        value_counts={1: 100, 2: 100},
+        null_value_counts={1: 0, 2: 0},
+        lower_bounds={1: event_day_bytes, 2: region_id_bytes},
+        upper_bounds={1: event_day_bytes, 2: region_id_bytes},
+    )
+
+
+def _metrics_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_file(manifest_number: int) -> ManifestFile:
+    return ManifestFile.from_args(
+        manifest_path=f"s3://bucket/manifest-{manifest_number}.avro",
+        manifest_length=1,
+        partition_spec_id=0,
+        content=ManifestContent.DATA,
+        sequence_number=1,
+        min_sequence_number=1,
+        added_snapshot_id=1,
+    )
+
+
+class _SerialExecutor:
+    def map(self, function: Callable[..., Any], *iterables: Any, **_kwargs: Any) -> Any:
+        return map(function, *iterables)
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    ("files_per_manifest", "partition_matches"),
+    [(1_000, True), (1, True), (1_000, False), (1, False)],
+    ids=["dense", "fragmented", "all-pruned-dense", "all-pruned-fragmented"],
+)
+def test_metrics_evaluator_reuse(
+    table_v2: Table,
+    monkeypatch: pytest.MonkeyPatch,
+    files_per_manifest: int,
+    partition_matches: bool,
+) -> None:
+    num_files = 1_000
+    schema = Schema(
+        NestedField(1, "event_day", LongType(), required=True),
+        NestedField(2, "region_id", LongType(), required=True),
+        *(NestedField(field_id, f"unused_{field_id}", LongType(), required=False) for field_id in range(3, 103)),
+        schema_id=table_v2.metadata.current_schema_id,
+    )
+    metadata = table_v2.metadata.model_copy(update={"schemas": [schema]})
+    planner = ManifestGroupPlanner(table_metadata=metadata, io=table_v2.io, row_filter=_metrics_filter())
+    data_files = [_data_file(file_number) for file_number in range(num_files)]
+    manifests = [_manifest_file(manifest_number) for manifest_number in range(0, num_files, files_per_manifest)]
+    files_by_manifest = {
+        manifest.manifest_path: data_files[start : start + files_per_manifest]
+        for manifest, start in zip(manifests, range(0, num_files, files_per_manifest), strict=True)
+    }
+
+    def open_manifest(
+        _io: Any,
+        manifest: ManifestFile,
+        partition_evaluator: Callable[[DataFile], bool],
+        metrics_evaluator: Callable[[DataFile], bool],
+    ) -> list[DataFile]:
+        return [
+            data_file
+            for data_file in files_by_manifest[manifest.manifest_path]
+            if partition_evaluator(data_file) and metrics_evaluator(data_file)
+        ]
+
+    monkeypatch.setattr(planner, "_build_manifest_evaluator", lambda _: lambda _: True)
+    monkeypatch.setattr(planner, "_build_partition_evaluator", lambda _: lambda _: partition_matches)
+    monkeypatch.setattr(table_module, "_open_manifest", open_manifest)
+    monkeypatch.setattr(ExecutorFactory, "get_or_create", staticmethod(lambda: _SerialExecutor()))
+
+    def evaluate_files() -> int:
+        return sum(len(entries) for entries in planner.plan_manifest_entries(manifests))
+
+    assert evaluate_files() == (67 if partition_matches else 0)
+    number = 1 if partition_matches else (1_000 if files_per_manifest == 1_000 else 10)
+    timings_ms = [timing * 1_000 / number for timing in timeit.repeat(evaluate_files, number=number, repeat=5)]
+    file_label = "file" if files_per_manifest == 1 else "files"
+    pruning_label = "all files pruned" if not partition_matches else "metrics evaluated"
+
+    print(
+        f"Evaluated metrics for {num_files} files with {files_per_manifest} {file_label} per manifest, "
+        f"a 102-column schema, and a 15-leaf predicate ({pruning_label}) in "
+        f"{statistics.mean(timings_ms):.3f}ms (best: {min(timings_ms):.3f}ms)"
+    )