paimon-python/pypaimon/execution/python/globalindex/global_index_evaluator.py

# 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.

"""Global index evaluator for filtering data using global indexes."""

import threading
import struct
from collections import deque
from concurrent.futures import Future
from typing import Callable, Collection, Dict, List, Optional

from pypaimon.globalindex.global_index_reader import GlobalIndexReader, FieldRef, _completed_future
from pypaimon.globalindex.global_index_result import GlobalIndexResult
from pypaimon.globalindex.global_index_evaluation import GlobalIndexEvaluation
from pypaimon.execution.python.globalindex.key_serializer import create_serializer
from pypaimon.common.predicate import Predicate
from pypaimon.schema.data_types import AtomicType, DataField


class GlobalIndexEvaluator:
    """Predicate evaluator for filtering data using global indexes.

    Reader visit methods return Future internally — the evaluator no longer
    dispatches to an executor.
    """

    def __init__(
        self,
        fields: List[DataField],
        readers_function: Callable[[DataField], Collection[GlobalIndexReader]],
    ):
        self._fields = fields
        self._field_by_name = {f.name: f for f in fields}
        self._readers_function = readers_function
        self._index_readers_cache: Dict[int, Collection[GlobalIndexReader]] = {}

    def evaluate(
        self,
        predicate: Optional[Predicate]
    ) -> Optional[GlobalIndexResult]:
        if predicate is None:
            return None
        evaluation = self._visit_async(predicate).result()
        return evaluation.result if evaluation is not None else None

    def evaluate_with_contributing_fields(
        self,
        predicate: Optional[Predicate]
    ) -> Optional[GlobalIndexEvaluation]:
        """Return matches and fields whose supported indexes contributed."""
        if predicate is None:
            return None
        return self._visit_async(predicate).result()

    def _visit_async(self, predicate) -> Future:
        if isinstance(predicate, Predicate) and predicate.method in ('and', 'or'):
            return self._visit_compound_async(predicate)
        return self._visit_leaf_async(predicate)

    def _visit_leaf_async(self, predicate: Predicate) -> Future:
        field = self._field_by_name.get(predicate.field)
        if field is None:
            f = Future()
            f.set_result(None)
            return f

        if not self._decimal_literals_supported(field, predicate):
            return _completed_future(None)

        field_id = field.id
        readers = self._index_readers_cache.get(field_id)
        if readers is None:
            readers = self._readers_function(field)
            self._index_readers_cache[field_id] = readers

        field_ref = FieldRef(predicate.index, predicate.field, str(field.type))

        reader_futures = []
        for reader in readers:
            reader_futures.append(
                self._visit_function(reader, predicate, field_ref)
            )

        all_done = Future()
        if not reader_futures:
            all_done.set_result(None)
            return all_done

        remaining = [len(reader_futures)]
        count_lock = threading.Lock()

        def on_done(_):
            with count_lock:
                remaining[0] -= 1
                if remaining[0] == 0:
                    try:
                        all_done.set_result(
                            self._combine_reader_results(
                                reader_futures, field_id)
                        )
                    except Exception as e:
                        all_done.set_exception(e)

        for rf in reader_futures:
            rf.add_done_callback(on_done)

        return all_done

    @staticmethod
    def _decimal_literals_supported(field: DataField, predicate: Predicate) -> bool:
        if not (isinstance(field.type, AtomicType)
                and field.type.type.upper().startswith(('DECIMAL', 'NUMERIC'))):
            return True
        # Like Native, decline a physical index lookup if conversion would
        # round a comparison literal. The original predicate must then filter
        # the data; a rounded bitmap lookup can silently discard matching rows.
        serializer = create_serializer(field.type)
        compare = serializer.create_comparator()
        try:
            return all(literal is None or compare(
                serializer.deserialize(serializer.serialize(literal)), literal) == 0
                for literal in (predicate.literals or []))
        except (ArithmeticError, ValueError, struct.error):
            return False

    def _combine_reader_results(
        self, reader_futures: List[Future], field_id: int,
    ) -> Optional[GlobalIndexEvaluation]:
        compound_result: Optional[GlobalIndexResult] = None
        for f in reader_futures:
            child_result = f.result()
            if child_result is None:
                continue
            if compound_result is not None:
                # Readers answer the same predicate: an exact result intersected
                # with a candidate superset remains exact.
                is_exact = compound_result.is_exact() or child_result.is_exact()
                compound_result = GlobalIndexResult.create(
                    compound_result.and_(child_result).results(), is_exact=is_exact)
            else:
                compound_result = child_result
            if compound_result.is_empty():
                break
        if compound_result is None:
            return None
        return GlobalIndexEvaluation(compound_result, frozenset([field_id]))

    def _visit_compound_async(self, predicate: Predicate) -> Future:
        children = self._flatten_children(predicate.method, predicate.literals)
        child_futures = [self._visit_async(child) for child in children]

        all_done = Future()
        if not child_futures:
            all_done.set_result(None)
            return all_done

        remaining = [len(child_futures)]
        lock = threading.Lock()

        def on_done(_):
            with lock:
                remaining[0] -= 1
                if remaining[0] == 0:
                    try:
                        results = [f.result() for f in child_futures]
                        all_done.set_result(
                            self._combine_results(results, predicate.method)
                        )
                    except Exception as e:
                        all_done.set_exception(e)

        for cf in child_futures:
            cf.add_done_callback(on_done)

        return all_done

    def _combine_results(
        self, results: List[Optional[GlobalIndexEvaluation]], method: str
    ) -> Optional[GlobalIndexEvaluation]:
        contributing_field_ids = set()
        if method == 'or':
            compound_result = GlobalIndexResult.create_empty()
            for child in results:
                if child is None:
                    return None
                compound_result = compound_result.or_(child.result)
                contributing_field_ids.update(child.contributing_field_ids)
            return GlobalIndexEvaluation(compound_result,
                                         frozenset(contributing_field_ids))
        else:
            compound_result: Optional[GlobalIndexResult] = None
            for child in results:
                if child is not None:
                    if compound_result is not None:
                        compound_result = compound_result.and_(child.result)
                    else:
                        compound_result = child.result
                    contributing_field_ids.update(
                        child.contributing_field_ids)
                if compound_result is not None and compound_result.is_empty():
                    break
            if compound_result is None:
                return None
            if any(child is None for child in results):
                # A dropped AND child can share a field with a supported child.
                # Contributing field ids alone therefore cannot prove exactness.
                compound_result = GlobalIndexResult.create(
                    compound_result.results(), is_exact=False)
            return GlobalIndexEvaluation(compound_result,
                                         frozenset(contributing_field_ids))

    def _flatten_children(self, method: str, children) -> list:
        result = []
        stack = deque(children)
        while stack:
            child = stack.popleft()
            if isinstance(child, Predicate) and child.method == method:
                for grandchild in reversed(child.literals):
                    stack.appendleft(grandchild)
            else:
                result.append(child)
        return result

    def _visit_function(
        self,
        reader: GlobalIndexReader,
        predicate: Predicate,
        field_ref: FieldRef
    ) -> 'Future[Optional[GlobalIndexResult]]':
        method = predicate.method
        literals = predicate.literals

        if method == 'equal':
            return reader.visit_equal(field_ref, literals[0])
        elif method == 'notEqual':
            return reader.visit_not_equal(field_ref, literals[0])
        elif method == 'lessThan':
            return reader.visit_less_than(field_ref, literals[0])
        elif method == 'lessOrEqual':
            return reader.visit_less_or_equal(field_ref, literals[0])
        elif method == 'greaterThan':
            return reader.visit_greater_than(field_ref, literals[0])
        elif method == 'greaterOrEqual':
            return reader.visit_greater_or_equal(field_ref, literals[0])
        elif method == 'isNull':
            return reader.visit_is_null(field_ref)
        elif method == 'isNotNull':
            return reader.visit_is_not_null(field_ref)
        elif method == 'in':
            return reader.visit_in(field_ref, literals)
        elif method == 'notIn':
            return reader.visit_not_in(field_ref, literals)
        elif method == 'startsWith':
            return reader.visit_starts_with(field_ref, literals[0])
        elif method == 'endsWith':
            return reader.visit_ends_with(field_ref, literals[0])
        elif method == 'contains':
            return reader.visit_contains(field_ref, literals[0])
        elif method == 'like':
            return reader.visit_like(field_ref, literals[0])
        elif method == 'between':
            return reader.visit_between(field_ref, literals[0], literals[1])
        elif method == 'notBetween':
            return reader.visit_not_between(field_ref, literals[0], literals[1])

        from pypaimon.globalindex.global_index_reader import _completed_future
        return _completed_future(None)

    def close(self) -> None:
        for readers in self._index_readers_cache.values():
            for reader in readers:
                try:
                    reader.close()
                except Exception:
                    pass
        self._index_readers_cache.clear()

    def __enter__(self) -> 'GlobalIndexEvaluator':
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()