[python] Read composite BTree indexes in data evolution scans

paimon-python/pypaimon/execution/python/globalindex/btree/composite_btree.py

@@ -0,0 +1,275 @@
+# 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.
+
+"""Read Java composite BTree keys with left-prefix predicate bounds."""
+
+from dataclasses import dataclass
+import datetime
+from decimal import Decimal
+import numbers
+import struct
+
+from pypaimon.execution.python.globalindex.btree.btree_index_reader import BTreeIndexReader
+from pypaimon.execution.python.globalindex.btree.btree_posting_list import add_row_ids
+from pypaimon.execution.python.globalindex.btree.composite_key import (
+    CompositeKeySerializer, compare_nullable,
+)
+from pypaimon.globalindex.global_index_result import GlobalIndexResult
+from pypaimon.utils.roaring_bitmap import RoaringBitmap64
+
+
+MAX_INTERVALS = 256
+_POINT = frozenset(('equal', 'in', 'isNull'))
+_RANGE = frozenset(('greaterThan', 'greaterOrEqual', 'lessThan',
+                    'lessOrEqual', 'between', 'isNotNull'))
+
+
+@dataclass(frozen=True)
+class Bound:
+    values: tuple
+    after: bool
+
+
+@dataclass
+class CompositePlan:
+    intervals: list
+    bound_columns: int
+    equal_columns: int
+    field_count: int
+    covered_predicates: tuple
+
+    @property
+    def point_lookup(self):
+        return self.equal_columns == self.field_count
+
+
+def _compare_values(left, right, comparators):
+    for index, (a, b) in enumerate(zip(left, right)):
+        result = compare_nullable(a, b, comparators[index])
+        if result:
+            return result
+    return 0
+
+
+def compare_bounds(left, right, comparators):
+    difference = _compare_values(left.values, right.values, comparators)
+    if difference:
+        return difference
+    if len(left.values) != len(right.values):
+        if len(left.values) < len(right.values):
+            return 1 if left.after else -1
+        return -1 if right.after else 1
+    return (left.after > right.after) - (left.after < right.after)
+
+
+def compare_key_bound(key, bound, comparators):
+    difference = _compare_values(key, bound.values, comparators)
+    return difference if difference else (-1 if bound.after else 1)
+
+
+def _conjuncts(predicate):
+    if predicate.method != 'and':
+        return [predicate]
+    result = []
+    for child in predicate.literals:
+        result.extend(_conjuncts(child))
+    return result
+
+
+def _as_leaf(predicate):
+    if predicate.method != 'or':
+        return predicate
+    children = predicate.literals or []
+    if not children or any(child.method != 'equal' or child.field != children[0].field
+                           for child in children):
+        return None
+    from pypaimon.common.predicate import Predicate
+    return Predicate('in', children[0].index, children[0].field,
+                     [child.literals[0] for child in children])
+
+
+def _matches(value, predicate, comparator):
+    method = predicate.method
+    values = predicate.literals or []
+    if method == 'isNull':
+        return value is None
+    if method == 'isNotNull':
+        return value is not None
+    if value is None:
+        return False
+    if method == 'in':
+        return any(literal is not None and comparator(value, literal) == 0
+                   for literal in values)
+    if any(literal is None for literal in values):
+        return False
+    if method == 'equal':
+        return comparator(value, values[0]) == 0
+    if method == 'between':
+        return (comparator(value, values[0]) >= 0
+                and comparator(value, values[1]) <= 0)
+    difference = comparator(value, values[0])
+    return {'greaterThan': difference > 0,
+            'greaterOrEqual': difference >= 0,
+            'lessThan': difference < 0,
+            'lessOrEqual': difference <= 0}[method]
+
+
+def _literal_supported(type_name, literal):
+    if type_name == 'BOOLEAN':
+        return isinstance(literal, bool)
+    if type_name in ('TINYINT', 'SMALLINT', 'INT', 'INTEGER', 'BIGINT'):
+        return isinstance(literal, numbers.Integral) and not isinstance(literal, bool)
+    if type_name.startswith(('CHAR', 'VARCHAR', 'STRING')):
+        return isinstance(literal, str)
+    if type_name in ('FLOAT', 'DOUBLE'):
+        # Java orders signed zeros separately, while Python row filtering does not.
+        return (isinstance(literal, numbers.Real)
+                and not isinstance(literal, bool) and literal != 0)
+    if type_name.startswith(('DECIMAL', 'NUMERIC')):
+        return isinstance(literal, (Decimal, numbers.Integral)) and not isinstance(literal, bool)
+    if type_name == 'DATE':
+        return ((isinstance(literal, datetime.date)
+                 and not isinstance(literal, datetime.datetime))
+                or isinstance(literal, numbers.Integral)
+                and not isinstance(literal, bool))
+    if type_name.startswith('TIME') and not type_name.startswith('TIMESTAMP'):
+        return (isinstance(literal, datetime.time)
+                and literal.microsecond % 1000 == 0
+                or isinstance(literal, numbers.Integral)
+                and not isinstance(literal, bool))
+    if type_name.startswith('TIMESTAMP'):
+        return (isinstance(literal, datetime.datetime)
+                and not getattr(literal, 'nanosecond', 0))
+    return False
+
+
+def plan_composite(fields, predicate):
+    """Return None when the index cannot safely bound this predicate."""
+    try:
+        serializer = CompositeKeySerializer(fields)
+    except ValueError:
+        return None
+    comparators = serializer._comparators
+    conjuncts = _conjuncts(predicate)
+    prefixes = [()]
+    equal_columns = 0
+    range_predicates = []
+    covered_predicates = []
+    for field, comparator, scalar in zip(
+            fields, comparators, serializer._serializers):
+        column = []
+        sources = []
+        for child in conjuncts:
+            leaf = _as_leaf(child)
+            if leaf is not None and leaf.field == field.name and (
+                    leaf.method in _POINT or leaf.method in _RANGE):
+                column.append(leaf)
+                sources.append(child)
+        if not column:
+            break
+        for leaf in column:
+            for literal in leaf.literals or []:
+                if literal is None:
+                    continue
+                if not _literal_supported(field.type.type.upper(), literal):
+                    return None
+                try:
+                    converted = scalar.deserialize(scalar.serialize(literal))
+                    if comparator(converted, literal) != 0:
+                        return None
+                except (ArithmeticError, ValueError, TypeError, struct.error):
+                    return None
+        point = [leaf for leaf in column if leaf.method in _POINT]
+        if not point:
+            range_predicates = column
+            covered_predicates.extend(sources)
+            break
+        selected = min(point, key=lambda leaf: len(leaf.literals or [None]))
+        candidates = ([None] if selected.method == 'isNull'
+                      else selected.literals or [])
+        values = []
+        for value in candidates:
+            if value is None and (selected.method != 'isNull' or not field.type.nullable):
+                continue
+            if all(_matches(value, leaf, comparator) for leaf in column):
+                if not any(compare_nullable(value, old, comparator) == 0 for old in values):
+                    values.append(value)
+        equal_columns += 1
+        covered_predicates.extend(sources)
+        if not values:
+            return CompositePlan([], equal_columns, equal_columns, len(fields),
+                                 tuple(covered_predicates))
+        if len(prefixes) * len(values) > MAX_INTERVALS:
+            return None
+        prefixes = [prefix + (value,) for prefix in prefixes for value in values]
+    bound_columns = equal_columns + bool(range_predicates)
+    if not bound_columns:
+        return None
+    intervals = []
+    for prefix in prefixes:
+        lower, upper = Bound(prefix, False), Bound(prefix, True)
+        if range_predicates:
+            lower = Bound(prefix + (None,), True)
+            empty = False
+            for leaf in range_predicates:
+                if leaf.method == 'isNotNull':
+                    continue
+                if None in (leaf.literals or []):
+                    empty = True
+                    break
+                if leaf.method in ('greaterThan', 'greaterOrEqual', 'between'):
+                    candidate = Bound(prefix + (leaf.literals[0],),
+                                      leaf.method == 'greaterThan')
+                    if compare_bounds(candidate, lower, comparators) > 0:
+                        lower = candidate
+                if leaf.method in ('lessThan', 'lessOrEqual', 'between'):
+                    candidate = Bound(prefix + (leaf.literals[-1],),
+                                      leaf.method != 'lessThan')
+                    if compare_bounds(candidate, upper, comparators) < 0:
+                        upper = candidate
+            if empty:
+                continue
+        if compare_bounds(lower, upper, comparators) < 0:
+            intervals.append((lower, upper))
+    return CompositePlan(intervals, bound_columns, equal_columns, len(fields),
+                         tuple(covered_predicates))
+
+
+class CompositeBTreeIndexReader(BTreeIndexReader):
+    def query(self, plan):
+        result = RoaringBitmap64()
+        if self.min_key is None:
+            return GlobalIndexResult.create(result, is_exact=False)
+        for lower, upper in plan.intervals:
+            if (compare_key_bound(self.max_key, lower, self.key_serializer._comparators) <= 0
+                    or compare_key_bound(self.min_key, upper,
+                                         self.key_serializer._comparators) >= 0):
+                continue
+            if plan.point_lookup:
+                posting = self.reader.lookup(
+                    self.key_serializer.serialize(lower.values))
+                if posting is not None:
+                    add_row_ids(posting, self.footer.version, result)
+                continue
+            seek = lower.values + (None,) * (plan.field_count - len(lower.values))
+            for entry in self._entries_from(self.key_serializer.serialize(seek)):
+                key = self.key_serializer.deserialize(entry.key)
+                if compare_key_bound(key, upper, self.key_serializer._comparators) >= 0:
+                    break
+                if compare_key_bound(key, lower, self.key_serializer._comparators) > 0:
+                    add_row_ids(entry.value, self.footer.version, result)
+        return GlobalIndexResult.create(result, is_exact=False)