| # 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. | |
| # pylint:disable=redefined-outer-name | |
| from concurrent.futures import ThreadPoolExecutor | |
| from threading import Event | |
| from typing import Any | |
| import pytest | |
| from pyiceberg.conversions import to_bytes | |
| from pyiceberg.expressions import ( | |
| AlwaysFalse, | |
| AlwaysTrue, | |
| And, | |
| BooleanExpression, | |
| BoundEqualTo, | |
| BoundGreaterThan, | |
| BoundGreaterThanOrEqual, | |
| BoundIn, | |
| BoundIsNaN, | |
| BoundIsNull, | |
| BoundLessThan, | |
| BoundLessThanOrEqual, | |
| BoundNotEqualTo, | |
| BoundNotIn, | |
| BoundNotNaN, | |
| BoundNotNull, | |
| BoundNotStartsWith, | |
| BoundPredicate, | |
| BoundReference, | |
| BoundStartsWith, | |
| BoundTerm, | |
| EqualTo, | |
| GreaterThan, | |
| GreaterThanOrEqual, | |
| In, | |
| IsNaN, | |
| IsNull, | |
| LessThan, | |
| LessThanOrEqual, | |
| Not, | |
| NotEqualTo, | |
| NotIn, | |
| NotNaN, | |
| NotNull, | |
| NotStartsWith, | |
| Or, | |
| Reference, | |
| StartsWith, | |
| UnboundPredicate, | |
| ) | |
| from pyiceberg.expressions.literals import Literal, literal | |
| from pyiceberg.expressions.visitors import ( | |
| BindVisitor, | |
| BooleanExpressionVisitor, | |
| BoundBooleanExpressionVisitor, | |
| _ExpressionEvaluator, | |
| _ManifestEvalVisitor, | |
| expression_evaluator, | |
| expression_to_plain_format, | |
| rewrite_not, | |
| rewrite_to_dnf, | |
| translate_column_names, | |
| visit, | |
| visit_bound_predicate, | |
| ) | |
| from pyiceberg.manifest import ManifestFile, PartitionFieldSummary | |
| from pyiceberg.schema import Accessor, Schema | |
| from pyiceberg.typedef import Record | |
| from pyiceberg.types import ( | |
| BinaryType, | |
| BooleanType, | |
| DoubleType, | |
| FloatType, | |
| IcebergType, | |
| IntegerType, | |
| NestedField, | |
| PrimitiveType, | |
| StringType, | |
| ) | |
| class ExampleVisitor(BooleanExpressionVisitor[list[str]]): | |
| """A test implementation of a BooleanExpressionVisitor. | |
| As this visitor visits each node, it appends an element to a `visit_history` list. | |
| This enables testing that a given expression is visited in an expected order by the `visit` method. | |
| """ | |
| def __init__(self) -> None: | |
| self.visit_history: list[str] = [] | |
| def visit_true(self) -> list[str]: | |
| self.visit_history.append("TRUE") | |
| return self.visit_history | |
| def visit_false(self) -> list[str]: | |
| self.visit_history.append("FALSE") | |
| return self.visit_history | |
| def visit_not(self, child_result: list[str]) -> list[str]: | |
| self.visit_history.append("NOT") | |
| return self.visit_history | |
| def visit_and(self, left_result: list[str], right_result: list[str]) -> list[str]: | |
| self.visit_history.append("AND") | |
| return self.visit_history | |
| def visit_or(self, left_result: list[str], right_result: list[str]) -> list[str]: | |
| self.visit_history.append("OR") | |
| return self.visit_history | |
| def visit_unbound_predicate(self, predicate: UnboundPredicate) -> list[str]: | |
| self.visit_history.append(str(predicate.__class__.__name__).upper()) | |
| return self.visit_history | |
| def visit_bound_predicate(self, predicate: BoundPredicate) -> list[str]: | |
| self.visit_history.append(str(predicate.__class__.__name__).upper()) | |
| return self.visit_history | |
| class FooBoundBooleanExpressionVisitor(BoundBooleanExpressionVisitor[list[str]]): | |
| """A test implementation of a BoundBooleanExpressionVisitor. | |
| As this visitor visits each node, it appends an element to a `visit_history` list. | |
| This enables testing that a given bound expression is visited in an expected order by the `visit` method. | |
| """ | |
| def __init__(self) -> None: | |
| self.visit_history: list[str] = [] | |
| def visit_in(self, term: BoundTerm, literals: set[Any]) -> list[str]: | |
| self.visit_history.append("IN") | |
| return self.visit_history | |
| def visit_not_in(self, term: BoundTerm, literals: set[Any]) -> list[str]: | |
| self.visit_history.append("NOT_IN") | |
| return self.visit_history | |
| def visit_is_nan(self, term: BoundTerm) -> list[str]: | |
| self.visit_history.append("IS_NAN") | |
| return self.visit_history | |
| def visit_not_nan(self, term: BoundTerm) -> list[str]: | |
| self.visit_history.append("NOT_NAN") | |
| return self.visit_history | |
| def visit_is_null(self, term: BoundTerm) -> list[str]: | |
| self.visit_history.append("IS_NULL") | |
| return self.visit_history | |
| def visit_not_null(self, term: BoundTerm) -> list[str]: | |
| self.visit_history.append("NOT_NULL") | |
| return self.visit_history | |
| def visit_equal(self, term: BoundTerm, literal: Literal[Any]) -> list[str]: # pylint: disable=redefined-outer-name | |
| self.visit_history.append("EQUAL") | |
| return self.visit_history | |
| def visit_not_equal(self, term: BoundTerm, literal: Literal[Any]) -> list[str]: # pylint: disable=redefined-outer-name | |
| self.visit_history.append("NOT_EQUAL") | |
| return self.visit_history | |
| def visit_greater_than_or_equal(self, term: BoundTerm, literal: Literal[Any]) -> list[str]: # pylint: disable=redefined-outer-name | |
| self.visit_history.append("GREATER_THAN_OR_EQUAL") | |
| return self.visit_history | |
| def visit_greater_than(self, term: BoundTerm, literal: Literal[Any]) -> list[str]: # pylint: disable=redefined-outer-name | |
| self.visit_history.append("GREATER_THAN") | |
| return self.visit_history | |
| def visit_less_than(self, term: BoundTerm, literal: Literal[Any]) -> list[str]: # pylint: disable=redefined-outer-name | |
| self.visit_history.append("LESS_THAN") | |
| return self.visit_history | |
| def visit_less_than_or_equal(self, term: BoundTerm, literal: Literal[Any]) -> list[str]: # pylint: disable=redefined-outer-name | |
| self.visit_history.append("LESS_THAN_OR_EQUAL") | |
| return self.visit_history | |
| def visit_true(self) -> list[str]: | |
| self.visit_history.append("TRUE") | |
| return self.visit_history | |
| def visit_false(self) -> list[str]: | |
| self.visit_history.append("FALSE") | |
| return self.visit_history | |
| def visit_not(self, child_result: list[str]) -> list[str]: | |
| self.visit_history.append("NOT") | |
| return self.visit_history | |
| def visit_and(self, left_result: list[str], right_result: list[str]) -> list[str]: | |
| self.visit_history.append("AND") | |
| return self.visit_history | |
| def visit_or(self, left_result: list[str], right_result: list[str]) -> list[str]: | |
| self.visit_history.append("OR") | |
| return self.visit_history | |
| def visit_starts_with(self, term: BoundTerm, literal: Literal[Any]) -> list[str]: | |
| self.visit_history.append("STARTS_WITH") | |
| return self.visit_history | |
| def visit_not_starts_with(self, term: BoundTerm, literal: Literal[Any]) -> list[str]: | |
| self.visit_history.append("NOT_STARTS_WITH") | |
| return self.visit_history | |
| def test_boolean_expression_visitor() -> None: | |
| """Test post-order traversal of boolean expression visit method""" | |
| expr = And( | |
| Or(Not(EqualTo("a", 1)), Not(NotEqualTo("b", 0)), EqualTo("a", 1), NotEqualTo("b", 0)), | |
| Not(EqualTo("a", 1)), | |
| NotEqualTo("b", 0), | |
| ) | |
| visitor = ExampleVisitor() | |
| result = visit(expr, visitor=visitor) | |
| assert result == [ | |
| "EQUALTO", | |
| "NOT", | |
| "NOTEQUALTO", | |
| "NOT", | |
| "OR", | |
| "EQUALTO", | |
| "NOTEQUALTO", | |
| "OR", | |
| "OR", | |
| "EQUALTO", | |
| "NOT", | |
| "NOTEQUALTO", | |
| "AND", | |
| "AND", | |
| ] | |
| def test_boolean_expression_visit_raise_not_implemented_error() -> None: | |
| """Test raise NotImplementedError when visiting an unsupported object type""" | |
| visitor = ExampleVisitor() | |
| with pytest.raises(NotImplementedError) as exc_info: | |
| visit("foo", visitor=visitor) # type: ignore | |
| assert str(exc_info.value) == "Cannot visit unsupported expression: foo" | |
| def test_bind_visitor_already_bound(table_schema_simple: Schema) -> None: | |
| bound = BoundEqualTo( | |
| term=BoundReference(table_schema_simple.find_field(1), table_schema_simple.accessor_for_field(1)), | |
| literal=literal("hello"), | |
| ) | |
| with pytest.raises(TypeError) as exc_info: | |
| visit(bound, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert ( | |
| "Found already bound predicate: BoundEqualTo(term=BoundReference(" | |
| "field=NestedField(field_id=1, name='foo', field_type=StringType(), required=False), " | |
| "accessor=Accessor(position=0,inner=None)), literal=literal('hello'))" | |
| ) == str(exc_info.value) | |
| def test_visit_bound_visitor_unknown_predicate() -> None: | |
| with pytest.raises(TypeError) as exc_info: | |
| visit_bound_predicate({"something"}, FooBoundBooleanExpressionVisitor()) # type: ignore | |
| assert "Unknown predicate: {'something'}" == str(exc_info.value) | |
| def test_always_true_expression_binding(table_schema_simple: Schema) -> None: | |
| """Test that visiting an always-true expression returns always-true""" | |
| unbound_expression = AlwaysTrue() | |
| bound_expression = visit(unbound_expression, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert bound_expression == AlwaysTrue() | |
| def test_always_false_expression_binding(table_schema_simple: Schema) -> None: | |
| """Test that visiting an always-false expression returns always-false""" | |
| unbound_expression = AlwaysFalse() | |
| bound_expression = visit(unbound_expression, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert bound_expression == AlwaysFalse() | |
| def test_always_false_and_always_true_expression_binding(table_schema_simple: Schema) -> None: | |
| """Test that visiting both an always-true AND always-false expression returns always-false""" | |
| unbound_expression = And(AlwaysTrue(), AlwaysFalse()) | |
| bound_expression = visit(unbound_expression, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert bound_expression == AlwaysFalse() | |
| def test_always_false_or_always_true_expression_binding(table_schema_simple: Schema) -> None: | |
| """Test that visiting always-true OR always-false expression returns always-true""" | |
| unbound_expression = Or(AlwaysTrue(), AlwaysFalse()) | |
| bound_expression = visit(unbound_expression, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert bound_expression == AlwaysTrue() | |
| @pytest.mark.parametrize( | |
| "unbound_and_expression,expected_bound_expression", | |
| [ | |
| ( | |
| And( | |
| In(Reference("foo"), {"foo", "bar"}), | |
| In(Reference("bar"), {1, 2, 3}), | |
| ), | |
| And( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("foo"), literal("bar")}, | |
| ), | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=2, name="bar", field_type=IntegerType(), required=True), | |
| accessor=Accessor(position=1, inner=None), | |
| ), | |
| {literal(1), literal(2), literal(3)}, | |
| ), | |
| ), | |
| ), | |
| ( | |
| And( | |
| In(Reference("foo"), ("bar", "baz")), | |
| In( | |
| Reference("bar"), | |
| (1,), | |
| ), | |
| In( | |
| Reference("foo"), | |
| ("baz",), | |
| ), | |
| ), | |
| And( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("bar"), literal("baz")}, | |
| ), | |
| And( | |
| BoundEqualTo( | |
| BoundReference( | |
| field=NestedField(field_id=2, name="bar", field_type=IntegerType(), required=True), | |
| accessor=Accessor(position=1, inner=None), | |
| ), | |
| literal(1), | |
| ), | |
| BoundEqualTo( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal("baz"), | |
| ), | |
| ), | |
| ), | |
| ), | |
| ], | |
| ) | |
| def test_and_expression_binding( | |
| unbound_and_expression: UnboundPredicate, expected_bound_expression: BoundPredicate, table_schema_simple: Schema | |
| ) -> None: | |
| """Test that visiting an unbound AND expression with a bind-visitor returns the expected bound expression""" | |
| bound_expression = visit(unbound_and_expression, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert bound_expression == expected_bound_expression | |
| @pytest.mark.parametrize( | |
| "unbound_or_expression,expected_bound_expression", | |
| [ | |
| ( | |
| Or( | |
| In(Reference("foo"), ("foo", "bar")), | |
| In(Reference("bar"), (1, 2, 3)), | |
| ), | |
| Or( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("foo"), literal("bar")}, | |
| ), | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=2, name="bar", field_type=IntegerType(), required=True), | |
| accessor=Accessor(position=1, inner=None), | |
| ), | |
| {literal(1), literal(2), literal(3)}, | |
| ), | |
| ), | |
| ), | |
| ( | |
| Or( | |
| In(Reference("foo"), ("bar", "baz")), | |
| In( | |
| Reference("foo"), | |
| ("bar",), | |
| ), | |
| In( | |
| Reference("foo"), | |
| ("baz",), | |
| ), | |
| ), | |
| Or( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("bar"), literal("baz")}, | |
| ), | |
| Or( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("bar")}, | |
| ), | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("baz")}, | |
| ), | |
| ), | |
| ), | |
| ), | |
| ( | |
| Or( | |
| AlwaysTrue(), | |
| AlwaysFalse(), | |
| ), | |
| AlwaysTrue(), | |
| ), | |
| ( | |
| Or( | |
| AlwaysTrue(), | |
| AlwaysTrue(), | |
| ), | |
| AlwaysTrue(), | |
| ), | |
| ( | |
| Or( | |
| AlwaysFalse(), | |
| AlwaysFalse(), | |
| ), | |
| AlwaysFalse(), | |
| ), | |
| ], | |
| ) | |
| def test_or_expression_binding( | |
| unbound_or_expression: UnboundPredicate, expected_bound_expression: BoundPredicate, table_schema_simple: Schema | |
| ) -> None: | |
| """Test that visiting an unbound OR expression with a bind-visitor returns the expected bound expression""" | |
| bound_expression = visit(unbound_or_expression, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert bound_expression == expected_bound_expression | |
| @pytest.mark.parametrize( | |
| "unbound_in_expression,expected_bound_expression", | |
| [ | |
| ( | |
| In(Reference("foo"), ("foo", "bar")), | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("foo"), literal("bar")}, | |
| ), | |
| ), | |
| ( | |
| In(Reference("foo"), ("bar", "baz")), | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("bar"), literal("baz")}, | |
| ), | |
| ), | |
| ( | |
| In( | |
| Reference("foo"), | |
| ("bar",), | |
| ), | |
| BoundEqualTo( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal("bar"), | |
| ), | |
| ), | |
| ], | |
| ) | |
| def test_in_expression_binding( | |
| unbound_in_expression: UnboundPredicate, expected_bound_expression: BoundPredicate, table_schema_simple: Schema | |
| ) -> None: | |
| """Test that visiting an unbound IN expression with a bind-visitor returns the expected bound expression""" | |
| bound_expression = visit(unbound_in_expression, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert bound_expression == expected_bound_expression | |
| @pytest.mark.parametrize( | |
| "unbound_not_expression,expected_bound_expression", | |
| [ | |
| ( | |
| Not(In(Reference("foo"), ("foo", "bar"))), | |
| Not( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("foo"), literal("bar")}, | |
| ) | |
| ), | |
| ), | |
| ( | |
| Not( | |
| Or( | |
| In(Reference("foo"), ("foo", "bar")), | |
| In(Reference("foo"), ("foo", "bar", "baz")), | |
| ) | |
| ), | |
| Not( | |
| Or( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("foo"), literal("bar")}, | |
| ), | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("foo"), literal("bar"), literal("baz")}, | |
| ), | |
| ), | |
| ), | |
| ), | |
| ], | |
| ) | |
| def test_not_expression_binding( | |
| unbound_not_expression: UnboundPredicate, expected_bound_expression: BoundPredicate, table_schema_simple: Schema | |
| ) -> None: | |
| """Test that visiting an unbound NOT expression with a bind-visitor returns the expected bound expression""" | |
| bound_expression = visit(unbound_not_expression, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| assert bound_expression == expected_bound_expression | |
| def test_bound_boolean_expression_visitor_and_in() -> None: | |
| """Test visiting an And and In expression with a bound boolean expression visitor""" | |
| bound_expression = And( | |
| BoundIn( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literals={literal("foo"), literal("bar")}, | |
| ), | |
| BoundIn( | |
| term=BoundReference( | |
| field=NestedField(field_id=2, name="bar", field_type=StringType(), required=False), | |
| accessor=Accessor(position=1, inner=None), | |
| ), | |
| literals={literal("baz"), literal("qux")}, | |
| ), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["IN", "IN", "AND"] | |
| def test_bound_boolean_expression_visitor_or() -> None: | |
| """Test visiting an Or expression with a bound boolean expression visitor""" | |
| bound_expression = Or( | |
| Not( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| {literal("foo"), literal("bar")}, | |
| ) | |
| ), | |
| Not( | |
| BoundIn( | |
| BoundReference( | |
| field=NestedField(field_id=2, name="bar", field_type=StringType(), required=False), | |
| accessor=Accessor(position=1, inner=None), | |
| ), | |
| {literal("baz"), literal("qux")}, | |
| ) | |
| ), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["IN", "NOT", "IN", "NOT", "OR"] | |
| def test_bound_boolean_expression_visitor_equal() -> None: | |
| bound_expression = BoundEqualTo( | |
| term=BoundReference( | |
| field=NestedField(field_id=2, name="bar", field_type=StringType(), required=False), | |
| accessor=Accessor(position=1, inner=None), | |
| ), | |
| literal=literal("foo"), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["EQUAL"] | |
| def test_bound_boolean_expression_visitor_not_equal() -> None: | |
| bound_expression = BoundNotEqualTo( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal=literal("foo"), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["NOT_EQUAL"] | |
| def test_bound_boolean_expression_visitor_always_true() -> None: | |
| bound_expression = AlwaysTrue() | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["TRUE"] | |
| def test_bound_boolean_expression_visitor_always_false() -> None: | |
| bound_expression = AlwaysFalse() | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["FALSE"] | |
| def test_bound_boolean_expression_visitor_in() -> None: | |
| bound_expression = BoundIn( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literals={literal("foo"), literal("bar")}, | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["IN"] | |
| def test_bound_boolean_expression_visitor_not_in() -> None: | |
| bound_expression = BoundNotIn( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literals={literal("foo"), literal("bar")}, | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["NOT_IN"] | |
| def test_bound_boolean_expression_visitor_is_nan() -> None: | |
| bound_expression = BoundIsNaN( | |
| term=BoundReference( | |
| field=NestedField(field_id=3, name="baz", field_type=FloatType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["IS_NAN"] | |
| def test_bound_boolean_expression_visitor_not_nan() -> None: | |
| bound_expression = BoundNotNaN( | |
| term=BoundReference( | |
| field=NestedField(field_id=3, name="baz", field_type=FloatType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["NOT_NAN"] | |
| def test_bound_boolean_expression_visitor_is_null() -> None: | |
| bound_expression = BoundIsNull( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["IS_NULL"] | |
| def test_bound_boolean_expression_visitor_not_null() -> None: | |
| bound_expression = BoundNotNull( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["NOT_NULL"] | |
| def test_bound_boolean_expression_visitor_greater_than() -> None: | |
| bound_expression = BoundGreaterThan( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal=literal("foo"), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["GREATER_THAN"] | |
| def test_bound_boolean_expression_visitor_greater_than_or_equal() -> None: | |
| bound_expression = BoundGreaterThanOrEqual( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal=literal("foo"), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["GREATER_THAN_OR_EQUAL"] | |
| def test_bound_boolean_expression_visitor_less_than() -> None: | |
| bound_expression = BoundLessThan( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal=literal("foo"), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["LESS_THAN"] | |
| def test_bound_boolean_expression_visitor_less_than_or_equal() -> None: | |
| bound_expression = BoundLessThanOrEqual( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal=literal("foo"), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["LESS_THAN_OR_EQUAL"] | |
| def test_bound_boolean_expression_visitor_raise_on_unbound_predicate() -> None: | |
| bound_expression = LessThanOrEqual( | |
| term=Reference("foo"), | |
| literal="foo", | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| with pytest.raises(TypeError) as exc_info: | |
| visit(bound_expression, visitor=visitor) | |
| assert "Not a bound predicate" in str(exc_info.value) | |
| def test_bound_boolean_expression_visitor_starts_with() -> None: | |
| bound_expression = BoundStartsWith( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal=literal("foo"), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["STARTS_WITH"] | |
| def test_bound_boolean_expression_visitor_not_starts_with() -> None: | |
| bound_expression = BoundNotStartsWith( | |
| term=BoundReference( | |
| field=NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ), | |
| literal=literal("foo"), | |
| ) | |
| visitor = FooBoundBooleanExpressionVisitor() | |
| result = visit(bound_expression, visitor=visitor) | |
| assert result == ["NOT_STARTS_WITH"] | |
| def _to_byte_buffer(field_type: IcebergType, val: Any) -> bytes: | |
| if not isinstance(field_type, PrimitiveType): | |
| raise ValueError(f"Expected a PrimitiveType, got: {type(field_type)}") | |
| return to_bytes(field_type, val) | |
| def _to_manifest_file(*partitions: PartitionFieldSummary) -> ManifestFile: | |
| """Helper to create a ManifestFile""" | |
| return ManifestFile.from_args(manifest_path="", manifest_length=0, partition_spec_id=0, partitions=partitions) | |
| INT_MIN_VALUE = 30 | |
| INT_MAX_VALUE = 79 | |
| INT_MIN = _to_byte_buffer(IntegerType(), INT_MIN_VALUE) | |
| INT_MAX = _to_byte_buffer(IntegerType(), INT_MAX_VALUE) | |
| STRING_MIN = _to_byte_buffer(StringType(), "a") | |
| STRING_MAX = _to_byte_buffer(StringType(), "z") | |
| @pytest.fixture | |
| def schema() -> Schema: | |
| return Schema( | |
| NestedField(1, "id", IntegerType(), required=True), | |
| NestedField(2, "all_nulls_missing_nan", StringType(), required=False), | |
| NestedField(3, "some_nulls", StringType(), required=False), | |
| NestedField(4, "no_nulls", StringType(), required=False), | |
| NestedField(5, "float", FloatType(), required=False), | |
| NestedField(6, "all_nulls_double", DoubleType(), required=False), | |
| NestedField(7, "all_nulls_no_nans", FloatType(), required=False), | |
| NestedField(8, "all_nans", DoubleType(), required=False), | |
| NestedField(9, "both_nan_and_null", FloatType(), required=False), | |
| NestedField(10, "no_nan_or_null", DoubleType(), required=False), | |
| NestedField(11, "all_nulls_missing_nan_float", FloatType(), required=False), | |
| NestedField(12, "all_same_value_or_null", StringType(), required=False), | |
| NestedField(13, "no_nulls_same_value_a", StringType(), required=False), | |
| ) | |
| @pytest.fixture | |
| def manifest_no_stats() -> ManifestFile: | |
| return _to_manifest_file() | |
| @pytest.fixture | |
| def manifest() -> ManifestFile: | |
| return _to_manifest_file( | |
| # id | |
| PartitionFieldSummary.from_args( | |
| contains_null=False, | |
| contains_nan=None, | |
| lower_bound=INT_MIN, | |
| upper_bound=INT_MAX, | |
| ), | |
| # all_nulls_missing_nan | |
| PartitionFieldSummary.from_args( | |
| contains_null=True, | |
| contains_nan=None, | |
| lower_bound=None, | |
| upper_bound=None, | |
| ), | |
| # some_nulls | |
| PartitionFieldSummary.from_args( | |
| contains_null=True, | |
| contains_nan=None, | |
| lower_bound=STRING_MIN, | |
| upper_bound=STRING_MAX, | |
| ), | |
| # no_nulls | |
| PartitionFieldSummary.from_args( | |
| contains_null=False, | |
| contains_nan=None, | |
| lower_bound=STRING_MIN, | |
| upper_bound=STRING_MAX, | |
| ), | |
| # float | |
| PartitionFieldSummary.from_args( | |
| contains_null=True, | |
| contains_nan=None, | |
| lower_bound=_to_byte_buffer(FloatType(), 0.0), | |
| upper_bound=_to_byte_buffer(FloatType(), 20.0), | |
| ), | |
| # all_nulls_double | |
| PartitionFieldSummary.from_args(contains_null=True, contains_nan=None, lower_bound=None, upper_bound=None), | |
| # all_nulls_no_nans | |
| PartitionFieldSummary.from_args( | |
| contains_null=True, | |
| contains_nan=False, | |
| lower_bound=None, | |
| upper_bound=None, | |
| ), | |
| # all_nans | |
| PartitionFieldSummary.from_args( | |
| contains_null=False, | |
| contains_nan=True, | |
| lower_bound=None, | |
| upper_bound=None, | |
| ), | |
| # both_nan_and_null | |
| PartitionFieldSummary.from_args( | |
| contains_null=True, | |
| contains_nan=True, | |
| lower_bound=None, | |
| upper_bound=None, | |
| ), | |
| # no_nan_or_null | |
| PartitionFieldSummary.from_args( | |
| contains_null=False, | |
| contains_nan=False, | |
| lower_bound=_to_byte_buffer(FloatType(), 0.0), | |
| upper_bound=_to_byte_buffer(FloatType(), 20.0), | |
| ), | |
| # all_nulls_missing_nan_float | |
| PartitionFieldSummary.from_args(contains_null=True, contains_nan=None, lower_bound=None, upper_bound=None), | |
| # all_same_value_or_null | |
| PartitionFieldSummary.from_args( | |
| contains_null=True, | |
| contains_nan=None, | |
| lower_bound=STRING_MIN, | |
| upper_bound=STRING_MIN, | |
| ), | |
| # no_nulls_same_value_a | |
| PartitionFieldSummary.from_args( | |
| contains_null=False, | |
| contains_nan=None, | |
| lower_bound=STRING_MIN, | |
| upper_bound=STRING_MIN, | |
| ), | |
| ) | |
| def test_all_nulls(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor(schema, NotNull(Reference("all_nulls_missing_nan")), case_sensitive=True).eval(manifest), ( | |
| "Should skip: all nulls column with non-floating type contains all null" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotNull(Reference("all_nulls_missing_nan_float")), case_sensitive=True).eval(manifest), ( | |
| "Should read: no NaN information may indicate presence of NaN value" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotNull(Reference("some_nulls")), case_sensitive=True).eval(manifest), ( | |
| "Should read: column with some nulls contains a non-null value" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotNull(Reference("no_nulls")), case_sensitive=True).eval(manifest), ( | |
| "Should read: non-null column contains a non-null value" | |
| ) | |
| def test_no_nulls(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, IsNull(Reference("all_nulls_missing_nan")), case_sensitive=True).eval(manifest), ( | |
| "Should read: at least one null value in all null column" | |
| ) | |
| assert _ManifestEvalVisitor(schema, IsNull(Reference("some_nulls")), case_sensitive=True).eval(manifest), ( | |
| "Should read: column with some nulls contains a null value" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, IsNull(Reference("no_nulls")), case_sensitive=True).eval(manifest), ( | |
| "Should skip: non-null column contains no null values" | |
| ) | |
| assert _ManifestEvalVisitor(schema, IsNull(Reference("both_nan_and_null")), case_sensitive=True).eval(manifest), ( | |
| "Should read: both_nan_and_null column contains no null values" | |
| ) | |
| def test_is_nan(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, IsNaN(Reference("float")), case_sensitive=True).eval(manifest), ( | |
| "Should read: no information on if there are nan value in float column" | |
| ) | |
| assert _ManifestEvalVisitor(schema, IsNaN(Reference("all_nulls_double")), case_sensitive=True).eval(manifest), ( | |
| "Should read: no NaN information may indicate presence of NaN value" | |
| ) | |
| assert _ManifestEvalVisitor(schema, IsNaN(Reference("all_nulls_missing_nan_float")), case_sensitive=True).eval(manifest), ( | |
| "Should read: no NaN information may indicate presence of NaN value" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, IsNaN(Reference("all_nulls_no_nans")), case_sensitive=True).eval(manifest), ( | |
| "Should skip: no nan column doesn't contain nan value" | |
| ) | |
| assert _ManifestEvalVisitor(schema, IsNaN(Reference("all_nans")), case_sensitive=True).eval(manifest), ( | |
| "Should read: all_nans column contains nan value" | |
| ) | |
| assert _ManifestEvalVisitor(schema, IsNaN(Reference("both_nan_and_null")), case_sensitive=True).eval(manifest), ( | |
| "Should read: both_nan_and_null column contains nan value" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, IsNaN(Reference("no_nan_or_null")), case_sensitive=True).eval(manifest), ( | |
| "Should skip: no_nan_or_null column doesn't contain nan value" | |
| ) | |
| def test_not_nan(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, NotNaN(Reference("float")), case_sensitive=True).eval(manifest), ( | |
| "Should read: no information on if there are nan value in float column" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotNaN(Reference("all_nulls_double")), case_sensitive=True).eval(manifest), ( | |
| "Should read: all null column contains non nan value" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotNaN(Reference("all_nulls_no_nans")), case_sensitive=True).eval(manifest), ( | |
| "Should read: no_nans column contains non nan value" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, NotNaN(Reference("all_nans")), case_sensitive=True).eval(manifest), ( | |
| "Should skip: all nans column doesn't contain non nan value" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotNaN(Reference("both_nan_and_null")), case_sensitive=True).eval(manifest), ( | |
| "Should read: both_nan_and_null nans column contains non nan value" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotNaN(Reference("no_nan_or_null")), case_sensitive=True).eval(manifest), ( | |
| "Should read: no_nan_or_null column contains non nan value" | |
| ) | |
| def test_missing_stats(schema: Schema, manifest_no_stats: ManifestFile) -> None: | |
| expressions: list[BooleanExpression] = [ | |
| LessThan(Reference("id"), 5), | |
| LessThanOrEqual(Reference("id"), 30), | |
| EqualTo(Reference("id"), 70), | |
| GreaterThan(Reference("id"), 78), | |
| GreaterThanOrEqual(Reference("id"), 90), | |
| NotEqualTo(Reference("id"), 101), | |
| IsNull(Reference("id")), | |
| NotNull(Reference("id")), | |
| IsNaN(Reference("float")), | |
| NotNaN(Reference("float")), | |
| ] | |
| for expr in expressions: | |
| assert _ManifestEvalVisitor(schema, expr, case_sensitive=True).eval(manifest_no_stats), ( | |
| f"Should read when missing stats for expr: {expr}" | |
| ) | |
| def test_not(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, Not(LessThan(Reference("id"), INT_MIN_VALUE - 25)), case_sensitive=True).eval(manifest), ( | |
| "Should read: not(false)" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, Not(GreaterThan(Reference("id"), INT_MIN_VALUE - 25)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should skip: not(true)" | |
| def test_and(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor( | |
| schema, | |
| And( | |
| LessThan(Reference("id"), INT_MIN_VALUE - 25), | |
| GreaterThanOrEqual(Reference("id"), INT_MIN_VALUE - 30), | |
| ), | |
| case_sensitive=True, | |
| ).eval(manifest), "Should skip: and(false, true)" | |
| assert not _ManifestEvalVisitor( | |
| schema, | |
| And( | |
| LessThan(Reference("id"), INT_MIN_VALUE - 25), | |
| GreaterThanOrEqual(Reference("id"), INT_MAX_VALUE + 1), | |
| ), | |
| case_sensitive=True, | |
| ).eval(manifest), "Should skip: and(false, false)" | |
| assert _ManifestEvalVisitor( | |
| schema, | |
| And( | |
| GreaterThan(Reference("id"), INT_MIN_VALUE - 25), | |
| LessThanOrEqual(Reference("id"), INT_MIN_VALUE), | |
| ), | |
| case_sensitive=True, | |
| ).eval(manifest), "Should read: and(true, true)" | |
| def test_or(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor( | |
| schema, | |
| Or( | |
| LessThan(Reference("id"), INT_MIN_VALUE - 25), | |
| GreaterThanOrEqual(Reference("id"), INT_MAX_VALUE + 1), | |
| ), | |
| case_sensitive=True, | |
| ).eval(manifest), "Should skip: or(false, false)" | |
| assert _ManifestEvalVisitor( | |
| schema, | |
| Or( | |
| LessThan(Reference("id"), INT_MIN_VALUE - 25), | |
| GreaterThanOrEqual(Reference("id"), INT_MAX_VALUE - 19), | |
| ), | |
| case_sensitive=True, | |
| ).eval(manifest), "Should read: or(false, true)" | |
| def test_integer_lt(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor(schema, LessThan(Reference("id"), INT_MIN_VALUE - 25), case_sensitive=True).eval(manifest), ( | |
| "Should not read: id range below lower bound (5 < 30)" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, LessThan(Reference("id"), INT_MIN_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should not read: id range below lower bound (30 is not < 30)" | |
| ) | |
| assert _ManifestEvalVisitor(schema, LessThan(Reference("id"), INT_MIN_VALUE + 1), case_sensitive=True).eval(manifest), ( | |
| "Should read: one possible id" | |
| ) | |
| assert _ManifestEvalVisitor(schema, LessThan(Reference("id"), INT_MAX_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: may possible ids" | |
| ) | |
| def test_integer_lt_eq(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor(schema, LessThanOrEqual(Reference("id"), INT_MIN_VALUE - 25), case_sensitive=True).eval( | |
| manifest | |
| ), "Should not read: id range below lower bound (5 < 30)" | |
| assert not _ManifestEvalVisitor(schema, LessThanOrEqual(Reference("id"), INT_MIN_VALUE - 1), case_sensitive=True).eval( | |
| manifest | |
| ), "Should not read: id range below lower bound (29 < 30)" | |
| assert _ManifestEvalVisitor(schema, LessThanOrEqual(Reference("id"), INT_MIN_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: one possible id" | |
| ) | |
| assert _ManifestEvalVisitor(schema, LessThanOrEqual(Reference("id"), INT_MAX_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: many possible ids" | |
| ) | |
| def test_integer_gt(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor(schema, GreaterThan(Reference("id"), INT_MAX_VALUE + 6), case_sensitive=True).eval( | |
| manifest | |
| ), "Should not read: id range above upper bound (85 < 79)" | |
| assert not _ManifestEvalVisitor(schema, GreaterThan(Reference("id"), INT_MAX_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should not read: id range above upper bound (79 is not > 79)" | |
| ) | |
| assert _ManifestEvalVisitor(schema, GreaterThan(Reference("id"), INT_MAX_VALUE - 1), case_sensitive=True).eval(manifest), ( | |
| "Should read: one possible id" | |
| ) | |
| assert _ManifestEvalVisitor(schema, GreaterThan(Reference("id"), INT_MAX_VALUE - 4), case_sensitive=True).eval(manifest), ( | |
| "Should read: may possible ids" | |
| ) | |
| def test_integer_gt_eq(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor(schema, GreaterThanOrEqual(Reference("id"), INT_MAX_VALUE + 6), case_sensitive=True).eval( | |
| manifest | |
| ), "Should not read: id range above upper bound (85 < 79)" | |
| assert not _ManifestEvalVisitor(schema, GreaterThanOrEqual(Reference("id"), INT_MAX_VALUE + 1), case_sensitive=True).eval( | |
| manifest | |
| ), "Should not read: id range above upper bound (80 > 79)" | |
| assert _ManifestEvalVisitor(schema, GreaterThanOrEqual(Reference("id"), INT_MAX_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: one possible id" | |
| ) | |
| assert _ManifestEvalVisitor(schema, GreaterThanOrEqual(Reference("id"), INT_MAX_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: may possible ids" | |
| ) | |
| def test_integer_eq(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor(schema, EqualTo(Reference("id"), INT_MIN_VALUE - 25), case_sensitive=True).eval(manifest), ( | |
| "Should not read: id below lower bound" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, EqualTo(Reference("id"), INT_MIN_VALUE - 1), case_sensitive=True).eval(manifest), ( | |
| "Should not read: id below lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, EqualTo(Reference("id"), INT_MIN_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: id equal to lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, EqualTo(Reference("id"), INT_MAX_VALUE - 4), case_sensitive=True).eval(manifest), ( | |
| "Should read: id between lower and upper bounds" | |
| ) | |
| assert _ManifestEvalVisitor(schema, EqualTo(Reference("id"), INT_MAX_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: id equal to upper bound" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, EqualTo(Reference("id"), INT_MAX_VALUE + 1), case_sensitive=True).eval(manifest), ( | |
| "Should not read: id above upper bound" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, EqualTo(Reference("id"), INT_MAX_VALUE + 6), case_sensitive=True).eval(manifest), ( | |
| "Should not read: id above upper bound" | |
| ) | |
| def test_integer_not_eq(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, NotEqualTo(Reference("id"), INT_MIN_VALUE - 25), case_sensitive=True).eval(manifest), ( | |
| "Should read: id below lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotEqualTo(Reference("id"), INT_MIN_VALUE - 1), case_sensitive=True).eval(manifest), ( | |
| "Should read: id below lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotEqualTo(Reference("id"), INT_MIN_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: id equal to lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotEqualTo(Reference("id"), INT_MAX_VALUE - 4), case_sensitive=True).eval(manifest), ( | |
| "Should read: id between lower and upper bounds" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotEqualTo(Reference("id"), INT_MAX_VALUE), case_sensitive=True).eval(manifest), ( | |
| "Should read: id equal to upper bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotEqualTo(Reference("id"), INT_MAX_VALUE + 1), case_sensitive=True).eval(manifest), ( | |
| "Should read: id above upper bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotEqualTo(Reference("id"), INT_MAX_VALUE + 6), case_sensitive=True).eval(manifest), ( | |
| "Should read: id above upper bound" | |
| ) | |
| def test_integer_not_eq_rewritten(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("id"), INT_MIN_VALUE - 25)), case_sensitive=True).eval(manifest), ( | |
| "Should read: id below lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("id"), INT_MIN_VALUE - 1)), case_sensitive=True).eval(manifest), ( | |
| "Should read: id below lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("id"), INT_MIN_VALUE)), case_sensitive=True).eval(manifest), ( | |
| "Should read: id equal to lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("id"), INT_MAX_VALUE - 4)), case_sensitive=True).eval(manifest), ( | |
| "Should read: id between lower and upper bounds" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("id"), INT_MAX_VALUE)), case_sensitive=True).eval(manifest), ( | |
| "Should read: id equal to upper bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("id"), INT_MAX_VALUE + 1)), case_sensitive=True).eval(manifest), ( | |
| "Should read: id above upper bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("id"), INT_MAX_VALUE + 6)), case_sensitive=True).eval(manifest), ( | |
| "Should read: id above upper bound" | |
| ) | |
| def test_integer_not_eq_rewritten_case_insensitive(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("ID"), INT_MIN_VALUE - 25)), case_sensitive=False).eval(manifest), ( | |
| "Should read: id below lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("ID"), INT_MIN_VALUE - 1)), case_sensitive=False).eval(manifest), ( | |
| "Should read: id below lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("ID"), INT_MIN_VALUE)), case_sensitive=False).eval(manifest), ( | |
| "Should read: id equal to lower bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("ID"), INT_MAX_VALUE - 4)), case_sensitive=False).eval(manifest), ( | |
| "Should read: id between lower and upper bounds" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("ID"), INT_MAX_VALUE)), case_sensitive=False).eval(manifest), ( | |
| "Should read: id equal to upper bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("ID"), INT_MAX_VALUE + 1)), case_sensitive=False).eval(manifest), ( | |
| "Should read: id above upper bound" | |
| ) | |
| assert _ManifestEvalVisitor(schema, Not(EqualTo(Reference("ID"), INT_MAX_VALUE + 6)), case_sensitive=False).eval(manifest), ( | |
| "Should read: id above upper bound" | |
| ) | |
| def test_integer_in(schema: Schema, manifest: ManifestFile) -> None: | |
| assert not _ManifestEvalVisitor( | |
| schema, In(Reference("id"), (INT_MIN_VALUE - 25, INT_MIN_VALUE - 24)), case_sensitive=True | |
| ).eval(manifest), "Should not read: id below lower bound (5 < 30, 6 < 30)" | |
| assert not _ManifestEvalVisitor( | |
| schema, In(Reference("id"), (INT_MIN_VALUE - 2, INT_MIN_VALUE - 1)), case_sensitive=True | |
| ).eval(manifest), "Should not read: id below lower bound (28 < 30, 29 < 30)" | |
| assert _ManifestEvalVisitor(schema, In(Reference("id"), (INT_MIN_VALUE - 1, INT_MIN_VALUE)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id equal to lower bound (30 == 30)" | |
| assert _ManifestEvalVisitor(schema, In(Reference("id"), (INT_MAX_VALUE - 4, INT_MAX_VALUE - 3)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id between lower and upper bounds (30 < 75 < 79, 30 < 76 < 79)" | |
| assert _ManifestEvalVisitor(schema, In(Reference("id"), (INT_MAX_VALUE, INT_MAX_VALUE + 1)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id equal to upper bound (79 == 79)" | |
| assert not _ManifestEvalVisitor( | |
| schema, In(Reference("id"), (INT_MAX_VALUE + 1, INT_MAX_VALUE + 2)), case_sensitive=True | |
| ).eval(manifest), "Should not read: id above upper bound (80 > 79, 81 > 79)" | |
| assert not _ManifestEvalVisitor( | |
| schema, In(Reference("id"), (INT_MAX_VALUE + 6, INT_MAX_VALUE + 7)), case_sensitive=True | |
| ).eval(manifest), "Should not read: id above upper bound (85 > 79, 86 > 79)" | |
| assert not _ManifestEvalVisitor(schema, In(Reference("all_nulls_missing_nan"), ("abc", "def")), case_sensitive=True).eval( | |
| manifest | |
| ), "Should skip: in on all nulls column" | |
| assert _ManifestEvalVisitor(schema, In(Reference("some_nulls"), ("abc", "def")), case_sensitive=True).eval(manifest), ( | |
| "Should read: in on some nulls column" | |
| ) | |
| assert _ManifestEvalVisitor(schema, In(Reference("no_nulls"), ("abc", "def")), case_sensitive=True).eval(manifest), ( | |
| "Should read: in on no nulls column" | |
| ) | |
| def test_integer_not_in(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor( | |
| schema, NotIn(Reference("id"), (INT_MIN_VALUE - 25, INT_MIN_VALUE - 24)), case_sensitive=True | |
| ).eval(manifest), "Should read: id below lower bound (5 < 30, 6 < 30)" | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("id"), (INT_MIN_VALUE - 2, INT_MIN_VALUE - 1)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id below lower bound (28 < 30, 29 < 30)" | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("id"), (INT_MIN_VALUE - 1, INT_MIN_VALUE)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id equal to lower bound (30 == 30)" | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("id"), (INT_MAX_VALUE - 4, INT_MAX_VALUE - 3)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id between lower and upper bounds (30 < 75 < 79, 30 < 76 < 79)" | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("id"), (INT_MAX_VALUE, INT_MAX_VALUE + 1)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id equal to upper bound (79 == 79)" | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("id"), (INT_MAX_VALUE + 1, INT_MAX_VALUE + 2)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id above upper bound (80 > 79, 81 > 79)" | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("id"), (INT_MAX_VALUE + 6, INT_MAX_VALUE + 7)), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: id above upper bound (85 > 79, 86 > 79)" | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("all_nulls_missing_nan"), ("abc", "def")), case_sensitive=True).eval( | |
| manifest | |
| ), "Should read: notIn on no nulls column" | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("some_nulls"), ("abc", "def")), case_sensitive=True).eval(manifest), ( | |
| "Should read: in on some nulls column" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotIn(Reference("no_nulls"), ("abc", "def")), case_sensitive=True).eval(manifest), ( | |
| "Should read: in on no nulls column" | |
| ) | |
| def test_string_starts_with(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, StartsWith(Reference("some_nulls"), "a"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, StartsWith(Reference("some_nulls"), "aa"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, StartsWith(Reference("some_nulls"), "dddd"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, StartsWith(Reference("some_nulls"), "z"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, StartsWith(Reference("no_nulls"), "a"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, StartsWith(Reference("some_nulls"), "zzzz"), case_sensitive=False).eval(manifest), ( | |
| "Should skip: range doesn't match" | |
| ) | |
| assert not _ManifestEvalVisitor(schema, StartsWith(Reference("some_nulls"), "1"), case_sensitive=False).eval(manifest), ( | |
| "Should skip: range doesn't match" | |
| ) | |
| def test_string_not_starts_with(schema: Schema, manifest: ManifestFile) -> None: | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("some_nulls"), "a"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("some_nulls"), "aa"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("some_nulls"), "dddd"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("some_nulls"), "z"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("no_nulls"), "a"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("some_nulls"), "zzzz"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("some_nulls"), "1"), case_sensitive=False).eval(manifest), ( | |
| "Should read: range matches" | |
| ) | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("all_same_value_or_null"), "a"), case_sensitive=False).eval( | |
| manifest | |
| ), "Should read: range matches" | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("all_same_value_or_null"), "aa"), case_sensitive=False).eval( | |
| manifest | |
| ), "Should read: range matches" | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("all_same_value_or_null"), "A"), case_sensitive=False).eval( | |
| manifest | |
| ), "Should read: range matches" | |
| # Iceberg does not implement SQL's 3-way boolean logic, so the choice of an all null column | |
| # matching is | |
| # by definition in order to surface more values to the query engine to allow it to make its own | |
| # decision. | |
| assert _ManifestEvalVisitor(schema, NotStartsWith(Reference("all_nulls_missing_nan"), "A"), case_sensitive=False).eval( | |
| manifest | |
| ), "Should read: range matches" | |
| assert not _ManifestEvalVisitor(schema, NotStartsWith(Reference("no_nulls_same_value_a"), "a"), case_sensitive=False).eval( | |
| manifest | |
| ), "Should not read: all values start with the prefix" | |
| def test_rewrite_not_equal_to() -> None: | |
| assert rewrite_not(Not(EqualTo(Reference("x"), 34.56))) == NotEqualTo(Reference("x"), 34.56) | |
| def test_rewrite_not_not_equal_to() -> None: | |
| assert rewrite_not(Not(NotEqualTo(Reference("x"), 34.56))) == EqualTo(Reference("x"), 34.56) | |
| def test_rewrite_not_in() -> None: | |
| assert rewrite_not(Not(In(Reference("x"), (34.56,)))) == NotIn(Reference("x"), (34.56,)) | |
| def test_rewrite_and() -> None: | |
| assert rewrite_not( | |
| Not( | |
| And( | |
| EqualTo(Reference("x"), 34.56), | |
| EqualTo(Reference("y"), 34.56), | |
| ) | |
| ) | |
| ) == Or( | |
| NotEqualTo(term=Reference(name="x"), literal=34.56), | |
| NotEqualTo(term=Reference(name="y"), literal=34.56), | |
| ) | |
| def test_rewrite_or() -> None: | |
| assert rewrite_not( | |
| Not( | |
| Or( | |
| EqualTo(Reference("x"), 34.56), | |
| EqualTo(Reference("y"), 34.56), | |
| ) | |
| ) | |
| ) == And( | |
| NotEqualTo(term=Reference(name="x"), literal=34.56), | |
| NotEqualTo(term=Reference(name="y"), literal=34.56), | |
| ) | |
| def test_rewrite_always_false() -> None: | |
| assert rewrite_not(Not(AlwaysFalse())) == AlwaysTrue() | |
| def test_rewrite_always_true() -> None: | |
| assert rewrite_not(Not(AlwaysTrue())) == AlwaysFalse() | |
| def test_rewrite_bound() -> None: | |
| schema = Schema(NestedField(2, "a", IntegerType(), required=False), schema_id=1) | |
| assert rewrite_not(IsNull(Reference("a")).bind(schema)) == BoundIsNull( | |
| term=BoundReference( | |
| field=NestedField(field_id=2, name="a", field_type=IntegerType(), required=False), | |
| accessor=Accessor(position=0, inner=None), | |
| ) | |
| ) | |
| def test_to_dnf() -> None: | |
| expr = Or(Not(EqualTo("P", "a")), And(EqualTo("Q", "b"), Not(Or(Not(EqualTo("R", "c")), EqualTo("S", "d"))))) | |
| assert rewrite_to_dnf(expr) == (NotEqualTo("P", "a"), And(EqualTo("Q", "b"), And(EqualTo("R", "c"), NotEqualTo("S", "d")))) | |
| def test_to_dnf_nested_or() -> None: | |
| expr = Or(EqualTo("P", "a"), And(EqualTo("Q", "b"), Or(EqualTo("R", "c"), EqualTo("S", "d")))) | |
| assert rewrite_to_dnf(expr) == ( | |
| EqualTo("P", "a"), | |
| And(EqualTo("Q", "b"), EqualTo("R", "c")), | |
| And(EqualTo("Q", "b"), EqualTo("S", "d")), | |
| ) | |
| def test_to_dnf_double_distribution() -> None: | |
| expr = And(Or(EqualTo("P", "a"), EqualTo("Q", "b")), Or(EqualTo("R", "c"), EqualTo("S", "d"))) | |
| assert rewrite_to_dnf(expr) == ( | |
| And( | |
| left=EqualTo(term=Reference(name="P"), literal=literal("a")), | |
| right=EqualTo(term=Reference(name="R"), literal=literal("c")), | |
| ), | |
| And( | |
| left=EqualTo(term=Reference(name="P"), literal=literal("a")), | |
| right=EqualTo(term=Reference(name="S"), literal=literal("d")), | |
| ), | |
| And( | |
| left=EqualTo(term=Reference(name="Q"), literal=literal("b")), | |
| right=EqualTo(term=Reference(name="R"), literal=literal("c")), | |
| ), | |
| And( | |
| left=EqualTo(term=Reference(name="Q"), literal=literal("b")), | |
| right=EqualTo(term=Reference(name="S"), literal=literal("d")), | |
| ), | |
| ) | |
| def test_to_dnf_double_negation() -> None: | |
| expr = rewrite_to_dnf(Not(Not(Not(Not(Not(Not(EqualTo("P", "a")))))))) | |
| assert expr == (EqualTo("P", "a"),) | |
| def test_to_dnf_and() -> None: | |
| expr = And(Not(EqualTo("Q", "b")), EqualTo("R", "c")) | |
| assert rewrite_to_dnf(expr) == (And(NotEqualTo("Q", "b"), EqualTo("R", "c")),) | |
| def test_to_dnf_not_and() -> None: | |
| expr = Not(And(Not(EqualTo("Q", "b")), EqualTo("R", "c"))) | |
| assert rewrite_to_dnf(expr) == (EqualTo("Q", "b"), NotEqualTo("R", "c")) | |
| def test_dnf_to_dask(table_schema_simple: Schema) -> None: | |
| expr = ( | |
| BoundGreaterThan( | |
| term=BoundReference(table_schema_simple.find_field(1), table_schema_simple.accessor_for_field(1)), | |
| literal=literal("hello"), | |
| ), | |
| And( | |
| BoundIn( | |
| term=BoundReference(table_schema_simple.find_field(2), table_schema_simple.accessor_for_field(2)), | |
| literals={literal(1), literal(2), literal(3)}, | |
| ), | |
| BoundEqualTo( | |
| term=BoundReference(table_schema_simple.find_field(3), table_schema_simple.accessor_for_field(3)), | |
| literal=literal(True), | |
| ), | |
| ), | |
| ) | |
| assert expression_to_plain_format(expr) == [[("foo", ">", "hello")], [("bar", "in", {1, 2, 3}), ("baz", "==", True)]] | |
| def test_expression_evaluator_null() -> None: | |
| struct = Record(None) | |
| schema = Schema(NestedField(1, "a", IntegerType(), required=False), schema_id=1) | |
| assert expression_evaluator(schema, In("a", {1, 2, 3}), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, NotIn("a", {1, 2, 3}), case_sensitive=True)(struct) is True | |
| assert expression_evaluator(schema, IsNaN("a"), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, NotNaN("a"), case_sensitive=True)(struct) is True | |
| assert expression_evaluator(schema, IsNull("a"), case_sensitive=True)(struct) is True | |
| assert expression_evaluator(schema, NotNull("a"), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, EqualTo("a", 1), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, NotEqualTo("a", 1), case_sensitive=True)(struct) is True | |
| assert expression_evaluator(schema, GreaterThanOrEqual("a", 1), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, GreaterThan("a", 1), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, LessThanOrEqual("a", 1), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, LessThan("a", 1), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, StartsWith("a", 1), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, NotStartsWith("a", 1), case_sensitive=True)(struct) is True | |
| def test_expression_evaluator_does_not_mutate_prepared_state() -> None: | |
| schema = Schema( | |
| NestedField(1, "a", IntegerType(), required=True), | |
| NestedField(2, "b", IntegerType(), required=True), | |
| ) | |
| evaluator = _ExpressionEvaluator(schema, And(EqualTo("a", 1), EqualTo("b", 1)), case_sensitive=True) | |
| initial_state = vars(evaluator).copy() | |
| assert evaluator.eval(Record(1, 1)) is True | |
| assert evaluator.eval(Record(0, 0)) is False | |
| assert evaluator.eval(Record(1, 1)) is True | |
| assert vars(evaluator) == initial_state | |
| def test_expression_evaluator_concurrent_calls_do_not_share_records() -> None: | |
| class BlockingRecord(Record): | |
| def __init__(self, first_read: Event, release_first_read: Event, *values: Any) -> None: | |
| super().__init__(*values) | |
| self.first_read = first_read | |
| self.release_first_read = release_first_read | |
| def __getitem__(self, pos: int) -> Any: | |
| value = super().__getitem__(pos) | |
| if pos == 0: | |
| self.first_read.set() | |
| if not self.release_first_read.wait(timeout=5): | |
| raise TimeoutError("Timed out waiting to interleave expression evaluations") | |
| return value | |
| schema = Schema( | |
| NestedField(1, "a", IntegerType(), required=True), | |
| NestedField(2, "b", IntegerType(), required=True), | |
| ) | |
| evaluator = expression_evaluator(schema, And(EqualTo("a", 1), EqualTo("b", 1)), case_sensitive=True) | |
| first_read = Event() | |
| release_first_read = Event() | |
| with ThreadPoolExecutor(max_workers=2) as executor: | |
| matching_result = executor.submit(evaluator, BlockingRecord(first_read, release_first_read, 1, 1)) | |
| assert first_read.wait(timeout=5) | |
| try: | |
| non_matching_result = executor.submit(evaluator, Record(0, 0)).result(timeout=5) | |
| finally: | |
| release_first_read.set() | |
| assert matching_result.result(timeout=5) is True | |
| assert non_matching_result is False | |
| def test_expression_evaluator_binary_starts_with() -> None: | |
| schema = Schema(NestedField(1, "x", BinaryType(), required=False), schema_id=1) | |
| struct = Record(b"aa") | |
| assert expression_evaluator(schema, StartsWith("x", b"a"), case_sensitive=True)(struct) is True | |
| assert expression_evaluator(schema, StartsWith("x", b"aa"), case_sensitive=True)(struct) is True | |
| assert expression_evaluator(schema, StartsWith("x", b"aaa"), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, StartsWith("x", b"b"), case_sensitive=True)(struct) is False | |
| def test_expression_evaluator_binary_not_starts_with() -> None: | |
| schema = Schema(NestedField(1, "x", BinaryType(), required=False), schema_id=1) | |
| struct = Record(b"aa") | |
| assert expression_evaluator(schema, NotStartsWith("x", b"a"), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, NotStartsWith("x", b"aa"), case_sensitive=True)(struct) is False | |
| assert expression_evaluator(schema, NotStartsWith("x", b"aaa"), case_sensitive=True)(struct) is True | |
| assert expression_evaluator(schema, NotStartsWith("x", b"b"), case_sensitive=True)(struct) is True | |
| def test_translate_column_names_simple_case(table_schema_simple: Schema) -> None: | |
| """Test translate_column_names with matching column names.""" | |
| # Create a bound expression using the original schema | |
| unbound_expr = EqualTo("foo", "test_value") | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=table_schema_simple, case_sensitive=True)) | |
| # File schema has the same column names | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="foo", field_type=StringType(), required=False), | |
| NestedField(field_id=2, name="bar", field_type=IntegerType(), required=True), | |
| NestedField(field_id=3, name="baz", field_type=BooleanType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Translate column names | |
| translated_expr = translate_column_names(bound_expr, file_schema, case_sensitive=True) | |
| # Should return an unbound expression with the same column name since they match | |
| assert isinstance(translated_expr, EqualTo) | |
| assert translated_expr.term == Reference("foo") | |
| assert translated_expr.literal == literal("test_value") | |
| def test_translate_column_names_different_column_names() -> None: | |
| """Test translate_column_names with different column names in file schema.""" | |
| # Original schema | |
| original_schema = Schema( | |
| NestedField(field_id=1, name="original_name", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Create bound expression | |
| unbound_expr = EqualTo("original_name", "test_value") | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema has different column name but same field ID | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="file_column_name", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Translate column names | |
| translated_expr = translate_column_names(bound_expr, file_schema, case_sensitive=True) | |
| # Should use the file schema's column name | |
| assert isinstance(translated_expr, EqualTo) | |
| assert translated_expr.term == Reference("file_column_name") | |
| assert translated_expr.literal == literal("test_value") | |
| def test_translate_column_names_missing_column() -> None: | |
| """Test translate_column_names when column is missing from file schema (such as in schema evolution).""" | |
| # Original schema | |
| original_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| NestedField(field_id=2, name="missing_col", field_type=IntegerType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Create bound expression for the missing column | |
| unbound_expr = EqualTo("missing_col", 42) | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema only has the existing column (field_id=1), missing field_id=2 | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Translate column names | |
| translated_expr = translate_column_names(bound_expr, file_schema, case_sensitive=True) | |
| # missing_col's default initial_default (None) does not match the expression literal (42) | |
| assert translated_expr == AlwaysFalse() | |
| def test_translate_column_names_missing_column_match_null() -> None: | |
| """Test translate_column_names when missing column matches null.""" | |
| # Original schema | |
| original_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| NestedField(field_id=2, name="missing_col", field_type=IntegerType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Create bound expression for the missing column | |
| unbound_expr = IsNull("missing_col") | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema only has the existing column (field_id=1), missing field_id=2 | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Translate column names | |
| translated_expr = translate_column_names(bound_expr, file_schema, case_sensitive=True) | |
| # Should evaluate to AlwaysTrue because the missing column is treated as null | |
| # missing_col's default initial_default (None) satisfies the IsNull predicate | |
| assert translated_expr == AlwaysTrue() | |
| def test_translate_column_names_missing_column_match_explicit_null() -> None: | |
| """Test translate_column_names when missing column matches null.""" | |
| # Original schema | |
| original_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| NestedField(field_id=2, name="missing_col", field_type=IntegerType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Create bound expression for the missing column | |
| unbound_expr = IsNull("missing_col") | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema only has the existing column (field_id=1), missing field_id=2 | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Translate column names | |
| translated_expr = translate_column_names(bound_expr, file_schema, projected_field_values={2: None}) | |
| # Should evaluate to AlwaysTrue because the missing column is treated as null | |
| # missing_col's default initial_default (None) satisfies the IsNull predicate | |
| assert translated_expr == AlwaysTrue() | |
| def test_translate_column_names_missing_column_with_initial_default() -> None: | |
| """Test translate_column_names when missing column's initial_default matches expression.""" | |
| # Original schema | |
| original_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| NestedField(field_id=2, name="missing_col", field_type=IntegerType(), required=False, initial_default=42), | |
| schema_id=1, | |
| ) | |
| # Create bound expression for the missing column | |
| unbound_expr = EqualTo("missing_col", 42) | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema only has the existing column (field_id=1), missing field_id=2 | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Translate column names | |
| translated_expr = translate_column_names(bound_expr, file_schema, case_sensitive=True) | |
| # Should evaluate to AlwaysTrue because the initial_default value (42) matches the literal (42) | |
| assert translated_expr == AlwaysTrue() | |
| def test_translate_column_names_missing_column_with_initial_default_mismatch() -> None: | |
| """Test translate_column_names when missing column's initial_default doesn't match expression.""" | |
| # Original schema | |
| original_schema = Schema( | |
| NestedField(field_id=2, name="missing_col", field_type=IntegerType(), required=False, initial_default=10), | |
| schema_id=1, | |
| ) | |
| # Create bound expression that won't match the default value | |
| unbound_expr = EqualTo("missing_col", 42) | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema doesn't have this column | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="other_col", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Translate column names | |
| translated_expr = translate_column_names(bound_expr, file_schema, case_sensitive=True) | |
| # Should evaluate to AlwaysFalse because initial_default value (10) doesn't match literal (42) | |
| assert translated_expr == AlwaysFalse() | |
| def test_translate_column_names_missing_column_with_projected_field_matches() -> None: | |
| """Test translate_column_names with projected field value that matches expression.""" | |
| # Original schema with a field that has no initial_default (defaults to None) | |
| original_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| NestedField(field_id=2, name="missing_col", field_type=IntegerType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Create bound expression for the missing column | |
| unbound_expr = EqualTo("missing_col", 42) | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema only has the existing column (field_id=1), missing field_id=2 | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Projected column that is missing in the file schema | |
| translated_expr = translate_column_names(bound_expr, file_schema, projected_field_values={2: 42}) | |
| # Should evaluate to AlwaysTrue since projected field value matches the expression literal | |
| # even though the field is missing in the file schema | |
| assert translated_expr == AlwaysTrue() | |
| def test_translate_column_names_missing_column_with_projected_field_mismatch() -> None: | |
| """Test translate_column_names with projected field value that doesn't match expression.""" | |
| # Original schema with a field that has no initial_default (defaults to None) | |
| original_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| NestedField(field_id=2, name="missing_col", field_type=IntegerType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Create bound expression for the missing column | |
| unbound_expr = EqualTo("missing_col", 42) | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema only has the existing column (field_id=1), missing field_id=2 | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Projected column that is missing in the file schema | |
| translated_expr = translate_column_names(bound_expr, file_schema, projected_field_values={2: 1}) | |
| # Should evaluate to AlwaysFalse since projected field value does not match the expression literal | |
| assert translated_expr == AlwaysFalse() | |
| def test_translate_column_names_missing_column_projected_field_ignores_initial_default() -> None: | |
| """Test translate_column_names when projected field value doesn't match but initial_default does.""" | |
| # Original schema with a field that has an initial_default | |
| original_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| NestedField(field_id=2, name="missing_col", field_type=IntegerType(), required=False, initial_default=42), | |
| schema_id=1, | |
| ) | |
| # Create bound expression for the missing column that would match initial_default | |
| unbound_expr = EqualTo("missing_col", 42) | |
| bound_expr = visit(unbound_expr, visitor=BindVisitor(schema=original_schema, case_sensitive=True)) | |
| # File schema only has the existing column (field_id=1), missing field_id=2 | |
| file_schema = Schema( | |
| NestedField(field_id=1, name="existing_col", field_type=StringType(), required=False), | |
| schema_id=1, | |
| ) | |
| # Projected field value that differs from both the expression literal and initial_default | |
| translated_expr = translate_column_names( | |
| bound_expr, | |
| file_schema, | |
| projected_field_values={2: 10}, # This doesn't match expression literal (42) | |
| ) | |
| # Should evaluate to AlwaysFalse since projected field value doesn't match the expression literal | |
| assert translated_expr == AlwaysFalse() |