refactor: implements compile_selection method (#1672)

Author: chelsea-lin

bigframes/core/compile/sqlglot/compiler.py

@@ -20,9 +20,11 @@
 
 from google.cloud import bigquery
 import pyarrow as pa
+import sqlglot.expressions as sge
 
 from bigframes.core import expression, identifiers, nodes, rewrite
 from bigframes.core.compile import configs
+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
 
@@ -38,7 +40,7 @@ def compile(
         ordered: bool = True,
         limit: typing.Optional[int] = None,
     ) -> str:
-        """Compile node into sql where rows are sorted with ORDER BY."""
+        """Compiles node into sql where rows are sorted with ORDER BY."""
         request = configs.CompileRequest(node, sort_rows=ordered, peek_count=limit)
         return self._compile_sql(request).sql
 
@@ -48,7 +50,7 @@ def compile_raw(
     ) -> typing.Tuple[
         str, typing.Sequence[bigquery.SchemaField], bf_ordering.RowOrdering
     ]:
-        """Compile node into sql that exposes all columns, including hidden
+        """Compiles node into sql that exposes all columns, including hidden
         ordering-only columns."""
         request = configs.CompileRequest(
             node, sort_rows=False, materialize_all_order_keys=True
@@ -163,6 +165,9 @@ def compile_readlocal(node: nodes.ReadLocalNode, *args) -> ir.SQLGlotIR:
 
 
 @_compile_node.register
-def compile_selection(node: nodes.SelectionNode, child: ir.SQLGlotIR):
-    # TODO: add support for selection
-    return child
+def compile_selection(node: nodes.SelectionNode, child: ir.SQLGlotIR) -> ir.SQLGlotIR:
+    select_cols: typing.Dict[str, sge.Expression] = {
+        id.name: scalar_compiler.compile_scalar_expression(expr)
+        for expr, id in node.input_output_pairs
+    }
+    return child.select(select_cols)

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -0,0 +1,33 @@
+# 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 functools
+
+import sqlglot.expressions as sge
+
+from bigframes.core import expression
+
+
+@functools.singledispatch
+def compile_scalar_expression(
+    expression: expression.Expression,
+) -> sge.Expression:
+    """Compiles BigFrames scalar expression into SQLGlot expression."""
+    raise ValueError(f"Can't compile unrecognized node: {expression}")
+
+
+@compile_scalar_expression.register
+def compile_deref_op(expr: expression.DerefOp):
+    return sge.ColumnDef(this=sge.to_identifier(expr.id.sql, quoted=True))

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -30,7 +30,7 @@
 class SQLGlotIR:
     """Helper class to build SQLGlot Query and generate SQL string."""
 
-    expr: sge.Expression = sge.Expression()
+    expr: sge.Select = sg.select()
     """The SQLGlot expression representing the query."""
 
     dialect = sqlglot.dialects.bigquery.BigQuery
@@ -90,6 +90,20 @@ def from_pandas(
         )
         return cls(expr=sg.select(sge.Star()).from_(expr))
 
+    def select(
+        self,
+        select_cols: typing.Dict[str, sge.Expression],
+    ) -> SQLGlotIR:
+        selected_cols = [
+            sge.Alias(
+                this=expr,
+                alias=sge.to_identifier(id, quoted=self.quoted),
+            )
+            for id, expr in select_cols.items()
+        ]
+        expr = self.expr.select(*selected_cols, append=False)
+        return SQLGlotIR(expr=expr)
+
 
 def _literal(value: typing.Any, dtype: str) -> sge.Expression:
     if value is None:

tests/unit/core/compile/sqlglot/snapshots/test_compile_readlocal/test_compile_readlocal/out.sql

@@ -1,3 +1,7 @@
 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 UNNEST(ARRAY<STRUCT<`bfcol_0` INT64, `bfcol_1` INT64, `bfcol_2` BOOLEAN, `bfcol_3` STRING, `bfcol_4` INT64>>[(1, -10, TRUE, 'b', 0), (2, 20, CAST(NULL AS BOOLEAN), 'aa', 1), (3, 30, FALSE, 'ccc', 2)])