refactor: Cache dtypes for scalar expressions for SQLGlot compiler (#1759)

* feat: include bq schema and query string in dry run results

* rename key

* fix tests

* refactor: cache dtypes for scalar expressions:

* fix deref expr type resolution bug

* add test

* remove dry_run changes from another branch

* remove more changes from dry_run PR

* rename DeferredDtype to AbsentDtype

* removed absentDtype and reuse bind_refs

* use a separate resolver for fields

* fix lint

* move field resolutions to a separate function

* update helper function name

* update doc and function names

* bind schema at compile time for SQLGlot compiler

* define a separate expression for field reference

Author: sycai

bigframes/core/bigframe_node.py

@@ -22,7 +22,7 @@
 import typing
 from typing import Callable, Dict, Generator, Iterable, Mapping, Sequence, Set, Tuple
 
-from bigframes.core import identifiers
+from bigframes.core import field, identifiers
 import bigframes.core.schema as schemata
 import bigframes.dtypes
 
@@ -34,23 +34,6 @@
 T = typing.TypeVar("T")
 
 
-@dataclasses.dataclass(frozen=True)
-class Field:
-    id: identifiers.ColumnId
-    dtype: bigframes.dtypes.Dtype
-    # Best effort, nullable=True if not certain
-    nullable: bool = True
-
-    def with_nullable(self) -> Field:
-        return Field(self.id, self.dtype, nullable=True)
-
-    def with_nonnull(self) -> Field:
-        return Field(self.id, self.dtype, nullable=False)
-
-    def with_id(self, id: identifiers.ColumnId) -> Field:
-        return Field(id, self.dtype, nullable=self.nullable)
-
-
 @dataclasses.dataclass(eq=False, frozen=True)
 class BigFrameNode:
     """
@@ -162,7 +145,7 @@ def roots(self) -> typing.Set[BigFrameNode]:
     # TODO: Store some local data lazily for select, aggregate nodes.
     @property
     @abc.abstractmethod
-    def fields(self) -> Sequence[Field]:
+    def fields(self) -> Sequence[field.Field]:
         ...
 
     @property
@@ -292,7 +275,7 @@ def _dtype_lookup(self) -> dict[identifiers.ColumnId, bigframes.dtypes.Dtype]:
         return {field.id: field.dtype for field in self.fields}
 
     @functools.cached_property
-    def field_by_id(self) -> Mapping[identifiers.ColumnId, Field]:
+    def field_by_id(self) -> Mapping[identifiers.ColumnId, field.Field]:
         return {field.id: field for field in self.fields}
 
     # Plan algorithms

bigframes/core/compile/sqlglot/compiler.py

@@ -26,6 +26,7 @@
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 import bigframes.core.ordering as bf_ordering
+from bigframes.core.rewrite import schema_binding
 
 
 class SQLGlotCompiler:
@@ -183,6 +184,6 @@ def compile_projection(
 
 def _replace_unsupported_ops(node: nodes.BigFrameNode):
     node = nodes.bottom_up(node, rewrite.rewrite_slice)
-    node = nodes.bottom_up(node, rewrite.rewrite_timedelta_expressions)
+    node = nodes.bottom_up(node, schema_binding.bind_schema_to_expressions)
     node = nodes.bottom_up(node, rewrite.rewrite_range_rolling)
     return node

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -35,6 +35,13 @@ 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

@@ -16,14 +16,16 @@
 
 import abc
 import dataclasses
+import functools
 import itertools
 import typing
 from typing import Generator, Mapping, TypeVar, Union
 
 import pandas as pd
 
+from bigframes import dtypes
+from bigframes.core import field
 import bigframes.core.identifiers as ids
-import bigframes.dtypes as dtypes
 import bigframes.operations
 import bigframes.operations.aggregations as agg_ops
 
@@ -50,7 +52,7 @@ class Aggregation(abc.ABC):
 
     @abc.abstractmethod
     def output_type(
-        self, input_types: dict[ids.ColumnId, dtypes.ExpressionType]
+        self, input_fields: Mapping[ids.ColumnId, field.Field]
     ) -> dtypes.ExpressionType:
         ...
 
@@ -72,7 +74,7 @@ class NullaryAggregation(Aggregation):
     op: agg_ops.NullaryWindowOp = dataclasses.field()
 
     def output_type(
-        self, input_types: dict[ids.ColumnId, bigframes.dtypes.Dtype]
+        self, input_fields: Mapping[ids.ColumnId, field.Field]
     ) -> dtypes.ExpressionType:
         return self.op.output_type()
 
@@ -86,13 +88,17 @@ def remap_column_refs(
 
 @dataclasses.dataclass(frozen=True)
 class UnaryAggregation(Aggregation):
-    op: agg_ops.UnaryWindowOp = dataclasses.field()
-    arg: Union[DerefOp, ScalarConstantExpression] = dataclasses.field()
+    op: agg_ops.UnaryWindowOp
+    arg: Union[DerefOp, ScalarConstantExpression]
 
     def output_type(
-        self, input_types: dict[ids.ColumnId, bigframes.dtypes.Dtype]
+        self, input_fields: Mapping[ids.ColumnId, field.Field]
     ) -> dtypes.ExpressionType:
-        return self.op.output_type(self.arg.output_type(input_types))
+        # TODO(b/419300717) Remove resolutions once defers are cleaned up.
+        resolved_expr = bind_schema_fields(self.arg, input_fields)
+        assert resolved_expr.is_resolved
+
+        return self.op.output_type(resolved_expr.output_type)
 
     @property
     def column_references(self) -> typing.Tuple[ids.ColumnId, ...]:
@@ -118,10 +124,16 @@ class BinaryAggregation(Aggregation):
     right: Union[DerefOp, ScalarConstantExpression] = dataclasses.field()
 
     def output_type(
-        self, input_types: dict[ids.ColumnId, bigframes.dtypes.Dtype]
+        self, input_fields: Mapping[ids.ColumnId, field.Field]
     ) -> dtypes.ExpressionType:
+        # TODO(b/419300717) Remove resolutions once defers are cleaned up.
+        left_resolved_expr = bind_schema_fields(self.left, input_fields)
+        assert left_resolved_expr.is_resolved
+        right_resolved_expr = bind_schema_fields(self.right, input_fields)
+        assert right_resolved_expr.is_resolved
+
         return self.op.output_type(
-            self.left.output_type(input_types), self.right.output_type(input_types)
+            left_resolved_expr.output_type, left_resolved_expr.output_type
         )
 
     @property
@@ -189,10 +201,17 @@ def remap_column_refs(
     def is_const(self) -> bool:
         ...
 
+    @property
     @abc.abstractmethod
-    def output_type(
-        self, input_types: dict[ids.ColumnId, dtypes.ExpressionType]
-    ) -> dtypes.ExpressionType:
+    def is_resolved(self) -> bool:
+        """
+        Returns true if and only if the expression's output type and nullability is available.
+        """
+        ...
+
+    @property
+    @abc.abstractmethod
+    def output_type(self) -> dtypes.ExpressionType:
         ...
 
     @abc.abstractmethod
@@ -256,9 +275,12 @@ def column_references(self) -> typing.Tuple[ids.ColumnId, ...]:
     def nullable(self) -> bool:
         return pd.isna(self.value)  # type: ignore
 
-    def output_type(
-        self, input_types: dict[ids.ColumnId, bigframes.dtypes.Dtype]
-    ) -> dtypes.ExpressionType:
+    @property
+    def is_resolved(self) -> bool:
+        return True
+
+    @property
+    def output_type(self) -> dtypes.ExpressionType:
         return self.dtype
 
     def bind_variables(
@@ -308,9 +330,12 @@ def is_const(self) -> bool:
     def column_references(self) -> typing.Tuple[ids.ColumnId, ...]:
         return ()
 
-    def output_type(
-        self, input_types: dict[ids.ColumnId, bigframes.dtypes.Dtype]
-    ) -> dtypes.ExpressionType:
+    @property
+    def is_resolved(self):
+        return False
+
+    @property
+    def output_type(self) -> dtypes.ExpressionType:
         raise ValueError(f"Type of variable {self.id} has not been fixed.")
 
     def bind_refs(
@@ -340,7 +365,7 @@ def is_identity(self) -> bool:
 
 @dataclasses.dataclass(frozen=True)
 class DerefOp(Expression):
-    """A variable expression representing an unbound variable."""
+    """An expression that refers to a column by ID."""
 
     id: ids.ColumnId
 
@@ -357,13 +382,13 @@ def nullable(self) -> bool:
         # Safe default, need to actually bind input schema to determine
         return True
 
-    def output_type(
-        self, input_types: dict[ids.ColumnId, bigframes.dtypes.Dtype]
-    ) -> dtypes.ExpressionType:
-        if self.id in input_types:
-            return input_types[self.id]
-        else:
-            raise ValueError(f"Type of variable {self.id} has not been fixed.")
+    @property
+    def is_resolved(self) -> bool:
+        return False
+
+    @property
+    def output_type(self) -> dtypes.ExpressionType:
+        raise ValueError(f"Type of variable {self.id} has not been fixed.")
 
     def bind_variables(
         self, bindings: Mapping[str, Expression], allow_partial_bindings: bool = False
@@ -390,6 +415,55 @@ def is_identity(self) -> bool:
         return True
 
 
+@dataclasses.dataclass(frozen=True)
+class SchemaFieldRefExpression(Expression):
+    """An expression representing a schema field. This is essentially a DerefOp with input schema bound."""
+
+    field: field.Field
+
+    @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
+
+    @property
+    def is_resolved(self) -> bool:
+        return True
+
+    @property
+    def output_type(self) -> dtypes.ExpressionType:
+        return self.field.dtype
+
+    def bind_variables(
+        self, bindings: Mapping[str, Expression], allow_partial_bindings: bool = False
+    ) -> Expression:
+        return self
+
+    def bind_refs(
+        self,
+        bindings: Mapping[ids.ColumnId, Expression],
+        allow_partial_bindings: bool = False,
+    ) -> Expression:
+        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
+
+
 @dataclasses.dataclass(frozen=True)
 class OpExpression(Expression):
     """An expression representing a scalar operation applied to 1 or more argument sub-expressions."""
@@ -429,13 +503,18 @@ def nullable(self) -> bool:
         )
         return not null_free
 
-    def output_type(
-        self, input_types: dict[ids.ColumnId, dtypes.ExpressionType]
-    ) -> dtypes.ExpressionType:
-        operand_types = tuple(
-            map(lambda x: x.output_type(input_types=input_types), self.inputs)
-        )
-        return self.op.output_type(*operand_types)
+    @functools.cached_property
+    def is_resolved(self) -> bool:
+        return all(input.is_resolved for input in self.inputs)
+
+    @functools.cached_property
+    def output_type(self) -> dtypes.ExpressionType:
+        if not self.is_resolved:
+            raise ValueError(f"Type of expression {self.op.name} has not been fixed.")
+
+        input_types = [input.output_type for input in self.inputs]
+
+        return self.op.output_type(*input_types)
 
     def bind_variables(
         self, bindings: Mapping[str, Expression], allow_partial_bindings: bool = False
@@ -475,4 +554,22 @@ def deterministic(self) -> bool:
         )
 
 
+def bind_schema_fields(
+    expr: Expression, field_by_id: Mapping[ids.ColumnId, field.Field]
+) -> Expression:
+    """
+    Updates `DerefOp` expressions by replacing column IDs with actual schema fields(columns).
+
+    We can only deduct an expression's output type and nullability after binding schema fields to
+    all its deref expressions.
+    """
+    if expr.is_resolved:
+        return expr
+
+    expr_by_id = {
+        id: SchemaFieldRefExpression(field) for id, field in field_by_id.items()
+    }
+    return expr.bind_refs(expr_by_id)
+
+
 RefOrConstant = Union[DerefOp, ScalarConstantExpression]

bigframes/core/field.py

@@ -0,0 +1,37 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import dataclasses
+
+from bigframes import dtypes
+from bigframes.core import identifiers
+
+
+@dataclasses.dataclass(frozen=True)
+class Field:
+    id: identifiers.ColumnId
+    dtype: dtypes.Dtype
+    # Best effort, nullable=True if not certain
+    nullable: bool = True
+
+    def with_nullable(self) -> Field:
+        return Field(self.id, self.dtype, nullable=True)
+
+    def with_nonnull(self) -> Field:
+        return Field(self.id, self.dtype, nullable=False)
+
+    def with_id(self, id: identifiers.ColumnId) -> Field:
+        return Field(id, self.dtype, nullable=self.nullable)