chore: inject dtypes to SQLGlot scalar expr compiler (#1821)

* chore: inject dtypes to SQLGlot scalar expr compiler

* fix format

Author: sycai

bigframes/core/compile/sqlglot/compiler.py

@@ -119,14 +119,16 @@ def _remap_variables(self, node: nodes.ResultNode) -> nodes.ResultNode:
         return typing.cast(nodes.ResultNode, result_node)
 
     def _compile_result_node(self, root: nodes.ResultNode) -> str:
-        sqlglot_ir = self.compile_node(root.child)
-
+        # Have to bind schema as the final step before compilation.
+        root = typing.cast(nodes.ResultNode, schema_binding.bind_schema_to_tree(root))
         selected_cols: tuple[tuple[str, sge.Expression], ...] = tuple(
             (name, scalar_compiler.compile_scalar_expression(ref))
             for ref, name in root.output_cols
         )
         # Skip squashing selections to ensure the right ordering and limit keys
-        sqlglot_ir = sqlglot_ir.select(selected_cols, squash_selections=False)
+        sqlglot_ir = self.compile_node(root.child).select(
+            selected_cols, squash_selections=False
+        )
 
         if root.order_by is not None:
             ordering_cols = tuple(
@@ -220,6 +222,5 @@ def compile_concat(
 
 def _replace_unsupported_ops(node: nodes.BigFrameNode):
     node = nodes.bottom_up(node, rewrite.rewrite_slice)
-    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

@@ -13,15 +13,25 @@
 # limitations under the License.
 from __future__ import annotations
 
+import dataclasses
 import functools
 
 import sqlglot.expressions as sge
 
+from bigframes import dtypes
 from bigframes.core import expression
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 import bigframes.operations as ops
 
 
+@dataclasses.dataclass(frozen=True)
+class TypedExpr:
+    """SQLGlot expression with type."""
+
+    expr: sge.Expression
+    dtype: dtypes.ExpressionType
+
+
 @functools.singledispatch
 def compile_scalar_expression(
     expression: expression.Expression,
@@ -50,9 +60,12 @@ def compile_constant_expression(
 
 
 @compile_scalar_expression.register
-def compile_op_expression(expr: expression.OpExpression):
+def compile_op_expression(expr: expression.OpExpression) -> sge.Expression:
     # Non-recursively compiles the children scalar expressions.
-    args = tuple(map(compile_scalar_expression, expr.inputs))
+    args = tuple(
+        TypedExpr(compile_scalar_expression(input), input.output_type)
+        for input in expr.inputs
+    )
 
     op = expr.op
     op_name = expr.op.__class__.__name__
@@ -79,8 +92,10 @@ def compile_op_expression(expr: expression.OpExpression):
 
 
 # TODO: add parenthesize for operators
-def compile_addop(
-    op: ops.AddOp, left: sge.Expression, right: sge.Expression
-) -> sge.Expression:
-    # TODO: support addop for string dtype.
-    return sge.Add(this=left, expression=right)
+def compile_addop(op: ops.AddOp, left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.dtype == dtypes.STRING_DTYPE and right.dtype == dtypes.STRING_DTYPE:
+        # String addition
+        return sge.Concat(expressions=[left.expr, right.expr])
+
+    # Numerical addition
+    return sge.Add(this=left.expr, expression=right.expr)

bigframes/core/rewrite/schema_binding.py

@@ -19,7 +19,13 @@
 from bigframes.core import nodes
 
 
-def bind_schema_to_expressions(
+def bind_schema_to_tree(
+    node: bigframe_node.BigFrameNode,
+) -> bigframe_node.BigFrameNode:
+    return nodes.bottom_up(node, bind_schema_to_node)
+
+
+def bind_schema_to_node(
     node: bigframe_node.BigFrameNode,
 ) -> bigframe_node.BigFrameNode:
     if isinstance(node, nodes.ProjectionNode):

tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_limit/out.sql

@@ -6,19 +6,14 @@ WITH `bfcte_0` AS (
     `float64_col` AS `bfcol_3`,
     `bool_col` AS `bfcol_4`
   FROM `test-project`.`test_dataset`.`test_table`
-), `bfcte_1` AS (
-  SELECT
-    *,
-    `bfcol_1` AS `bfcol_5`
-  FROM `bfcte_0`
 )
 SELECT
   `bfcol_0` AS `rowindex`,
   `bfcol_1` AS `int64_col`,
   `bfcol_2` AS `string_col`,
   `bfcol_3` AS `float64_col`,
   `bfcol_4` AS `bool_col`
-FROM `bfcte_1`
+FROM `bfcte_0`
 ORDER BY
-  `bfcol_5` ASC NULLS LAST
+  `bfcol_1` ASC NULLS LAST
 LIMIT 10

tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_ordering/out.sql

@@ -6,35 +6,13 @@ WITH `bfcte_0` AS (
     `float64_col` AS `bfcol_3`,
     `bool_col` AS `bfcol_4`
   FROM `test-project`.`test_dataset`.`test_table`
-), `bfcte_1` AS (
-  SELECT
-    `bfcol_0` AS `bfcol_5`,
-    `bfcol_1` AS `bfcol_6`,
-    `bfcol_2` AS `bfcol_7`,
-    `bfcol_3` AS `bfcol_8`,
-    `bfcol_4` AS `bfcol_9`
-  FROM `bfcte_0`
-), `bfcte_2` AS (
-  SELECT
-    *,
-    `bfcol_5` AS `bfcol_10`
-  FROM `bfcte_1`
-), `bfcte_3` AS (
-  SELECT
-    `bfcol_5` AS `bfcol_11`,
-    `bfcol_6` AS `bfcol_12`,
-    `bfcol_7` AS `bfcol_13`,
-    `bfcol_8` AS `bfcol_14`,
-    `bfcol_9` AS `bfcol_15`,
-    `bfcol_10` AS `bfcol_16`
-  FROM `bfcte_2`
 )
 SELECT
-  `bfcol_11` AS `rowindex`,
-  `bfcol_12` AS `int64_col`,
-  `bfcol_13` AS `string_col`,
-  `bfcol_14` AS `float64_col`,
-  `bfcol_15` AS `bool_col`
-FROM `bfcte_3`
+  `bfcol_0` AS `rowindex`,
+  `bfcol_1` AS `int64_col`,
+  `bfcol_2` AS `string_col`,
+  `bfcol_3` AS `float64_col`,
+  `bfcol_4` AS `bool_col`
+FROM `bfcte_0`
 ORDER BY
-  `bfcol_16` ASC NULLS LAST
+  `bfcol_0` ASC NULLS LAST

tests/unit/core/compile/sqlglot/snapshots/test_compile_scalar_expr/test_compile_numerical_add/out.sql

@@ -0,0 +1,33 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `rowindex` AS `bfcol_0`,
+    `int64_col` AS `bfcol_1`,
+    `string_col` AS `bfcol_2`,
+    `float64_col` AS `bfcol_3`,
+    `bool_col` AS `bfcol_4`
+  FROM `test-project`.`test_dataset`.`test_table`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_0` AS `bfcol_5`,
+    `bfcol_2` AS `bfcol_6`,
+    `bfcol_3` AS `bfcol_7`,
+    `bfcol_4` AS `bfcol_8`,
+    `bfcol_1` + `bfcol_1` AS `bfcol_9`
+  FROM `bfcte_0`
+), `bfcte_2` AS (
+  SELECT
+    `bfcol_5` AS `bfcol_10`,
+    `bfcol_9` AS `bfcol_11`,
+    `bfcol_6` AS `bfcol_12`,
+    `bfcol_7` AS `bfcol_13`,
+    `bfcol_8` AS `bfcol_14`
+  FROM `bfcte_1`
+)
+SELECT
+  `bfcol_10` AS `rowindex`,
+  `bfcol_11` AS `int64_col`,
+  `bfcol_12` AS `string_col`,
+  `bfcol_13` AS `float64_col`,
+  `bfcol_14` AS `bool_col`
+FROM `bfcte_2`

tests/unit/core/compile/sqlglot/snapshots/test_compile_scalar_expr/test_compile_string_add/out.sql

@@ -0,0 +1,33 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `rowindex` AS `bfcol_0`,
+    `int64_col` AS `bfcol_1`,
+    `string_col` AS `bfcol_2`,
+    `float64_col` AS `bfcol_3`,
+    `bool_col` AS `bfcol_4`
+  FROM `test-project`.`test_dataset`.`test_table`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_0` AS `bfcol_5`,
+    `bfcol_1` AS `bfcol_6`,
+    `bfcol_3` AS `bfcol_7`,
+    `bfcol_4` AS `bfcol_8`,
+    CONCAT(`bfcol_2`, 'a') AS `bfcol_9`
+  FROM `bfcte_0`
+), `bfcte_2` AS (
+  SELECT
+    `bfcol_5` AS `bfcol_10`,
+    `bfcol_6` AS `bfcol_11`,
+    `bfcol_9` AS `bfcol_12`,
+    `bfcol_7` AS `bfcol_13`,
+    `bfcol_8` AS `bfcol_14`
+  FROM `bfcte_1`
+)
+SELECT
+  `bfcol_10` AS `rowindex`,
+  `bfcol_11` AS `int64_col`,
+  `bfcol_12` AS `string_col`,
+  `bfcol_13` AS `float64_col`,
+  `bfcol_14` AS `bool_col`
+FROM `bfcte_2`

tests/unit/core/compile/sqlglot/test_compile_scalar_expr.py

@@ -0,0 +1,31 @@
+# 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.
+
+import pytest
+
+import bigframes
+
+pytest.importorskip("pytest_snapshot")
+
+
+def test_compile_numerical_add(compiler_session: bigframes.Session, snapshot):
+    bf_df = compiler_session.read_gbq_table("test-project.test_dataset.test_table")
+    bf_df["int64_col"] = bf_df["int64_col"] + bf_df["int64_col"]
+    snapshot.assert_match(bf_df.sql, "out.sql")
+
+
+def test_compile_string_add(compiler_session: bigframes.Session, snapshot):
+    bf_df = compiler_session.read_gbq_table("test-project.test_dataset.test_table")
+    bf_df["string_col"] = bf_df["string_col"] + "a"
+    snapshot.assert_match(bf_df.sql, "out.sql")