refactor: fix test_concat_dataframe by de-deuplicating merged ctes

Author: chelsea-lin

bigframes/core/compile/sqlglot/compiler.py

@@ -41,8 +41,6 @@
 def compile_sql(request: configs.CompileRequest) -> configs.CompileResult:
     """Compiles a BigFrameNode according to the request into SQL using SQLGlot."""
 
-    # Generator for unique identifiers.
-    uid_gen = guid.SequentialUIDGenerator()
     output_names = tuple((expression.DerefOp(id), id.sql) for id in request.node.ids)
     result_node = nodes.ResultNode(
         request.node,
@@ -61,22 +59,16 @@ def compile_sql(request: configs.CompileRequest) -> configs.CompileResult:
     )
     if request.sort_rows:
         result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node))
-        result_node = _remap_variables(result_node, uid_gen)
-        result_node = typing.cast(
-            nodes.ResultNode, rewrite.defer_selection(result_node)
-        )
-        sql = _compile_result_node(result_node, uid_gen)
+        sql = _compile_result_node(result_node)
         return configs.CompileResult(
             sql, result_node.schema.to_bigquery(), result_node.order_by
         )
 
     ordering: typing.Optional[bf_ordering.RowOrdering] = result_node.order_by
     result_node = dataclasses.replace(result_node, order_by=None)
     result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node))
+    sql = _compile_result_node(result_node)
 
-    result_node = _remap_variables(result_node, uid_gen)
-    result_node = typing.cast(nodes.ResultNode, rewrite.defer_selection(result_node))
-    sql = _compile_result_node(result_node, uid_gen)
     # Return the ordering iff no extra columns are needed to define the row order
     if ordering is not None:
         output_order = (
@@ -97,11 +89,16 @@ def _remap_variables(
     return typing.cast(nodes.ResultNode, result_node)
 
 
-def _compile_result_node(
-    root: nodes.ResultNode, uid_gen: guid.SequentialUIDGenerator
-) -> str:
+def _compile_result_node(root: nodes.ResultNode) -> str:
+    # Create UIDs to standardize variable names and ensure consistent compilation
+    # of nodes using the same generator.
+    uid_gen = guid.SequentialUIDGenerator()
+    root = _remap_variables(root, uid_gen)
+    root = typing.cast(nodes.ResultNode, rewrite.defer_selection(root))
+
     # Have to bind schema as the final step before compilation.
     root = typing.cast(nodes.ResultNode, schema_binding.bind_schema_to_tree(root))
+
     selected_cols: tuple[tuple[str, sge.Expression], ...] = tuple(
         (name, scalar_compiler.scalar_op_compiler.compile_expression(ref))
         for ref, name in root.output_cols
@@ -127,7 +124,6 @@ def _compile_result_node(
     return sqlglot_ir.sql
 
 
-@functools.lru_cache(maxsize=5000)
 def compile_node(
     node: nodes.BigFrameNode, uid_gen: guid.SequentialUIDGenerator
 ) -> ir.SQLGlotIR:

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -206,7 +206,7 @@ def from_union(
 
             select_expr = select.copy()
             select_expr, select_ctes = _pop_query_ctes(select_expr)
-            existing_ctes = [*existing_ctes, *select_ctes]
+            existing_ctes = _merge_ctes(existing_ctes, select_ctes)
             union_selects.append(select_expr)
 
         union_expr: sge.Query = union_selects[0].subquery()
@@ -337,7 +337,7 @@ def join(
 
         left_select, left_ctes = _pop_query_ctes(left_select)
         right_select, right_ctes = _pop_query_ctes(right_select)
-        merged_ctes = [*left_ctes, *right_ctes]
+        merged_ctes = _merge_ctes(left_ctes, right_ctes)
 
         join_on = _and(
             tuple(
@@ -374,7 +374,7 @@ def isin_join(
 
         left_select, left_ctes = _pop_query_ctes(left_select)
         right_select, right_ctes = _pop_query_ctes(right_select)
-        merged_ctes = [*left_ctes, *right_ctes]
+        merged_ctes = _merge_ctes(left_ctes, right_ctes)
 
         left_condition = typed_expr.TypedExpr(
             sge.Column(this=conditions[0].expr, table=left_cte_name),
@@ -827,6 +827,15 @@ def _set_query_ctes(
     return new_expr
 
 
+def _merge_ctes(ctes1: list[sge.CTE], ctes2: list[sge.CTE]) -> list[sge.CTE]:
+    """Merges two lists of CTEs, de-duplicating by alias name."""
+    seen = {cte.alias: cte for cte in ctes1}
+    for cte in ctes2:
+        if cte.alias not in seen:
+            seen[cte.alias] = cte
+    return list(seen.values())
+
+
 def _pop_query_ctes(
     expr: sge.Select,
 ) -> tuple[sge.Select, list[sge.CTE]]: