non-recursively implement remap_variables

Author: chelsea-lin

bigframes/core/rewrite/identifiers.py

@@ -30,28 +30,47 @@ def remap_variables(
 
     Note: this will convert a DAG to a tree.
     """
-    child_replacement_map = dict()
-    ref_mapping = dict()
+
+    def _remap_node(
+        node: nodes.BigFrameNode,
+        child_results: typing.Tuple[
+            typing.Tuple[
+                nodes.BigFrameNode,
+                dict[identifiers.ColumnId, identifiers.ColumnId],
+            ],
+            ...,
+        ],
+    ) -> typing.Tuple[
+        nodes.BigFrameNode,
+        dict[identifiers.ColumnId, identifiers.ColumnId],
+    ]:
+        child_replacement_map = {
+            original_child: result[0]
+            for original_child, result in zip(node.child_nodes, child_results)
+        }
+        ref_mapping = {}
+        for _, child_var_mapping in child_results:
+            ref_mapping.update(child_var_mapping)
+
+        # This is actually invalid until we've replaced all of children, refs and var defs
+        with_new_children = node.transform_children(
+            lambda child_node: child_replacement_map[child_node]
+        )
+
+        with_new_refs = with_new_children.remap_refs(ref_mapping)
+
+        node_var_mapping = {
+            old_id: next(id_generator) for old_id in node.node_defined_ids
+        }
+        with_new_vars = with_new_refs.remap_vars(node_var_mapping)
+        with_new_vars._validate()
+
+        return (
+            with_new_vars,
+            node_var_mapping
+            if node.defines_namespace
+            else (ref_mapping | node_var_mapping),
+        )
+
     # Sequential ids are assigned bottom-up left-to-right
-    for child in root.child_nodes:
-        new_child, child_var_mapping = remap_variables(child, id_generator=id_generator)
-        child_replacement_map[child] = new_child
-        ref_mapping.update(child_var_mapping)
-
-    # This is actually invalid until we've replaced all of children, refs and var defs
-    with_new_children = root.transform_children(
-        lambda node: child_replacement_map[node]
-    )
-
-    with_new_refs = with_new_children.remap_refs(ref_mapping)
-
-    node_var_mapping = {old_id: next(id_generator) for old_id in root.node_defined_ids}
-    with_new_vars = with_new_refs.remap_vars(node_var_mapping)
-    with_new_vars._validate()
-
-    return (
-        with_new_vars,
-        node_var_mapping
-        if root.defines_namespace
-        else (ref_mapping | node_var_mapping),
-    )
+    return root.reduce_up(_remap_node)

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

@@ -1,78 +1,78 @@
-WITH `bfcte_1` AS (
+WITH `bfcte_0` AS (
   SELECT
     *
   FROM UNNEST(ARRAY<STRUCT<`bfcol_0` INT64, `bfcol_1` INT64, `bfcol_2` INT64, `bfcol_3` STRING, `bfcol_4` INT64>>[STRUCT(0, 123456789, 0, 'Hello, World!', 0), STRUCT(1, -987654321, 1, 'こんにちは', 1), STRUCT(2, 314159, 2, '  ¡Hola Mundo!  ', 2), STRUCT(3, CAST(NULL AS INT64), 3, CAST(NULL AS STRING), 3), STRUCT(4, -234892, 4, 'Hello, World!', 4), STRUCT(5, 55555, 5, 'Güten Tag!', 5), STRUCT(6, 101202303, 6, 'capitalize, This ', 6), STRUCT(7, -214748367, 7, ' سلام', 7), STRUCT(8, 2, 8, 'T', 8)])
-), `bfcte_3` AS (
+), `bfcte_2` AS (
   SELECT
     *,
-    `bfcol_4` AS `bfcol_10`
-  FROM `bfcte_1`
-), `bfcte_5` AS (
+    `bfcol_4` AS `bfcol_11`
+  FROM `bfcte_0`
+), `bfcte_4` AS (
   SELECT
     *,
-    0 AS `bfcol_16`
-  FROM `bfcte_3`
+    0 AS `bfcol_23`
+  FROM `bfcte_2`
 ), `bfcte_6` AS (
   SELECT
-    `bfcol_0` AS `bfcol_17`,
-    `bfcol_2` AS `bfcol_18`,
-    `bfcol_1` AS `bfcol_19`,
-    `bfcol_3` AS `bfcol_20`,
-    `bfcol_16` AS `bfcol_21`,
-    `bfcol_10` AS `bfcol_22`
-  FROM `bfcte_5`
-), `bfcte_0` AS (
+    `bfcol_0` AS `bfcol_30`,
+    `bfcol_2` AS `bfcol_31`,
+    `bfcol_1` AS `bfcol_32`,
+    `bfcol_3` AS `bfcol_33`,
+    `bfcol_23` AS `bfcol_34`,
+    `bfcol_11` AS `bfcol_35`
+  FROM `bfcte_4`
+), `bfcte_1` AS (
   SELECT
     *
-  FROM UNNEST(ARRAY<STRUCT<`bfcol_23` INT64, `bfcol_24` INT64, `bfcol_25` INT64, `bfcol_26` STRING, `bfcol_27` INT64>>[STRUCT(0, 123456789, 0, 'Hello, World!', 0), STRUCT(1, -987654321, 1, 'こんにちは', 1), STRUCT(2, 314159, 2, '  ¡Hola Mundo!  ', 2), STRUCT(3, CAST(NULL AS INT64), 3, CAST(NULL AS STRING), 3), STRUCT(4, -234892, 4, 'Hello, World!', 4), STRUCT(5, 55555, 5, 'Güten Tag!', 5), STRUCT(6, 101202303, 6, 'capitalize, This ', 6), STRUCT(7, -214748367, 7, ' سلام', 7), STRUCT(8, 2, 8, 'T', 8)])
-), `bfcte_2` AS (
+  FROM UNNEST(ARRAY<STRUCT<`bfcol_0` INT64, `bfcol_1` INT64, `bfcol_2` INT64, `bfcol_3` STRING, `bfcol_4` INT64>>[STRUCT(0, 123456789, 0, 'Hello, World!', 0), STRUCT(1, -987654321, 1, 'こんにちは', 1), STRUCT(2, 314159, 2, '  ¡Hola Mundo!  ', 2), STRUCT(3, CAST(NULL AS INT64), 3, CAST(NULL AS STRING), 3), STRUCT(4, -234892, 4, 'Hello, World!', 4), STRUCT(5, 55555, 5, 'Güten Tag!', 5), STRUCT(6, 101202303, 6, 'capitalize, This ', 6), STRUCT(7, -214748367, 7, ' سلام', 7), STRUCT(8, 2, 8, 'T', 8)])
+), `bfcte_3` AS (
   SELECT
     *,
-    `bfcol_27` AS `bfcol_33`
-  FROM `bfcte_0`
-), `bfcte_4` AS (
+    `bfcol_4` AS `bfcol_10`
+  FROM `bfcte_1`
+), `bfcte_5` AS (
   SELECT
     *,
-    1 AS `bfcol_39`
-  FROM `bfcte_2`
+    1 AS `bfcol_22`
+  FROM `bfcte_3`
 ), `bfcte_7` AS (
   SELECT
-    `bfcol_23` AS `bfcol_40`,
-    `bfcol_25` AS `bfcol_41`,
-    `bfcol_24` AS `bfcol_42`,
-    `bfcol_26` AS `bfcol_43`,
-    `bfcol_39` AS `bfcol_44`,
-    `bfcol_33` AS `bfcol_45`
-  FROM `bfcte_4`
+    `bfcol_0` AS `bfcol_24`,
+    `bfcol_2` AS `bfcol_25`,
+    `bfcol_1` AS `bfcol_26`,
+    `bfcol_3` AS `bfcol_27`,
+    `bfcol_22` AS `bfcol_28`,
+    `bfcol_10` AS `bfcol_29`
+  FROM `bfcte_5`
 ), `bfcte_8` AS (
   SELECT
     *
   FROM (
     SELECT
-      `bfcol_17` AS `bfcol_46`,
-      `bfcol_18` AS `bfcol_47`,
-      `bfcol_19` AS `bfcol_48`,
-      `bfcol_20` AS `bfcol_49`,
-      `bfcol_21` AS `bfcol_50`,
-      `bfcol_22` AS `bfcol_51`
+      `bfcol_30` AS `bfcol_36`,
+      `bfcol_31` AS `bfcol_37`,
+      `bfcol_32` AS `bfcol_38`,
+      `bfcol_33` AS `bfcol_39`,
+      `bfcol_34` AS `bfcol_40`,
+      `bfcol_35` AS `bfcol_41`
     FROM `bfcte_6`
     UNION ALL
     SELECT
-      `bfcol_40` AS `bfcol_46`,
-      `bfcol_41` AS `bfcol_47`,
-      `bfcol_42` AS `bfcol_48`,
-      `bfcol_43` AS `bfcol_49`,
-      `bfcol_44` AS `bfcol_50`,
-      `bfcol_45` AS `bfcol_51`
+      `bfcol_24` AS `bfcol_36`,
+      `bfcol_25` AS `bfcol_37`,
+      `bfcol_26` AS `bfcol_38`,
+      `bfcol_27` AS `bfcol_39`,
+      `bfcol_28` AS `bfcol_40`,
+      `bfcol_29` AS `bfcol_41`
     FROM `bfcte_7`
   )
 )
 SELECT
-  `bfcol_46` AS `rowindex`,
-  `bfcol_47` AS `rowindex_1`,
-  `bfcol_48` AS `int64_col`,
-  `bfcol_49` AS `string_col`
+  `bfcol_36` AS `rowindex`,
+  `bfcol_37` AS `rowindex_1`,
+  `bfcol_38` AS `int64_col`,
+  `bfcol_39` AS `string_col`
 FROM `bfcte_8`
 ORDER BY
-  `bfcol_50` ASC NULLS LAST,
-  `bfcol_51` ASC NULLS LAST
+  `bfcol_40` ASC NULLS LAST,
+  `bfcol_41` ASC NULLS LAST

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

@@ -1,32 +1,32 @@
 WITH `bfcte_1` AS (
   SELECT
-    `int64_col` AS `bfcol_0`,
-    `rowindex` AS `bfcol_1`
+    `int64_col` AS `bfcol_2`,
+    `rowindex` AS `bfcol_3`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_2` AS (
   SELECT
-    `bfcol_1` AS `bfcol_2`,
-    `bfcol_0` AS `bfcol_3`
+    `bfcol_3` AS `bfcol_6`,
+    `bfcol_2` AS `bfcol_7`
   FROM `bfcte_1`
 ), `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_4`,
-    `int64_too` AS `bfcol_5`
+    `int64_col` AS `bfcol_0`,
+    `int64_too` AS `bfcol_1`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_3` AS (
   SELECT
-    `bfcol_4` AS `bfcol_6`,
-    `bfcol_5` AS `bfcol_7`
+    `bfcol_0` AS `bfcol_4`,
+    `bfcol_1` AS `bfcol_5`
   FROM `bfcte_0`
 ), `bfcte_4` AS (
   SELECT
     *
   FROM `bfcte_2`
   LEFT JOIN `bfcte_3`
-    ON COALESCE(`bfcol_2`, 0) = COALESCE(`bfcol_6`, 0)
-    AND COALESCE(`bfcol_2`, 1) = COALESCE(`bfcol_6`, 1)
+    ON COALESCE(`bfcol_6`, 0) = COALESCE(`bfcol_4`, 0)
+    AND COALESCE(`bfcol_6`, 1) = COALESCE(`bfcol_4`, 1)
 )
 SELECT
-  `bfcol_3` AS `int64_col`,
-  `bfcol_7` AS `int64_too`
+  `bfcol_7` AS `int64_col`,
+  `bfcol_5` AS `int64_too`
 FROM `bfcte_4`

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

@@ -10,24 +10,24 @@ WITH `bfcte_1` AS (
   FROM `bfcte_1`
 ), `bfcte_0` AS (
   SELECT
-    `bool_col` AS `bfcol_4`,
-    `rowindex` AS `bfcol_5`
+    `bool_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_3` AS (
   SELECT
-    `bfcol_5` AS `bfcol_6`,
-    `bfcol_4` AS `bfcol_7`
+    `bfcol_1` AS `bfcol_4`,
+    `bfcol_0` AS `bfcol_5`
   FROM `bfcte_0`
 ), `bfcte_4` AS (
   SELECT
     *
   FROM `bfcte_2`
   INNER JOIN `bfcte_3`
-    ON COALESCE(CAST(`bfcol_3` AS STRING), '0') = COALESCE(CAST(`bfcol_7` AS STRING), '0')
-    AND COALESCE(CAST(`bfcol_3` AS STRING), '1') = COALESCE(CAST(`bfcol_7` AS STRING), '1')
+    ON COALESCE(CAST(`bfcol_3` AS STRING), '0') = COALESCE(CAST(`bfcol_5` AS STRING), '0')
+    AND COALESCE(CAST(`bfcol_3` AS STRING), '1') = COALESCE(CAST(`bfcol_5` AS STRING), '1')
 )
 SELECT
   `bfcol_2` AS `rowindex_x`,
   `bfcol_3` AS `bool_col`,
-  `bfcol_6` AS `rowindex_y`
+  `bfcol_4` AS `rowindex_y`
 FROM `bfcte_4`

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

@@ -10,24 +10,24 @@ WITH `bfcte_1` AS (
   FROM `bfcte_1`
 ), `bfcte_0` AS (
   SELECT
-    `float64_col` AS `bfcol_4`,
-    `rowindex` AS `bfcol_5`
+    `float64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_3` AS (
   SELECT
-    `bfcol_5` AS `bfcol_6`,
-    `bfcol_4` AS `bfcol_7`
+    `bfcol_1` AS `bfcol_4`,
+    `bfcol_0` AS `bfcol_5`
   FROM `bfcte_0`
 ), `bfcte_4` AS (
   SELECT
     *
   FROM `bfcte_2`
   INNER JOIN `bfcte_3`
-    ON IF(IS_NAN(`bfcol_3`), 2, COALESCE(`bfcol_3`, 0)) = IF(IS_NAN(`bfcol_7`), 2, COALESCE(`bfcol_7`, 0))
-    AND IF(IS_NAN(`bfcol_3`), 3, COALESCE(`bfcol_3`, 1)) = IF(IS_NAN(`bfcol_7`), 3, COALESCE(`bfcol_7`, 1))
+    ON IF(IS_NAN(`bfcol_3`), 2, COALESCE(`bfcol_3`, 0)) = IF(IS_NAN(`bfcol_5`), 2, COALESCE(`bfcol_5`, 0))
+    AND IF(IS_NAN(`bfcol_3`), 3, COALESCE(`bfcol_3`, 1)) = IF(IS_NAN(`bfcol_5`), 3, COALESCE(`bfcol_5`, 1))
 )
 SELECT
   `bfcol_2` AS `rowindex_x`,
   `bfcol_3` AS `float64_col`,
-  `bfcol_6` AS `rowindex_y`
+  `bfcol_4` AS `rowindex_y`
 FROM `bfcte_4`

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

@@ -10,24 +10,24 @@ WITH `bfcte_1` AS (
   FROM `bfcte_1`
 ), `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_4`,
-    `rowindex` AS `bfcol_5`
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_3` AS (
   SELECT
-    `bfcol_5` AS `bfcol_6`,
-    `bfcol_4` AS `bfcol_7`
+    `bfcol_1` AS `bfcol_4`,
+    `bfcol_0` AS `bfcol_5`
   FROM `bfcte_0`
 ), `bfcte_4` AS (
   SELECT
     *
   FROM `bfcte_2`
   INNER JOIN `bfcte_3`
-    ON COALESCE(`bfcol_3`, 0) = COALESCE(`bfcol_7`, 0)
-    AND COALESCE(`bfcol_3`, 1) = COALESCE(`bfcol_7`, 1)
+    ON COALESCE(`bfcol_3`, 0) = COALESCE(`bfcol_5`, 0)
+    AND COALESCE(`bfcol_3`, 1) = COALESCE(`bfcol_5`, 1)
 )
 SELECT
   `bfcol_2` AS `rowindex_x`,
   `bfcol_3` AS `int64_col`,
-  `bfcol_6` AS `rowindex_y`
+  `bfcol_4` AS `rowindex_y`
 FROM `bfcte_4`

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

@@ -1,4 +1,4 @@
-WITH `bfcte_1` AS (
+WITH `bfcte_0` AS (
   SELECT
     `numeric_col` AS `bfcol_0`,
     `rowindex` AS `bfcol_1`
@@ -7,27 +7,27 @@ WITH `bfcte_1` AS (
   SELECT
     `bfcol_1` AS `bfcol_2`,
     `bfcol_0` AS `bfcol_3`
-  FROM `bfcte_1`
-), `bfcte_0` AS (
+  FROM `bfcte_0`
+), `bfcte_1` AS (
   SELECT
-    `numeric_col` AS `bfcol_4`,
-    `rowindex` AS `bfcol_5`
+    `numeric_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_3` AS (
   SELECT
-    `bfcol_5` AS `bfcol_6`,
-    `bfcol_4` AS `bfcol_7`
-  FROM `bfcte_0`
+    `bfcol_1` AS `bfcol_4`,
+    `bfcol_0` AS `bfcol_5`
+  FROM `bfcte_1`
 ), `bfcte_4` AS (
   SELECT
     *
   FROM `bfcte_2`
   INNER JOIN `bfcte_3`
-    ON COALESCE(`bfcol_3`, CAST(0 AS NUMERIC)) = COALESCE(`bfcol_7`, CAST(0 AS NUMERIC))
-    AND COALESCE(`bfcol_3`, CAST(1 AS NUMERIC)) = COALESCE(`bfcol_7`, CAST(1 AS NUMERIC))
+    ON COALESCE(`bfcol_3`, CAST(0 AS NUMERIC)) = COALESCE(`bfcol_5`, CAST(0 AS NUMERIC))
+    AND COALESCE(`bfcol_3`, CAST(1 AS NUMERIC)) = COALESCE(`bfcol_5`, CAST(1 AS NUMERIC))
 )
 SELECT
   `bfcol_2` AS `rowindex_x`,
   `bfcol_3` AS `numeric_col`,
-  `bfcol_6` AS `rowindex_y`
+  `bfcol_4` AS `rowindex_y`
 FROM `bfcte_4`