Merge branch 'main' into shuowei-anywidget-cell-visual

t pu# Please enter a commit message to explain why this merge is necessary,

Author: shuoweil

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]]:

bigframes/formatting_helpers.py

@@ -336,7 +336,7 @@ def get_job_url(
     """
     if project_id is None or location is None or job_id is None:
         return None
-    return f"""https://console.cloud. google.com/bigquery?project={project_id}&j=bq:{location}:{job_id}&page=queryresults"""
+    return f"""https://console.cloud.google.com/bigquery?project={project_id}&j=bq:{location}:{job_id}&page=queryresults"""
 
 
 def render_bqquery_sent_event_html(

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

@@ -50,12 +50,6 @@ WITH `bfcte_2` AS (
     *,
     2 AS `bfcol_26`
   FROM `bfcte_5`
-), `bfcte_0` AS (
-  SELECT
-    `bool_col`,
-    `float64_col`,
-    `int64_too`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_3` AS (
   SELECT
     *

tests/unit/test_formatting_helpers.py

@@ -197,3 +197,18 @@ def test_render_bqquery_finished_event_plaintext():
     assert "finished" in text
     assert "1.0 kB processed" in text
     assert "Slot time: 2 seconds" in text
+
+
+def test_get_job_url():
+    job_id = "my-job-id"
+    location = "us-central1"
+    project_id = "my-project"
+    expected_url = (
+        f"https://console.cloud.google.com/bigquery?project={project_id}"
+        f"&j=bq:{location}:{job_id}&page=queryresults"
+    )
+
+    actual_url = formatting_helpers.get_job_url(
+        job_id=job_id, location=location, project_id=project_id
+    )
+    assert actual_url == expected_url