wrap more nodes in ctes

Author: TrevorBergeron

bigframes/core/compile/sqlglot/compiler.py

@@ -113,7 +113,8 @@ def _compile_result_node(root: nodes.ResultNode) -> str:
     # Probably, should defer even further
     root = typing.cast(nodes.ResultNode, schema_binding.bind_schema_to_tree(root))
 
-    sqlglot_ir_obj = compile_node(rewrite.as_sql_nodes(root), uid_gen)
+    # TODO: Bake all IDs in tree, stop passing uid_gen to emitters
+    sqlglot_ir_obj = compile_node(rewrite.as_sql_nodes(root, uid_gen), uid_gen)
     return sqlglot_ir_obj.sql
 
 

bigframes/core/rewrite/as_sql.py

@@ -19,6 +19,7 @@
 from bigframes.core import (
     agg_expressions,
     expression,
+    guid,
     identifiers,
     nodes,
     ordering,
@@ -222,31 +223,59 @@ def _as_sql_node(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         return node
 
 
-def _extract_ctes(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
-    topological_ctes = list(
-        filter(lambda n: isinstance(n, nodes.CteNode), root.iter_nodes_topo())
+# In the future, we will have sql nodes for each of these node types.
+_LOGICAL_NODE_TYPES_TO_WRAP = (
+    nodes.ReadLocalNode,
+    nodes.ExplodeNode,
+    nodes.InNode,
+    nodes.AggregateNode,
+    nodes.FromRangeNode,
+    nodes.ConcatNode,
+)
+
+
+def _extract_nodes_as_ctes(
+    root: nodes.BigFrameNode, uid_gen: guid.SequentialUIDGenerator
+) -> nodes.BigFrameNode:
+    # CTE nodes we must wrap, and logical nodes we want to wrap for style
+    topological_extracted_nodes = list(
+        filter(
+            lambda n: isinstance(n, (nodes.CteNode, *_LOGICAL_NODE_TYPES_TO_WRAP)),
+            root.iter_nodes_topo(),
+        )
+    )
+    cte_names = tuple(
+        next(uid_gen.get_uid_stream("bfcte_"))
+        for _ in range(len(topological_extracted_nodes))
     )
-    cte_names = tuple(f"bfcte_{i}" for i in range(len(topological_ctes)))
 
-    if len(topological_ctes) == 0:
+    if len(topological_extracted_nodes) == 0:
         return root
 
     mapping = {
-        cte_node: sql_nodes.SqlCteRefNode(cte_name, tuple(cte_node.fields))
-        for cte_node, cte_name in zip(topological_ctes, cte_names)
+        extracted_node: sql_nodes.SqlCteRefNode(cte_name, tuple(extracted_node.fields))
+        for extracted_node, cte_name in zip(topological_extracted_nodes, cte_names)
     }
 
+    def _maybe_unwrap_cte(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+        """Unwrap CTE nodes, but not logical nodes."""
+        if isinstance(node, nodes.CteNode):
+            return node.child
+        return node
+
     # Replace all CTEs with CTE references and wrap the new root in a WITH clause
     return sql_nodes.SqlWithCtesNode(
         root.top_down(lambda x: mapping.get(x, x)),
         cte_names,
         tuple(
-            cte.child.top_down(lambda x: mapping.get(x, x)) for cte in topological_ctes  # type: ignore
+            _maybe_unwrap_cte(extracted_node).top_down(lambda x: mapping.get(x, x) if x != extracted_node else x) for extracted_node in topological_extracted_nodes  # type: ignore
         ),
     )
 
 
-def as_sql_nodes(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
-    # TODO: Aggregations, Unions, Joins, raw data sources
-    root = _extract_ctes(root)
-    return nodes.bottom_up(root, _as_sql_node)
+def as_sql_nodes(
+    root: nodes.BigFrameNode, uid_gen: guid.SequentialUIDGenerator
+) -> nodes.BigFrameNode:
+    root = _extract_nodes_as_ctes(root, uid_gen)
+    root = nodes.bottom_up(root, _as_sql_node)
+    return root

bigframes/session/bq_caching_executor.py

@@ -308,10 +308,10 @@ def _export_gbq(
 
             ir = sqlglot_ir.SQLGlotIR.from_unparsed_query(sql)
             if spec.if_exists == "append":
-                sql = sg_sql.to_sql(sg_sql.insert(ir.expr, spec.table))
+                sql = sg_sql.to_sql(sg_sql.insert(ir.expr.as_select_all(), spec.table))
             else:  # for "replace"
                 assert spec.if_exists == "replace"
-                sql = sg_sql.to_sql(sg_sql.replace(ir.expr, spec.table))
+                sql = sg_sql.to_sql(sg_sql.replace(ir.expr.as_select_all(), spec.table))
         else:
             dispositions = {
                 "fail": bigquery.WriteDisposition.WRITE_EMPTY,

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_corr/out.sql

@@ -1,6 +1,4 @@
-SELECT
-  `bfcol_2` AS `corr_col`
-FROM (
+WITH `bfcte_0` AS (
   SELECT
     CORR(`int64_col`, `float64_col`) AS `bfcol_2`
   FROM (
@@ -9,4 +7,7 @@ FROM (
       `float64_col`
     FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
   )
-)
+)
+SELECT
+  `bfcol_2` AS `corr_col`
+FROM `bfcte_0`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_cov/out.sql

@@ -1,6 +1,4 @@
-SELECT
-  `bfcol_2` AS `cov_col`
-FROM (
+WITH `bfcte_0` AS (
   SELECT
     COVAR_SAMP(`int64_col`, `float64_col`) AS `bfcol_2`
   FROM (
@@ -9,4 +7,7 @@ FROM (
       `float64_col`
     FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
   )
-)
+)
+SELECT
+  `bfcol_2` AS `cov_col`
+FROM `bfcte_0`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_size/out.sql

@@ -1,11 +1,12 @@
-SELECT
-  `bfcol_32` AS `size`
-FROM (
+WITH `bfcte_0` AS (
   SELECT
     COUNT(1) AS `bfcol_32`
   FROM (
     SELECT
       *
     FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
   )
-)
+)
+SELECT
+  `bfcol_32` AS `size`
+FROM `bfcte_0`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_ordered_unary_compiler/test_array_agg/out.sql

@@ -1,11 +1,12 @@
-SELECT
-  `bfcol_1` AS `int64_col`
-FROM (
+WITH `bfcte_0` AS (
   SELECT
     ARRAY_AGG(`int64_col` IGNORE NULLS ORDER BY `int64_col` IS NULL ASC, `int64_col` ASC) AS `bfcol_1`
   FROM (
     SELECT
       `int64_col`
     FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
   )
-)
+)
+SELECT
+  `bfcol_1` AS `int64_col`
+FROM `bfcte_0`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_ordered_unary_compiler/test_string_agg/out.sql

@@ -1,6 +1,4 @@
-SELECT
-  `bfcol_1` AS `string_col`
-FROM (
+WITH `bfcte_0` AS (
   SELECT
     COALESCE(
       STRING_AGG(`string_col`, ','
@@ -14,4 +12,7 @@ FROM (
       `string_col`
     FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
   )
-)
+)
+SELECT
+  `bfcol_1` AS `string_col`
+FROM `bfcte_0`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all/out.sql

@@ -1,7 +1,4 @@
-SELECT
-  `bfcol_2` AS `bool_col`,
-  `bfcol_3` AS `int64_col`
-FROM (
+WITH `bfcte_0` AS (
   SELECT
     COALESCE(LOGICAL_AND(`bool_col`), TRUE) AS `bfcol_2`,
     COALESCE(LOGICAL_AND(`int64_col` <> 0), TRUE) AS `bfcol_3`
@@ -11,4 +8,8 @@ FROM (
       `int64_col`
     FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
   )
-)
+)
+SELECT
+  `bfcol_2` AS `bool_col`,
+  `bfcol_3` AS `int64_col`
+FROM `bfcte_0`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_any/out.sql

@@ -1,7 +1,4 @@
-SELECT
-  `bfcol_2` AS `bool_col`,
-  `bfcol_3` AS `int64_col`
-FROM (
+WITH `bfcte_0` AS (
   SELECT
     COALESCE(LOGICAL_OR(`bool_col`), FALSE) AS `bfcol_2`,
     COALESCE(LOGICAL_OR(`int64_col` <> 0), FALSE) AS `bfcol_3`
@@ -11,4 +8,8 @@ FROM (
       `int64_col`
     FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
   )
-)
+)
+SELECT
+  `bfcol_2` AS `bool_col`,
+  `bfcol_3` AS `int64_col`
+FROM `bfcte_0`