| # 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. | |
| """ | |
| Predicate-driven bucket pruning for HASH_FIXED tables. | |
| Mirrors Java's ``org.apache.paimon.operation.BucketSelectConverter``: | |
| walk the predicate, isolate AND clauses that constrain bucket-key fields | |
| with Equal/In, take the cartesian product of literal values, hash each | |
| combination using the writer's hash routine, and produce the set of | |
| buckets the query can possibly hit. All other entries are safely dropped. | |
| Hard correctness contract: the bucket set this returns is a *superset* of | |
| the buckets that contain any matching rows. False-positive (over-keep) | |
| allowed; false-negative (drop a bucket that has matching rows) MUST never | |
| happen — that would be silent data loss. | |
| The hashing routine reuses ``_hash_bytes_by_words`` / | |
| ``_bucket_from_hash`` from ``pypaimon.execution.python.write.input_routing`` — the same | |
| code path the writer uses to assign rows to buckets. Reusing it (rather | |
| than copying) is what guarantees read/write hash agreement in the face of | |
| future routine changes. | |
| Conservative scope (deliberately narrower than Java's general flexibility): | |
| * Only HASH_FIXED tables (caller's responsibility to gate; this module | |
| does not look at the bucket mode itself). | |
| * All bucket-key fields must be constrained, with Equal or In, in a | |
| single AND-of-OR-of-literals shape. If any bucket-key column is | |
| unconstrained, return None — the caller must scan all buckets. | |
| * Repeated constraints on the same bucket-key column under top-level | |
| AND (e.g. ``id IN (1,2,3) AND id IN (2,3,4)``) intersect their | |
| literal sets (mirrors Java ``BucketSelector.retainAll``). An empty | |
| intersection means the predicate is unsatisfiable, and we return | |
| None. | |
| * Total cartesian product capped at MAX_VALUES (1000), again matching | |
| Java; above that, fall back to a full scan. | |
| Returns a callable ``selector(partition, bucket: int, total_buckets: int) | |
| -> bool``. The callable is cached per ``(partition, total_buckets)`` to | |
| handle (a) bucket count variation across snapshots (rescale) and (b) | |
| per-partition predicate specialisation: predicates of the form | |
| ``(part='a' AND bk IN (1,2)) OR (part='b' AND bk IN (3,4))`` are | |
| simplified per concrete partition value before bucket selection, so each | |
| partition gets its own tight bucket set. | |
| When ``partition`` is ``None`` (early manifest filter that has not yet | |
| deserialised the entry), the selector falls back to a partition-agnostic | |
| result — sound but possibly wider than the per-partition tight set. | |
| """ | |
| from itertools import product | |
| from typing import Any, Dict, FrozenSet, List, Optional, Set, Tuple, Union | |
| from pypaimon.common.predicate import Predicate | |
| from pypaimon.schema.data_types import DataField | |
| from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer | |
| from pypaimon.table.row.internal_row import InternalRow, RowKind | |
| from pypaimon.execution.python.write.input_routing import ( | |
| _bucket_from_hash, | |
| _hash_bytes_by_words, | |
| ) | |
| MAX_VALUES = 1000 | |
| # Bucket-key column types where the Python serializer is not byte-aligned | |
| # with the writer's logical value, or with Java's ``BinaryRow`` byte layout. | |
| # A divergent hash is silent data loss (false-negative), so the selector | |
| # refuses to build at all when a bucket-key field has one of these types. | |
| # | |
| # Two reasons something gets blacklisted: | |
| # | |
| # 1. Locale / precision drift between writer and reader for equal logical | |
| # values (DECIMAL via float-vs-Decimal, TIMESTAMP via naive datetime | |
| # timezone interpretation). | |
| # 2. Composite / nested types whose ``GenericRowSerializer`` byte layout | |
| # hasn't been cross-validated against Java's ``BinaryRow`` (ARRAY, | |
| # MAP, ROW, MULTISET, VARIANT, BLOB). Until that validation lands, | |
| # treating them as safe risks a hash divergence. | |
| _UNSAFE_BUCKET_KEY_TYPES = frozenset({ | |
| 'DECIMAL', | |
| 'TIMESTAMP', | |
| 'TIMESTAMP_WITH_LOCAL_TIME_ZONE', | |
| 'ARRAY', | |
| 'MAP', | |
| 'ROW', | |
| 'MULTISET', | |
| 'VARIANT', | |
| 'BLOB', | |
| }) | |
| def _has_unsafe_bucket_key_type(bucket_key_fields: List[DataField]) -> bool: | |
| for f in bucket_key_fields: | |
| type_name = getattr(getattr(f, 'type', None), 'type', '') | |
| if not type_name: | |
| continue | |
| head = type_name.split('(')[0].split('<')[0].strip().upper() | |
| if head in _UNSAFE_BUCKET_KEY_TYPES: | |
| return True | |
| return False | |
| def _split_and(p: Predicate) -> List[Predicate]: | |
| if p.method == 'and': | |
| out: List[Predicate] = [] | |
| for child in (p.literals or []): | |
| out.extend(_split_and(child)) | |
| return out | |
| return [p] | |
| def _split_or(p: Predicate) -> List[Predicate]: | |
| if p.method == 'or': | |
| out: List[Predicate] = [] | |
| for child in (p.literals or []): | |
| out.extend(_split_or(child)) | |
| return out | |
| return [p] | |
| def _extract_or_clause(or_pred: Predicate, | |
| bk_name_to_slot: Dict[str, int]) -> Optional[List[Any]]: | |
| """For one AND-child predicate, return either: | |
| * ``[slot_index, [literal, ...]]`` — the OR/leaf is a pure | |
| Equal-or-In list on a single bucket-key field; or | |
| * ``None`` — the clause is not a bucket-key constraint we can | |
| safely use; the caller skips it. | |
| All disjuncts must hit the same bucket-key column. Mixed columns or | |
| non-Equal/In operators disqualify the entire AND clause. | |
| """ | |
| slot: Optional[int] = None | |
| values: List[Any] = [] | |
| for clause in _split_or(or_pred): | |
| if clause.method not in ('equal', 'in'): | |
| return None | |
| if clause.field is None or clause.field not in bk_name_to_slot: | |
| return None | |
| this_slot = bk_name_to_slot[clause.field] | |
| if slot is not None and slot != this_slot: | |
| return None | |
| slot = this_slot | |
| for lit in (clause.literals or []): | |
| # Java filters nulls; null literals are degenerate (NULL = NULL | |
| # is UNKNOWN in SQL). Producing zero values for a slot will | |
| # cascade through the cartesian product to "match nothing", | |
| # which is the same observable behaviour as Java. | |
| if lit is None: | |
| continue | |
| values.append(lit) | |
| return None if slot is None else [slot, values] | |
| def _build_combinations( | |
| predicate: Predicate, | |
| bucket_key_fields: List[DataField]) -> Optional[List[List[Any]]]: | |
| """Walk ``predicate`` for top-level AND clauses constraining bucket-key | |
| columns by Equal/In, intersect repeated constraints, and return the | |
| cartesian product of literal values (one row per combination). | |
| Returns None when the predicate doesn't pin down every bucket-key | |
| column or when the cartesian product exceeds ``MAX_VALUES`` — the | |
| caller treats that as "no pruning, all buckets accept". | |
| """ | |
| bk_name_to_slot: Dict[str, int] = { | |
| f.name: i for i, f in enumerate(bucket_key_fields) | |
| } | |
| n_slots = len(bucket_key_fields) | |
| slot_values: List[Optional[List[Any]]] = [None] * n_slots | |
| for and_child in _split_and(predicate): | |
| extracted = _extract_or_clause(and_child, bk_name_to_slot) | |
| if extracted is None: | |
| # Not a bucket-key constraint — that's fine, just skip it. The | |
| # remaining predicate still describes a SUPERSET of matching | |
| # rows; bucket pruning stays sound as long as we don't *add* | |
| # constraints that aren't actually true. | |
| continue | |
| slot, values = extracted | |
| if slot_values[slot] is not None: | |
| # Same bucket-key column constrained twice in top-level AND | |
| # (e.g. ``id IN (1,2,3) AND id IN (2,3,4)``). Mirror Java's | |
| # ``retainAll``: keep the intersection, bail only when it is | |
| # empty (the predicate is unsatisfiable). | |
| new_values_set = set(values) | |
| intersection = [v for v in slot_values[slot] | |
| if v in new_values_set] | |
| if not intersection: | |
| return None | |
| slot_values[slot] = intersection | |
| else: | |
| slot_values[slot] = values | |
| # Every bucket-key column must be constrained. | |
| for v in slot_values: | |
| if v is None: | |
| return None | |
| # Cartesian-product cap. Above the cap the bucket set is essentially | |
| # all buckets anyway; punting saves the hash computation. | |
| total = 1 | |
| for v in slot_values: | |
| # An empty slot (e.g. all literals were null) collapses the | |
| # product to 0 — observable behaviour: empty bucket set, drop | |
| # everything. Mirrors Java. | |
| total *= len(v) | |
| if total > MAX_VALUES: | |
| return None | |
| return [list(combo) for combo in product(*slot_values)] | |
| def _hash_combinations(combinations: List[List[Any]], | |
| bucket_key_fields: List[DataField], | |
| total_buckets: int) -> FrozenSet[int]: | |
| result = set() | |
| for combo in combinations: | |
| row = GenericRow(list(combo), bucket_key_fields, RowKind.INSERT) | |
| serialized = GenericRowSerializer.to_bytes(row) | |
| # Skip the 4-byte length prefix — matches the writer's hash | |
| # input exactly (see RowKeyExtractor._binary_row_hash_code). | |
| h = _hash_bytes_by_words(serialized[4:]) | |
| result.add(_bucket_from_hash(h, total_buckets)) | |
| return frozenset(result) | |
| def _predicate_touches_partition(predicate: Predicate, | |
| partition_field_names: Set[str]) -> bool: | |
| """True if ``predicate`` references any partition column directly or | |
| inside an AND/OR/NOT subtree.""" | |
| if predicate.method in ('and', 'or', 'not'): | |
| return any(_predicate_touches_partition(c, partition_field_names) | |
| for c in (predicate.literals or [])) | |
| return predicate.field is not None and predicate.field in partition_field_names | |
| def _evaluate_partition_leaf(predicate: Predicate, | |
| partition_values: Dict[str, Any]) -> Optional[bool]: | |
| """Evaluate ``predicate`` (a leaf on a partition column) against the | |
| concrete partition values. Returns True / False, or ``None`` if the | |
| leaf isn't safely evaluable here (caller should keep the leaf | |
| unchanged — bucket selection stays sound as long as we don't fold | |
| away an evaluable False). | |
| """ | |
| field_value = partition_values.get(predicate.field) | |
| tester = Predicate.testers.get(predicate.method) | |
| if tester is None: | |
| return None | |
| try: | |
| return tester.test_by_value(field_value, predicate.literals) | |
| except Exception: | |
| return None | |
| _AlwaysFalse = False # sentinel: predicate always evaluates to False | |
| _AlwaysTrue = None # sentinel: predicate cleared (always True) | |
| def replace_partition_predicate( | |
| predicate: Predicate, | |
| partition_field_names: Set[str], | |
| partition_values: Dict[str, Any]) -> Optional[Union[bool, Predicate]]: | |
| """Substitute partition-column leaves with their concrete values and | |
| fold away always-true / always-false sub-expressions. | |
| Three-way return: | |
| * ``None`` — predicate is unconditionally True after substitution | |
| (no constraint left for this partition). | |
| * ``False`` — predicate is unconditionally False (this partition | |
| cannot contain matching rows). | |
| * ``Predicate`` — the simplified predicate; partition leaves are | |
| gone. The caller continues bucket-key extraction on this. | |
| """ | |
| if predicate.method == 'and': | |
| new_children: List[Predicate] = [] | |
| for child in (predicate.literals or []): | |
| simplified = replace_partition_predicate( | |
| child, partition_field_names, partition_values) | |
| if simplified is _AlwaysFalse: | |
| return _AlwaysFalse | |
| if simplified is _AlwaysTrue: | |
| continue | |
| new_children.append(simplified) | |
| if not new_children: | |
| return _AlwaysTrue | |
| if len(new_children) == 1: | |
| return new_children[0] | |
| return Predicate(method='and', index=None, field=None, | |
| literals=new_children) | |
| if predicate.method == 'or': | |
| new_children = [] | |
| for child in (predicate.literals or []): | |
| simplified = replace_partition_predicate( | |
| child, partition_field_names, partition_values) | |
| if simplified is _AlwaysTrue: | |
| return _AlwaysTrue | |
| if simplified is _AlwaysFalse: | |
| continue | |
| new_children.append(simplified) | |
| if not new_children: | |
| return _AlwaysFalse | |
| if len(new_children) == 1: | |
| return new_children[0] | |
| return Predicate(method='or', index=None, field=None, | |
| literals=new_children) | |
| # Leaf predicate. | |
| if predicate.field is not None and predicate.field in partition_field_names: | |
| truth = _evaluate_partition_leaf(predicate, partition_values) | |
| if truth is True: | |
| return _AlwaysTrue | |
| if truth is False: | |
| return _AlwaysFalse | |
| # Couldn't safely evaluate — keep the leaf. Bucket selection | |
| # stays sound: the leaf still gets ANDed in, just doesn't help | |
| # narrow buckets for this partition. | |
| return predicate | |
| # Non-partition leaf: keep as-is. | |
| return predicate | |
| def _partition_to_dict(partition: Optional[InternalRow], | |
| partition_fields: List[DataField]) -> Dict[str, Any]: | |
| """Pull each partition column's value out of ``partition`` keyed by | |
| field name. Returns an empty dict when ``partition`` is None.""" | |
| if partition is None: | |
| return {} | |
| out: Dict[str, Any] = {} | |
| for i, field in enumerate(partition_fields): | |
| try: | |
| out[field.name] = partition.get_field(i) | |
| except Exception: | |
| out[field.name] = None | |
| return out | |
| def _partition_to_cache_key(partition: Optional[InternalRow], | |
| partition_fields: List[DataField] | |
| ) -> Optional[Tuple[Any, ...]]: | |
| if partition is None or not partition_fields: | |
| return None | |
| try: | |
| return tuple(partition.get_field(i) for i in range(len(partition_fields))) | |
| except Exception: | |
| return None | |
| class _Selector: | |
| """Callable bucket filter, lazy + cached per ``(partition, total_buckets)``.""" | |
| __slots__ = ('_predicate', '_bucket_key_fields', '_partition_fields', | |
| '_cache') | |
| def __init__(self, predicate: Predicate, | |
| bucket_key_fields: List[DataField], | |
| partition_fields: Optional[List[DataField]] = None): | |
| self._predicate = predicate | |
| self._bucket_key_fields = bucket_key_fields | |
| self._partition_fields = list(partition_fields or []) | |
| self._cache: Dict[Tuple[Optional[Tuple[Any, ...]], int], FrozenSet[int]] = {} | |
| def __call__(self, *args) -> bool: | |
| # Accept ``(bucket, total_buckets)`` (early manifest filter that | |
| # hasn't deserialised the entry yet — partition unknown) or | |
| # ``(partition, bucket, total_buckets)`` (late filter on a fully | |
| # decoded ``ManifestEntry``). The two-arg form is partition- | |
| # agnostic; partition substitution is skipped. | |
| if len(args) == 2: | |
| partition = None | |
| bucket, total_buckets = args | |
| elif len(args) == 3: | |
| partition, bucket, total_buckets = args | |
| else: | |
| raise TypeError( | |
| "_Selector expects 2 or 3 positional args, got %d" % len(args)) | |
| # ``total_buckets <= 0`` shows up for postpone / legacy / special | |
| # entries and must NOT be pruned: returning False here would drop | |
| # rows the writer hashed under a different convention. Fail open. | |
| if bucket < 0 or total_buckets <= 0: | |
| return True | |
| try: | |
| return bucket in self._compute(partition, total_buckets) | |
| except Exception: | |
| # Fail open on any hashing / serialization / specialisation | |
| # error (e.g. a literal type that doesn't match the bucket-key | |
| # column's atomic type). Crashing the entire scan here would | |
| # be worse than skipping pruning; the soundness contract still | |
| # forbids false-negatives. | |
| return True | |
| def may_contain(self, min_bucket: int, max_bucket: int, total_buckets: int) -> bool: | |
| """Conservatively test an inclusive manifest range without a known partition.""" | |
| if min_bucket < 0 or max_bucket < min_bucket or total_buckets <= 0: | |
| return True | |
| try: | |
| return any(min_bucket <= bucket <= max_bucket | |
| for bucket in self._compute(None, total_buckets)) | |
| except Exception: | |
| # Use the same fail-open behavior as entry-level bucket selection. | |
| return True | |
| def _compute(self, partition, total_buckets: int) -> FrozenSet[int]: | |
| cache_key = (_partition_to_cache_key(partition, self._partition_fields), | |
| total_buckets) | |
| cached = self._cache.get(cache_key) | |
| if cached is not None: | |
| return cached | |
| effective_predicate: Optional[Union[bool, Predicate]] = self._predicate | |
| if partition is not None and self._partition_fields: | |
| partition_values = _partition_to_dict(partition, self._partition_fields) | |
| partition_field_names = {f.name for f in self._partition_fields} | |
| effective_predicate = replace_partition_predicate( | |
| self._predicate, partition_field_names, partition_values) | |
| if effective_predicate is _AlwaysFalse: | |
| # Java cannot extract Equal/In values from AlwaysFalse. Leave | |
| # rejection to the partition predicate rather than the key hash. | |
| frozen: FrozenSet[int] = frozenset(range(total_buckets)) | |
| self._cache[cache_key] = frozen | |
| return frozen | |
| if effective_predicate is _AlwaysTrue: | |
| # Predicate cleared after partition substitution — accept all | |
| # buckets for this partition. | |
| frozen = frozenset(range(total_buckets)) | |
| self._cache[cache_key] = frozen | |
| return frozen | |
| combinations = _build_combinations(effective_predicate, | |
| self._bucket_key_fields) | |
| if combinations is None: | |
| # Couldn't pin down all bucket keys (or above MAX_VALUES) — | |
| # fall back to "all buckets accept" for soundness. | |
| frozen = frozenset(range(total_buckets)) | |
| self._cache[cache_key] = frozen | |
| return frozen | |
| frozen = _hash_combinations(combinations, self._bucket_key_fields, | |
| total_buckets) | |
| self._cache[cache_key] = frozen | |
| return frozen | |
| def create_bucket_selector( | |
| predicate: Optional[Predicate], | |
| bucket_key_fields: List[DataField], | |
| partition_fields: Optional[List[DataField]] = None, | |
| ) -> Optional[_Selector]: | |
| """Try to derive a bucket selector from ``predicate`` constrained to | |
| ``bucket_key_fields``. | |
| Returns: | |
| A callable ``(partition, bucket, total_buckets) -> bool``. When | |
| ``partition_fields`` is given and the predicate references those | |
| partition columns, the selector specialises the predicate per | |
| partition value before hashing — this catches mixed forms like | |
| ``(part='a' AND bk IN (1,2)) OR (part='b' AND bk IN (3,4))`` that | |
| would otherwise be unprunable. ``partition=None`` callsites | |
| (early manifest filter that hasn't deserialised the entry yet) | |
| simply get the partition-agnostic result. | |
| Returns None when the predicate carries no usable bucket-key | |
| constraint at all (caller must NOT prune by bucket). | |
| ``may_contain(min_bucket, max_bucket, total_buckets)`` tests a | |
| manifest's inclusive bucket range without partition specialisation. | |
| """ | |
| if predicate is None or not bucket_key_fields: | |
| return None | |
| # See ``_UNSAFE_BUCKET_KEY_TYPES``: refuse pruning when the bucket-key | |
| # column types are prone to writer/reader byte-level disagreement on | |
| # equal logical values. Fail open rather than risk false-negatives. | |
| if _has_unsafe_bucket_key_type(bucket_key_fields): | |
| return None | |
| # Sanity gate: if the predicate without any partition substitution | |
| # already fails to pin down bucket keys AND it doesn't touch any | |
| # partition columns, there's no point handing the caller a selector | |
| # that always returns "all buckets" — preserve the original "return | |
| # None for unprunable" contract so the caller can skip the wrap. | |
| partition_names = {f.name for f in (partition_fields or [])} | |
| touches_partition = ( | |
| bool(partition_names) | |
| and _predicate_touches_partition(predicate, partition_names) | |
| ) | |
| if not touches_partition: | |
| if _build_combinations(predicate, bucket_key_fields) is None: | |
| return None | |
| return _Selector(predicate, bucket_key_fields, partition_fields) |