refactor: Add cte factoring to new compiler

Author: TrevorBergeron

bigframes/core/compile/sqlglot/compiler.py

@@ -48,6 +48,8 @@ def compile_sql(request: configs.CompileRequest) -> configs.CompileResult:
         output_cols=output_names,
         limit=request.peek_count,
     )
+    # Extract CTEs early, as later rewriters could otherwise make common subtrees diverge
+    result_node = typing.cast(nodes.ResultNode, rewrite.extract_ctes(result_node))
     if request.sort_rows:
         # Can only pullup slice if we are doing ORDER BY in outermost SELECT
         # Need to do this before replacing unsupported ops, as that will rewrite slice ops
@@ -103,11 +105,21 @@ def _compile_result_node(root: nodes.ResultNode) -> str:
     root = _remap_variables(root, uid_gen)
     root = typing.cast(nodes.ResultNode, rewrite.defer_selection(root))
 
+    # TODO: Extract out CTEs to a with_ctes node?
+    cte_nodes = _get_ctes(root)
+
     # Have to bind schema as the final step before compilation.
     # Probably, should defer even further
     root = typing.cast(nodes.ResultNode, schema_binding.bind_schema_to_tree(root))
 
     sqlglot_ir = compile_node(rewrite.as_sql_nodes(root), uid_gen)
+    sqlglot_ir = sqlglot_ir.with_ctes(
+        tuple(
+            (compile_node(cte_node, uid_gen)._as_select(), cte_node.name)
+            for cte_node in cte_nodes
+        )
+    )
+
     return sqlglot_ir.sql
 
 
@@ -247,6 +259,31 @@ def compile_isin_join(
     )
 
 
+@_compile_node.register
+def compile_cte_node(node: nodes.CteRefNode, _)
+    table = node.name
+    return ir.SQLGlotIR.from_table(
+        table.project_id,
+        table.dataset_id,
+        table.table_id,
+        uid_gen=child.uid_gen,
+        sql_predicate=node.source.sql_predicate,
+        system_time=node.source.at_time,
+    )
+
+@_compile_node.register
+def compile_cte_node(node: nodes.CteNode, _)
+    table = node.source.table
+    return ir.SQLGlotIR.from_table(
+        table.project_id,
+        table.dataset_id,
+        table.table_id,
+        uid_gen=child.uid_gen,
+        sql_predicate=node.source.sql_predicate,
+        system_time=node.source.at_time,
+    )
+
+
 @_compile_node.register
 def compile_concat(node: nodes.ConcatNode, *children: ir.SQLGlotIR) -> ir.SQLGlotIR:
     assert len(children) >= 1
@@ -312,3 +349,16 @@ def _replace_unsupported_ops(node: nodes.BigFrameNode):
     node = nodes.bottom_up(node, rewrite.rewrite_slice)
     node = nodes.bottom_up(node, rewrite.rewrite_range_rolling)
     return node
+
+
+def _get_ctes(root: nodes.ResultNode) -> typing.Sequence[nodes.CteNode]:
+    """
+    Get ctes from plan in topological order.
+    """
+
+    def merge_list(node, cte_list):
+        if isinstance(node, nodes.CteNode):
+            return (*cte_list, node)
+        return cte_list
+
+    return root.reduce_up(merge_list)

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -165,7 +165,7 @@ def select(
     ) -> SQLGlotIR:
         # TODO: Explicitly insert CTEs into plan
         if isinstance(self.expr, sge.Select):
-            new_expr, _ = self._select_to_cte()
+            new_expr, _ = self._as_from_item()
         else:
             new_expr = sge.Select().from_(self.expr)
 
