| # 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 math |
| from abc import ABC, abstractmethod |
| from functools import singledispatch |
| from typing import ( |
| Any, |
| Callable, |
| Dict, |
| Generic, |
| List, |
| Set, |
| SupportsFloat, |
| Tuple, |
| TypeVar, |
| Union, |
| ) |
| |
| from pyiceberg.conversions import from_bytes |
| from pyiceberg.expressions import ( |
| AlwaysFalse, |
| AlwaysTrue, |
| And, |
| BooleanExpression, |
| BoundEqualTo, |
| BoundGreaterThan, |
| BoundGreaterThanOrEqual, |
| BoundIn, |
| BoundIsNaN, |
| BoundIsNull, |
| BoundLessThan, |
| BoundLessThanOrEqual, |
| BoundLiteralPredicate, |
| BoundNotEqualTo, |
| BoundNotIn, |
| BoundNotNaN, |
| BoundNotNull, |
| BoundNotStartsWith, |
| BoundPredicate, |
| BoundSetPredicate, |
| BoundStartsWith, |
| BoundTerm, |
| BoundUnaryPredicate, |
| Not, |
| Or, |
| UnboundPredicate, |
| ) |
| from pyiceberg.expressions.literals import Literal |
| from pyiceberg.manifest import DataFile, ManifestFile, PartitionFieldSummary |
| from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec |
| from pyiceberg.schema import Schema |
| from pyiceberg.typedef import EMPTY_DICT, L, Record, StructProtocol |
| from pyiceberg.types import ( |
| DoubleType, |
| FloatType, |
| IcebergType, |
| NestedField, |
| PrimitiveType, |
| StructType, |
| TimestampType, |
| TimestamptzType, |
| ) |
| from pyiceberg.utils.datetime import micros_to_timestamp, micros_to_timestamptz |
| |
| T = TypeVar("T") |
| |
| |
| class BooleanExpressionVisitor(Generic[T], ABC): |
| @abstractmethod |
| def visit_true(self) -> T: |
| """Visit method for an AlwaysTrue boolean expression. |
| |
| Note: This visit method has no arguments since AlwaysTrue instances have no context. |
| """ |
| |
| @abstractmethod |
| def visit_false(self) -> T: |
| """Visit method for an AlwaysFalse boolean expression. |
| |
| Note: This visit method has no arguments since AlwaysFalse instances have no context. |
| """ |
| |
| @abstractmethod |
| def visit_not(self, child_result: T) -> T: |
| """Visit method for a Not boolean expression. |
| |
| Args: |
| child_result (T): The result of visiting the child of the Not boolean expression. |
| """ |
| |
| @abstractmethod |
| def visit_and(self, left_result: T, right_result: T) -> T: |
| """Visit method for an And boolean expression. |
| |
| Args: |
| left_result (T): The result of visiting the left side of the expression. |
| right_result (T): The result of visiting the right side of the expression. |
| """ |
| |
| @abstractmethod |
| def visit_or(self, left_result: T, right_result: T) -> T: |
| """Visit method for an Or boolean expression. |
| |
| Args: |
| left_result (T): The result of visiting the left side of the expression. |
| right_result (T): The result of visiting the right side of the expression. |
| """ |
| |
| @abstractmethod |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> T: |
| """Visit method for an unbound predicate in an expression tree. |
| |
| Args: |
| predicate (UnboundPredicate[L): An instance of an UnboundPredicate. |
| """ |
| |
| @abstractmethod |
| def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> T: |
| """Visit method for a bound predicate in an expression tree. |
| |
| Args: |
| predicate (BoundPredicate[L]): An instance of a BoundPredicate. |
| """ |
| |
| |
| @singledispatch |
| def visit(obj: BooleanExpression, visitor: BooleanExpressionVisitor[T]) -> T: |
| """Apply a boolean expression visitor to any point within an expression. |
| |
| The function traverses the expression in post-order fashion. |
| |
| Args: |
| obj (BooleanExpression): An instance of a BooleanExpression. |
| visitor (BooleanExpressionVisitor[T]): An instance of an implementation of the generic BooleanExpressionVisitor base class. |
| |
| Raises: |
| NotImplementedError: If attempting to visit an unsupported expression. |
| """ |
| raise NotImplementedError(f"Cannot visit unsupported expression: {obj}") |
| |
| |
| @visit.register(AlwaysTrue) |
| def _(_: AlwaysTrue, visitor: BooleanExpressionVisitor[T]) -> T: |
| """Visit an AlwaysTrue boolean expression with a concrete BooleanExpressionVisitor.""" |
| return visitor.visit_true() |
| |
| |
| @visit.register(AlwaysFalse) |
| def _(_: AlwaysFalse, visitor: BooleanExpressionVisitor[T]) -> T: |
| """Visit an AlwaysFalse boolean expression with a concrete BooleanExpressionVisitor.""" |
| return visitor.visit_false() |
| |
| |
| @visit.register(Not) |
| def _(obj: Not, visitor: BooleanExpressionVisitor[T]) -> T: |
| """Visit a Not boolean expression with a concrete BooleanExpressionVisitor.""" |
| child_result: T = visit(obj.child, visitor=visitor) |
| return visitor.visit_not(child_result=child_result) |
| |
| |
| @visit.register(And) |
| def _(obj: And, visitor: BooleanExpressionVisitor[T]) -> T: |
| """Visit an And boolean expression with a concrete BooleanExpressionVisitor.""" |
| left_result: T = visit(obj.left, visitor=visitor) |
| right_result: T = visit(obj.right, visitor=visitor) |
| return visitor.visit_and(left_result=left_result, right_result=right_result) |
| |
| |
| @visit.register(UnboundPredicate) |
| def _(obj: UnboundPredicate[L], visitor: BooleanExpressionVisitor[T]) -> T: |
| """Visit an unbound boolean expression with a concrete BooleanExpressionVisitor.""" |
| return visitor.visit_unbound_predicate(predicate=obj) |
| |
| |
| @visit.register(BoundPredicate) |
| def _(obj: BoundPredicate[L], visitor: BooleanExpressionVisitor[T]) -> T: |
| """Visit a bound boolean expression with a concrete BooleanExpressionVisitor.""" |
| return visitor.visit_bound_predicate(predicate=obj) |
| |
| |
| @visit.register(Or) |
| def _(obj: Or, visitor: BooleanExpressionVisitor[T]) -> T: |
| """Visit an Or boolean expression with a concrete BooleanExpressionVisitor.""" |
| left_result: T = visit(obj.left, visitor=visitor) |
| right_result: T = visit(obj.right, visitor=visitor) |
| return visitor.visit_or(left_result=left_result, right_result=right_result) |
| |
| |
| def bind(schema: Schema, expression: BooleanExpression, case_sensitive: bool) -> BooleanExpression: |
| """Travers over an expression to bind the predicates to the schema. |
| |
| Args: |
| schema (Schema): A schema to use when binding the expression. |
| expression (BooleanExpression): An expression containing UnboundPredicates that can be bound. |
| case_sensitive (bool): Whether to consider case when binding a reference to a field in a schema, defaults to True. |
| |
| Raises: |
| TypeError: In the case a predicate is already bound. |
| """ |
| return visit(expression, BindVisitor(schema, case_sensitive)) |
| |
| |
| class BindVisitor(BooleanExpressionVisitor[BooleanExpression]): |
| """Rewrites a boolean expression by replacing unbound references with references to fields in a struct schema. |
| |
| Args: |
| schema (Schema): A schema to use when binding the expression. |
| case_sensitive (bool): Whether to consider case when binding a reference to a field in a schema, defaults to True. |
| |
| Raises: |
| TypeError: In the case a predicate is already bound. |
| """ |
| |
| schema: Schema |
| case_sensitive: bool |
| |
| def __init__(self, schema: Schema, case_sensitive: bool) -> None: |
| self.schema = schema |
| self.case_sensitive = case_sensitive |
| |
| def visit_true(self) -> BooleanExpression: |
| return AlwaysTrue() |
| |
| def visit_false(self) -> BooleanExpression: |
| return AlwaysFalse() |
| |
| def visit_not(self, child_result: BooleanExpression) -> BooleanExpression: |
| return Not(child=child_result) |
| |
| def visit_and(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return And(left=left_result, right=right_result) |
| |
| def visit_or(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return Or(left=left_result, right=right_result) |
| |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> BooleanExpression: |
| return predicate.bind(self.schema, case_sensitive=self.case_sensitive) |
| |
| def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> BooleanExpression: |
| raise TypeError(f"Found already bound predicate: {predicate}") |
| |
| |
| class BoundBooleanExpressionVisitor(BooleanExpressionVisitor[T], ABC): |
| @abstractmethod |
| def visit_in(self, term: BoundTerm[L], literals: Set[L]) -> T: |
| """Visit a bound In predicate.""" |
| |
| @abstractmethod |
| def visit_not_in(self, term: BoundTerm[L], literals: Set[L]) -> T: |
| """Visit a bound NotIn predicate.""" |
| |
| @abstractmethod |
| def visit_is_nan(self, term: BoundTerm[L]) -> T: |
| """Visit a bound IsNan predicate.""" |
| |
| @abstractmethod |
| def visit_not_nan(self, term: BoundTerm[L]) -> T: |
| """Visit a bound NotNan predicate.""" |
| |
| @abstractmethod |
| def visit_is_null(self, term: BoundTerm[L]) -> T: |
| """Visit a bound IsNull predicate.""" |
| |
| @abstractmethod |
| def visit_not_null(self, term: BoundTerm[L]) -> T: |
| """Visit a bound NotNull predicate.""" |
| |
| @abstractmethod |
| def visit_equal(self, term: BoundTerm[L], literal: Literal[L]) -> T: |
| """Visit a bound Equal predicate.""" |
| |
| @abstractmethod |
| def visit_not_equal(self, term: BoundTerm[L], literal: Literal[L]) -> T: |
| """Visit a bound NotEqual predicate.""" |
| |
| @abstractmethod |
| def visit_greater_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> T: |
| """Visit a bound GreaterThanOrEqual predicate.""" |
| |
| @abstractmethod |
| def visit_greater_than(self, term: BoundTerm[L], literal: Literal[L]) -> T: |
| """Visit a bound GreaterThan predicate.""" |
| |
| @abstractmethod |
| def visit_less_than(self, term: BoundTerm[L], literal: Literal[L]) -> T: |
| """Visit a bound LessThan predicate.""" |
| |
| @abstractmethod |
| def visit_less_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> T: |
| """Visit a bound LessThanOrEqual predicate.""" |
| |
| @abstractmethod |
| def visit_true(self) -> T: |
| """Visit a bound True predicate.""" |
| |
| @abstractmethod |
| def visit_false(self) -> T: |
| """Visit a bound False predicate.""" |
| |
| @abstractmethod |
| def visit_not(self, child_result: T) -> T: |
| """Visit a bound Not predicate.""" |
| |
| @abstractmethod |
| def visit_and(self, left_result: T, right_result: T) -> T: |
| """Visit a bound And predicate.""" |
| |
| @abstractmethod |
| def visit_or(self, left_result: T, right_result: T) -> T: |
| """Visit a bound Or predicate.""" |
| |
| @abstractmethod |
| def visit_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> T: |
| """Visit bound StartsWith predicate.""" |
| |
| @abstractmethod |
| def visit_not_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> T: |
| """Visit bound NotStartsWith predicate.""" |
| |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> T: |
| """Visit an unbound predicate. |
| |
| Args: |
| predicate (UnboundPredicate[L]): An unbound predicate. |
| Raises: |
| TypeError: This always raises since an unbound predicate is not expected in a bound boolean expression. |
| """ |
| raise TypeError(f"Not a bound predicate: {predicate}") |
| |
| def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> T: |
| """Visit a bound predicate. |
| |
| Args: |
| predicate (BoundPredicate[L]): A bound predicate. |
| """ |
| return visit_bound_predicate(predicate, self) |
| |
| |
| @singledispatch |
| def visit_bound_predicate(expr: BoundPredicate[L], _: BooleanExpressionVisitor[T]) -> T: |
| raise TypeError(f"Unknown predicate: {expr}") |
| |
| |
| @visit_bound_predicate.register(BoundIn) |
| def _(expr: BoundIn[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_in(term=expr.term, literals=expr.value_set) |
| |
| |
| @visit_bound_predicate.register(BoundNotIn) |
| def _(expr: BoundNotIn[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_not_in(term=expr.term, literals=expr.value_set) |
| |
| |
| @visit_bound_predicate.register(BoundIsNaN) |
| def _(expr: BoundIsNaN[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_is_nan(term=expr.term) |
| |
| |
| @visit_bound_predicate.register(BoundNotNaN) |
| def _(expr: BoundNotNaN[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_not_nan(term=expr.term) |
| |
| |
| @visit_bound_predicate.register(BoundIsNull) |
| def _(expr: BoundIsNull[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_is_null(term=expr.term) |
| |
| |
| @visit_bound_predicate.register(BoundNotNull) |
| def _(expr: BoundNotNull[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_not_null(term=expr.term) |
| |
| |
| @visit_bound_predicate.register(BoundEqualTo) |
| def _(expr: BoundEqualTo[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_equal(term=expr.term, literal=expr.literal) |
| |
| |
| @visit_bound_predicate.register(BoundNotEqualTo) |
| def _(expr: BoundNotEqualTo[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_not_equal(term=expr.term, literal=expr.literal) |
| |
| |
| @visit_bound_predicate.register(BoundGreaterThanOrEqual) |
| def _(expr: BoundGreaterThanOrEqual[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| """Visit a bound GreaterThanOrEqual predicate.""" |
| return visitor.visit_greater_than_or_equal(term=expr.term, literal=expr.literal) |
| |
| |
| @visit_bound_predicate.register(BoundGreaterThan) |
| def _(expr: BoundGreaterThan[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_greater_than(term=expr.term, literal=expr.literal) |
| |
| |
| @visit_bound_predicate.register(BoundLessThan) |
| def _(expr: BoundLessThan[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_less_than(term=expr.term, literal=expr.literal) |
| |
| |
| @visit_bound_predicate.register(BoundLessThanOrEqual) |
| def _(expr: BoundLessThanOrEqual[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_less_than_or_equal(term=expr.term, literal=expr.literal) |
| |
| |
| @visit_bound_predicate.register(BoundStartsWith) |
| def _(expr: BoundStartsWith[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_starts_with(term=expr.term, literal=expr.literal) |
| |
| |
| @visit_bound_predicate.register(BoundNotStartsWith) |
| def _(expr: BoundNotStartsWith[L], visitor: BoundBooleanExpressionVisitor[T]) -> T: |
| return visitor.visit_not_starts_with(term=expr.term, literal=expr.literal) |
| |
| |
| def rewrite_not(expr: BooleanExpression) -> BooleanExpression: |
| return visit(expr, _RewriteNotVisitor()) |
| |
| |
| class _RewriteNotVisitor(BooleanExpressionVisitor[BooleanExpression]): |
| """Inverts the negations.""" |
| |
| def visit_true(self) -> BooleanExpression: |
| return AlwaysTrue() |
| |
| def visit_false(self) -> BooleanExpression: |
| return AlwaysFalse() |
| |
| def visit_not(self, child_result: BooleanExpression) -> BooleanExpression: |
| return ~child_result |
| |
| def visit_and(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return And(left=left_result, right=right_result) |
| |
| def visit_or(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return Or(left=left_result, right=right_result) |
| |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> BooleanExpression: |
| return predicate |
| |
| def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> BooleanExpression: |
| return predicate |
| |
| |
| def expression_evaluator(schema: Schema, unbound: BooleanExpression, case_sensitive: bool) -> Callable[[StructProtocol], bool]: |
| return _ExpressionEvaluator(schema, unbound, case_sensitive).eval |
| |
| |
| class _ExpressionEvaluator(BoundBooleanExpressionVisitor[bool]): |
| bound: BooleanExpression |
| struct: StructProtocol |
| |
| def __init__(self, schema: Schema, unbound: BooleanExpression, case_sensitive: bool): |
| self.bound = bind(schema, unbound, case_sensitive) |
| |
| def eval(self, struct: StructProtocol) -> bool: |
| self.struct = struct |
| return visit(self.bound, self) |
| |
| def visit_in(self, term: BoundTerm[L], literals: Set[L]) -> bool: |
| return term.eval(self.struct) in literals |
| |
| def visit_not_in(self, term: BoundTerm[L], literals: Set[L]) -> bool: |
| return term.eval(self.struct) not in literals |
| |
| def visit_is_nan(self, term: BoundTerm[L]) -> bool: |
| val = term.eval(self.struct) |
| return val != val |
| |
| def visit_not_nan(self, term: BoundTerm[L]) -> bool: |
| val = term.eval(self.struct) |
| return val == val |
| |
| def visit_is_null(self, term: BoundTerm[L]) -> bool: |
| return term.eval(self.struct) is None |
| |
| def visit_not_null(self, term: BoundTerm[L]) -> bool: |
| return term.eval(self.struct) is not None |
| |
| def visit_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| return term.eval(self.struct) == literal.value |
| |
| def visit_not_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| return term.eval(self.struct) != literal.value |
| |
| def visit_greater_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| value = term.eval(self.struct) |
| return value is not None and value >= literal.value |
| |
| def visit_greater_than(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| value = term.eval(self.struct) |
| return value is not None and value > literal.value |
| |
| def visit_less_than(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| value = term.eval(self.struct) |
| return value is not None and value < literal.value |
| |
| def visit_less_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| value = term.eval(self.struct) |
| return value is not None and value <= literal.value |
| |
| def visit_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| eval_res = term.eval(self.struct) |
| return eval_res is not None and str(eval_res).startswith(str(literal.value)) |
| |
| def visit_not_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| return not self.visit_starts_with(term, literal) |
| |
| def visit_true(self) -> bool: |
| return True |
| |
| def visit_false(self) -> bool: |
| return False |
| |
| def visit_not(self, child_result: bool) -> bool: |
| return not child_result |
| |
| def visit_and(self, left_result: bool, right_result: bool) -> bool: |
| return left_result and right_result |
| |
| def visit_or(self, left_result: bool, right_result: bool) -> bool: |
| return left_result or right_result |
| |
| |
| ROWS_MIGHT_MATCH = True |
| ROWS_MUST_MATCH = True |
| ROWS_CANNOT_MATCH = False |
| ROWS_MIGHT_NOT_MATCH = False |
| IN_PREDICATE_LIMIT = 200 |
| |
| |
| def _from_byte_buffer(field_type: IcebergType, val: bytes) -> Any: |
| if not isinstance(field_type, PrimitiveType): |
| raise ValueError(f"Expected a PrimitiveType, got: {type(field_type)}") |
| return from_bytes(field_type, val) |
| |
| |
| class _ManifestEvalVisitor(BoundBooleanExpressionVisitor[bool]): |
| partition_fields: List[PartitionFieldSummary] |
| partition_filter: BooleanExpression |
| |
| def __init__(self, partition_struct_schema: Schema, partition_filter: BooleanExpression, case_sensitive: bool) -> None: |
| self.partition_filter = bind(partition_struct_schema, rewrite_not(partition_filter), case_sensitive) |
| |
| def eval(self, manifest: ManifestFile) -> bool: |
| if partitions := manifest.partitions: |
| self.partition_fields = partitions |
| return visit(self.partition_filter, self) |
| |
| # No partition information |
| return ROWS_MIGHT_MATCH |
| |
| def visit_in(self, term: BoundTerm[L], literals: Set[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| |
| if field.lower_bound is None: |
| return ROWS_CANNOT_MATCH |
| |
| if len(literals) > IN_PREDICATE_LIMIT: |
| return ROWS_MIGHT_MATCH |
| |
| lower = _from_byte_buffer(term.ref().field.field_type, field.lower_bound) |
| |
| if all(lower > val for val in literals): |
| return ROWS_CANNOT_MATCH |
| |
| if field.upper_bound is not None: |
| upper = _from_byte_buffer(term.ref().field.field_type, field.upper_bound) |
| if all(upper < val for val in literals): |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_in(self, term: BoundTerm[L], literals: Set[L]) -> bool: |
| # because the bounds are not necessarily a min or max value, this cannot be answered using |
| # them. notIn(col, {X, ...}) with (X, Y) doesn't guarantee that X is a value in col. |
| return ROWS_MIGHT_MATCH |
| |
| def visit_is_nan(self, term: BoundTerm[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| |
| if field.contains_nan is False: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_nan(self, term: BoundTerm[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| |
| if field.contains_nan is True and field.contains_null is False and field.lower_bound is None: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_is_null(self, term: BoundTerm[L]) -> bool: |
| pos = term.ref().accessor.position |
| |
| if self.partition_fields[pos].contains_null is False: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_null(self, term: BoundTerm[L]) -> bool: |
| pos = term.ref().accessor.position |
| |
| # contains_null encodes whether at least one partition value is null, |
| # lowerBound is null if all partition values are null |
| all_null = self.partition_fields[pos].contains_null is True and self.partition_fields[pos].lower_bound is None |
| |
| if all_null and isinstance(term.ref().field.field_type, (DoubleType, FloatType)): |
| # floating point types may include NaN values, which we check separately. |
| # In case bounds don't include NaN value, contains_nan needs to be checked against. |
| all_null = self.partition_fields[pos].contains_nan is False |
| |
| if all_null: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| |
| if field.lower_bound is None or field.upper_bound is None: |
| # values are all null and literal cannot contain null |
| return ROWS_CANNOT_MATCH |
| |
| lower = _from_byte_buffer(term.ref().field.field_type, field.lower_bound) |
| |
| if lower > literal.value: |
| return ROWS_CANNOT_MATCH |
| |
| upper = _from_byte_buffer(term.ref().field.field_type, field.upper_bound) |
| |
| if literal.value > upper: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| # because the bounds are not necessarily a min or max value, this cannot be answered using |
| # them. notEq(col, X) with (X, Y) doesn't guarantee that X is a value in col. |
| return ROWS_MIGHT_MATCH |
| |
| def visit_greater_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| |
| if field.upper_bound is None: |
| return ROWS_CANNOT_MATCH |
| |
| upper = _from_byte_buffer(term.ref().field.field_type, field.upper_bound) |
| |
| if literal.value > upper: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_greater_than(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| |
| if field.upper_bound is None: |
| return ROWS_CANNOT_MATCH |
| |
| upper = _from_byte_buffer(term.ref().field.field_type, field.upper_bound) |
| |
| if literal.value >= upper: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_less_than(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| |
| if field.lower_bound is None: |
| return ROWS_CANNOT_MATCH |
| |
| lower = _from_byte_buffer(term.ref().field.field_type, field.lower_bound) |
| |
| if literal.value <= lower: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_less_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| |
| if field.lower_bound is None: |
| return ROWS_CANNOT_MATCH |
| |
| lower = _from_byte_buffer(term.ref().field.field_type, field.lower_bound) |
| |
| if literal.value < lower: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| prefix = str(literal.value) |
| len_prefix = len(prefix) |
| |
| if field.lower_bound is None: |
| return ROWS_CANNOT_MATCH |
| |
| lower = _from_byte_buffer(term.ref().field.field_type, field.lower_bound) |
| # truncate lower bound so that its length is not greater than the length of prefix |
| if lower is not None and lower[:len_prefix] > prefix: |
| return ROWS_CANNOT_MATCH |
| |
| if field.upper_bound is None: |
| return ROWS_CANNOT_MATCH |
| |
| upper = _from_byte_buffer(term.ref().field.field_type, field.upper_bound) |
| # truncate upper bound so that its length is not greater than the length of prefix |
| if upper is not None and upper[:len_prefix] < prefix: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| pos = term.ref().accessor.position |
| field = self.partition_fields[pos] |
| prefix = str(literal.value) |
| len_prefix = len(prefix) |
| |
| if field.contains_null or field.lower_bound is None or field.upper_bound is None: |
| return ROWS_MIGHT_MATCH |
| |
| # not_starts_with will match unless all values must start with the prefix. This happens when |
| # the lower and upper bounds both start with the prefix. |
| lower = _from_byte_buffer(term.ref().field.field_type, field.lower_bound) |
| upper = _from_byte_buffer(term.ref().field.field_type, field.upper_bound) |
| |
| if lower is not None and upper is not None: |
| # if lower is shorter than the prefix then lower doesn't start with the prefix |
| if len(lower) < len_prefix: |
| return ROWS_MIGHT_MATCH |
| |
| if lower[:len_prefix] == prefix: |
| # if upper is shorter than the prefix then upper can't start with the prefix |
| if len(upper) < len_prefix: |
| return ROWS_MIGHT_MATCH |
| |
| if upper[:len_prefix] == prefix: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_true(self) -> bool: |
| return ROWS_MIGHT_MATCH |
| |
| def visit_false(self) -> bool: |
| return ROWS_CANNOT_MATCH |
| |
| def visit_not(self, child_result: bool) -> bool: |
| return not child_result |
| |
| def visit_and(self, left_result: bool, right_result: bool) -> bool: |
| return left_result and right_result |
| |
| def visit_or(self, left_result: bool, right_result: bool) -> bool: |
| return left_result or right_result |
| |
| |
| def manifest_evaluator( |
| partition_spec: PartitionSpec, schema: Schema, partition_filter: BooleanExpression, case_sensitive: bool = True |
| ) -> Callable[[ManifestFile], bool]: |
| partition_type = partition_spec.partition_type(schema) |
| partition_schema = Schema(*partition_type.fields) |
| evaluator = _ManifestEvalVisitor(partition_schema, partition_filter, case_sensitive) |
| return evaluator.eval |
| |
| |
| class ProjectionEvaluator(BooleanExpressionVisitor[BooleanExpression], ABC): |
| schema: Schema |
| spec: PartitionSpec |
| case_sensitive: bool |
| |
| def __init__(self, schema: Schema, spec: PartitionSpec, case_sensitive: bool): |
| self.schema = schema |
| self.spec = spec |
| self.case_sensitive = case_sensitive |
| |
| def project(self, expr: BooleanExpression) -> BooleanExpression: |
| # projections assume that there are no NOT nodes in the expression tree. to ensure that this |
| # is the case, the expression is rewritten to push all NOT nodes down to the expression |
| # leaf nodes. |
| # this is necessary to ensure that the default expression returned when a predicate can't be |
| # projected is correct. |
| return visit(bind(self.schema, rewrite_not(expr), self.case_sensitive), self) |
| |
| def visit_true(self) -> BooleanExpression: |
| return AlwaysTrue() |
| |
| def visit_false(self) -> BooleanExpression: |
| return AlwaysFalse() |
| |
| def visit_not(self, child_result: BooleanExpression) -> BooleanExpression: |
| raise ValueError(f"Cannot project not expression, should be rewritten: {child_result}") |
| |
| def visit_and(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return And(left_result, right_result) |
| |
| def visit_or(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return Or(left_result, right_result) |
| |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> BooleanExpression: |
| raise ValueError(f"Cannot project unbound predicate: {predicate}") |
| |
| |
| class InclusiveProjection(ProjectionEvaluator): |
| def visit_bound_predicate(self, predicate: BoundPredicate[Any]) -> BooleanExpression: |
| parts = self.spec.fields_by_source_id(predicate.term.ref().field.field_id) |
| |
| result: BooleanExpression = AlwaysTrue() |
| for part in parts: |
| # consider (d = 2019-01-01) with bucket(7, d) and bucket(5, d) |
| # projections: b1 = bucket(7, '2019-01-01') = 5, b2 = bucket(5, '2019-01-01') = 0 |
| # any value where b1 != 5 or any value where b2 != 0 cannot be the '2019-01-01' |
| # |
| # similarly, if partitioning by day(ts) and hour(ts), the more restrictive |
| # projection should be used. ts = 2019-01-01T01:00:00 produces day=2019-01-01 and |
| # hour=2019-01-01-01. the value will be in 2019-01-01-01 and not in 2019-01-01-02. |
| incl_projection = part.transform.project(name=part.name, pred=predicate) |
| if incl_projection is not None: |
| result = And(result, incl_projection) |
| |
| return result |
| |
| |
| def inclusive_projection( |
| schema: Schema, spec: PartitionSpec, case_sensitive: bool = True |
| ) -> Callable[[BooleanExpression], BooleanExpression]: |
| return InclusiveProjection(schema, spec, case_sensitive).project |
| |
| |
| class _ColumnNameTranslator(BooleanExpressionVisitor[BooleanExpression]): |
| """Converts the column names with the ones in the actual file. |
| |
| Args: |
| file_schema (Schema): The schema of the file. |
| case_sensitive (bool): Whether to consider case when binding a reference to a field in a schema, defaults to True. |
| |
| Raises: |
| TypeError: In the case of an UnboundPredicate. |
| ValueError: When a column name cannot be found. |
| """ |
| |
| file_schema: Schema |
| case_sensitive: bool |
| |
| def __init__(self, file_schema: Schema, case_sensitive: bool) -> None: |
| self.file_schema = file_schema |
| self.case_sensitive = case_sensitive |
| |
| def visit_true(self) -> BooleanExpression: |
| return AlwaysTrue() |
| |
| def visit_false(self) -> BooleanExpression: |
| return AlwaysFalse() |
| |
| def visit_not(self, child_result: BooleanExpression) -> BooleanExpression: |
| return Not(child=child_result) |
| |
| def visit_and(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return And(left=left_result, right=right_result) |
| |
| def visit_or(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return Or(left=left_result, right=right_result) |
| |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> BooleanExpression: |
| raise TypeError(f"Expected Bound Predicate, got: {predicate.term}") |
| |
| def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> BooleanExpression: |
| file_column_name = self.file_schema.find_column_name(predicate.term.ref().field.field_id) |
| |
| if file_column_name is None: |
| # In the case of schema evolution, the column might not be present |
| # in the file schema when reading older data |
| if isinstance(predicate, BoundIsNull): |
| return AlwaysTrue() |
| else: |
| return AlwaysFalse() |
| |
| if isinstance(predicate, BoundUnaryPredicate): |
| return predicate.as_unbound(file_column_name) |
| elif isinstance(predicate, BoundLiteralPredicate): |
| return predicate.as_unbound(file_column_name, predicate.literal) |
| elif isinstance(predicate, BoundSetPredicate): |
| return predicate.as_unbound(file_column_name, predicate.literals) |
| else: |
| raise ValueError(f"Unsupported predicate: {predicate}") |
| |
| |
| def translate_column_names(expr: BooleanExpression, file_schema: Schema, case_sensitive: bool) -> BooleanExpression: |
| return visit(expr, _ColumnNameTranslator(file_schema, case_sensitive)) |
| |
| |
| class _ExpressionFieldIDs(BooleanExpressionVisitor[Set[int]]): |
| """Extracts the field IDs used in the BooleanExpression.""" |
| |
| def visit_true(self) -> Set[int]: |
| return set() |
| |
| def visit_false(self) -> Set[int]: |
| return set() |
| |
| def visit_not(self, child_result: Set[int]) -> Set[int]: |
| return child_result |
| |
| def visit_and(self, left_result: Set[int], right_result: Set[int]) -> Set[int]: |
| return left_result.union(right_result) |
| |
| def visit_or(self, left_result: Set[int], right_result: Set[int]) -> Set[int]: |
| return left_result.union(right_result) |
| |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> Set[int]: |
| raise ValueError("Only works on bound records") |
| |
| def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> Set[int]: |
| return {predicate.term.ref().field.field_id} |
| |
| |
| def extract_field_ids(expr: BooleanExpression) -> Set[int]: |
| return visit(expr, _ExpressionFieldIDs()) |
| |
| |
| class _RewriteToDNF(BooleanExpressionVisitor[Tuple[BooleanExpression, ...]]): |
| def visit_true(self) -> Tuple[BooleanExpression, ...]: |
| return (AlwaysTrue(),) |
| |
| def visit_false(self) -> Tuple[BooleanExpression, ...]: |
| return (AlwaysFalse(),) |
| |
| def visit_not(self, child_result: Tuple[BooleanExpression, ...]) -> Tuple[BooleanExpression, ...]: |
| raise ValueError(f"Not expressions are not allowed: {child_result}") |
| |
| def visit_and( |
| self, left_result: Tuple[BooleanExpression, ...], right_result: Tuple[BooleanExpression, ...] |
| ) -> Tuple[BooleanExpression, ...]: |
| # Distributive law: |
| # ((P OR Q) AND (R OR S)) AND (((P AND R) OR (P AND S)) OR ((Q AND R) OR ((Q AND S))) |
| # A AND (B OR C) = (A AND B) OR (A AND C) |
| # (A OR B) AND C = (A AND C) OR (B AND C) |
| return tuple(And(le, re) for le in left_result for re in right_result) |
| |
| def visit_or( |
| self, left_result: Tuple[BooleanExpression, ...], right_result: Tuple[BooleanExpression, ...] |
| ) -> Tuple[BooleanExpression, ...]: |
| return left_result + right_result |
| |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> Tuple[BooleanExpression, ...]: |
| return (predicate,) |
| |
| def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> Tuple[BooleanExpression, ...]: |
| return (predicate,) |
| |
| |
| def rewrite_to_dnf(expr: BooleanExpression) -> Tuple[BooleanExpression, ...]: |
| # Rewrites an arbitrary boolean expression to disjunctive normal form (DNF): |
| # (A AND NOT(B) AND C) OR (NOT(D) AND E AND F) OR (G) |
| expr_without_not = rewrite_not(expr) |
| return visit(expr_without_not, _RewriteToDNF()) |
| |
| |
| class ExpressionToPlainFormat(BoundBooleanExpressionVisitor[List[Tuple[str, str, Any]]]): |
| cast_int_to_date: bool |
| |
| def __init__(self, cast_int_to_date: bool = False) -> None: |
| self.cast_int_to_date = cast_int_to_date |
| |
| def _cast_if_necessary(self, iceberg_type: IcebergType, literal: Union[L, Set[L]]) -> Union[L, Set[L]]: |
| if self.cast_int_to_date: |
| iceberg_type_class = type(iceberg_type) |
| conversions = {TimestampType: micros_to_timestamp, TimestamptzType: micros_to_timestamptz} |
| if iceberg_type_class in conversions: |
| conversion_function = conversions[iceberg_type_class] |
| if isinstance(literal, set): |
| return {conversion_function(lit) for lit in literal} # type: ignore |
| else: |
| return conversion_function(literal) # type: ignore |
| return literal |
| |
| def visit_in(self, term: BoundTerm[L], literals: Set[L]) -> List[Tuple[str, str, Any]]: |
| field = term.ref().field |
| return [(term.ref().field.name, "in", self._cast_if_necessary(field.field_type, literals))] |
| |
| def visit_not_in(self, term: BoundTerm[L], literals: Set[L]) -> List[Tuple[str, str, Any]]: |
| field = term.ref().field |
| return [(field.name, "not in", self._cast_if_necessary(field.field_type, literals))] |
| |
| def visit_is_nan(self, term: BoundTerm[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, "==", float("nan"))] |
| |
| def visit_not_nan(self, term: BoundTerm[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, "!=", float("nan"))] |
| |
| def visit_is_null(self, term: BoundTerm[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, "==", None)] |
| |
| def visit_not_null(self, term: BoundTerm[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, "!=", None)] |
| |
| def visit_equal(self, term: BoundTerm[L], literal: Literal[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, "==", self._cast_if_necessary(term.ref().field.field_type, literal.value))] |
| |
| def visit_not_equal(self, term: BoundTerm[L], literal: Literal[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, "!=", self._cast_if_necessary(term.ref().field.field_type, literal.value))] |
| |
| def visit_greater_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, ">=", self._cast_if_necessary(term.ref().field.field_type, literal.value))] |
| |
| def visit_greater_than(self, term: BoundTerm[L], literal: Literal[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, ">", self._cast_if_necessary(term.ref().field.field_type, literal.value))] |
| |
| def visit_less_than(self, term: BoundTerm[L], literal: Literal[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, "<", self._cast_if_necessary(term.ref().field.field_type, literal.value))] |
| |
| def visit_less_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> List[Tuple[str, str, Any]]: |
| return [(term.ref().field.name, "<=", self._cast_if_necessary(term.ref().field.field_type, literal.value))] |
| |
| def visit_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> List[Tuple[str, str, Any]]: |
| return [] |
| |
| def visit_not_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> List[Tuple[str, str, Any]]: |
| return [] |
| |
| def visit_true(self) -> List[Tuple[str, str, Any]]: |
| return [] # Not supported |
| |
| def visit_false(self) -> List[Tuple[str, str, Any]]: |
| raise ValueError("Not supported: AlwaysFalse") |
| |
| def visit_not(self, child_result: List[Tuple[str, str, Any]]) -> List[Tuple[str, str, Any]]: |
| raise ValueError(f"Not allowed: {child_result}") |
| |
| def visit_and( |
| self, left_result: List[Tuple[str, str, Any]], right_result: List[Tuple[str, str, Any]] |
| ) -> List[Tuple[str, str, Any]]: |
| return left_result + right_result |
| |
| def visit_or( |
| self, left_result: List[Tuple[str, str, Any]], right_result: List[Tuple[str, str, Any]] |
| ) -> List[Tuple[str, str, Any]]: |
| raise ValueError(f"Not allowed: {left_result} || {right_result}") |
| |
| |
| def expression_to_plain_format( |
| expressions: Tuple[BooleanExpression, ...], cast_int_to_datetime: bool = False |
| ) -> List[List[Tuple[str, str, Any]]]: |
| """Format a Disjunctive Normal Form expression. |
| |
| These are the formats that the expression can be fed into: |
| |
| - https://arrow.apache.org/docs/python/generated/pyarrow.parquet.read_table.html |
| - https://docs.dask.org/en/stable/generated/dask.dataframe.read_parquet.html |
| |
| Contrary to normal DNF that may contain Not expressions, but here they should have |
| been rewritten. This can be done using ``rewrite_not(...)``. |
| |
| Keep in mind that this is only used for page skipping, and still needs to filter |
| on a row level. |
| |
| Args: |
| expressions: Expression in Disjunctive Normal Form. |
| |
| Returns: |
| Formatter filter compatible with Dask and PyArrow. |
| """ |
| # In the form of expr1 ∨ expr2 ∨ ... ∨ exprN |
| visitor = ExpressionToPlainFormat(cast_int_to_datetime) |
| return [visit(expression, visitor) for expression in expressions] |
| |
| |
| class _MetricsEvaluator(BoundBooleanExpressionVisitor[bool], ABC): |
| value_counts: Dict[int, int] |
| null_counts: Dict[int, int] |
| nan_counts: Dict[int, int] |
| lower_bounds: Dict[int, bytes] |
| upper_bounds: Dict[int, bytes] |
| |
| def visit_true(self) -> bool: |
| # all rows match |
| return ROWS_MIGHT_MATCH |
| |
| def visit_false(self) -> bool: |
| # all rows fail |
| return ROWS_CANNOT_MATCH |
| |
| def visit_not(self, child_result: bool) -> bool: |
| raise ValueError(f"NOT should be rewritten: {child_result}") |
| |
| def visit_and(self, left_result: bool, right_result: bool) -> bool: |
| return left_result and right_result |
| |
| def visit_or(self, left_result: bool, right_result: bool) -> bool: |
| return left_result or right_result |
| |
| def _contains_nulls_only(self, field_id: int) -> bool: |
| if (value_count := self.value_counts.get(field_id)) and (null_count := self.null_counts.get(field_id)): |
| return value_count == null_count |
| return False |
| |
| def _contains_nans_only(self, field_id: int) -> bool: |
| if (nan_count := self.nan_counts.get(field_id)) and (value_count := self.value_counts.get(field_id)): |
| return nan_count == value_count |
| return False |
| |
| def _is_nan(self, val: Any) -> bool: |
| try: |
| return math.isnan(val) |
| except TypeError: |
| # In the case of None or other non-numeric types |
| return False |
| |
| |
| class _InclusiveMetricsEvaluator(_MetricsEvaluator): |
| struct: StructType |
| expr: BooleanExpression |
| |
| def __init__( |
| self, schema: Schema, expr: BooleanExpression, case_sensitive: bool = True, include_empty_files: bool = False |
| ) -> None: |
| self.struct = schema.as_struct() |
| self.include_empty_files = include_empty_files |
| self.expr = bind(schema, rewrite_not(expr), case_sensitive) |
| |
| def eval(self, file: DataFile) -> bool: |
| """Test whether the file may contain records that match the expression.""" |
| if not self.include_empty_files and file.record_count == 0: |
| return ROWS_CANNOT_MATCH |
| |
| if file.record_count < 0: |
| # Older version don't correctly implement record count from avro file and thus |
| # set record count -1 when importing avro tables to iceberg tables. This should |
| # be updated once we implemented and set correct record count. |
| return ROWS_MIGHT_MATCH |
| |
| self.value_counts = file.value_counts or EMPTY_DICT |
| self.null_counts = file.null_value_counts or EMPTY_DICT |
| self.nan_counts = file.nan_value_counts or EMPTY_DICT |
| self.lower_bounds = file.lower_bounds or EMPTY_DICT |
| self.upper_bounds = file.upper_bounds or EMPTY_DICT |
| |
| return visit(self.expr, self) |
| |
| def _may_contain_null(self, field_id: int) -> bool: |
| return self.null_counts is None or (field_id in self.null_counts and self.null_counts.get(field_id) is not None) |
| |
| def _contains_nans_only(self, field_id: int) -> bool: |
| if (nan_count := self.nan_counts.get(field_id)) and (value_count := self.value_counts.get(field_id)): |
| return nan_count == value_count |
| return False |
| |
| def visit_is_null(self, term: BoundTerm[L]) -> bool: |
| field_id = term.ref().field.field_id |
| |
| if self.null_counts.get(field_id) == 0: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_null(self, term: BoundTerm[L]) -> bool: |
| # no need to check whether the field is required because binding evaluates that case |
| # if the column has no non-null values, the expression cannot match |
| field_id = term.ref().field.field_id |
| |
| if self._contains_nulls_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_is_nan(self, term: BoundTerm[L]) -> bool: |
| field_id = term.ref().field.field_id |
| |
| if self.nan_counts.get(field_id) == 0: |
| return ROWS_CANNOT_MATCH |
| |
| # when there's no nanCounts information, but we already know the column only contains null, |
| # it's guaranteed that there's no NaN value |
| if self._contains_nulls_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_nan(self, term: BoundTerm[L]) -> bool: |
| field_id = term.ref().field.field_id |
| |
| if self._contains_nans_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_less_than(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| field = term.ref().field |
| field_id = field.field_id |
| |
| if self._contains_nulls_only(field_id) or self._contains_nans_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| if not isinstance(field.field_type, PrimitiveType): |
| raise ValueError(f"Expected PrimitiveType: {field.field_type}") |
| |
| if lower_bound_bytes := self.lower_bounds.get(field_id): |
| lower_bound = from_bytes(field.field_type, lower_bound_bytes) |
| |
| if self._is_nan(lower_bound): |
| # NaN indicates unreliable bounds. See the InclusiveMetricsEvaluator docs for more. |
| return ROWS_MIGHT_MATCH |
| |
| if lower_bound >= literal.value: # type: ignore[operator] |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_less_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| field = term.ref().field |
| field_id = field.field_id |
| |
| if self._contains_nulls_only(field_id) or self._contains_nans_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| if not isinstance(field.field_type, PrimitiveType): |
| raise ValueError(f"Expected PrimitiveType: {field.field_type}") |
| |
| if lower_bound_bytes := self.lower_bounds.get(field_id): |
| lower_bound = from_bytes(field.field_type, lower_bound_bytes) |
| if self._is_nan(lower_bound): |
| # NaN indicates unreliable bounds. See the InclusiveMetricsEvaluator docs for more. |
| return ROWS_MIGHT_MATCH |
| |
| if lower_bound > literal.value: # type: ignore[operator] |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_greater_than(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| field = term.ref().field |
| field_id = field.field_id |
| |
| if self._contains_nulls_only(field_id) or self._contains_nans_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| if not isinstance(field.field_type, PrimitiveType): |
| raise ValueError(f"Expected PrimitiveType: {field.field_type}") |
| |
| if upper_bound_bytes := self.upper_bounds.get(field_id): |
| upper_bound = from_bytes(field.field_type, upper_bound_bytes) |
| if upper_bound <= literal.value: # type: ignore[operator] |
| if self._is_nan(upper_bound): |
| # NaN indicates unreliable bounds. See the InclusiveMetricsEvaluator docs for more. |
| return ROWS_MIGHT_MATCH |
| |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_greater_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| field = term.ref().field |
| field_id = field.field_id |
| |
| if self._contains_nulls_only(field_id) or self._contains_nans_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| if not isinstance(field.field_type, PrimitiveType): |
| raise ValueError(f"Expected PrimitiveType: {field.field_type}") |
| |
| if upper_bound_bytes := self.upper_bounds.get(field_id): |
| upper_bound = from_bytes(field.field_type, upper_bound_bytes) |
| if upper_bound < literal.value: # type: ignore[operator] |
| if self._is_nan(upper_bound): |
| # NaN indicates unreliable bounds. See the InclusiveMetricsEvaluator docs for more. |
| return ROWS_MIGHT_MATCH |
| |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| field = term.ref().field |
| field_id = field.field_id |
| |
| if self._contains_nulls_only(field_id) or self._contains_nans_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| if not isinstance(field.field_type, PrimitiveType): |
| raise ValueError(f"Expected PrimitiveType: {field.field_type}") |
| |
| if lower_bound_bytes := self.lower_bounds.get(field_id): |
| lower_bound = from_bytes(field.field_type, lower_bound_bytes) |
| if self._is_nan(lower_bound): |
| # NaN indicates unreliable bounds. See the InclusiveMetricsEvaluator docs for more. |
| return ROWS_MIGHT_MATCH |
| |
| if lower_bound > literal.value: # type: ignore[operator] |
| return ROWS_CANNOT_MATCH |
| |
| if upper_bound_bytes := self.upper_bounds.get(field_id): |
| upper_bound = from_bytes(field.field_type, upper_bound_bytes) |
| if self._is_nan(upper_bound): |
| # NaN indicates unreliable bounds. See the InclusiveMetricsEvaluator docs for more. |
| return ROWS_MIGHT_MATCH |
| |
| if upper_bound < literal.value: # type: ignore[operator] |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| return ROWS_MIGHT_MATCH |
| |
| def visit_in(self, term: BoundTerm[L], literals: Set[L]) -> bool: |
| field = term.ref().field |
| field_id = field.field_id |
| |
| if self._contains_nulls_only(field_id) or self._contains_nans_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| if len(literals) > IN_PREDICATE_LIMIT: |
| # skip evaluating the predicate if the number of values is too big |
| return ROWS_MIGHT_MATCH |
| |
| if not isinstance(field.field_type, PrimitiveType): |
| raise ValueError(f"Expected PrimitiveType: {field.field_type}") |
| |
| if lower_bound_bytes := self.lower_bounds.get(field_id): |
| lower_bound = from_bytes(field.field_type, lower_bound_bytes) |
| if self._is_nan(lower_bound): |
| # NaN indicates unreliable bounds. See the InclusiveMetricsEvaluator docs for more. |
| return ROWS_MIGHT_MATCH |
| |
| literals = {lit for lit in literals if lower_bound <= lit} # type: ignore[operator] |
| if len(literals) == 0: |
| return ROWS_CANNOT_MATCH |
| |
| if upper_bound_bytes := self.upper_bounds.get(field_id): |
| upper_bound = from_bytes(field.field_type, upper_bound_bytes) |
| # this is different from Java, here NaN is always larger |
| if self._is_nan(upper_bound): |
| return ROWS_MIGHT_MATCH |
| |
| literals = {lit for lit in literals if upper_bound >= lit} # type: ignore[operator] |
| if len(literals) == 0: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_in(self, term: BoundTerm[L], literals: Set[L]) -> bool: |
| # because the bounds are not necessarily a min or max value, this cannot be answered using |
| # them. notIn(col, {X, ...}) with (X, Y) doesn't guarantee that X is a value in col. |
| return ROWS_MIGHT_MATCH |
| |
| def visit_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| field = term.ref().field |
| field_id: int = field.field_id |
| |
| if self._contains_nulls_only(field_id): |
| return ROWS_CANNOT_MATCH |
| |
| if not isinstance(field.field_type, PrimitiveType): |
| raise ValueError(f"Expected PrimitiveType: {field.field_type}") |
| |
| prefix = str(literal.value) |
| len_prefix = len(prefix) |
| |
| if lower_bound_bytes := self.lower_bounds.get(field_id): |
| lower_bound = str(from_bytes(field.field_type, lower_bound_bytes)) |
| |
| # truncate lower bound so that its length is not greater than the length of prefix |
| if lower_bound and lower_bound[:len_prefix] > prefix: |
| return ROWS_CANNOT_MATCH |
| |
| if upper_bound_bytes := self.upper_bounds.get(field_id): |
| upper_bound = str(from_bytes(field.field_type, upper_bound_bytes)) |
| |
| # truncate upper bound so that its length is not greater than the length of prefix |
| if upper_bound is not None and upper_bound[:len_prefix] < prefix: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| def visit_not_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| field = term.ref().field |
| field_id: int = field.field_id |
| |
| if self._may_contain_null(field_id): |
| return ROWS_MIGHT_MATCH |
| |
| if not isinstance(field.field_type, PrimitiveType): |
| raise ValueError(f"Expected PrimitiveType: {field.field_type}") |
| |
| prefix = str(literal.value) |
| len_prefix = len(prefix) |
| |
| # not_starts_with will match unless all values must start with the prefix. This happens when |
| # the lower and upper bounds both start with the prefix. |
| if (lower_bound_bytes := self.lower_bounds.get(field_id)) and (upper_bound_bytes := self.upper_bounds.get(field_id)): |
| lower_bound = str(from_bytes(field.field_type, lower_bound_bytes)) |
| upper_bound = str(from_bytes(field.field_type, upper_bound_bytes)) |
| |
| # if lower is shorter than the prefix then lower doesn't start with the prefix |
| if len(lower_bound) < len_prefix: |
| return ROWS_MIGHT_MATCH |
| |
| if lower_bound[:len_prefix] == prefix: |
| # if upper is shorter than the prefix then upper can't start with the prefix |
| if len(upper_bound) < len_prefix: |
| return ROWS_MIGHT_MATCH |
| |
| if upper_bound[:len_prefix] == prefix: |
| return ROWS_CANNOT_MATCH |
| |
| return ROWS_MIGHT_MATCH |
| |
| |
| def strict_projection( |
| schema: Schema, spec: PartitionSpec, case_sensitive: bool = True |
| ) -> Callable[[BooleanExpression], BooleanExpression]: |
| return StrictProjection(schema, spec, case_sensitive).project |
| |
| |
| class StrictProjection(ProjectionEvaluator): |
| def visit_bound_predicate(self, predicate: BoundPredicate[Any]) -> BooleanExpression: |
| parts = self.spec.fields_by_source_id(predicate.term.ref().field.field_id) |
| |
| result: BooleanExpression = AlwaysFalse() |
| for part in parts: |
| # consider (ts > 2019-01-01T01:00:00) with day(ts) and hour(ts) |
| # projections: d >= 2019-01-02 and h >= 2019-01-01-02 (note the inclusive bounds). |
| # any timestamp where either projection predicate is true must match the original |
| # predicate. For example, ts = 2019-01-01T03:00:00 matches the hour projection but not |
| # the day, but does match the original predicate. |
| strict_projection = part.transform.strict_project(name=part.name, pred=predicate) |
| if strict_projection is not None: |
| result = Or(result, strict_projection) |
| |
| return result |
| |
| |
| class _StrictMetricsEvaluator(_MetricsEvaluator): |
| struct: StructType |
| expr: BooleanExpression |
| |
| def __init__( |
| self, schema: Schema, expr: BooleanExpression, case_sensitive: bool = True, include_empty_files: bool = False |
| ) -> None: |
| self.struct = schema.as_struct() |
| self.include_empty_files = include_empty_files |
| self.expr = bind(schema, rewrite_not(expr), case_sensitive) |
| |
| def eval(self, file: DataFile) -> bool: |
| """Test whether all records within the file match the expression. |
| |
| Args: |
| file: A data file |
| |
| Returns: false if the file may contain any row that doesn't match |
| the expression, true otherwise. |
| """ |
| if file.record_count <= 0: |
| # Older version don't correctly implement record count from avro file and thus |
| # set record count -1 when importing avro tables to iceberg tables. This should |
| # be updated once we implemented and set correct record count. |
| return ROWS_MUST_MATCH |
| |
| self.value_counts = file.value_counts or EMPTY_DICT |
| self.null_counts = file.null_value_counts or EMPTY_DICT |
| self.nan_counts = file.nan_value_counts or EMPTY_DICT |
| self.lower_bounds = file.lower_bounds or EMPTY_DICT |
| self.upper_bounds = file.upper_bounds or EMPTY_DICT |
| |
| return visit(self.expr, self) |
| |
| def visit_is_null(self, term: BoundTerm[L]) -> bool: |
| # no need to check whether the field is required because binding evaluates that case |
| # if the column has any non-null values, the expression does not match |
| field_id = term.ref().field.field_id |
| |
| if self._contains_nulls_only(field_id): |
| return ROWS_MUST_MATCH |
| else: |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_not_null(self, term: BoundTerm[L]) -> bool: |
| # no need to check whether the field is required because binding evaluates that case |
| # if the column has any non-null values, the expression does not match |
| field_id = term.ref().field.field_id |
| |
| if (null_count := self.null_counts.get(field_id)) is not None and null_count == 0: |
| return ROWS_MUST_MATCH |
| else: |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_is_nan(self, term: BoundTerm[L]) -> bool: |
| field_id = term.ref().field.field_id |
| |
| if self._contains_nans_only(field_id): |
| return ROWS_MUST_MATCH |
| else: |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_not_nan(self, term: BoundTerm[L]) -> bool: |
| field_id = term.ref().field.field_id |
| |
| if (nan_count := self.nan_counts.get(field_id)) is not None and nan_count == 0: |
| return ROWS_MUST_MATCH |
| |
| if self._contains_nulls_only(field_id): |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_less_than(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| # Rows must match when: <----------Min----Max---X-------> |
| |
| field_id = term.ref().field.field_id |
| |
| if self._can_contain_nulls(field_id) or self._can_contain_nans(field_id): |
| return ROWS_MIGHT_NOT_MATCH |
| |
| if upper_bytes := self.upper_bounds.get(field_id): |
| field = self._get_field(field_id) |
| upper = _from_byte_buffer(field.field_type, upper_bytes) |
| |
| if upper < literal.value: |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_less_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| # Rows must match when: <----------Min----Max---X-------> |
| |
| field_id = term.ref().field.field_id |
| |
| if self._can_contain_nulls(field_id) or self._can_contain_nans(field_id): |
| return ROWS_MIGHT_NOT_MATCH |
| |
| if upper_bytes := self.upper_bounds.get(field_id): |
| field = self._get_field(field_id) |
| upper = _from_byte_buffer(field.field_type, upper_bytes) |
| |
| if upper <= literal.value: |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_greater_than(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| # Rows must match when: <-------X---Min----Max----------> |
| |
| field_id = term.ref().field.field_id |
| |
| if self._can_contain_nulls(field_id) or self._can_contain_nans(field_id): |
| return ROWS_MIGHT_NOT_MATCH |
| |
| if lower_bytes := self.lower_bounds.get(field_id): |
| field = self._get_field(field_id) |
| lower = _from_byte_buffer(field.field_type, lower_bytes) |
| |
| if self._is_nan(lower): |
| # NaN indicates unreliable bounds. |
| # See the _StrictMetricsEvaluator docs for more. |
| return ROWS_MIGHT_NOT_MATCH |
| |
| if lower > literal.value: |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_greater_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| # Rows must match when: <-------X---Min----Max----------> |
| field_id = term.ref().field.field_id |
| |
| if self._can_contain_nulls(field_id) or self._can_contain_nans(field_id): |
| return ROWS_MIGHT_NOT_MATCH |
| |
| if lower_bytes := self.lower_bounds.get(field_id): |
| field = self._get_field(field_id) |
| lower = _from_byte_buffer(field.field_type, lower_bytes) |
| |
| if self._is_nan(lower): |
| # NaN indicates unreliable bounds. |
| # See the _StrictMetricsEvaluator docs for more. |
| return ROWS_MIGHT_NOT_MATCH |
| |
| if lower >= literal.value: |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| # Rows must match when Min == X == Max |
| field_id = term.ref().field.field_id |
| |
| if self._can_contain_nulls(field_id) or self._can_contain_nans(field_id): |
| return ROWS_MIGHT_NOT_MATCH |
| |
| if (lower_bytes := self.lower_bounds.get(field_id)) and (upper_bytes := self.upper_bounds.get(field_id)): |
| field = self._get_field(field_id) |
| lower = _from_byte_buffer(field.field_type, lower_bytes) |
| upper = _from_byte_buffer(field.field_type, upper_bytes) |
| |
| if lower != literal.value or upper != literal.value: |
| return ROWS_MIGHT_NOT_MATCH |
| else: |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_not_equal(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| # Rows must match when X < Min or Max < X because it is not in the range |
| field_id = term.ref().field.field_id |
| |
| if self._can_contain_nulls(field_id) or self._can_contain_nans(field_id): |
| return ROWS_MUST_MATCH |
| |
| field = self._get_field(field_id) |
| |
| if lower_bytes := self.lower_bounds.get(field_id): |
| lower = _from_byte_buffer(field.field_type, lower_bytes) |
| |
| if self._is_nan(lower): |
| # NaN indicates unreliable bounds. |
| # See the _StrictMetricsEvaluator docs for more. |
| return ROWS_MIGHT_NOT_MATCH |
| |
| if lower > literal.value: |
| return ROWS_MUST_MATCH |
| |
| if upper_bytes := self.upper_bounds.get(field_id): |
| upper = _from_byte_buffer(field.field_type, upper_bytes) |
| |
| if upper < literal.value: |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_in(self, term: BoundTerm[L], literals: Set[L]) -> bool: |
| field_id = term.ref().field.field_id |
| |
| if self._can_contain_nulls(field_id) or self._can_contain_nans(field_id): |
| return ROWS_MIGHT_NOT_MATCH |
| |
| field = self._get_field(field_id) |
| |
| if (lower_bytes := self.lower_bounds.get(field_id)) and (upper_bytes := self.upper_bounds.get(field_id)): |
| # similar to the implementation in eq, first check if the lower bound is in the set |
| lower = _from_byte_buffer(field.field_type, lower_bytes) |
| if lower not in literals: |
| return ROWS_MIGHT_NOT_MATCH |
| |
| # check if the upper bound is in the set |
| upper = _from_byte_buffer(field.field_type, upper_bytes) |
| if upper not in literals: |
| return ROWS_MIGHT_NOT_MATCH |
| |
| # finally check if the lower bound and the upper bound are equal |
| if lower != upper: |
| return ROWS_MIGHT_NOT_MATCH |
| |
| # All values must be in the set if the lower bound and the upper bound are |
| # in the set and are equal. |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_not_in(self, term: BoundTerm[L], literals: Set[L]) -> bool: |
| field_id = term.ref().field.field_id |
| |
| if self._can_contain_nulls(field_id) or self._can_contain_nans(field_id): |
| return ROWS_MUST_MATCH |
| |
| field = self._get_field(field_id) |
| |
| if lower_bytes := self.lower_bounds.get(field_id): |
| lower = _from_byte_buffer(field.field_type, lower_bytes) |
| |
| if self._is_nan(lower): |
| # NaN indicates unreliable bounds. |
| # See the StrictMetricsEvaluator docs for more. |
| return ROWS_MIGHT_NOT_MATCH |
| |
| literals = {val for val in literals if lower <= val} |
| if len(literals) == 0: |
| return ROWS_MUST_MATCH |
| |
| if upper_bytes := self.upper_bounds.get(field_id): |
| upper = _from_byte_buffer(field.field_type, upper_bytes) |
| |
| literals = {val for val in literals if upper >= val} |
| |
| if len(literals) == 0: |
| return ROWS_MUST_MATCH |
| |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def visit_not_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> bool: |
| return ROWS_MIGHT_NOT_MATCH |
| |
| def _get_field(self, field_id: int) -> NestedField: |
| field = self.struct.field(field_id=field_id) |
| if field is None: |
| raise ValueError(f"Cannot find field, might be nested or missing: {field_id}") |
| |
| return field |
| |
| def _can_contain_nulls(self, field_id: int) -> bool: |
| return (null_count := self.null_counts.get(field_id)) is not None and null_count > 0 |
| |
| def _can_contain_nans(self, field_id: int) -> bool: |
| return (nan_count := self.nan_counts.get(field_id)) is not None and nan_count > 0 |
| |
| |
| class ResidualVisitor(BoundBooleanExpressionVisitor[BooleanExpression], ABC): |
| """Finds the residuals for an Expression the partitions in the given PartitionSpec. |
| |
| A residual expression is made by partially evaluating an expression using partition values. |
| For example, if a table is partitioned by day(utc_timestamp) and is read with a filter expression |
| utc_timestamp > a and utc_timestamp < b, then there are 4 possible residuals expressions |
| for the partition data, d: |
| |
| |
| 1. If d > day(a) and d < day(b), the residual is always true |
| 2. If d == day(a) and d != day(b), the residual is utc_timestamp > a |
| 3. if d == day(b) and d != day(a), the residual is utc_timestamp < b |
| 4. If d == day(a) == day(b), the residual is utc_timestamp > a and utc_timestamp < b |
| Partition data is passed using StructLike. Residuals are returned by residualFor(StructLike). |
| """ |
| |
| schema: Schema |
| spec: PartitionSpec |
| case_sensitive: bool |
| expr: BooleanExpression |
| |
| def __init__(self, schema: Schema, spec: PartitionSpec, case_sensitive: bool, expr: BooleanExpression) -> None: |
| self.schema = schema |
| self.spec = spec |
| self.case_sensitive = case_sensitive |
| self.expr = expr |
| |
| def eval(self, partition_data: Record) -> BooleanExpression: |
| self.struct = partition_data |
| return visit(self.expr, visitor=self) |
| |
| def visit_true(self) -> BooleanExpression: |
| return AlwaysTrue() |
| |
| def visit_false(self) -> BooleanExpression: |
| return AlwaysFalse() |
| |
| def visit_not(self, child_result: BooleanExpression) -> BooleanExpression: |
| return Not(child_result) |
| |
| def visit_and(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return And(left_result, right_result) |
| |
| def visit_or(self, left_result: BooleanExpression, right_result: BooleanExpression) -> BooleanExpression: |
| return Or(left_result, right_result) |
| |
| def visit_is_null(self, term: BoundTerm[L]) -> BooleanExpression: |
| if term.eval(self.struct) is None: |
| return AlwaysTrue() |
| else: |
| return AlwaysFalse() |
| |
| def visit_not_null(self, term: BoundTerm[L]) -> BooleanExpression: |
| if term.eval(self.struct) is not None: |
| return AlwaysTrue() |
| else: |
| return AlwaysFalse() |
| |
| def visit_is_nan(self, term: BoundTerm[L]) -> BooleanExpression: |
| val = term.eval(self.struct) |
| if isinstance(val, SupportsFloat) and math.isnan(val): |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_not_nan(self, term: BoundTerm[L]) -> BooleanExpression: |
| val = term.eval(self.struct) |
| if isinstance(val, SupportsFloat) and not math.isnan(val): |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_less_than(self, term: BoundTerm[L], literal: Literal[L]) -> BooleanExpression: |
| if term.eval(self.struct) < literal.value: |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_less_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> BooleanExpression: |
| if term.eval(self.struct) <= literal.value: |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_greater_than(self, term: BoundTerm[L], literal: Literal[L]) -> BooleanExpression: |
| if term.eval(self.struct) > literal.value: |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_greater_than_or_equal(self, term: BoundTerm[L], literal: Literal[L]) -> BooleanExpression: |
| if term.eval(self.struct) >= literal.value: |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_equal(self, term: BoundTerm[L], literal: Literal[L]) -> BooleanExpression: |
| if term.eval(self.struct) == literal.value: |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_not_equal(self, term: BoundTerm[L], literal: Literal[L]) -> BooleanExpression: |
| if term.eval(self.struct) != literal.value: |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_in(self, term: BoundTerm[L], literals: Set[L]) -> BooleanExpression: |
| if term.eval(self.struct) in literals: |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_not_in(self, term: BoundTerm[L], literals: Set[L]) -> BooleanExpression: |
| if term.eval(self.struct) not in literals: |
| return self.visit_true() |
| else: |
| return self.visit_false() |
| |
| def visit_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> BooleanExpression: |
| eval_res = term.eval(self.struct) |
| if eval_res is not None and str(eval_res).startswith(str(literal.value)): |
| return AlwaysTrue() |
| else: |
| return AlwaysFalse() |
| |
| def visit_not_starts_with(self, term: BoundTerm[L], literal: Literal[L]) -> BooleanExpression: |
| if not self.visit_starts_with(term, literal): |
| return AlwaysTrue() |
| else: |
| return AlwaysFalse() |
| |
| def visit_bound_predicate(self, predicate: BoundPredicate[Any]) -> BooleanExpression: |
| """ |
| If there is no strict projection or if it evaluates to false, then return the predicate. |
| |
| Get the strict projection and inclusive projection of this predicate in partition data, |
| then use them to determine whether to return the original predicate. The strict projection |
| returns true iff the original predicate would have returned true, so the predicate can be |
| eliminated if the strict projection evaluates to true. Similarly the inclusive projection |
| returns false iff the original predicate would have returned false, so the predicate can |
| also be eliminated if the inclusive projection evaluates to false. |
| |
| """ |
| parts = self.spec.fields_by_source_id(predicate.term.ref().field.field_id) |
| if parts == []: |
| return predicate |
| |
| def struct_to_schema(struct: StructType) -> Schema: |
| return Schema(*struct.fields) |
| |
| for part in parts: |
| strict_projection = part.transform.strict_project(part.name, predicate) |
| strict_result = None |
| |
| if strict_projection is not None: |
| bound = strict_projection.bind( |
| struct_to_schema(self.spec.partition_type(self.schema)), case_sensitive=self.case_sensitive |
| ) |
| if isinstance(bound, BoundPredicate): |
| strict_result = super().visit_bound_predicate(bound) |
| else: |
| # if the result is not a predicate, then it must be a constant like alwaysTrue or alwaysFalse |
| strict_result = bound |
| |
| if isinstance(strict_result, AlwaysTrue): |
| return AlwaysTrue() |
| |
| inclusive_projection = part.transform.project(part.name, predicate) |
| inclusive_result = None |
| if inclusive_projection is not None: |
| bound_inclusive = inclusive_projection.bind( |
| struct_to_schema(self.spec.partition_type(self.schema)), case_sensitive=self.case_sensitive |
| ) |
| if isinstance(bound_inclusive, BoundPredicate): |
| # using predicate method specific to inclusive |
| inclusive_result = super().visit_bound_predicate(bound_inclusive) |
| else: |
| # if the result is not a predicate, then it must be a constant like alwaysTrue or |
| # alwaysFalse |
| inclusive_result = bound_inclusive |
| if isinstance(inclusive_result, AlwaysFalse): |
| return AlwaysFalse() |
| |
| return predicate |
| |
| def visit_unbound_predicate(self, predicate: UnboundPredicate[L]) -> BooleanExpression: |
| bound = predicate.bind(self.schema, case_sensitive=self.case_sensitive) |
| |
| if isinstance(bound, BoundPredicate): |
| bound_residual = self.visit_bound_predicate(predicate=bound) |
| if not isinstance(bound_residual, (AlwaysFalse, AlwaysTrue)): |
| # replace inclusive original unbound predicate |
| return predicate |
| |
| # use the non-predicate residual (e.g. alwaysTrue) |
| return bound_residual |
| |
| # if binding didn't result in a Predicate, return the expression |
| return bound |
| |
| |
| class ResidualEvaluator(ResidualVisitor): |
| def residual_for(self, partition_data: Record) -> BooleanExpression: |
| return self.eval(partition_data) |
| |
| |
| class UnpartitionedResidualEvaluator(ResidualEvaluator): |
| # Finds the residuals for an Expression the partitions in the given PartitionSpec |
| def __init__(self, schema: Schema, expr: BooleanExpression): |
| super().__init__(schema=schema, spec=UNPARTITIONED_PARTITION_SPEC, expr=expr, case_sensitive=False) |
| self.expr = expr |
| |
| def residual_for(self, partition_data: Record) -> BooleanExpression: |
| return self.expr |
| |
| |
| def residual_evaluator_of( |
| spec: PartitionSpec, expr: BooleanExpression, case_sensitive: bool, schema: Schema |
| ) -> ResidualEvaluator: |
| return ( |
| UnpartitionedResidualEvaluator(schema=schema, expr=expr) |
| if spec.is_unpartitioned() |
| else ResidualEvaluator(spec=spec, expr=expr, schema=schema, case_sensitive=case_sensitive) |
| ) |
