separate logical cte nodes from concrete sql ones

Author: TrevorBergeron

bigframes/core/compile/sqlglot/compiler.py

@@ -106,21 +106,11 @@ 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_obj = compile_node(rewrite.as_sql_nodes(root), uid_gen)
-    sqlglot_ir_obj = sqlglot_ir_obj.with_ctes(
-        tuple(
-            (compile_node(cte_node, uid_gen)._as_select(), cte_node.name)
-            for cte_node in cte_nodes
-        )
-    )
-
     return sqlglot_ir_obj.sql
 
 
@@ -269,16 +259,20 @@ def compile_isin_join(
 
 
 @_compile_node.register
-def compile_cte_ref_node(node: nodes.CteRefNode, child: sqlglot_ir.SQLGlotIR):
+def compile_cte_ref_node(node: sql_nodes.SqlCteRefNode, child: sqlglot_ir.SQLGlotIR):
     return sqlglot_ir.SQLGlotIR.from_cte_ref(
-        node.child.name,  # type: ignore
+        node.cte_name,
         uid_gen=child.uid_gen,
     )
 
 
 @_compile_node.register
-def compile_cte_node(node: nodes.CteNode, child: sqlglot_ir.SQLGlotIR):
-    raise ValueError("CTE definitions should not be directly compiled")
+def compile_with_ctes_node(
+    node: sql_nodes.SqlWithCtesNode,
+    child: sqlglot_ir.SQLGlotIR,
+    *ctes: sqlglot_ir.SQLGlotIR,
+):
+    return child.with_ctes(tuple(zip(node.cte_names, ctes)))
 
 
 @_compile_node.register

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -386,11 +386,11 @@ def aggregate(
 
     def with_ctes(
         self,
-        ctes: tuple[tuple[str, sge.Select], ...],
+        ctes: tuple[tuple[str, SQLGlotIR], ...],
     ) -> SQLGlotIR:
         sge_ctes = [
             sge.CTE(
-                this=cte,
+                this=cte._as_select(),
                 alias=cte_name,
             )
             for cte_name, cte in ctes

bigframes/core/nodes.py

@@ -1713,52 +1713,18 @@ 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
+    """
+    Semantically a no-op, used to indicate shared subtrees and act as optimization boundary.
+    """
 
     @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

bigframes/core/rewrite/as_sql.py

@@ -222,6 +222,26 @@ def _as_sql_node(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         return node
 
 
+def _extract_ctes(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
+    topological_ctes = list(
+        filter(lambda n: isinstance(n, nodes.CteNode), root.iter_nodes_topo())
+    )
+    cte_names = tuple(f"cte_{i}" for i in range(len(topological_ctes)))
+
+    mapping = {
+        cte_node: sql_nodes.SqlCteRefNode(cte_name, tuple(cte_node.fields))
+        for cte_node, cte_name in zip(topological_ctes, cte_names)
+    }
+
+    # Replace all CTEs with CTE references and wrap the new root in a WITH clause
+    return sql_nodes.SqlWithCtesNode(
+        root.top_down(lambda x: mapping.get(x, x)),
+        cte_names,
+        tuple(cte.top_down(lambda x: mapping.get(x, x)) for cte in topological_ctes),
+    )
+
+
 def as_sql_nodes(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
     # TODO: Aggregations, Unions, Joins, raw data sources
+    root = _extract_ctes(root)
     return nodes.bottom_up(root, _as_sql_node)

bigframes/core/rewrite/ctes.py

@@ -15,29 +15,24 @@
 
 from collections import defaultdict
 
-from bigframes.core import expression, nodes
+from bigframes.core import nodes
 
 
 def extract_ctes(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
     # identify candidates
-    # candidates
     node_parents: dict[nodes.BigFrameNode, int] = defaultdict(int)
     for parent in root.unique_nodes():
         for child in parent.child_nodes:
             node_parents[child] += 1
 
     counter = 0
-    # ok time to replace via extract
+
     # we just mark in place, rather than pull out of the tree.
-    # if we did pull out of tree, we'd want to make sure to extract bottom-up
     def insert_cte_markers(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
         nonlocal counter
         if node_parents[node] > 1:
             counter += 1
-            return nodes.CteRefNode(
-                nodes.CteNode(node, name=f"cte_{counter})"),
-                cols=tuple(expression.DerefOp(id) for id in node.ids),
-            )
+            return nodes.CteNode(node)
         return node
 
     return root.top_down(insert_cte_markers)

bigframes/core/rewrite/pruning.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import collections
 import dataclasses
 import functools
 import itertools
@@ -23,39 +22,7 @@
 def column_pruning(
     root: nodes.BigFrameNode,
 ) -> nodes.BigFrameNode:
-    # We wrap the entire process in a fixed-point iteration to ensure
-    # the global push-down through CTEs fully settles.
-    @to_fixed(max_iterations=100)
-    def prune_tree(current_root: nodes.BigFrameNode) -> nodes.BigFrameNode:
-        # Apply local top-down pruning rules (pushes selections to CteRefNodes)
-        pushed_root = nodes.top_down(current_root, prune_columns)
-
-        # Gather the union of required columns globally across all CTE refs
-        cte_reqs: typing.DefaultDict[
-            nodes.BigFrameNode, set[identifiers.ColumnId]
-        ] = collections.defaultdict(set)
-
-        def gather_reqs(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
-            if isinstance(node, nodes.CteRefNode):
-                # node.child is the referenced CTE definition (CteNode)
-                for col in node.cols:
-                    cte_reqs[node.child].add(col.id)
-            return node
-
-        nodes.top_down(pushed_root, gather_reqs)
-
-        # Apply the unioned required columns to the CTE definitions
-        def apply_cte_reqs(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
-            if isinstance(node, nodes.CteNode) and node in cte_reqs:
-                needed_ids = frozenset(cte_reqs[node])
-                pruned_child = prune_node(node.child, needed_ids)
-                if pruned_child is not node.child:
-                    return node.replace_child(pruned_child)
-            return node
-
-        return nodes.top_down(pushed_root, apply_cte_reqs)
-
-    return prune_tree(root)
+    return nodes.top_down(root, prune_columns)
 
 
 def to_fixed(max_iterations: int = 100):
@@ -110,24 +77,6 @@ def prune_selection_child(
             {id: ref.id for ref, id in child.input_output_pairs}
         ).replace_child(child.child)
 
-    elif isinstance(child, nodes.CteRefNode):
-        # Push selection locally into the CTE Reference
-        needed_ids = selection.consumed_ids
-        new_cols = (
-            tuple(col for col in child.cols if col.id in needed_ids) or child.cols[0:1]
-        )
-
-        if new_cols == child.cols:
-            return selection
-        return selection.replace_child(dataclasses.replace(child, cols=new_cols))
-
-    elif isinstance(child, nodes.CteNode):
-        # Pass-through selection to the CTE Definition just in case it's wrapped locally
-        pruned_child = prune_node(child.child, selection.consumed_ids)
-        if pruned_child is child.child:
-            return selection
-        return selection.replace_child(child.replace_child(pruned_child))
-
     elif isinstance(child, nodes.AdditiveNode):
         if not set(field.id for field in child.added_fields) & selection.consumed_ids:
             return selection.replace_child(child.additive_base)

bigframes/core/sql_nodes.py

@@ -16,13 +16,18 @@
 
 import dataclasses
 import functools
-from typing import Mapping, Optional, Sequence, Tuple
+from typing import Callable, Mapping, Optional, Sequence, Tuple
 
 from bigframes.core import bq_data, identifiers, nodes
 import bigframes.core.expression as ex
 from bigframes.core.ordering import OrderingExpression
 import bigframes.dtypes
 
+# SQL Nodes are generally terminal, so don't support rich transformation methods
+# like remap_vars, remap_refs, etc.
+# Still, fields should be defined on them, as typing info is still used for
+# dispatching some operators in the emitter, and for validation.
+
 
 # TODO: Join node, union node
 @dataclasses.dataclass(frozen=True)
@@ -84,6 +89,127 @@ def remap_refs(
         raise NotImplementedError()  # type: ignore
 
 
+@dataclasses.dataclass(frozen=True)
+class SqlWithCtesNode(nodes.BigFrameNode):
+    # def, name pairs
+    child: nodes.BigFrameNode
+    cte_names: tuple[str, ...]
+    cte_defs: tuple[nodes.BigFrameNode, ...]
+
+    @property
+    def child_nodes(self) -> Sequence[nodes.BigFrameNode]:
+        return (self.child, *self.cte_defs)
+
+    @property
+    def fields(self) -> Sequence[nodes.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 defines_namespace(self) -> bool:
+        return True
+
+    @property
+    def explicitly_ordered(self) -> bool:
+        return False
+
+    @property
+    def order_ambiguous(self) -> bool:
+        return True
+
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
+    @property
+    def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
+        return tuple(self.ids)
+
+    @property
+    def consumed_ids(self):
+        return ()
+
+    @property
+    def _node_expressions(self):
+        return ()
+
+    def remap_vars(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> SqlWithCtesNode:
+        raise NotImplementedError()
+
+    def remap_refs(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> SqlWithCtesNode:
+        raise NotImplementedError()  # type: ignore
+
+    def transform_children(
+        self, transform: Callable[[nodes.BigFrameNode], nodes.BigFrameNode]
+    ) -> SqlWithCtesNode:
+        return SqlWithCtesNode(
+            transform(self.child),
+            self.cte_names,
+            tuple(transform(cte) for cte in self.cte_defs),
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class SqlCteRefNode(nodes.LeafNode):
+    cte_name: str
+    cte_schema: tuple[nodes.Field, ...]
+
+    @property
+    def fields(self) -> Sequence[nodes.Field]:
+        return self.cte_schema
+
+    @property
+    def variables_introduced(self) -> int:
+        # This operation only renames variables, doesn't actually create new ones
+        return 0
+
+    @property
+    def defines_namespace(self) -> bool:
+        return True
+
+    @property
+    def explicitly_ordered(self) -> bool:
+        return False
+
+    @property
+    def order_ambiguous(self) -> bool:
+        return True
+
+    @property
+    def row_count(self) -> Optional[int]:
+        raise NotImplementedError()
+
+    @property
+    def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
+        return tuple(self.ids)
+
+    @property
+    def consumed_ids(self):
+        return ()
+
+    @property
+    def _node_expressions(self):
+        return ()
+
+    def remap_vars(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> SqlCteRefNode:
+        raise NotImplementedError()
+
+    def remap_refs(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> SqlCteRefNode:
+        raise NotImplementedError()  # type: ignore
+
+
 @dataclasses.dataclass(frozen=True)
 class SqlSelectNode(nodes.UnaryNode):
     selections: tuple[nodes.ColumnDef, ...] = ()