clean up select vs table expr cases in ir

Author: TrevorBergeron

bigframes/core/compile/sqlglot/compiler.py

@@ -108,7 +108,6 @@ def _compile_result_node(root: nodes.ResultNode) -> str:
     root = typing.cast(nodes.ResultNode, schema_binding.bind_schema_to_tree(root))
 
     sqlglot_ir = compile_node(rewrite.as_sql_nodes(root), uid_gen)
-    print(sqlglot_ir.sql)
     return sqlglot_ir.sql
 
 
@@ -261,7 +260,7 @@ def compile_concat(node: nodes.ConcatNode, *children: ir.SQLGlotIR) -> ir.SQLGlo
     ]
 
     return ir.SQLGlotIR.from_union(
-        [child.expr for child in children],
+        [child._as_select() for child in children],
         output_aliases=output_aliases,
         uid_gen=uid_gen,
     )

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -44,7 +44,7 @@
 class SQLGlotIR:
     """Helper class to build SQLGlot Query and generate SQL string."""
 
-    expr: sge.Select = sg.select()
+    expr: typing.Union[sge.Select, sge.Table] = sg.select()
     """The SQLGlot expression representing the query."""
 
     dialect = sg.dialects.bigquery.BigQuery
@@ -163,15 +163,9 @@ def select(
         sorting: tuple[sge.Ordered, ...] = (),
         limit: typing.Optional[int] = None,
     ) -> SQLGlotIR:
-
         # TODO: Explicitly insert CTEs into plan
         if isinstance(self.expr, sge.Select):
-            new_expr = _select_to_cte(
-                self.expr,
-                sge.to_identifier(
-                    next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-                ),
-            )
+            new_expr, _ = self._select_to_cte()
         else:
             new_expr = sge.Select().from_(self.expr)
 
@@ -272,15 +266,8 @@ def join(
         joins_nulls: bool = True,
     ) -> SQLGlotIR:
         """Joins the current query with another SQLGlotIR instance."""
-        left_cte_name = sge.to_identifier(
-            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-        )
-        right_cte_name = sge.to_identifier(
-            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-        )
-
-        left_select = _select_to_cte(self.expr, left_cte_name)
-        right_select = _select_to_cte(right.expr, right_cte_name)
+        left_select, left_cte_name = self._select_to_cte()
+        right_select, right_cte_name = self._select_to_cte()
 
         left_select, left_ctes = _pop_query_ctes(left_select)
         right_select, right_ctes = _pop_query_ctes(right_select)
@@ -311,13 +298,9 @@ def isin_join(
         joins_nulls: bool = True,
     ) -> SQLGlotIR:
         """Joins the current query with another SQLGlotIR instance."""
-        left_cte_name = sge.to_identifier(
-            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-        )
-
-        left_select = _select_to_cte(self.expr, left_cte_name)
+        left_select, left_cte_name = self._select_to_cte()
         # Prefer subquery over CTE for the IN clause's right side to improve SQL readability.
-        right_select = right.expr
+        right_select = right._as_select()
 
         left_select, left_ctes = _pop_query_ctes(left_select)
         right_select, right_ctes = _pop_query_ctes(right_select)
