fix id issues preventing cte factoring

Author: TrevorBergeron

bigframes/core/bigframe_node.py

@@ -330,22 +330,12 @@ def top_down(
         """
         Perform a top-down transformation of the BigFrameNode tree.
         """
-        to_process = [self]
-        results: Dict[BigFrameNode, BigFrameNode] = {}
 
-        while to_process:
-            item = to_process.pop()
-            if item not in results.keys():
-                item_result = transform(item)
-                results[item] = item_result
-                to_process.extend(item_result.child_nodes)
+        @functools.cache
+        def recursive_transform(node: BigFrameNode) -> BigFrameNode:
+            return transform(node).transform_children(recursive_transform)
 
-        to_process = [self]
-        # for each processed item, replace its children
-        for item in reversed(list(results.keys())):
-            results[item] = results[item].transform_children(lambda x: results[x])
-
-        return results[self]
+        return recursive_transform(self)
 
     def bottom_up(
         self: BigFrameNode,

bigframes/core/compile/sqlglot/compiler.py

@@ -98,6 +98,7 @@ def _remap_variables(
     result_node, _ = rewrite.remap_variables(
         node, map(identifiers.ColumnId, uid_gen.get_uid_stream("bfcol_"))
     )
+    result_node.validate_tree()
     return typing.cast(nodes.ResultNode, result_node)
 
 

bigframes/core/rewrite/ctes.py

@@ -25,10 +25,17 @@ def extract_ctes(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
         for child in parent.child_nodes:
             node_parents[child] += 1
 
-    # we just mark in place, rather than pull out of the tree.
+    # everywhere a multi-parent node is referenced, wrap it in a CTE node
     def insert_cte_markers(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
-        if node_parents[node] > 1:
-            return nodes.CteNode(node)
-        return node
+        def _add_cte_if_needed(child: nodes.BigFrameNode) -> nodes.BigFrameNode:
+            if node_parents[child] > 1:
+                return nodes.CteNode(child)
+            return child
+
+        if isinstance(node, nodes.CteNode):
+            # don't re-wrap CTE nodes
+            return node
+
+        return node.transform_children(_add_cte_if_needed)
 
     return root.top_down(insert_cte_markers)

bigframes/core/rewrite/identifiers.py

@@ -38,13 +38,13 @@ def _create_mapping_operator(
 
     def _mapping_operator(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         # Step 1: Get the local remapping for the current node.
-        local_def_remaps = id_def_remapping_by_node.get(node, {})
-        local_ref_remaps = id_ref_remapping_by_node.get(node, {})
+        local_def_remaps = id_def_remapping_by_node[node]
+        local_ref_remaps = id_ref_remapping_by_node[node]
 
-        node = node.remap_vars(local_def_remaps)
-        node = node.remap_refs(local_ref_remaps)
+        result = node.remap_vars(local_def_remaps)
+        result = result.remap_refs(local_ref_remaps)
 
-        return node
+        return result
 
     return _mapping_operator
 
@@ -80,19 +80,33 @@ def remap_variables(
     id_ref_remaps: dict[
         nodes.BigFrameNode, dict[identifiers.ColumnId, identifiers.ColumnId]
     ] = {}
-    for node in root.iter_nodes_topo():  # bottom_up
+    for node in root.iter_nodes_topo():  # bottom up
         local_def_remaps = {
             col_id: next(id_generator) for col_id in node.node_defined_ids
         }
         id_def_remaps[node] = local_def_remaps
 
         local_ref_remaps = {}
-        for child in node.child_nodes:  # inherit ref and def mappings from children
+
+        # InNode is special case as ID scope inherited purely from left side
+        inheriting_nodes = (
+            [node.child_nodes[0]]
+            if isinstance(node, nodes.InNode)
+            else node.child_nodes
+        )
+        for child in inheriting_nodes:  # inherit ref and def mappings from children
+            if not child.defines_namespace:  # these nodes represent new id spaces
+                local_ref_remaps.update(
+                    {
+                        old_id: new_id
+                        for old_id, new_id in id_ref_remaps[child].items()
+                        if old_id in child.ids
+                    }
+                )
             local_ref_remaps.update(id_def_remaps[child])
-            if not child.defines_namespace:
-                local_ref_remaps.update(id_ref_remaps[child])
         id_ref_remaps[node] = local_ref_remaps
 
+    # have to do top down to preserve node identities
     return (
         root.top_down(_create_mapping_operator(id_def_remaps, id_ref_remaps)),
         id_def_remaps[root],

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

@@ -1,59 +1,48 @@
 WITH `bfcte_0` AS (
-  SELECT
-    `rowindex` AS `bfcol_18`,
-    `rowindex` AS `bfcol_19`,
-    `int64_col` AS `bfcol_20`,
-    `string_col` AS `bfcol_21`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
-), `bfcte_1` AS (
   SELECT
     `rowindex` AS `bfcol_3`,
     `rowindex` AS `bfcol_4`,
     `int64_col` AS `bfcol_5`,
     `string_col` AS `bfcol_6`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
-), `bfcte_2` AS (
-  SELECT
-    *
-  FROM `bfcte_0`
 )
 SELECT
-  `bfcol_30` AS `rowindex`,
-  `bfcol_31` AS `rowindex_1`,
-  `bfcol_32` AS `int64_col`,
-  `bfcol_33` AS `string_col`
+  `bfcol_23` AS `rowindex`,
+  `bfcol_24` AS `rowindex_1`,
+  `bfcol_25` AS `int64_col`,
+  `bfcol_26` AS `string_col`
 FROM (
   SELECT
-    `bfcol_9` AS `bfcol_30`,
-    `bfcol_10` AS `bfcol_31`,
-    `bfcol_11` AS `bfcol_32`,
-    `bfcol_12` AS `bfcol_33`,
-    `bfcol_13` AS `bfcol_34`,
-    `bfcol_14` AS `bfcol_35`
+    `bfcol_17` AS `bfcol_23`,
+    `bfcol_18` AS `bfcol_24`,
+    `bfcol_19` AS `bfcol_25`,
+    `bfcol_20` AS `bfcol_26`,
+    `bfcol_21` AS `bfcol_27`,
+    `bfcol_22` AS `bfcol_28`
   FROM (
     (
       SELECT
-        `bfcol_3` AS `bfcol_9`,
-        `bfcol_4` AS `bfcol_10`,
-        `bfcol_5` AS `bfcol_11`,
-        `bfcol_6` AS `bfcol_12`,
-        0 AS `bfcol_13`,
-        ROW_NUMBER() OVER () - 1 AS `bfcol_14`
-      FROM `bfcte_1`
+        `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
-        `bfcol_18` AS `bfcol_24`,
-        `bfcol_19` AS `bfcol_25`,
-        `bfcol_20` AS `bfcol_26`,
-        `bfcol_21` AS `bfcol_27`,
-        1 AS `bfcol_28`,
-        ROW_NUMBER() OVER () - 1 AS `bfcol_29`
-      FROM `bfcte_2`
+        `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`
     )
   )
 )
 ORDER BY
-  `bfcol_34` ASC NULLS LAST,
-  `bfcol_35` ASC NULLS LAST
+  `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

@@ -1,83 +1,63 @@
 WITH `bfcte_0` AS (
   SELECT
-    `float64_col` AS `bfcol_23`,
-    `int64_col` AS `bfcol_24`
+    `float64_col` AS `bfcol_5`,
+    `int64_col` AS `bfcol_6`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_1` AS (
   SELECT
-    `float64_col` AS `bfcol_2`,
-    `int64_col` AS `bfcol_3`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
-), `bfcte_2` AS (
-  SELECT
-    `float64_col` AS `bfcol_34`,
-    `int64_too` AS `bfcol_35`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
-  WHERE
-    `bool_col`
-), `bfcte_3` AS (
-  SELECT
-    `float64_col` AS `bfcol_13`,
-    `int64_too` AS `bfcol_14`
+    `float64_col` AS `bfcol_7`,
+    `int64_too` AS `bfcol_8`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
   WHERE
     `bool_col`
-), `bfcte_4` AS (
-  SELECT
-    *
-  FROM `bfcte_0`
-), `bfcte_5` AS (
-  SELECT
-    *
-  FROM `bfcte_2`
 )
 SELECT
-  `bfcol_42` AS `float64_col`,
-  `bfcol_43` AS `int64_col`
+  `bfcol_33` AS `float64_col`,
+  `bfcol_34` AS `int64_col`
 FROM (
   SELECT
-    `bfcol_6` AS `bfcol_42`,
-    `bfcol_7` AS `bfcol_43`,
-    `bfcol_8` AS `bfcol_44`,
-    `bfcol_9` AS `bfcol_45`
+    `bfcol_21` AS `bfcol_33`,
+    `bfcol_22` AS `bfcol_34`,
+    `bfcol_23` AS `bfcol_35`,
+    `bfcol_24` AS `bfcol_36`
   FROM (
     (
       SELECT
-        `bfcol_2` AS `bfcol_6`,
-        `bfcol_3` AS `bfcol_7`,
-        0 AS `bfcol_8`,
-        ROW_NUMBER() OVER (ORDER BY `bfcol_3` ASC NULLS LAST) - 1 AS `bfcol_9`
-      FROM `bfcte_1`
+        `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_0`
     )
     UNION ALL
     (
       SELECT
-        `bfcol_13` AS `bfcol_17`,
-        `bfcol_14` AS `bfcol_18`,
-        1 AS `bfcol_19`,
-        ROW_NUMBER() OVER () - 1 AS `bfcol_20`
-      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_1`
     )
     UNION ALL
     (
       SELECT
-        `bfcol_23` AS `bfcol_27`,
-        `bfcol_24` AS `bfcol_28`,
-        2 AS `bfcol_29`,
-        ROW_NUMBER() OVER (ORDER BY `bfcol_24` ASC NULLS LAST) - 1 AS `bfcol_30`
-      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_0`
     )
     UNION ALL
     (
       SELECT
-        `bfcol_34` AS `bfcol_38`,
-        `bfcol_35` AS `bfcol_39`,
-        3 AS `bfcol_40`,
-        ROW_NUMBER() OVER () - 1 AS `bfcol_41`
-      FROM `bfcte_5`
+        `bfcol_7` AS `bfcol_25`,
+        `bfcol_8` AS `bfcol_26`,
+        3 AS `bfcol_27`,
+        ROW_NUMBER() OVER () - 1 AS `bfcol_28`
+      FROM `bfcte_1`
     )
   )
 )
 ORDER BY
-  `bfcol_44` ASC NULLS LAST,
-  `bfcol_45` ASC NULLS LAST
+  `bfcol_35` ASC NULLS LAST,
+  `bfcol_36` ASC NULLS LAST