refactor: add compile_window to the sqlglot compiler

Author: chelsea-lin

bigframes/core/compile/sqlglot/aggregate_compiler.py

@@ -15,7 +15,7 @@
 
 import sqlglot.expressions as sge
 
-from bigframes.core import expression
+from bigframes.core import expression, window_spec
 from bigframes.core.compile.sqlglot.aggregations import (
     binary_compiler,
     nullary_compiler,
@@ -56,3 +56,21 @@ def compile_aggregate(
         return binary_compiler.compile(aggregate.op, left, right)
     else:
         raise ValueError(f"Unexpected aggregation: {aggregate}")
+
+
+def compile_analytic(
+    aggregate: expression.Aggregation,
+    window: window_spec.WindowSpec,
+) -> sge.Expression:
+    if isinstance(aggregate, expression.NullaryAggregation):
+        return nullary_compiler.compile(aggregate.op)
+    if isinstance(aggregate, expression.UnaryAggregation):
+        column = typed_expr.TypedExpr(
+            scalar_compiler.compile_scalar_expression(aggregate.arg),
+            aggregate.arg.output_type,
+        )
+        return unary_compiler.compile(aggregate.op, column, window)
+    elif isinstance(aggregate, expression.BinaryAggregation):
+        raise NotImplementedError("binary analytic operations not yet supported")
+    else:
+        raise ValueError(f"Unexpected analytic operation: {aggregate}")

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

@@ -20,7 +20,7 @@
 
 from bigframes.core import window_spec
 import bigframes.core.compile.sqlglot.aggregations.op_registration as reg
-from bigframes.core.compile.sqlglot.aggregations.utils import apply_window_if_present
+from bigframes.core.compile.sqlglot.aggregations.window import apply_window_if_present
 from bigframes.operations import aggregations as agg_ops
 
 NULLARY_OP_REGISTRATION = reg.OpRegistration()

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

@@ -18,9 +18,10 @@
 
 import sqlglot.expressions as sge
 
+from bigframes import dtypes
 from bigframes.core import window_spec
 import bigframes.core.compile.sqlglot.aggregations.op_registration as reg
-from bigframes.core.compile.sqlglot.aggregations.utils import apply_window_if_present
+from bigframes.core.compile.sqlglot.aggregations.window import apply_window_if_present
 import bigframes.core.compile.sqlglot.expressions.typed_expr as typed_expr
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 from bigframes.operations import aggregations as agg_ops
@@ -42,8 +43,11 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=column.expr, to="INT64")
     # Will be null if all inputs are null. Pandas defaults to zero sum though.
-    expr = apply_window_if_present(sge.func("SUM", column.expr), window)
+    expr = apply_window_if_present(sge.func("SUM", expr), window)
     return sge.func("IFNULL", expr, ir._literal(0, column.dtype))
 
 

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

@@ -0,0 +1,148 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+from __future__ import annotations
+
+import typing
+
+import sqlglot.expressions as sge
+
+from bigframes.core import utils, window_spec
+import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
+import bigframes.core.ordering as ordering_spec
+
+
+def apply_window_if_present(
+    value: sge.Expression,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    if window is None:
+        return value
+
+    kind = (
+        "ROWS" if isinstance(window.bounds, window_spec.RowsWindowBounds) else "RANGE"
+    )
+
+    start: typing.Union[int, float, None] = None
+    end: typing.Union[int, float, None] = None
+    if isinstance(window.bounds, window_spec.RangeWindowBounds):
+        if window.bounds.start is not None:
+            start = utils.timedelta_to_micros(window.bounds.start)
+        if window.bounds.end is not None:
+            end = utils.timedelta_to_micros(window.bounds.end)
+    elif window.bounds:
+        start = window.bounds.start
+        end = window.bounds.end
+
+    start_value, start_side = _get_window_bounds(start, is_start_bounds=True)
+    end_value, end_side = _get_window_bounds(end, is_start_bounds=False)
+
+    spec = sge.WindowSpec(
+        kind=kind,
+        start=start_value,
+        start_side=start_side,
+        end=end_value,
+        end_side=end_side,
+        over="OVER",
+    )
+
+    if window.is_row_bounded and not window.ordering:
+        raise ValueError("No ordering provided for ordered analytic function")
+    elif (
+        not window.is_row_bounded
+        and not window.is_range_bounded
+        and not window.ordering
+    ):
+        # Unbound grouping window.
+        order_by = None
+    elif window.is_range_bounded:
+        # Note that, when the window is range-bounded, we only need one ordering key.
+        # There are two reasons:
+        # 1. Manipulating null positions requires more than one ordering key, which
+        #  is forbidden by SQL window syntax for range rolling.
+        # 2. Pandas does not allow range rolling on timeseries with nulls.
+        order_by = get_window_order_by((window.ordering[0],), override_null_order=False)
+    else:
+        order_by = get_window_order_by(window.ordering, override_null_order=True)
+
+    order = sge.Order(expressions=order_by) if order_by else None
+
+    group_by = (
+        [scalar_compiler.compile_scalar_expression(key) for key in window.grouping_keys]
+        if window.grouping_keys
+        else None
+    )
+
+    return sge.Window(this=value, partition_by=group_by, order=order, spec=spec)
+
+
+def get_window_order_by(
+    ordering: typing.Tuple[ordering_spec.OrderingExpression, ...],
+    override_null_order: bool = False,
+) -> typing.Optional[typing.List[sge.Ordered]]:
+    """Returns the SQL order by clause for a window specification."""
+    if not ordering:
+        return None
+
+    order_by = []
+    for ordering_spec_item in ordering:
+        expr = scalar_compiler.compile_scalar_expression(
+            ordering_spec_item.scalar_expression
+        )
+        desc = not ordering_spec_item.direction.is_ascending
+        nulls_first = not ordering_spec_item.na_last
+
+        if override_null_order:
+            # Bigquery SQL considers NULLS to be "smallest" values, but we need
+            # to override in these cases.
+            is_null_expr = sge.Is(this=expr, expression=sge.Null())
+            if nulls_first and desc:
+                order_by.append(
+                    sge.Ordered(
+                        this=is_null_expr,
+                        desc=desc,
+                        nulls_first=nulls_first,
+                    )
+                )
+            elif not nulls_first and not desc:
+                order_by.append(
+                    sge.Ordered(
+                        this=is_null_expr,
+                        desc=desc,
+                        nulls_first=nulls_first,
+                    )
+                )
+
+        order_by.append(
+            sge.Ordered(
+                this=expr,
+                desc=desc,
+                nulls_first=nulls_first,
+            )
+        )
+    return order_by
+
+
+def _get_window_bounds(
+    value, is_start_bounds: bool
+) -> tuple[typing.Union[str, sge.Expression], typing.Optional[str]]:
+    """Compiles a single boundary value into its SQL components."""
+    if value is None:
+        side = "PRECEDING" if is_start_bounds else "FOLLOWING"
+        return "UNBOUNDED", side
+
+    if value == 0:
+        return "CURRENT ROW", None
+
+    side = "PRECEDING" if value < 0 else "FOLLOWING"
+    return sge.convert(abs(value)), side

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

@@ -23,6 +23,8 @@
 from bigframes.core import expression, guid, identifiers, nodes, pyarrow_utils, rewrite
 from bigframes.core.compile import configs
 import bigframes.core.compile.sqlglot.aggregate_compiler as aggregate_compiler
+from bigframes.core.compile.sqlglot.aggregations import window
+from bigframes.core.compile.sqlglot.aggregations.window import apply_window_if_present
 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
@@ -272,18 +274,16 @@ def compile_random_sample(
     def compile_aggregate(
         self, node: nodes.AggregateNode, child: ir.SQLGlotIR
     ) -> ir.SQLGlotIR:
-        ordering_cols = tuple(
-            sge.Ordered(
-                this=scalar_compiler.compile_scalar_expression(
-                    ordering.scalar_expression
-                ),
-                desc=ordering.direction.is_ascending is False,
-                nulls_first=ordering.na_last is False,
-            )
-            for ordering in node.order_by
+        ordering_cols = window.get_window_order_by(
+            node.order_by, override_null_order=True
         )
         aggregations: tuple[tuple[str, sge.Expression], ...] = tuple(
-            (id.sql, aggregate_compiler.compile_aggregate(agg, order_by=ordering_cols))
+            (
+                id.sql,
+                aggregate_compiler.compile_aggregate(
+                    agg, order_by=tuple(ordering_cols) if ordering_cols else ()
+                ),
+            )
             for agg, id in node.aggregations
         )
         by_cols: tuple[sge.Expression, ...] = tuple(
@@ -299,6 +299,72 @@ def compile_aggregate(
 
         return child.aggregate(aggregations, by_cols, tuple(dropna_cols))
 
+    @_compile_node.register
+    def compile_window(
+        self, node: nodes.WindowOpNode, child: ir.SQLGlotIR
+    ) -> ir.SQLGlotIR:
+        window_spec = node.window_spec
+        if node.expression.op.order_independent and window_spec.is_unbounded:
+            # notably percentile_cont does not support ordering clause
+            window_spec = window_spec.without_order()
+
+        window_op = aggregate_compiler.compile_analytic(node.expression, window_spec)
+
+        inputs: tuple[sge.Expression, ...] = tuple(
+            scalar_compiler.compile_scalar_expression(expression.DerefOp(column))
+            for column in node.expression.column_references
+        )
+
+        clauses: list[tuple[sge.Expression, sge.Expression]] = []
+        if node.expression.op.skips_nulls and not node.never_skip_nulls:
+            for column in inputs:
+                clauses.append((sge.Is(this=column, expression=sge.Null()), sge.Null()))
+
+        if window_spec.min_periods and len(inputs) > 0:
+            if node.expression.op.skips_nulls:
+                # Most operations do not count NULL values towards min_periods
+                not_null_columns = [
+                    sge.Not(this=sge.Is(this=column, expression=sge.Null()))
+                    for column in inputs
+                ]
+                # All inputs must be non-null for observation to count
+                if not not_null_columns:
+                    is_observation_expr: sge.Expression = sge.convert(True)
+                else:
+                    is_observation_expr = not_null_columns[0]
+                    for expr in not_null_columns[1:]:
+                        is_observation_expr = sge.And(
+                            this=is_observation_expr, expression=expr
+                        )
+                is_observation = ir._cast(is_observation_expr, "INT64")
+            else:
+                # Operations like count treat even NULLs as valid observations
+                # for the sake of min_periods notnull is just used to convert
+                # null values to non-null (FALSE) values to be counted.
+                is_observation = ir._cast(
+                    sge.Not(this=sge.Is(this=inputs[0], expression=sge.Null())),
+                    "INT64",
+                )
+
+            observation_count = apply_window_if_present(
+                sge.func("SUM", is_observation), window_spec
+            )
+            clauses.append(
+                (
+                    observation_count < sge.convert(window_spec.min_periods),
+                    sge.Null(),
+                )
+            )
+        if clauses:
+            when_expressions = [sge.When(this=cond, true=res) for cond, res in clauses]
+            window_op = sge.Case(ifs=when_expressions, default=window_op)
+
+        # TODO: check if we can directly window the expression.
+        return child.window(
+            window_op=window_op,
+            output_column_id=node.output_name.sql,
+        )
+
 
 def _replace_unsupported_ops(node: nodes.BigFrameNode):
     node = nodes.bottom_up(node, rewrite.rewrite_slice)

bigframes/core/compile/sqlglot/compiler.py

@@ -394,3 +394,8 @@ def _(op: ops.ParseJSON, expr: TypedExpr) -> sge.Expression:
 @UNARY_OP_REGISTRATION.register(ops.ToJSONString)
 def _(op: ops.ToJSONString, expr: TypedExpr) -> sge.Expression:
     return sge.func("TO_JSON_STRING", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.UnixMicros)
+def _(op: ops.UnixMicros, expr: TypedExpr) -> sge.Expression:
+    return sge.func("UNIX_MICROS", expr.expr)

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

@@ -31,7 +31,7 @@
 
 @functools.singledispatch
 def compile_scalar_expression(
-    expression: expression.Expression,
+    expr: expression.Expression,
 ) -> sge.Expression:
     """Compiles BigFrames scalar expression into SQLGlot expression."""
     raise ValueError(f"Can't compile unrecognized node: {expression}")

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -409,6 +409,13 @@ def aggregate(
             new_expr = new_expr.where(condition, append=False)
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
+    def window(
+        self,
+        window_op: sge.Expression,
+        output_column_id: str,
+    ) -> SQLGlotIR:
+        return self.project(((output_column_id, window_op),))
+
     def insert(
         self,
         destination: bigquery.TableReference,