dont extract union all children as ctes

Author: TrevorBergeron

bigframes/core/rewrite/as_sql.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import dataclasses
+import itertools
 from typing import Optional, Sequence, Union
 
 from bigframes.core import (
@@ -240,11 +241,23 @@ def _insert_cte_markers(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
     wrapped_nodes = set(
         node.child for node in root.unique_nodes() if isinstance(node, nodes.CteNode)
     )
+    # don't wrap child nodes of ConcatNode
+    union_child_nodes = set(
+        itertools.chain.from_iterable(
+            node.child_nodes
+            for node in root.unique_nodes()
+            if isinstance(node, nodes.ConcatNode)
+        )
+    )
 
     def maybe_insert_cte_marker(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         if node == root:
             return node
-        if isinstance(node, _LOGICAL_NODE_TYPES_TO_WRAP) and node not in wrapped_nodes:
+        if (
+            isinstance(node, _LOGICAL_NODE_TYPES_TO_WRAP)
+            and node not in wrapped_nodes
+            and node not in union_child_nodes
+        ):
             wrapped_nodes.add(node)
             return nodes.CteNode(node)
         return node

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

@@ -6,24 +6,6 @@ WITH `bfcte_0` AS (
     `string_col` AS `bfcol_6`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_1` AS (
-  SELECT
-    `bfcol_3` AS `bfcol_11`,
-    `bfcol_4` AS `bfcol_12`,
-    `bfcol_5` AS `bfcol_13`,
-    `bfcol_6` AS `bfcol_14`,
-    1 AS `bfcol_15`,
-    ROW_NUMBER() OVER () - 1 AS `bfcol_16`
-  FROM `bfcte_0`
-), `bfcte_2` AS (
-  SELECT
-    `bfcol_3` AS `bfcol_17`,
-    `bfcol_4` AS `bfcol_18`,
-    `bfcol_5` AS `bfcol_19`,
-    `bfcol_6` AS `bfcol_20`,
-    0 AS `bfcol_21`,
-    ROW_NUMBER() OVER () - 1 AS `bfcol_22`
-  FROM `bfcte_0`
-), `bfcte_3` AS (
   SELECT
     `bfcol_17` AS `bfcol_23`,
     `bfcol_18` AS `bfcol_24`,
@@ -34,14 +16,24 @@ WITH `bfcte_0` AS (
   FROM (
     (
       SELECT
-        *
-      FROM `bfcte_2`
+        `bfcol_3` AS `bfcol_17`,
+        `bfcol_4` AS `bfcol_18`,
+        `bfcol_5` AS `bfcol_19`,
+        `bfcol_6` AS `bfcol_20`,
+        0 AS `bfcol_21`,
+        ROW_NUMBER() OVER () - 1 AS `bfcol_22`
+      FROM `bfcte_0`
     )
     UNION ALL
     (
       SELECT
-        *
-      FROM `bfcte_1`
+        `bfcol_3` AS `bfcol_11`,
+        `bfcol_4` AS `bfcol_12`,
+        `bfcol_5` AS `bfcol_13`,
+        `bfcol_6` AS `bfcol_14`,
+        1 AS `bfcol_15`,
+        ROW_NUMBER() OVER () - 1 AS `bfcol_16`
+      FROM `bfcte_0`
     )
   )
 )
@@ -50,7 +42,7 @@ SELECT
   `bfcol_24` AS `rowindex_1`,
   `bfcol_25` AS `int64_col`,
   `bfcol_26` AS `string_col`
-FROM `bfcte_3`
+FROM `bfcte_1`
 ORDER BY
   `bfcol_27` ASC NULLS LAST,
   `bfcol_28` ASC NULLS LAST

tests/unit/core/compile/sqlglot/snapshots/test_compile_concat/test_compile_concat_filter_sorted/out.sql

@@ -11,34 +11,6 @@ WITH `bfcte_0` AS (
     `int64_col` AS `bfcol_6`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_2` AS (
-  SELECT
-    `bfcol_7` AS `bfcol_25`,
-    `bfcol_8` AS `bfcol_26`,
-    3 AS `bfcol_27`,
-    ROW_NUMBER() OVER () - 1 AS `bfcol_28`
-  FROM `bfcte_0`
-), `bfcte_3` AS (
-  SELECT
-    `bfcol_7` AS `bfcol_29`,
-    `bfcol_8` AS `bfcol_30`,
-    1 AS `bfcol_31`,
-    ROW_NUMBER() OVER () - 1 AS `bfcol_32`
-  FROM `bfcte_0`
-), `bfcte_4` AS (
-  SELECT
-    `bfcol_5` AS `bfcol_17`,
-    `bfcol_6` AS `bfcol_18`,
-    2 AS `bfcol_19`,
-    ROW_NUMBER() OVER (ORDER BY `bfcol_6` ASC NULLS LAST) - 1 AS `bfcol_20`
-  FROM `bfcte_1`
-), `bfcte_5` AS (
-  SELECT
-    `bfcol_5` AS `bfcol_21`,
-    `bfcol_6` AS `bfcol_22`,
-    0 AS `bfcol_23`,
-    ROW_NUMBER() OVER (ORDER BY `bfcol_6` ASC NULLS LAST) - 1 AS `bfcol_24`
-  FROM `bfcte_1`
-), `bfcte_6` AS (
   SELECT
     `bfcol_21` AS `bfcol_33`,
     `bfcol_22` AS `bfcol_34`,
@@ -47,33 +19,45 @@ WITH `bfcte_0` AS (
   FROM (
     (
       SELECT
-        *
-      FROM `bfcte_5`
+        `bfcol_5` AS `bfcol_21`,
+        `bfcol_6` AS `bfcol_22`,
+        0 AS `bfcol_23`,
+        ROW_NUMBER() OVER (ORDER BY `bfcol_6` ASC NULLS LAST) - 1 AS `bfcol_24`
+      FROM `bfcte_1`
     )
     UNION ALL
     (
       SELECT
-        *
-      FROM `bfcte_3`
+        `bfcol_7` AS `bfcol_29`,
+        `bfcol_8` AS `bfcol_30`,
+        1 AS `bfcol_31`,
+        ROW_NUMBER() OVER () - 1 AS `bfcol_32`
+      FROM `bfcte_0`
     )
     UNION ALL
     (
       SELECT
-        *
-      FROM `bfcte_4`
+        `bfcol_5` AS `bfcol_17`,
+        `bfcol_6` AS `bfcol_18`,
+        2 AS `bfcol_19`,
+        ROW_NUMBER() OVER (ORDER BY `bfcol_6` ASC NULLS LAST) - 1 AS `bfcol_20`
+      FROM `bfcte_1`
     )
     UNION ALL
     (
       SELECT
-        *
-      FROM `bfcte_2`
+        `bfcol_7` AS `bfcol_25`,
+        `bfcol_8` AS `bfcol_26`,
+        3 AS `bfcol_27`,
+        ROW_NUMBER() OVER () - 1 AS `bfcol_28`
+      FROM `bfcte_0`
     )
   )
 )
 SELECT
   `bfcol_33` AS `float64_col`,
   `bfcol_34` AS `int64_col`
-FROM `bfcte_6`
+FROM `bfcte_2`
 ORDER BY
   `bfcol_35` ASC NULLS LAST,
   `bfcol_36` ASC NULLS LAST