extract projected columns evaluator into separate visitor

Author: Roman Shanin

pyiceberg/expressions/visitors.py

@@ -860,9 +860,7 @@ class _ColumnNameTranslator(BooleanExpressionVisitor[BooleanExpression]):
 
     Args:
       file_schema (Schema): The schema of the file.
-      projected_schema (Schema): The schema to project onto the data files.
       case_sensitive (bool): Whether to consider case when binding a reference to a field in a schema, defaults to True.
-      projected_missing_fields(dict[str, Any]): Map of fields missing in file_schema, but present as partition values.
 
     Raises:
         TypeError: In the case of an UnboundPredicate.
@@ -872,13 +870,9 @@ class _ColumnNameTranslator(BooleanExpressionVisitor[BooleanExpression]):
     file_schema: Schema
     case_sensitive: bool
 
-    def __init__(
-        self, file_schema: Schema, projected_schema: Schema, case_sensitive: bool, projected_missing_fields: dict[str, Any]
-    ) -> None:
+    def __init__(self, file_schema: Schema, case_sensitive: bool) -> None:
         self.file_schema = file_schema
-        self.projected_schema = projected_schema
         self.case_sensitive = case_sensitive
-        self.projected_missing_fields = projected_missing_fields
 
     def visit_true(self) -> BooleanExpression:
         return AlwaysTrue()
@@ -906,24 +900,6 @@ def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> BooleanExpressi
             # in the file schema when reading older data
             if isinstance(predicate, BoundIsNull):
                 return AlwaysTrue()
-            # Evaluate projected field by value extracted from partition
-            elif (field_name := predicate.term.ref().field.name) in self.projected_missing_fields:
-                unbound_predicate: BooleanExpression
-                if isinstance(predicate, BoundUnaryPredicate):
-                    unbound_predicate = predicate.as_unbound(field_name)
-                elif isinstance(predicate, BoundLiteralPredicate):
-                    unbound_predicate = predicate.as_unbound(field_name, predicate.literal)
-                elif isinstance(predicate, BoundSetPredicate):
-                    unbound_predicate = predicate.as_unbound(field_name, predicate.literals)
-                else:
-                    raise ValueError(f"Unsupported predicate: {predicate}")
-                field = self.projected_schema.find_field(field_name)
-                schema = Schema(field)
-                evaluator = expression_evaluator(schema, unbound_predicate, self.case_sensitive)
-                if evaluator(Record(self.projected_missing_fields[field_name])):
-                    return AlwaysTrue()
-                else:
-                    return AlwaysFalse()
             else:
                 return AlwaysFalse()
 
@@ -937,14 +913,84 @@ def visit_bound_predicate(self, predicate: BoundPredicate[L]) -> BooleanExpressi
             raise ValueError(f"Unsupported predicate: {predicate}")
 
 
