paimon-python/pypaimon/common/predicate.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.

import re
from abc import ABC, ABCMeta, abstractmethod
from dataclasses import dataclass
from decimal import Decimal
from typing import Any, ClassVar, Dict, List, Optional, Sequence, cast

import pyarrow
from pyarrow import compute as pyarrow_compute
from pyarrow import dataset as pyarrow_dataset

from pypaimon.common.decimal_arrow_predicate import (
    DecimalLiteral,
    decimal_arrow_expression,
)
from pypaimon.manifest.schema.simple_stats import SimpleStats
from pypaimon.table.row.internal_row import InternalRow


def _combine_arrow_expressions(expressions, combine):
    while len(expressions) > 1:
        next_level = []
        for index in range(0, len(expressions), 2):
            if index + 1 == len(expressions):
                next_level.append(expressions[index])
            else:
                next_level.append(combine(
                    expressions[index], expressions[index + 1],
                ))
        expressions = next_level
    return expressions[0]


@dataclass
class Predicate:
    method: str
    index: Optional[int]
    field: Optional[str]
    literals: Optional[List[Any]] = None

    testers: ClassVar[Dict[str, Any]] = {}

    def new_index(self, index: int):
        return Predicate(
            method=self.method,
            index=index,
            field=self.field,
            literals=self.literals)

    def new_literals(self, literals: List[Any]):
        return Predicate(
            method=self.method,
            index=self.index,
            field=self.field,
            literals=literals)

    def test(self, record: InternalRow) -> bool:
        if self.method == 'and':
            return all(p.test(record) for p in self.literals)
        if self.method == 'or':
            t = any(p.test(record) for p in self.literals)
            return t
        field_value = record.get_field(self.index)
        tester = Predicate.testers.get(self.method)
        if tester:
            return tester.test_by_value(field_value, self.literals)
        raise ValueError(f"Unsupported predicate method: {self.method}")

    def test_by_simple_stats(self, stat: SimpleStats, row_count: int) -> bool:
        """Test predicate against BinaryRow stats with denseIndexMapping like Java implementation."""
        if self.method == 'and':
            return all(p.test_by_simple_stats(stat, row_count) for p in self.literals)
        if self.method == 'or':
            return any(p.test_by_simple_stats(stat, row_count) for p in self.literals)

        index = self.index
        if index is None or index < 0:
            # Missing stats cannot prove that the file does not match.
            return True

        null_count = (
            stat.null_counts[index]
            if stat.null_counts is not None and index < len(stat.null_counts)
            else None
        )

        if self.method == 'isNull':
            return null_count is None or null_count > 0
        if self.method == 'isNotNull':
            return null_count is None or row_count is None or null_count < row_count

        try:
            min_value = (
                stat.min_values.get_field(index)
                if index < len(stat.min_values)
                else None
            )
            max_value = (
                stat.max_values.get_field(index)
                if index < len(stat.max_values)
                else None
            )
        except IndexError:
            return True

        if min_value is None or max_value is None or (null_count is not None and null_count == row_count):
            # invalid stats, skip validation
            return True

        tester = Predicate.testers.get(self.method)
        if tester:
            return tester.test_by_stats(min_value, max_value, self.literals)
        raise ValueError(f"Unsupported predicate method: {self.method}")

    def to_arrow(self, schema: Optional[pyarrow.Schema] = None) -> Any:
        """Convert to Arrow, using the input schema for exact decimal comparisons."""
        if self.method == 'and':
            return _combine_arrow_expressions(
                [
                    p.to_arrow(schema) if schema is not None else p.to_arrow()
                    for p in self.literals
                ],
                lambda left, right: left & right,
            )
        if self.method == 'or':
            return _combine_arrow_expressions(
                [
                    p.to_arrow(schema) if schema is not None else p.to_arrow()
                    for p in self.literals
                ],
                lambda left, right: left | right,
            )

        if self.method == 'startsWith':
            pattern = self.literals[0]
            # For PyArrow compatibility - improved approach
            try:
                field_ref = pyarrow_dataset.field(self.field)
                # Ensure the field is cast to string type
                string_field = field_ref.cast(pyarrow.string())
                result = pyarrow_compute.starts_with(string_field, pattern)
                return result
            except Exception:
                # Fallback to True
                return pyarrow_dataset.field(self.field).is_valid() | pyarrow_dataset.field(self.field).is_null()
        if self.method == 'endsWith':
            pattern = self.literals[0]
            # For PyArrow compatibility
            try:
                field_ref = pyarrow_dataset.field(self.field)
                # Ensure the field is cast to string type
                string_field = field_ref.cast(pyarrow.string())
                result = pyarrow_compute.ends_with(string_field, pattern)
                return result
            except Exception:
                # Fallback to True
                return pyarrow_dataset.field(self.field).is_valid() | pyarrow_dataset.field(self.field).is_null()
        if self.method == 'contains':
            pattern = self.literals[0]
            # For PyArrow compatibility
            try:
                field_ref = pyarrow_dataset.field(self.field)
                # Ensure the field is cast to string type
                string_field = field_ref.cast(pyarrow.string())
                result = pyarrow_compute.match_substring(string_field, pattern)
                return result
            except Exception:
                # Fallback to True
                return pyarrow_dataset.field(self.field).is_valid() | pyarrow_dataset.field(self.field).is_null()
        if self.method == 'like':
            pattern = self.literals[0]
            try:
                field_ref = pyarrow_dataset.field(self.field)
                # Ensure the field is cast to string type
                string_field = field_ref.cast(pyarrow.string())
                return pyarrow_compute.match_like(string_field, pattern)
            except Exception:
                return None

        field = pyarrow_dataset.field(self.field)
        literals = self.literals
        if schema is not None and self.field in schema.names:
            field_type = schema.field(self.field).type
            if (
                pyarrow.types.is_decimal(field_type)
                and literals
                and all(
                    value is None or isinstance(value, (Decimal, int))
                    for value in literals
                )
                and self.method
                in (
                    'equal',
                    'notEqual',
                    'lessThan',
                    'lessOrEqual',
                    'greaterThan',
                    'greaterOrEqual',
                    'in',
                    'notIn',
                    'between',
                    'notBetween',
                )
            ):
                return decimal_arrow_expression(
                    field,
                    field_type,
                    self.method,
                    cast(Sequence[Optional[DecimalLiteral]], literals),
                )
        tester = Predicate.testers.get(self.method)
        if tester:
            return tester.test_by_arrow(field, literals)

        raise ValueError("Unsupported predicate method: {}".format(self.method))


