68
68
BooleanExpressionVisitor ,
69
69
BoundBooleanExpressionVisitor ,
70
70
_ManifestEvalVisitor ,
71
+ bind ,
71
72
expression_evaluator ,
72
73
expression_to_plain_format ,
73
74
rewrite_not ,
74
75
rewrite_to_dnf ,
76
+ translate_column_names ,
75
77
visit ,
76
78
visit_bound_predicate ,
77
79
)
@@ -1623,3 +1625,22 @@ def test_expression_evaluator_null() -> None:
1623
1625
assert expression_evaluator (schema , LessThan ("a" , 1 ), case_sensitive = True )(struct ) is False
1624
1626
assert expression_evaluator (schema , StartsWith ("a" , 1 ), case_sensitive = True )(struct ) is False
1625
1627
assert expression_evaluator (schema , NotStartsWith ("a" , 1 ), case_sensitive = True )(struct ) is True
1628
+
1629
+ @pytest .mark .parametrize (
1630
+ "before_expression,after_expression" ,
1631
+ [
1632
+ (In ("id" , {1 , 2 , 3 }), AlwaysTrue ()),
1633
+ (EqualTo ("id" , 3 ), AlwaysFalse ()),
1634
+ (And (EqualTo ("id" , 1 ), EqualTo ("all_same_value_or_null" , "string" )), And (AlwaysTrue (), EqualTo ("all_same_value_or_null" , "string" ))),
1635
+ (And (EqualTo ("all_same_value_or_null" , "string" ), GreaterThan ("id" , 2 )), And (EqualTo ("all_same_value_or_null" , "string" ), AlwaysFalse ())),
1636
+ (
1637
+ Or (And (EqualTo ("id" , 1 ), EqualTo ("all_same_value_or_null" , "string" )), And (EqualTo ("all_same_value_or_null" , "string" ), GreaterThan ("id" , 2 ))),
1638
+ Or (And (AlwaysTrue (), EqualTo ("all_same_value_or_null" , "string" )), And (EqualTo ("all_same_value_or_null" , "string" ), AlwaysFalse ())),
1639
+ )
1640
+ ]
1641
+ )
1642
+ def test_translate_column_names_eval_projected_fields (schema : Schema , before_expression : BooleanExpression , after_expression : BooleanExpression ) -> None :
1643
+ # exclude id from file_schema pretending that it's part of partition values
1644
+ file_schema = Schema (* [field for field in schema .columns if field .name != "id" ])
1645
+ projected_missing_fields = { "id" : 1 }
1646
+ assert translate_column_names (bind (schema , before_expression , True ), file_schema , schema , True , projected_missing_fields ) == after_expression
0 commit comments