-def translate_column_names(
+def translate_column_names(expr: BooleanExpression, file_schema: Schema, case_sensitive: bool) -> BooleanExpression:
+    return visit(expr, _ColumnNameTranslator(file_schema, case_sensitive))
+
+
+class _ProjectedColumnsEvaluator(BooleanExpressionVisitor[BooleanExpression]):
+    """Evaluated predicates which involve projected columns missing from the file.
+
+    Args:
+      file_schema (Schema): The schema of the file.
+      projected_schema (Schema): The schema to project onto the data files.
+      case_sensitive (bool): Whether to consider case when binding a reference to a field in a schema, defaults to True.
+      projected_missing_fields(dict[str, Any]): Map of fields missing in file_schema, but present as partition values.
+
+    Raises:
+        TypeError: In the case of an UnboundPredicate.
+    """
+
+    file_schema: Schema
+    case_sensitive: bool
+
+    def __init__(
+        self, file_schema: Schema, projected_schema: Schema, case_sensitive: bool, projected_missing_fields: dict[str, Any]
+    ) -> None:
+        self.file_schema = file_schema
+        self.projected_schema = projected_schema
+        self.case_sensitive = case_sensitive
+        self.projected_missing_fields = projected_missing_fields
+
+    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 and (field_name := predicate.term.ref().field.name) in self.projected_missing_fields:
+            unbound_predicate: BooleanExpression
+            if isinstance(predicate, BoundUnaryPredicate):
+                unbound_predicate = predicate.as_unbound(field_name)
+            elif isinstance(predicate, BoundLiteralPredicate):
+                unbound_predicate = predicate.as_unbound(field_name, predicate.literal)
+            elif isinstance(predicate, BoundSetPredicate):
+                unbound_predicate = predicate.as_unbound(field_name, predicate.literals)
+            else:
+                raise ValueError(f"Unsupported predicate: {predicate}")
+            field = self.projected_schema.find_field(field_name)
+            schema = Schema(field)
+            evaluator = expression_evaluator(schema, unbound_predicate, self.case_sensitive)
+            if evaluator(Record(self.projected_missing_fields[field_name])):
+                return AlwaysTrue()
+            else:
+                return AlwaysFalse()
+
+        return predicate
+
+
+def evaluate_projected_columns(
     expr: BooleanExpression,
     file_schema: Schema,
     projected_schema: Schema,
     case_sensitive: bool,
     projected_missing_fields: dict[str, Any],
 ) -> BooleanExpression:
-    return visit(expr, _ColumnNameTranslator(file_schema, projected_schema, case_sensitive, projected_missing_fields))
+    return visit(expr, _ProjectedColumnsEvaluator(file_schema, projected_schema, case_sensitive, projected_missing_fields))
 
 
 class _ExpressionFieldIDs(BooleanExpressionVisitor[Set[int]]):

pyiceberg/io/pyarrow.py

@@ -78,6 +78,7 @@
 from pyiceberg.expressions.visitors import (
     BoundBooleanExpressionVisitor,
     bind,
+    evaluate_projected_columns,
     extract_field_ids,
     translate_column_names,
 )
@@ -1466,13 +1467,18 @@ def _task_to_record_batches(
 
         pyarrow_filter = None
         if bound_row_filter is not AlwaysTrue():
-            translated_row_filter = translate_column_names(
+            evaluated_projected_columns_filter = evaluate_projected_columns(
                 bound_row_filter,
                 file_schema,
                 projected_schema,
                 case_sensitive=case_sensitive,
                 projected_missing_fields=projected_missing_fields,
             )
+            translated_row_filter = translate_column_names(
+                evaluated_projected_columns_filter,
+                file_schema,
+                case_sensitive=case_sensitive,
+            )
             bound_file_filter = bind(file_schema, translated_row_filter, case_sensitive=case_sensitive)
             pyarrow_filter = expression_to_pyarrow(bound_file_filter)
 

tests/expressions/test_visitors.py

@@ -69,11 +69,11 @@
     BoundBooleanExpressionVisitor,
     _ManifestEvalVisitor,
     bind,
+    evaluate_projected_columns,
     expression_evaluator,
     expression_to_plain_format,
     rewrite_not,
     rewrite_to_dnf,
-    translate_column_names,
     visit,
     visit_bound_predicate,
 )
@@ -1652,13 +1652,10 @@ def test_expression_evaluator_null() -> None:
         ),
     ],
 )
-def test_translate_column_names_eval_projected_fields(
-    schema: Schema, before_expression: BooleanExpression, after_expression: BooleanExpression
-) -> None:
+def test_eval_projected_fields(schema: Schema, before_expression: BooleanExpression, after_expression: BooleanExpression) -> None:
     # exclude id from file_schema pretending that it's part of partition values
     file_schema = Schema(*[field for field in schema.columns if field.name != "id"])
     projected_missing_fields = {"id": 1}
-    assert (
-        translate_column_names(bind(schema, before_expression, True), file_schema, schema, True, projected_missing_fields)
-        == after_expression
-    )
+    assert evaluate_projected_columns(
+        bind(schema, before_expression, True), file_schema, schema, True, projected_missing_fields
+    ) == bind(schema, after_expression, True)