class RegisterMeta(ABCMeta):
    def __init__(cls, name, bases, dct):
        super().__init__(name, bases, dct)
        if not bool(cls.__abstractmethods__):
            Predicate.testers[cls.name] = cls()


class Tester(ABC, metaclass=RegisterMeta):
    name = None

    @abstractmethod
    def test_by_value(self, val, literals) -> bool:
        """
        Test based on the specific val and literals.
        """

    @abstractmethod
    def test_by_stats(self, min_v, max_v, literals) -> bool:
        """
        Test based on the specific min_value and max_value and literals.
        """

    @abstractmethod
    def test_by_arrow(self, val, literals) -> bool:
        """
        Test based on the specific arrow value and literals.
        """


class Equal(Tester):
    name = 'equal'

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return val == literals[0]

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return min_v <= literals[0] <= max_v

    def test_by_arrow(self, val, literals) -> bool:
        return val == literals[0]


class NotEqual(Tester):
    name = "notEqual"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return val != literals[0]

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return not (min_v == literals[0] == max_v)

    def test_by_arrow(self, val, literals) -> bool:
        return val != literals[0]


class LessThan(Tester):
    name = "lessThan"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return val < literals[0]

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return literals[0] > min_v

    def test_by_arrow(self, val, literals) -> bool:
        return val < literals[0]


class LessOrEqual(Tester):
    name = "lessOrEqual"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return val <= literals[0]

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return literals[0] >= min_v

    def test_by_arrow(self, val, literals) -> bool:
        return val <= literals[0]


class GreaterThan(Tester):
    name = "greaterThan"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return val > literals[0]

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return literals[0] < max_v

    def test_by_arrow(self, val, literals) -> bool:
        return val > literals[0]


class GreaterOrEqual(Tester):
    name = "greaterOrEqual"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return val >= literals[0]

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return literals[0] <= max_v

    def test_by_arrow(self, val, literals) -> bool:
        return val >= literals[0]


