refactor: subclass DerefOp as ResolvedDerefOp

Author: chelsea-lin

bigframes/core/compile/polars/compiler.py

@@ -120,7 +120,7 @@ def _(
         @compile_expression.register
         def _(
             self,
-            expression: ex.SchemaFieldRefExpression,
+            expression: ex.ResolvedDerefOp,
         ) -> pl.Expr:
             return pl.col(expression.field.id.sql)
 

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -42,13 +42,6 @@ def compile_deref_expression(expr: expression.DerefOp) -> sge.Expression:
     return sge.ColumnDef(this=sge.to_identifier(expr.id.sql, quoted=True))
 
 
-@compile_scalar_expression.register
-def compile_field_ref_expression(
-    expr: expression.SchemaFieldRefExpression,
-) -> sge.Expression:
-    return sge.ColumnDef(this=sge.to_identifier(expr.field.id.sql, quoted=True))
-
-
 @compile_scalar_expression.register
 def compile_constant_expression(
     expr: expression.ScalarConstantExpression,

bigframes/core/expression.py

@@ -429,19 +429,23 @@ def transform_children(self, t: Callable[[Expression], Expression]) -> Expressio
 
 
 @dataclasses.dataclass(frozen=True)
-class SchemaFieldRefExpression(Expression):
-    """An expression representing a schema field. This is essentially a DerefOp with input schema bound."""
+class ResolvedDerefOp(DerefOp):
+    """An expression that refers to a column by ID and resolved with schema bound."""
 
     field: field.Field
 
+    # Re-declare 'id' from the parent to remove it from the __init__ method
+    id: ids.ColumnId = dataclasses.field(init=False)
+
+    def __post_init__(self):
+        # Initialize the parent's 'id' field after the object is created.
+        # We must use object.__setattr__ because the dataclass is frozen.
+        object.__setattr__(self, "id", self.field.id)
+
     @property
     def column_references(self) -> typing.Tuple[ids.ColumnId, ...]:
         return (self.field.id,)
 
-    @property
-    def is_const(self) -> bool:
-        return False
-
     @property
     def nullable(self) -> bool:
         return self.field.nullable
@@ -464,21 +468,11 @@ def bind_refs(
         bindings: Mapping[ids.ColumnId, Expression],
         allow_partial_bindings: bool = False,
     ) -> Expression:
+        # TODO: Check if we can remove.
         if self.field.id in bindings.keys():
             return bindings[self.field.id]
         return self
 
-    @property
-    def is_bijective(self) -> bool:
-        return True
-
-    @property
-    def is_identity(self) -> bool:
-        return True
-
-    def transform_children(self, t: Callable[[Expression], Expression]) -> Expression:
-        return self
-
 
 @dataclasses.dataclass(frozen=True)
 class OpExpression(Expression):
@@ -588,9 +582,7 @@ def bind_schema_fields(
     if expr.is_resolved:
         return expr
 
-    expr_by_id = {
-        id: SchemaFieldRefExpression(field) for id, field in field_by_id.items()
-    }
+    expr_by_id = {id: ResolvedDerefOp(field) for id, field in field_by_id.items()}
     return expr.bind_refs(expr_by_id)
 
 

bigframes/core/rewrite/schema_binding.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import dataclasses
+import typing
 
 from bigframes.core import bigframe_node
 from bigframes.core import expression as ex
@@ -52,4 +53,23 @@ def bind_schema_to_node(
 
         return dataclasses.replace(node, by=tuple(bound_bys))
 
+    if isinstance(node, nodes.JoinNode):
+        conditions = tuple(
+            (
+                typing.cast(
+                    ex.ResolvedDerefOp,
+                    ex.bind_schema_fields(left, node.left_child.field_by_id),
+                ),
+                typing.cast(
+                    ex.ResolvedDerefOp,
+                    ex.bind_schema_fields(right, node.right_child.field_by_id),
+                ),
+            )
+            for left, right in node.conditions
+        )
+        return dataclasses.replace(
+            node,
+            conditions=conditions,
+        )
+
     return node

tests/unit/core/test_expression.py

@@ -77,7 +77,7 @@ def test_deref_op_dtype_resolution():
 
 
 def test_field_ref_expr_dtype_resolution_short_circuit():
-    expression = ex.SchemaFieldRefExpression(
+    expression = ex.ResolvedDerefOp(
         field.Field(ids.ColumnId("mycol"), dtype=dtypes.INT_DTYPE)
     )
     field_bindings = _create_field_bindings({"anotherCol": dtypes.STRING_DTYPE})