| # 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', |
| ]) |
| |
| |
| 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 |
| if predicate.method == 'and' or predicate.method == 'or': |
| return all( |
| predicate_supports_arrow_filter(p) |
| for p in (predicate.literals or []) |
| ) |
| return predicate.method not in _UNSAFE_ARROW_FILTER_METHODS |
| |
| |
| 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 |