@@ -380,21 +363,12 @@ def explode(
 
     def sample(self, fraction: float) -> SQLGlotIR:
         """Uniform samples a fraction of the rows."""
-        uuid_col = sge.to_identifier(
-            next(self.uid_gen.get_uid_stream("bfcol_")), quoted=self.quoted
-        )
-        uuid_expr = sge.Alias(this=sge.func("RAND"), alias=uuid_col)
         condition = sge.LT(
-            this=uuid_col,
+            this=sge.func("RAND"),
             expression=_literal(fraction, dtypes.FLOAT_DTYPE),
         )
 
-        new_cte_name = sge.to_identifier(
-            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-        )
-        new_expr = _select_to_cte(
-            self.expr.select(uuid_expr, append=True), new_cte_name
-        ).where(condition, append=False)
+        new_expr = self._select_to_cte()[0].where(condition, append=False)
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
     def aggregate(
@@ -418,12 +392,7 @@ def aggregate(
             for id, expr in aggregations
         ]
 
-        new_expr = _select_to_cte(
-            self.expr,
-            sge.to_identifier(
-                next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-            ),
-        )
+        new_expr, _ = self._select_to_cte()
         new_expr = new_expr.group_by(*by_cols).select(
             *[*by_cols, *aggregations_expr], append=False
         )
@@ -443,7 +412,7 @@ def insert(
         destination: bigquery.TableReference,
     ) -> str:
         """Generates an INSERT INTO SQL statement from the current SELECT clause."""
-        return sge.insert(self.expr.subquery(), _table(destination)).sql(
+        return sge.insert(self._as_from_item(), _table(destination)).sql(
             dialect=self.dialect, pretty=self.pretty
         )
 
@@ -467,7 +436,7 @@ def replace(
 
         merge_str = sge.Merge(
             this=_table(destination),
-            using=self.expr.subquery(),
+            using=self._as_from_item(),
             on=_literal(False, dtypes.BOOL_DTYPE),
         ).sql(dialect=self.dialect, pretty=self.pretty)
         return f"{merge_str}\n{whens_str}"
@@ -490,12 +459,7 @@ def _explode_single_column(
         )
         selection = sge.Star(replace=[unnested_column_alias.as_(column)])
 
-        new_expr = _select_to_cte(
-            self.expr,
-            sge.to_identifier(
-                next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-            ),
-        )
+        new_expr, _ = self._select_to_cte()
         # Use LEFT JOIN to preserve rows when unnesting empty arrays.
         new_expr = new_expr.select(selection, append=False).join(
             unnest_expr, join_type="LEFT"
@@ -546,32 +510,46 @@ def _explode_multiple_columns(
                 for column in columns
             ]
         )
-        new_expr = _select_to_cte(
-            self.expr,
-            sge.to_identifier(
-                next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-            ),
-        )
+        new_expr, _ = self._select_to_cte()
         # Use LEFT JOIN to preserve rows when unnesting empty arrays.
         new_expr = new_expr.select(selection, append=False).join(
             unnest_expr, join_type="LEFT"
         )
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
+    def _as_from_item(self) -> typing.Union[sge.Table, sge.Subquery]:
+        if isinstance(self.expr, sge.Select):
+            return self.expr.subquery()
+        else:  # table
+            return self.expr
 
-def _select_to_cte(expr: sge.Select, cte_name: sge.Identifier) -> sge.Select:
-    """Transforms a given sge.Select query by pushing its main SELECT statement
-    into a new CTE and then generates a 'SELECT * FROM new_cte_name'
-    for the new query."""
-    select_expr = expr.copy()
-    select_expr, existing_ctes = _pop_query_ctes(select_expr)
-    new_cte = sge.CTE(
-        this=select_expr,
-        alias=cte_name,
-    )
-    new_select_expr = sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
-    new_select_expr = _set_query_ctes(new_select_expr, [*existing_ctes, new_cte])
-    return new_select_expr
+    def _as_select(self) -> sge.Select:
+        if isinstance(self.expr, sge.Select):
+            return self.expr
+        else:  # table
+            return sge.Select().from_(self.expr)
+
+    def _as_subquery(self) -> sge.Subquery:
+        return self._as_select().subquery()
+
+    def _select_to_cte(self) -> tuple[sge.Select, sge.Identifier]:
+        """Transforms a given sge.Select query by pushing its main SELECT statement
+        into a new CTE and then generates a 'SELECT * FROM new_cte_name'
+        for the new query."""
+        cte_name = sge.to_identifier(
+            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
+        )
+        select_expr = self.expr._as_select().copy()
+        select_expr, existing_ctes = _pop_query_ctes(select_expr)
+        new_cte = sge.CTE(
+            this=select_expr,
+            alias=cte_name,
+        )
+        new_select_expr = (
+            sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
+        )
+        new_select_expr = _set_query_ctes(new_select_expr, [*existing_ctes, new_cte])
+        return new_select_expr, cte_name
 
 
 def _is_null_literal(expr: sge.Expression) -> bool:

bigframes/core/sql_nodes.py

@@ -91,12 +91,6 @@ class SqlSelectNode(nodes.UnaryNode):
     sorting: tuple[OrderingExpression, ...] = ()
     limit: Optional[int] = None
 
-    def __post_init__(self):
-        try:
-            self.fields
-        except Exception:
-            ...
-
     @functools.cached_property
     def fields(self) -> Sequence[nodes.Field]:
         fields = []

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

@@ -5,8 +5,8 @@ WITH `bfcte_0` AS (
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    `int64_col` AS `bfcol_6`,
-    `int64_too` AS `bfcol_7`
+    `rowindex` AS `bfcol_2`,
+    `int64_col` AS `bfcol_3`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_2` AS (
   SELECT

tests/unit/core/compile/sqlglot/snapshots/test_compile_join/test_compile_join_w_on/bool_col/out.sql

@@ -5,8 +5,8 @@ WITH `bfcte_0` AS (
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    `rowindex` AS `bfcol_6`,
-    `bool_col` AS `bfcol_7`
+    `rowindex` AS `bfcol_2`,
+    `bool_col` AS `bfcol_3`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_2` AS (
   SELECT

tests/unit/core/compile/sqlglot/snapshots/test_compile_join/test_compile_join_w_on/float64_col/out.sql

@@ -5,8 +5,8 @@ WITH `bfcte_0` AS (
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    `rowindex` AS `bfcol_6`,
-    `float64_col` AS `bfcol_7`
+    `rowindex` AS `bfcol_2`,
+    `float64_col` AS `bfcol_3`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_2` AS (
   SELECT

tests/unit/core/compile/sqlglot/snapshots/test_compile_join/test_compile_join_w_on/int64_col/out.sql

@@ -5,8 +5,8 @@ WITH `bfcte_0` AS (
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    `rowindex` AS `bfcol_6`,
-    `int64_col` AS `bfcol_7`
+    `rowindex` AS `bfcol_2`,
+    `int64_col` AS `bfcol_3`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_2` AS (
   SELECT

tests/unit/core/compile/sqlglot/snapshots/test_compile_join/test_compile_join_w_on/numeric_col/out.sql

@@ -5,8 +5,8 @@ WITH `bfcte_0` AS (
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    `rowindex` AS `bfcol_6`,
-    `numeric_col` AS `bfcol_7`
+    `rowindex` AS `bfcol_2`,
+    `numeric_col` AS `bfcol_3`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_2` AS (
   SELECT

tests/unit/core/compile/sqlglot/snapshots/test_compile_join/test_compile_join_w_on/string_col/out.sql

@@ -5,8 +5,8 @@ WITH `bfcte_0` AS (
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    `rowindex` AS `bfcol_4`,
-    `string_col` AS `bfcol_5`
+    `rowindex` AS `bfcol_0`,
+    `string_col` AS `bfcol_1`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_2` AS (
   SELECT

tests/unit/core/compile/sqlglot/snapshots/test_compile_join/test_compile_join_w_on/time_col/out.sql

@@ -5,8 +5,8 @@ WITH `bfcte_0` AS (
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    `rowindex` AS `bfcol_4`,
-    `time_col` AS `bfcol_5`
+    `rowindex` AS `bfcol_0`,
+    `time_col` AS `bfcol_1`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_2` AS (
   SELECT