Merge branch 'main' into tswast-patch-1

Author: tswast

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.windows 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

@@ -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.windows 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

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

@@ -0,0 +1,153 @@
+# 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
+
+    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
+    )
+
+    # This is the key change. Don't create a spec for the default window frame
+    # if there's no ordering. This avoids generating an `ORDER BY NULL` clause.
+    if not window.bounds and not order:
+        return sge.Window(this=value, partition_by=group_by)
+
+    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_preceding=True)
+    end_value, end_side = _get_window_bounds(end, is_preceding=False)
+
+    spec = sge.WindowSpec(
+        kind=kind,
+        start=start_value,
+        start_side=start_side,
+        end=end_value,
+        end_side=end_side,
+        over="OVER",
+    )
+
+    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[tuple[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 tuple(order_by)
+
+
+def _get_window_bounds(
+    value, is_preceding: 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_preceding 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/windows.py

@@ -23,6 +23,7 @@
 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 windows
 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 +273,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 = windows.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=ordering_cols if ordering_cols else ()
+                ),
+            )
             for agg, id in node.aggregations
         )
         by_cols: tuple[sge.Expression, ...] = tuple(

bigframes/core/compile/sqlglot/compiler.py

@@ -16,6 +16,7 @@
 
 import functools
 import inspect
+import os
 import threading
 import typing
 from typing import (
@@ -56,6 +57,7 @@
 from bigframes.session import dry_runs
 import bigframes.session._io.bigquery
 import bigframes.session.clients
+import bigframes.session.metrics
 
 # Note: the following methods are duplicated from Session. This duplication
 # enables the following:
@@ -625,6 +627,11 @@ def _get_bqclient() -> bigquery.Client:
 
 def _dry_run(query, bqclient) -> bigquery.QueryJob:
     job = bqclient.query(query, bigquery.QueryJobConfig(dry_run=True))
+
+    # Fix for b/435183833. Log metrics even if a Session isn't available.
+    if bigframes.session.metrics.LOGGING_NAME_ENV_VAR in os.environ:
+        metrics = bigframes.session.metrics.ExecutionMetrics()
+        metrics.count_job_stats(job)
     return job
 
 

bigframes/pandas/io/api.py

@@ -40,32 +40,54 @@ def count_job_stats(
     ):
         if query_job is None:
             assert row_iterator is not None
-            total_bytes_processed = getattr(row_iterator, "total_bytes_processed", None)
-            query = getattr(row_iterator, "query", None)
-            if total_bytes_processed is None or query is None:
-                return
+
+            # TODO(tswast): Pass None after making benchmark publishing robust to missing data.
+            bytes_processed = getattr(row_iterator, "total_bytes_processed", 0)
+            query_char_count = len(getattr(row_iterator, "query", ""))
+            slot_millis = getattr(row_iterator, "slot_millis", 0)
+            exec_seconds = 0.0
 
             self.execution_count += 1
-            self.query_char_count += len(query)
-            self.bytes_processed += total_bytes_processed
-            write_stats_to_disk(len(query), total_bytes_processed)
-            return
+            self.query_char_count += query_char_count
+            self.bytes_processed += bytes_processed
+            self.slot_millis += slot_millis
+
+        elif query_job.configuration.dry_run:
+            query_char_count = len(query_job.query)
 
-        if query_job.configuration.dry_run:
-            write_stats_to_disk(len(query_job.query), 0, 0, 0)
+            # TODO(tswast): Pass None after making benchmark publishing robust to missing data.
+            bytes_processed = 0
+            slot_millis = 0
+            exec_seconds = 0.0
 
-        stats = get_performance_stats(query_job)
-        if stats is not None:
-            query_char_count, bytes_processed, slot_millis, execution_secs = stats
+        elif (stats := get_performance_stats(query_job)) is not None:
+            query_char_count, bytes_processed, slot_millis, exec_seconds = stats
             self.execution_count += 1
             self.query_char_count += query_char_count
             self.bytes_processed += bytes_processed
             self.slot_millis += slot_millis
-            self.execution_secs += execution_secs
+            self.execution_secs += exec_seconds
             write_stats_to_disk(
-                query_char_count, bytes_processed, slot_millis, execution_secs
+                query_char_count=query_char_count,
+                bytes_processed=bytes_processed,
+                slot_millis=slot_millis,
+                exec_seconds=exec_seconds,
             )
 
+        else:
+            # TODO(tswast): Pass None after making benchmark publishing robust to missing data.
+            bytes_processed = 0
+            query_char_count = 0
+            slot_millis = 0
+            exec_seconds = 0
+
+        write_stats_to_disk(
+            query_char_count=query_char_count,
+            bytes_processed=bytes_processed,
+            slot_millis=slot_millis,
+            exec_seconds=exec_seconds,
+        )
+
 
 def get_performance_stats(
     query_job: bigquery.QueryJob,
@@ -103,10 +125,11 @@ def get_performance_stats(
 
 
 def write_stats_to_disk(
+    *,
     query_char_count: int,
     bytes_processed: int,
-    slot_millis: Optional[int] = None,
-    exec_seconds: Optional[float] = None,
+    slot_millis: int,
+    exec_seconds: float,
 ):
     """For pytest runs only, log information about the query job
     to a file in order to create a performance report.
@@ -118,18 +141,17 @@ def write_stats_to_disk(
     test_name = os.environ[LOGGING_NAME_ENV_VAR]
     current_directory = os.getcwd()
 
-    if (slot_millis is not None) and (exec_seconds is not None):
-        # store slot milliseconds
-        slot_file = os.path.join(current_directory, test_name + ".slotmillis")
-        with open(slot_file, "a") as f:
-            f.write(str(slot_millis) + "\n")
+    # store slot milliseconds
+    slot_file = os.path.join(current_directory, test_name + ".slotmillis")
+    with open(slot_file, "a") as f:
+        f.write(str(slot_millis) + "\n")
 
-        # store execution time seconds
-        exec_time_file = os.path.join(
-            current_directory, test_name + ".bq_exec_time_seconds"
-        )
-        with open(exec_time_file, "a") as f:
-            f.write(str(exec_seconds) + "\n")
+    # store execution time seconds
+    exec_time_file = os.path.join(
+        current_directory, test_name + ".bq_exec_time_seconds"
+    )
+    with open(exec_time_file, "a") as f:
+        f.write(str(exec_seconds) + "\n")
 
     # store length of query
     query_char_count_file = os.path.join(