refactor: Define CTE-related SQL nodes for emitter (#2495)

Author: TrevorBergeron

bigframes/core/bigframe_node.py

@@ -330,22 +330,12 @@ def top_down(
         """
         Perform a top-down transformation of the BigFrameNode tree.
         """
-        to_process = [self]
-        results: Dict[BigFrameNode, BigFrameNode] = {}
 
-        while to_process:
-            item = to_process.pop()
-            if item not in results.keys():
-                item_result = transform(item)
-                results[item] = item_result
-                to_process.extend(item_result.child_nodes)
+        @functools.cache
+        def recursive_transform(node: BigFrameNode) -> BigFrameNode:
+            return transform(node).transform_children(recursive_transform)
 
-        to_process = [self]
-        # for each processed item, replace its children
-        for item in reversed(list(results.keys())):
-            results[item] = results[item].transform_children(lambda x: results[x])
-
-        return results[self]
+        return recursive_transform(self)
 
     def bottom_up(
         self: BigFrameNode,

bigframes/core/compile/sqlglot/compiler.py

@@ -62,6 +62,8 @@ def compile_sql(request: configs.CompileRequest) -> configs.CompileResult:
     if request.sort_rows:
         result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node))
         encoded_type_refs = data_type_logger.encode_type_refs(result_node)
+        # TODO: Extract CTEs earlier
+        result_node = typing.cast(nodes.ResultNode, rewrite.extract_ctes(result_node))
         sql = _compile_result_node(result_node)
         return configs.CompileResult(
             sql,
@@ -74,6 +76,8 @@ def compile_sql(request: configs.CompileRequest) -> configs.CompileResult:
     result_node = dataclasses.replace(result_node, order_by=None)
     result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node))
     encoded_type_refs = data_type_logger.encode_type_refs(result_node)
+    # TODO: Extract CTEs earlier
+    result_node = typing.cast(nodes.ResultNode, rewrite.extract_ctes(result_node))
     sql = _compile_result_node(result_node)
     # Return the ordering iff no extra columns are needed to define the row order
     if ordering is not None:
@@ -94,6 +98,7 @@ def _remap_variables(
     result_node, _ = rewrite.remap_variables(
         node, map(identifiers.ColumnId, uid_gen.get_uid_stream("bfcol_"))
     )
+    result_node.validate_tree()
     return typing.cast(nodes.ResultNode, result_node)
 
 
@@ -102,13 +107,16 @@ def _compile_result_node(root: nodes.ResultNode) -> str:
     # of nodes using the same generator.
     uid_gen = guid.SequentialUIDGenerator()
     root = _remap_variables(root, uid_gen)
+    # Remap variables creates too mayn new
+    # root = rewrite.select_pullup(root, prefer_source_names=False)
     root = typing.cast(nodes.ResultNode, rewrite.defer_selection(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)
+    # TODO: Bake all IDs in tree, stop passing uid_gen to emitters
+    sqlglot_ir_obj = compile_node(rewrite.as_sql_nodes(root, uid_gen), uid_gen)
     return sqlglot_ir_obj.sql
 
 
@@ -121,7 +129,7 @@ def compile_node(
     for current_node in list(node.iter_nodes_topo()):
         if current_node.child_nodes == ():
             # For leaf node, generates a dumpy child to pass the UID generator.
-            child_results = tuple([sqlglot_ir.SQLGlotIR(uid_gen=uid_gen)])
+            child_results = tuple([sqlglot_ir.SQLGlotIR.empty(uid_gen=uid_gen)])
         else:
             # Child nodes should have been compiled in the reverse topological order.
             child_results = tuple(
@@ -256,6 +264,23 @@ def compile_isin_join(
     )
 
 
+@_compile_node.register
+def compile_cte_ref_node(node: sql_nodes.SqlCteRefNode, child: sqlglot_ir.SQLGlotIR):
+    return sqlglot_ir.SQLGlotIR.from_cte_ref(
+        node.cte_name,
+        uid_gen=child.uid_gen,
+    )
+
+
+@_compile_node.register
+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
 def compile_concat(
     node: nodes.ConcatNode, *children: sqlglot_ir.SQLGlotIR
@@ -271,7 +296,7 @@ def compile_concat(
     ]
 
     return sqlglot_ir.SQLGlotIR.from_union(
-        [child._as_select() for child in children],
+        [child.expr.as_select_all() for child in children],
         output_aliases=output_aliases,
         uid_gen=uid_gen,
     )

bigframes/core/compile/sqlglot/expressions/comparison_ops.py

@@ -33,27 +33,39 @@
 @register_unary_op(ops.IsInOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.IsInOp) -> sge.Expression:
     values = []
+    # bools are not comparable to non-bools in SQL, so we need to cast the expression to INT64 if the values contain non-bools.
+    must_upcast_bools = dtypes.is_numeric(expr.dtype, include_bool=False) or any(
+        dtypes.is_numeric(dtypes.bigframes_type(type(value)), include_bool=False)
+        for value in op.values
+        if not _is_null(value)
+    )
     for value in op.values:
         if _is_null(value):
             continue
         dtype = dtypes.bigframes_type(type(value))
         if dtypes.can_compare(expr.dtype, dtype):
+            if must_upcast_bools and dtype == dtypes.BOOL_DTYPE:
+                value = int(value)
             values.append(sge.convert(value))
 
+    sg_lexpr: sge.Expression = expr.expr
+    if expr.dtype == dtypes.BOOL_DTYPE and must_upcast_bools:
+        sg_lexpr = sge.cast(expr.expr, "INT64")
+
     if op.match_nulls:
         contains_nulls = any(_is_null(value) for value in op.values)
         if contains_nulls:
             if len(values) == 0:
-                return sge.Is(this=expr.expr, expression=sge.Null())
-            return sge.Is(this=expr.expr, expression=sge.Null()) | sge.In(
-                this=expr.expr, expressions=values
+                return sge.Is(this=sg_lexpr, expression=sge.Null())
+            return sge.Is(this=sg_lexpr, expression=sge.Null()) | sge.In(
+                this=sg_lexpr, expressions=values
             )
 
     if len(values) == 0:
         return sge.convert(False)
 
     return sge.func(
-        "COALESCE", sge.In(this=expr.expr, expressions=values), sge.convert(False)
+        "COALESCE", sge.In(this=sg_lexpr, expressions=values), sge.convert(False)
     )