syun64 commented on code in PR #955: URL: https://github.com/apache/iceberg-python/pull/955#discussion_r1688622860
########## 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 + + +@pytest.fixture +def bound_reference_str() -> BoundReference[Any]: + return BoundReference( + field=NestedField(1, "field_str_unmentioned", StringType(), required=False), accessor=Accessor(position=0, inner=None) + ) + + +@pytest.fixture +def bound_reference_float() -> BoundReference[Any]: + return BoundReference( + field=NestedField(2, "field_float_mentioned_nan", FloatType(), required=False), accessor=Accessor(position=1, inner=None) + ) + + +@pytest.fixture +def bound_reference_double() -> BoundReference[Any]: + return BoundReference( + field=NestedField(3, "field_double_mentioned_null", DoubleType(), required=False), + accessor=Accessor(position=2, inner=None), + ) + + +@pytest.fixture +def bound_eq_str_field(bound_reference_str: BoundReference[Any]) -> BoundEqualTo[Any]: + return BoundEqualTo(term=bound_reference_str, literal=literal("hello")) + + +@pytest.fixture +def bound_greater_than_float_field(bound_reference_float: BoundReference[Any]) -> BoundGreaterThan[Any]: + return BoundGreaterThan(term=bound_reference_float, literal=literal(100)) + + +@pytest.fixture +def bound_is_nan_float_field(bound_reference_float: BoundReference[Any]) -> BoundIsNaN[Any]: + return BoundIsNaN(bound_reference_float) + + +@pytest.fixture +def bound_eq_double_field(bound_reference_double: BoundReference[Any]) -> BoundEqualTo[Any]: + return BoundEqualTo(term=bound_reference_double, literal=literal(False)) + + +@pytest.fixture +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: Hi @jqin61 - this looks almost ready to merge. Thank you again for all your effort in fixing this critical issue in this feature for this release. Personally, I'm finding these field names difficult to understand. I feel as though we may have named the fields this way for convenience in understanding the test case, but I'm actually finding it more difficult to understand because I can get over the fact that a field name should be agnostic of the predicate expression that is applied on it - is_nan or a predicate without nan or null can be applied on the same field. I'm curious to source other reviewers' opinions on this as well! -- 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: issues-unsubscr...@iceberg.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org For additional commands, e-mail: issues-h...@iceberg.apache.org