paimon-python/pypaimon/read/push_down_utils.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.

from typing import Dict, List, Optional, Set

from pypaimon.common.predicate import Predicate
from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon.schema.data_types import DataField

_UNSAFE_ARROW_FILTER_METHODS = frozenset([
    'startsWith',
    'endsWith',
    'contains',
    'like',
])

# Large boolean trees can overflow or crash native Dataset scanners even when
# balanced. Keep complex predicates on Paimon's exact row-filter path instead.
_MAX_ARROW_FILTER_LEAVES = 256


def extract_partition_spec_from_predicate(
    predicate: Predicate, partition_keys: List[str]
) -> Optional[Dict[str, str]]:
    if not predicate or not partition_keys:
        return None
    parts = _split_and(predicate)
    spec: Dict[str, str] = {}
    for p in parts:
        if p.method != "equal" or p.field is None or p.literals is None or len(p.literals) != 1:
            continue
        if p.field in partition_keys:
            spec[p.field] = str(p.literals[0])
    if set(spec.keys()) == set(partition_keys):
        return spec
    return None


def trim_and_transform_predicate(input_predicate: Predicate, all_fields: List[str], trimmed_keys: List[str]):
    new_predicate = trim_predicate_by_fields(input_predicate, trimmed_keys)
    part_to_index = {element: idx for idx, element in enumerate(trimmed_keys)}
    mapping: Dict[int, int] = {
        i: part_to_index.get(all_fields[i], -1)
        for i in range(len(all_fields))
    }
    return _change_index(new_predicate, mapping)


def trim_predicate_by_fields(input_predicate: Predicate, trimmed_keys: List[str]):
    if not input_predicate or not trimmed_keys:
        return None

    predicates: list[Predicate] = _split_and(input_predicate)
    predicates = [element for element in predicates if _get_all_fields(element).issubset(trimmed_keys)]
    return PredicateBuilder.and_predicates(predicates)


def _split_and(input_predicate: Predicate):
    if not input_predicate:
        return list()

    if input_predicate.method == 'and':
        return [p for element in (input_predicate.literals or []) for p in _split_and(element)]

    return [input_predicate]


def rewrite_predicate_indices(
    input_predicate: Optional[Predicate],
    read_fields: List[DataField],
) -> Optional[Predicate]:
    """Rewrite predicate leaf indices to match positions in ``read_fields``.

    Predicate leaves are built against the original table schema (via
    PredicateBuilder), so their ``index`` field encodes that schema's column
    order. When the same predicate is later evaluated row-by-row against a
    projected scan (read_type narrower or reordered), those indices no longer
    match the OffsetRow layout the reader hands to FilterRecordReader, and
    ``OffsetRow.get_field(idx)`` raises IndexError.

    Returns a new predicate where every leaf's ``index`` is rebound to its
    column's position in ``read_fields``. The caller is responsible for
    ensuring that every leaf field is present in ``read_fields``.
    """
    if input_predicate is None:
        return None
    name_to_pos = {f.name: i for i, f in enumerate(read_fields)}
    return _rewrite_by_name(input_predicate, name_to_pos)


def _rewrite_by_name(p: Predicate, name_to_pos: Dict[str, int]) -> Predicate:
    if p.method == 'and' or p.method == 'or':
        return p.new_literals(
            [_rewrite_by_name(c, name_to_pos) for c in (p.literals or [])]
        )
    if p.field is None or p.field not in name_to_pos:
        raise ValueError(
            "Cannot rewrite predicate index for leaf {!r}: field {!r} is not "
            "in read fields {}. The caller must ensure all referenced columns "
            "are projected.".format(p, p.field, list(name_to_pos))
        )
    return p.new_index(name_to_pos[p.field])


def _change_index(input_predicate: Predicate, mapping: Dict[int, int]):
    if not input_predicate:
        return None

    if input_predicate.method == 'and' or input_predicate.method == 'or':
        predicates: list[Predicate] = input_predicate.literals
        new_predicates = [_change_index(element, mapping) for element in predicates]
        return input_predicate.new_literals(new_predicates)

    return input_predicate.new_index(mapping[input_predicate.index])


def predicate_field_names(predicate: Predicate) -> Set[str]:
    """Return all column names referenced by predicate leaves."""
    if predicate.field is not None:
        return {predicate.field}
    involved_fields = set()
    if predicate.literals:
        for sub_predicate in predicate.literals:
            involved_fields.update(predicate_field_names(sub_predicate))
    return involved_fields


def _get_all_fields(predicate: Predicate) -> Set[str]:
    return predicate_field_names(predicate)


def predicate_supports_arrow_filter(predicate: Optional[Predicate]) -> bool:
    """Whether ``predicate.to_arrow()`` is safe for batch filtering.

    PyArrow 6 accepts dataset expressions for comparisons, null checks, and
    isin, but string match compute functions do not accept dataset expressions.
    Predicate.to_arrow() currently falls back to a truthy expression or None for
    those methods, which is safe for file pruning but not for final row filters.
    """
    if predicate is None:
        return True

    leaves = 0
    pending = [predicate]
    while pending:
        current = pending.pop()
        if current.method == 'and' or current.method == 'or':
            pending.extend(current.literals or [])
            continue
        if current.method in _UNSAFE_ARROW_FILTER_METHODS:
            return False
        leaves += 1
        if leaves > _MAX_ARROW_FILTER_LEAVES:
            return False
    return True


def remove_row_id_filter(predicate: Predicate) -> Optional[Predicate]:
    from pypaimon.table.special_fields import SpecialFields

    if not predicate:
        return None
    if predicate.field == SpecialFields.ROW_ID.name:
        return None
    if predicate.method == "and":
        parts = _split_and(predicate)
        non_row_id = [
            p for p in parts
            if _get_all_fields(p) != {SpecialFields.ROW_ID.name}
        ]
        if not non_row_id:
            return None
        filtered = []
        for p in non_row_id:
            r = remove_row_id_filter(p)
            if r is None:
                return None
            filtered.append(r)
        return PredicateBuilder.and_predicates(filtered)
    if predicate.method == "or":
        fields = _get_all_fields(predicate)
        if (
            SpecialFields.ROW_ID.name in fields
            and fields != {SpecialFields.ROW_ID.name}
        ):
            return predicate
        new_children = []
        for c in predicate.literals or []:
            r = remove_row_id_filter(c)
            if r is not None:
                new_children.append(r)
        if not new_children:
            return None
        if len(new_children) == 1:
            return new_children[0]
        return PredicateBuilder.or_predicates(new_children)
    return predicate


def exclude_predicate_with_fields(predicate: Optional[Predicate], fields: Set[str]) -> Optional[Predicate]:
    """Drop predicate parts referencing any of ``fields`` (mirrors Java
    PredicateBuilder.excludePredicateWithFields)."""
    if not predicate or not fields:
        return predicate
    if predicate.method == "and":
        kept = []
        for p in _split_and(predicate):
            r = exclude_predicate_with_fields(p, fields)
            if r is not None:
                kept.append(r)
        return PredicateBuilder.and_predicates(kept) if kept else None
    # leaf or OR: drop the whole thing if it touches any field (OR isn't split apart)
    if _get_all_fields(predicate) & fields:
        return None
    return predicate