cte logic, cleanup

Author: TrevorBergeron

bigframes/core/compile/sqlglot/aggregate_compiler.py

@@ -22,8 +22,8 @@
     ordered_unary_compiler,
     unary_compiler,
 )
+import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
 from bigframes.core.compile.sqlglot.expressions import typed_expr
-import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
 
 def compile_aggregate(
@@ -35,7 +35,7 @@ def compile_aggregate(
         return nullary_compiler.compile(aggregate.op)
     if isinstance(aggregate, agg_expressions.UnaryAggregation):
         column = typed_expr.TypedExpr(
-            scalar_compiler.scalar_op_compiler.compile_expression(aggregate.arg),
+            expression_compiler.expression_compiler.compile_expression(aggregate.arg),
             aggregate.arg.output_type,
         )
         if not aggregate.op.order_independent:
@@ -46,11 +46,11 @@ def compile_aggregate(
             return unary_compiler.compile(aggregate.op, column)
     elif isinstance(aggregate, agg_expressions.BinaryAggregation):
         left = typed_expr.TypedExpr(
-            scalar_compiler.scalar_op_compiler.compile_expression(aggregate.left),
+            expression_compiler.expression_compiler.compile_expression(aggregate.left),
             aggregate.left.output_type,
         )
         right = typed_expr.TypedExpr(
-            scalar_compiler.scalar_op_compiler.compile_expression(aggregate.right),
+            expression_compiler.expression_compiler.compile_expression(aggregate.right),
             aggregate.right.output_type,
         )
         return binary_compiler.compile(aggregate.op, left, right)
@@ -66,7 +66,7 @@ def compile_analytic(
         return nullary_compiler.compile(aggregate.op, window)
     if isinstance(aggregate, agg_expressions.UnaryAggregation):
         column = typed_expr.TypedExpr(
-            scalar_compiler.scalar_op_compiler.compile_expression(aggregate.arg),
+            expression_compiler.expression_compiler.compile_expression(aggregate.arg),
             aggregate.arg.output_type,
         )
         return unary_compiler.compile(aggregate.op, column, window)

bigframes/core/compile/sqlglot/aggregations/windows.py

@@ -18,7 +18,7 @@
 import bigframes_vendored.sqlglot.expressions as sge
 
 from bigframes.core import utils, window_spec
-import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
 import bigframes.core.expression as ex
 import bigframes.core.ordering as ordering_spec
 import bigframes.dtypes as dtypes
@@ -116,7 +116,7 @@ def get_window_order_by(
 
     order_by = []
     for ordering_spec_item in ordering:
-        expr = scalar_compiler.scalar_op_compiler.compile_expression(
+        expr = expression_compiler.expression_compiler.compile_expression(
             ordering_spec_item.scalar_expression
         )
         desc = not ordering_spec_item.direction.is_ascending
@@ -191,7 +191,7 @@ def _get_window_bounds(
 
 
 def _compile_group_by_key(key: ex.Expression) -> sge.Expression:
-    expr = scalar_compiler.scalar_op_compiler.compile_expression(key)
+    expr = expression_compiler.expression_compiler.compile_expression(key)
     # The group_by keys has been rewritten by bind_schema_to_node
     assert isinstance(key, ex.ResolvedDerefOp)
 

bigframes/core/compile/sqlglot/compiler.py

@@ -31,8 +31,8 @@
 from bigframes.core.compile import configs
 import bigframes.core.compile.sqlglot.aggregate_compiler as aggregate_compiler
 from bigframes.core.compile.sqlglot.aggregations import windows
+import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
 from bigframes.core.compile.sqlglot.expressions import typed_expr
-import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 from bigframes.core.logging import data_types as data_type_logger
 import bigframes.core.ordering as bf_ordering
@@ -143,38 +143,32 @@ def _compile_node(
 @_compile_node.register
 def compile_sql_select(node: sql_nodes.SqlSelectNode, child: ir.SQLGlotIR):
     sqlglot_ir = child
-    if node.sorting is not None:
-        ordering_cols = tuple(
-            sge.Ordered(
-                this=scalar_compiler.scalar_op_compiler.compile_expression(
-                    ordering.scalar_expression
-                ),
-                desc=ordering.direction.is_ascending is False,
-                nulls_first=ordering.na_last is False,
-            )
-            for ordering in node.sorting
+    ordering_cols = tuple(
+        sge.Ordered(
+            this=expression_compiler.expression_compiler.compile_expression(
+                ordering.scalar_expression
+            ),
+            desc=ordering.direction.is_ascending is False,
+            nulls_first=ordering.na_last is False,
         )
-        sqlglot_ir = sqlglot_ir.order_by(ordering_cols)
+        for ordering in node.sorting
+    )
 
     projected_cols: tuple[tuple[str, sge.Expression], ...] = tuple(
         (
             cdef.id.sql,
-            scalar_compiler.scalar_op_compiler.compile_expression(cdef.expression),
+            expression_compiler.expression_compiler.compile_expression(cdef.expression),
         )
         for cdef in node.selections
     )
     sqlglot_ir = sqlglot_ir.select(projected_cols)
 
-    if len(node.predicates) > 0:
-        sge_predicates = tuple(
-            scalar_compiler.scalar_op_compiler.compile_expression(expression)
-            for expression in node.predicates
-        )
-        sqlglot_ir = sqlglot_ir.filter(sge_predicates)
-    if node.limit is not None:
-        sqlglot_ir = sqlglot_ir.limit(node.limit)
+    sge_predicates = tuple(
+        expression_compiler.expression_compiler.compile_expression(expression)
+        for expression in node.predicates
+    )
 
-    return sqlglot_ir
+    return child.select(projected_cols, sge_predicates, ordering_cols, node.limit)
 
 
 @_compile_node.register
@@ -210,11 +204,11 @@ def compile_join(
     conditions = tuple(
         (
             typed_expr.TypedExpr(
-                scalar_compiler.scalar_op_compiler.compile_expression(left),
+                expression_compiler.expression_compiler.compile_expression(left),
                 left.output_type,
             ),
             typed_expr.TypedExpr(
-                scalar_compiler.scalar_op_compiler.compile_expression(right),
+                expression_compiler.expression_compiler.compile_expression(right),
                 right.output_type,
             ),
         )
@@ -236,11 +230,11 @@ def compile_isin_join(
     right_field = node.right_child.fields[0]
     conditions = (
         typed_expr.TypedExpr(
-            scalar_compiler.scalar_op_compiler.compile_expression(node.left_col),
+            expression_compiler.expression_compiler.compile_expression(node.left_col),
             node.left_col.output_type,
         ),
         typed_expr.TypedExpr(
-            scalar_compiler.scalar_op_compiler.compile_expression(
+            expression_compiler.expression_compiler.compile_expression(
                 expression.DerefOp(right_field.id)
             ),
             right_field.dtype,
@@ -303,7 +297,7 @@ def compile_aggregate(node: nodes.AggregateNode, child: ir.SQLGlotIR) -> ir.SQLG
         for agg, id in node.aggregations
     )
     by_cols: tuple[sge.Expression, ...] = tuple(
-        scalar_compiler.scalar_op_compiler.compile_expression(by_col)
+        expression_compiler.expression_compiler.compile_expression(by_col)
         for by_col in node.by_column_ids
     )
 

bigframes/core/compile/sqlglot/expression_compiler.py

@@ -1,4 +1,4 @@
-# Copyright 2026 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,13 +18,14 @@
 
 import bigframes_vendored.sqlglot.expressions as sge
 
+import bigframes.core.agg_expressions as agg_exprs
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 import bigframes.core.expression as ex
 import bigframes.operations as ops
 
 
-class ScalarOpCompiler:
+class ExpressionCompiler:
     # Mapping of operation name to implemenations
     _registry: dict[
         str,
@@ -78,6 +79,15 @@ def _(self, expr: ex.DerefOp) -> sge.Expression:
     def _(self, expr: ex.ScalarConstantExpression) -> sge.Expression:
         return ir._literal(expr.value, expr.dtype)
 
+    @compile_expression.register
+    def _(self, expr: agg_exprs.WindowExpression) -> sge.Expression:
+        import bigframes.core.compile.sqlglot.aggregate_compiler as agg_compile
+
+        return agg_compile.compile_analytic(
+            expr.analytic_expr,
+            expr.window,
+        )
+
     @compile_expression.register
     def _(self, expr: ex.OpExpression) -> sge.Expression:
         # Non-recursively compiles the children scalar expressions.
@@ -218,4 +228,4 @@ def _add_parentheses(cls, expr: TypedExpr) -> TypedExpr:
 
 
 # Singleton compiler
-scalar_op_compiler = ScalarOpCompiler()
+expression_compiler = ExpressionCompiler()

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

@@ -19,10 +19,10 @@
 import bigframes_vendored.sqlglot.expressions as sge
 
 from bigframes import operations as ops
-from bigframes.core.compile.sqlglot import scalar_compiler
+from bigframes.core.compile.sqlglot import expression_compiler
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 
-register_nary_op = scalar_compiler.scalar_op_compiler.register_nary_op
+register_nary_op = expression_compiler.expression_compiler.register_nary_op
 
 
 @register_nary_op(ops.AIGenerate, pass_op=True)

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

@@ -20,16 +20,16 @@
 import bigframes_vendored.sqlglot.expressions as sge
 
 from bigframes import operations as ops
+import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
 from bigframes.core.compile.sqlglot.expressions.string_ops import (
     string_index,
     string_slice,
 )
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
-import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 import bigframes.dtypes as dtypes
 
-register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
-register_nary_op = scalar_compiler.scalar_op_compiler.register_nary_op
+register_unary_op = expression_compiler.expression_compiler.register_unary_op
+register_nary_op = expression_compiler.expression_compiler.register_nary_op
 
 
 @register_unary_op(ops.ArrayIndexOp, pass_op=True)

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

@@ -17,11 +17,11 @@
 import bigframes_vendored.sqlglot.expressions as sge
 
 from bigframes import operations as ops
+import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
-import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
-register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
-register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
+register_unary_op = expression_compiler.expression_compiler.register_unary_op
+register_binary_op = expression_compiler.expression_compiler.register_binary_op
 
 
 @register_unary_op(ops.obj_fetch_metadata_op)

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

@@ -18,10 +18,10 @@
 
 from bigframes import dtypes
 from bigframes import operations as ops
+import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
-import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
-register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
+register_binary_op = expression_compiler.expression_compiler.register_binary_op
 
 
 @register_binary_op(ops.and_op)

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

@@ -23,11 +23,11 @@
 from bigframes import dtypes
 from bigframes import operations as ops
 from bigframes.core.compile.sqlglot import sqlglot_ir
+import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
-import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
-register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
-register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
+register_unary_op = expression_compiler.expression_compiler.register_unary_op
+register_binary_op = expression_compiler.expression_compiler.register_binary_op
 
 
 @register_unary_op(ops.IsInOp, pass_op=True)

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

@@ -17,10 +17,10 @@
 import bigframes_vendored.sqlglot.expressions as sge
 
 from bigframes import operations as ops
+import bigframes.core.compile.sqlglot.expression_compiler as expression_compiler
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
-import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
-register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_unary_op = expression_compiler.expression_compiler.register_unary_op
 
 
 @register_unary_op(ops.date_op)