class In(Tester):
    name = "in"

    def test_by_value(self, val, literals) -> bool:
        if val is None:
            return False
        return val in literals

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return any(
            min_v <= literal <= max_v
            for literal in literals
            if literal is not None
        )

    def test_by_arrow(self, val, literals) -> bool:
        # Arrow treats null as a set member, while SQL IN never returns true
        # solely because both the field and an IN literal are null.
        non_null_literals = [literal for literal in literals if literal is not None]
        if not non_null_literals:
            return val.is_valid() & val.is_null()
        return val.isin(non_null_literals) & val.is_valid()


class NotIn(Tester):
    name = "notIn"

    def test_by_value(self, val, literals) -> bool:
        if val is None or any(literal is None for literal in literals):
            return False
        return val not in literals

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        if any(literal is None for literal in literals):
            return False
        return not any(min_v == l == max_v for l in literals)

    def test_by_arrow(self, val, literals) -> bool:
        # Any null literal makes SQL NOT IN unknown for every non-matching row.
        if any(literal is None for literal in literals):
            return val.is_valid() & val.is_null()
        return (~val.isin(literals)) & val.is_valid()


class Between(Tester):
    name = "between"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals or len(literals) < 2:
            return False
        return literals[0] <= val <= literals[1]

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return literals[0] <= max_v and literals[1] >= min_v

    def test_by_arrow(self, val, literals) -> bool:
        return (val >= literals[0]) & (val <= literals[1])


class StartsWith(Tester):
    name = "startsWith"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return isinstance(val, str) and val.startswith(literals[0])

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return ((isinstance(min_v, str) and isinstance(max_v, str)) and
                ((min_v.startswith(literals[0]) or min_v < literals[0]) and
                 (max_v.startswith(literals[0]) or max_v > literals[0])))

    def test_by_arrow(self, val, literals) -> bool:
        return True


class EndsWith(Tester):
    name = "endsWith"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return isinstance(val, str) and val.endswith(literals[0])

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return True

    def test_by_arrow(self, val, literals) -> bool:
        return True


class Contains(Tester):
    name = "contains"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        return isinstance(val, str) and literals[0] in val

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return True

    def test_by_arrow(self, val, literals) -> bool:
        return True


class IsNull(Tester):
    name = "isNull"

    def test_by_value(self, val, literals) -> bool:
        return val is None

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return True

    def test_by_arrow(self, val, literals) -> bool:
        return val.is_null()


class IsNotNull(Tester):
    name = "isNotNull"

    def test_by_value(self, val, literals) -> bool:
        return val is not None

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return True

    def test_by_arrow(self, val, literals) -> bool:
        return val.is_valid()


class NotBetween(Tester):
    name = "notBetween"

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals or len(literals) < 2:
            return False
        return val < literals[0] or val > literals[1]

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return literals[0] > min_v or literals[1] < max_v

    def test_by_arrow(self, val, literals) -> bool:
        return (val < literals[0]) | (val > literals[1])


class Like(Tester):
    name = "like"

    @staticmethod
    def _sql_like_to_regex(pattern: str, escape_char: str = '\\') -> str:
        """Convert a SQL LIKE pattern to a Python regex pattern."""
        regex_parts = []
        index = 0
        length = len(pattern)
        while index < length:
            char = pattern[index]
            if char == escape_char:
                if index + 1 < length:
                    next_char = pattern[index + 1]
                    if next_char in ('_', '%', escape_char):
                        regex_parts.append(re.escape(next_char))
                        index += 2
                        continue
                    else:
                        raise ValueError(
                            f"Invalid escape sequence '{pattern}' at position {index}")
                else:
                    raise ValueError(
                        f"Invalid escape sequence '{pattern}' at position {index}")
            elif char == '_':
                regex_parts.append('.')
            elif char == '%':
                regex_parts.append('(?s:.*)')
            else:
                regex_parts.append(re.escape(char))
            index += 1
        return ''.join(regex_parts)

    def test_by_value(self, val, literals) -> bool:
        if val is None or not literals:
            return False
        if not isinstance(val, str):
            return False
        pattern = self._sql_like_to_regex(str(literals[0]))
        return bool(re.fullmatch(pattern, val))

    def test_by_stats(self, min_v, max_v, literals) -> bool:
        return True

    def test_by_arrow(self, val, literals) -> bool:
        return True