jqin61 commented on code in PR #955:
URL: https://github.com/apache/iceberg-python/pull/955#discussion_r1688678348
##########
tests/io/test_pyarrow_visitor.py:
##########
@@ -580,3 +594,114 @@ def
test_pyarrow_schema_ensure_large_types(pyarrow_schema_nested_without_ids: pa
),
])
assert
_pyarrow_schema_ensure_large_types(pyarrow_schema_nested_without_ids) ==
expected_schema
+
+
[email protected]
+def bound_reference_str() -> BoundReference[Any]:
+ return BoundReference(
+ field=NestedField(1, "field_str_unmentioned", StringType(),
required=False), accessor=Accessor(position=0, inner=None)
+ )
+
+
[email protected]
+def bound_reference_float() -> BoundReference[Any]:
+ return BoundReference(
+ field=NestedField(2, "field_float_mentioned_nan", FloatType(),
required=False), accessor=Accessor(position=1, inner=None)
+ )
+
+
[email protected]
+def bound_reference_double() -> BoundReference[Any]:
+ return BoundReference(
+ field=NestedField(3, "field_double_mentioned_null", DoubleType(),
required=False),
+ accessor=Accessor(position=2, inner=None),
+ )
+
+
[email protected]
+def bound_eq_str_field(bound_reference_str: BoundReference[Any]) ->
BoundEqualTo[Any]:
+ return BoundEqualTo(term=bound_reference_str, literal=literal("hello"))
+
+
[email protected]
+def bound_greater_than_float_field(bound_reference_float: BoundReference[Any])
-> BoundGreaterThan[Any]:
+ return BoundGreaterThan(term=bound_reference_float, literal=literal(100))
+
+
[email protected]
+def bound_is_nan_float_field(bound_reference_float: BoundReference[Any]) ->
BoundIsNaN[Any]:
+ return BoundIsNaN(bound_reference_float)
+
+
[email protected]
+def bound_eq_double_field(bound_reference_double: BoundReference[Any]) ->
BoundEqualTo[Any]:
+ return BoundEqualTo(term=bound_reference_double, literal=literal(False))
+
+
[email protected]
+def bound_is_null_double_field(bound_reference_double: BoundReference[Any]) ->
BoundIsNull[Any]:
+ return BoundIsNull(bound_reference_double)
+
+
+def test_collect_null_nan_unmentioned_terms(
+ bound_eq_str_field: BoundEqualTo[Any], bound_is_nan_float_field:
BoundIsNaN[Any], bound_is_null_double_field: BoundIsNull[Any]
+) -> None:
+ bound_expr = And(
+ Or(And(bound_eq_str_field, bound_is_nan_float_field),
bound_is_null_double_field), Not(bound_is_nan_float_field)
+ )
+ categorized_terms = _get_null_nan_refs(bound_expr)
+
+ assert {f.field.name for f in categorized_terms[0]} ==
{"field_float_mentioned_nan", "field_str_unmentioned"}
+ assert {f.field.name for f in categorized_terms[1]} ==
{"field_str_unmentioned", "field_double_mentioned_null"}
+ assert {f.field.name for f in categorized_terms[2]} == {
+ "field_double_mentioned_null",
+ }
+ assert {f.field.name for f in categorized_terms[3]} ==
{"field_float_mentioned_nan"}
+
+
+def
test_collect_null_nan_unmentioned_terms_with_multiple_predicates_on_the_same_term(
+ bound_eq_str_field: BoundEqualTo[Any],
+ bound_greater_than_float_field: BoundGreaterThan[Any],
+ bound_is_nan_float_field: BoundIsNaN[Any],
+ bound_eq_double_field: BoundEqualTo[Any],
+ bound_is_null_double_field: BoundIsNull[Any],
+) -> None:
+ """Test a single term appears multiple places in the expression tree"""
+ bound_expr = And(
+ Or(
+ And(bound_eq_str_field, bound_greater_than_float_field),
+ And(bound_is_nan_float_field, bound_eq_double_field),
+ bound_greater_than_float_field,
+ ),
+ Not(bound_is_null_double_field),
+ )
+ categorized_terms = _get_null_nan_refs(bound_expr)
+ assert {f.field.name for f in categorized_terms[0]} ==
{"field_float_mentioned_nan", "field_str_unmentioned"}
+ assert {f.field.name for f in categorized_terms[1]} ==
{"field_str_unmentioned", "field_double_mentioned_null"}
+ assert {f.field.name for f in categorized_terms[2]} == {
+ "field_double_mentioned_null",
+ }
+ assert {f.field.name for f in categorized_terms[3]} ==
{"field_float_mentioned_nan"}
+
+
+def test__expression_to_complementary_pyarrow(
+ bound_eq_str_field: BoundEqualTo[Any],
+ bound_greater_than_float_field: BoundGreaterThan[Any],
+ bound_is_nan_float_field: BoundIsNaN[Any],
+ bound_eq_double_field: BoundEqualTo[Any],
+ bound_is_null_double_field: BoundIsNull[Any],
+) -> None:
+ bound_expr = And(
+ Or(
+ And(bound_eq_str_field, bound_greater_than_float_field),
+ And(bound_is_nan_float_field, bound_eq_double_field),
+ bound_greater_than_float_field,
+ ),
+ Not(bound_is_null_double_field),
+ )
+ result = _expression_to_complementary_pyarrow(bound_expr)
+ # Notice an isNan predicate on a str column is automatically converted to
always false and removed from Or and thus will not appear in the pc.expr.
+ assert (
+ repr(result)
+ == """<pyarrow.compute.Expression (((invert((((((field_str_unmentioned
== "hello") and (field_float_mentioned_nan > 100)) or
(is_nan(field_float_mentioned_nan) and (field_double_mentioned_null == 0))) or
(field_float_mentioned_nan > 100)) and
invert(is_null(field_double_mentioned_null, {nan_is_null=false})))) or
is_null(field_float_mentioned_nan, {nan_is_null=false})) or
is_null(field_str_unmentioned, {nan_is_null=false})) or
is_nan(field_double_mentioned_null))>"""
Review Comment:
gotcha, will change it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