@@ -222,21 +222,8 @@ def from_union(
         assert (
             len(list(selects)) >= 2
         ), f"At least two select expressions must be provided, but got {selects}."
-
-        existing_ctes: list[sge.CTE] = []
-        union_selects: list[sge.Select] = []
-        for select in selects:
-            assert isinstance(
-                select, sge.Select
-            ), f"All provided expressions must be of type sge.Select, but got {type(select)}"
-
-            select_expr = select.copy()
-            select_expr, select_ctes = _pop_query_ctes(select_expr)
-            existing_ctes = _merge_ctes(existing_ctes, select_ctes)
-            union_selects.append(select_expr)
-
-        union_expr: sge.Query = union_selects[0].subquery()
-        for select in union_selects[1:]:
+        union_expr: sge.Query = selects[0].subquery()
+        for select in selects[1:]:
             union_expr = sge.Union(
                 this=union_expr,
                 expression=select.subquery(),
@@ -254,7 +241,6 @@ def from_union(
         final_select_expr = (
             sge.Select().select(*selections).from_(union_expr.subquery())
         )
-        final_select_expr = _set_query_ctes(final_select_expr, existing_ctes)
         return cls(expr=final_select_expr, uid_gen=uid_gen)
 
     def join(
@@ -266,12 +252,8 @@ def join(
         joins_nulls: bool = True,
     ) -> SQLGlotIR:
         """Joins the current query with another SQLGlotIR instance."""
-        left_select, left_cte_name = self._select_to_cte()
-        right_select, right_cte_name = right._select_to_cte()
-
-        left_select, left_ctes = _pop_query_ctes(left_select)
-        right_select, right_ctes = _pop_query_ctes(right_select)
-        merged_ctes = _merge_ctes(left_ctes, right_ctes)
+        left_from = self._as_from_item()
+        right_from = right._as_from_item()
 
         join_on = _and(
             tuple(
@@ -283,10 +265,9 @@ def join(
         new_expr = (
             sge.Select()
             .select(sge.Star())
-            .from_(sge.Table(this=left_cte_name))
-            .join(sge.Table(this=right_cte_name), on=join_on, join_type=join_type_str)
+            .from_(left_from)
+            .join(right_from, on=join_on, join_type=join_type_str)
         )
-        new_expr = _set_query_ctes(new_expr, merged_ctes)
 
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
@@ -298,16 +279,12 @@ def isin_join(
         joins_nulls: bool = True,
     ) -> SQLGlotIR:
         """Joins the current query with another SQLGlotIR instance."""
-        left_select, left_cte_name = self._select_to_cte()
+        left_from = self._as_from_item()
         # Prefer subquery over CTE for the IN clause's right side to improve SQL readability.
         right_select = right._as_select()
 
-        left_select, left_ctes = _pop_query_ctes(left_select)
-        right_select, right_ctes = _pop_query_ctes(right_select)
-        merged_ctes = _merge_ctes(left_ctes, right_ctes)
-
         left_condition = typed_expr.TypedExpr(
-            sge.Column(this=conditions[0].expr, table=left_cte_name),
+            sge.Column(this=conditions[0].expr, table=left_from),
             conditions[0].dtype,
         )
 
@@ -341,10 +318,9 @@ def isin_join(
 
         new_expr = (
             sge.Select()
-            .select(sge.Column(this=sge.Star(), table=left_cte_name), new_column)
-            .from_(sge.Table(this=left_cte_name))
+            .select(sge.Column(this=sge.Star(), table=left_from), new_column)
+            .from_(left_from)
         )
-        new_expr = _set_query_ctes(new_expr, merged_ctes)
 
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
@@ -368,7 +344,7 @@ def sample(self, fraction: float) -> SQLGlotIR:
             expression=_literal(fraction, dtypes.FLOAT_DTYPE),
         )
 
-        new_expr = self._select_to_cte()[0].where(condition, append=False)
+        new_expr = self._as_select().where(condition, append=False)
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
     def aggregate(
@@ -392,7 +368,7 @@ def aggregate(
             for id, expr in aggregations
         ]
 
-        new_expr, _ = self._select_to_cte()
+        new_expr = self._as_select()
         new_expr = new_expr.group_by(*by_cols).select(
             *[*by_cols, *aggregations_expr], append=False
         )
@@ -407,12 +383,26 @@ def aggregate(
             new_expr = new_expr.where(condition, append=False)
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
+    def with_ctes(
+        self,
+        ctes: tuple[tuple[str, sge.Select], ...],
+    ) -> SQLGlotIR:
+        sge_ctes = [
+            sge.CTE(
+                this=cte,
+                alias=cte_name,
+            )
+            for cte_name, cte in ctes
+        ]
+        select_expr = _set_query_ctes(self._as_select(), sge_ctes)
+        return SQLGlotIR(expr=select_expr, uid_gen=self.uid_gen)
+
     def insert(
         self,
         destination: bigquery.TableReference,
     ) -> str:
         """Generates an INSERT INTO SQL statement from the current SELECT clause."""
-        return sge.insert(self._as_from_item(), _table(destination)).sql(
+        return sge.insert(self._as_select(), _table(destination)).sql(
             dialect=self.dialect, pretty=self.pretty
         )
 
@@ -436,7 +426,7 @@ def replace(
 
         merge_str = sge.Merge(
             this=_table(destination),
-            using=self._as_from_item(),
+            using=self._as_select(),
             on=_literal(False, dtypes.BOOL_DTYPE),
         ).sql(dialect=self.dialect, pretty=self.pretty)
         return f"{merge_str}\n{whens_str}"
@@ -459,7 +449,7 @@ def _explode_single_column(
         )
         selection = sge.Star(replace=[unnested_column_alias.as_(column)])
 
-        new_expr, _ = self._select_to_cte()
+        new_expr = self._as_select()
         # Use LEFT JOIN to preserve rows when unnesting empty arrays.
         new_expr = new_expr.select(selection, append=False).join(
             unnest_expr, join_type="LEFT"
@@ -510,7 +500,7 @@ def _explode_multiple_columns(
                 for column in columns
             ]
         )
-        new_expr, _ = self._select_to_cte()
+        new_expr = self._as_select()
         # Use LEFT JOIN to preserve rows when unnesting empty arrays.
         new_expr = new_expr.select(selection, append=False).join(
             unnest_expr, join_type="LEFT"
@@ -532,25 +522,6 @@ def _as_select(self) -> sge.Select:
     def _as_subquery(self) -> sge.Subquery:
         return self._as_select().subquery()
 
-    def _select_to_cte(self) -> tuple[sge.Select, sge.Identifier]:
-        """Transforms a given sge.Select query by pushing its main SELECT statement
-        into a new CTE and then generates a 'SELECT * FROM new_cte_name'
-        for the new query."""
-        cte_name = sge.to_identifier(
-            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
-        )
-        select_expr = self._as_select().copy()
-        select_expr, existing_ctes = _pop_query_ctes(select_expr)
-        new_cte = sge.CTE(
-            this=select_expr,
-            alias=cte_name,
-        )
-        new_select_expr = (
-            sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
-        )
-        new_select_expr = _set_query_ctes(new_select_expr, [*existing_ctes, new_cte])
-        return new_select_expr, cte_name
-
 
 def _is_null_literal(expr: sge.Expression) -> bool:
     """Checks if the given expression is a NULL literal."""
@@ -743,26 +714,3 @@ def _set_query_ctes(
     else:
         raise ValueError("The expression does not support 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]]:
-    """Pops the CTEs of a given sge.Select expression."""
-    if "with" in expr.arg_types.keys():
-        expr_ctes = expr.args.pop("with", [])
-        return expr, expr_ctes
-    elif "with_" in expr.arg_types.keys():
-        expr_ctes = expr.args.pop("with_", [])
-        return expr, expr_ctes
-    else:
-        raise ValueError("The expression does not support CTEs.")

bigframes/core/nodes.py

@@ -1715,6 +1715,73 @@ def _node_expressions(self):
         return tuple(ref for ref, _ in self.output_cols)
 
 
+@dataclasses.dataclass(frozen=True, eq=False)
+class CteRefNode(UnaryNode):
+    cols: tuple[ex.DerefOp, ...]
+
+    @property
+    def fields(self) -> Sequence[Field]:
+        # Fields property here is for output schema, not to be consumed by a parent node.
+        input_fields_by_id = {field.id: field for field in self.child.fields}
+        return tuple(input_fields_by_id[ref.id] for ref in self.cols)
+
+    @property
+    def variables_introduced(self) -> int:
+        # This operation only renames variables, doesn't actually create new ones
+        return 0
+
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
+    @property
+    def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
+        return ()
+
+    def remap_vars(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> CteRefNode:
+        return self
+
+    def remap_refs(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> CteRefNode:
+        new_cols = tuple(id.remap_column_refs(mappings) for id in self.cols)
+        return dataclasses.replace(self, cols=new_cols)
+
+
+@dataclasses.dataclass(frozen=True, eq=False)
+class CteNode(UnaryNode):
+    name: str
+
+    @property
+    def fields(self) -> Sequence[Field]:
+        return self.child.fields
+
+    @property
+    def variables_introduced(self) -> int:
+        # This operation only renames variables, doesn't actually create new ones
+        return 0
+
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
+    @property
+    def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
+        return ()
+
+    def remap_vars(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> CteNode:
+        return self
+
+    def remap_refs(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> CteNode:
+        return self
+
+
 # Tree operators
 def top_down(
     root: BigFrameNode,

bigframes/core/rewrite/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from bigframes.core.rewrite.as_sql import as_sql_nodes
+from bigframes.core.rewrite.ctes import extract_ctes
 from bigframes.core.rewrite.fold_row_count import fold_row_counts
 from bigframes.core.rewrite.identifiers import remap_variables
 from bigframes.core.rewrite.implicit_align import try_row_join
@@ -34,6 +35,7 @@
 
 __all__ = [
     "as_sql_nodes",
+    "extract_ctes",
     "legacy_join_as_projection",
     "try_row_join",
     "rewrite_slice",