Merge branch 'main' into shuowei-anywidget-fix-empty-index

Author: shuoweil

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

@@ -44,6 +44,7 @@ def apply_window_if_present(
         order_by = None
     elif window.is_range_bounded:
         order_by = get_window_order_by((window.ordering[0],))
+        order_by = remove_null_ordering_for_range_windows(order_by)
     else:
         order_by = get_window_order_by(window.ordering)
 
@@ -150,6 +151,30 @@ def get_window_order_by(
     return tuple(order_by)
 
 
+def remove_null_ordering_for_range_windows(
+    order_by: typing.Optional[tuple[sge.Ordered, ...]],
+) -> typing.Optional[tuple[sge.Ordered, ...]]:
+    """Removes NULL FIRST/LAST from ORDER BY expressions in RANGE windows.
+    Here's the support matrix:
+    ✅ sum(x) over (order by y desc nulls last)
+    🚫 sum(x) over (order by y asc nulls last)
+    ✅ sum(x) over (order by y asc nulls first)
+    🚫 sum(x) over (order by y desc nulls first)
+    """
+    if order_by is None:
+        return None
+
+    new_order_by = []
+    for key in order_by:
+        kargs = key.args
+        if kargs.get("desc") is True and kargs.get("nulls_first", False):
+            kargs["nulls_first"] = False
+        elif kargs.get("desc") is False and not kargs.setdefault("nulls_first", True):
+            kargs["nulls_first"] = True
+        new_order_by.append(sge.Ordered(**kargs))
+    return tuple(new_order_by)
+
+
 def _get_window_bounds(
     value, is_preceding: bool
 ) -> tuple[typing.Union[str, sge.Expression], typing.Optional[str]]:

bigframes/core/compile/sqlglot/compiler.py

@@ -356,6 +356,9 @@ def compile_window(node: nodes.WindowOpNode, child: ir.SQLGlotIR) -> ir.SQLGlotI
                 observation_count = windows.apply_window_if_present(
                     sge.func("SUM", is_observation), window_spec
                 )
+                observation_count = sge.func(
+                    "COALESCE", observation_count, sge.convert(0)
+                )
             else:
                 # Operations like count treat even NULLs as valid observations
                 # for the sake of min_periods notnull is just used to convert

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

@@ -89,27 +89,39 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.ge_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.GTE(this=left_expr, expression=right_expr)
 
 
 @register_binary_op(ops.gt_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.GT(this=left_expr, expression=right_expr)
 
 
 @register_binary_op(ops.lt_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.LT(this=left_expr, expression=right_expr)
 
 
 @register_binary_op(ops.le_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.LTE(this=left_expr, expression=right_expr)

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

@@ -140,6 +140,19 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.Coalesce(this=left.expr, expressions=[right.expr])
 
 
+@register_binary_op(ops.BinaryRemoteFunctionOp, pass_op=True)
+def _(
+    left: TypedExpr, right: TypedExpr, op: ops.BinaryRemoteFunctionOp
+) -> sge.Expression:
+    routine_ref = op.function_def.routine_ref
+    # Quote project, dataset, and routine IDs to avoid keyword clashes.
+    func_name = (
+        f"`{routine_ref.project}`.`{routine_ref.dataset_id}`.`{routine_ref.routine_id}`"
+    )
+
+    return sge.func(func_name, left.expr, right.expr)
+
+
 @register_nary_op(ops.case_when_op)
 def _(*cases_and_outputs: TypedExpr) -> sge.Expression:
     # Need to upcast BOOL to INT if any output is numeric

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

@@ -388,6 +388,9 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.add_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     if left.dtype == dtypes.STRING_DTYPE and right.dtype == dtypes.STRING_DTYPE:
         # String addition
         return sge.Concat(expressions=[left.expr, right.expr])
@@ -442,6 +445,9 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.floordiv_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
 
@@ -525,6 +531,9 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.mul_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
 
@@ -548,6 +557,9 @@ def _(expr: TypedExpr, n_digits: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.sub_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if left.expr == sge.null() or right.expr == sge.null():
+        return sge.null()
+
     if dtypes.is_numeric(left.dtype) and dtypes.is_numeric(right.dtype):
         left_expr = _coerce_bool_to_int(left)
         right_expr = _coerce_bool_to_int(right)

docs/conf.py

@@ -267,6 +267,12 @@
 
 # https://sphinx-sitemap.readthedocs.io/en/latest/getting-started.html#usage
 html_baseurl = "https://dataframes.bigquery.dev/"
+sitemap_locales = [None]
+
+# We don't have any immediate plans to translate the API reference, so omit the
+# language from the URLs.
+# https://sphinx-sitemap.readthedocs.io/en/latest/advanced-configuration.html#configuration-customizing-url-scheme
+sitemap_url_scheme = "{link}"
 
 # -- Options for warnings ------------------------------------------------------
 

tests/unit/core/compile/sqlglot/aggregations/test_windows.py

@@ -127,7 +127,7 @@ def test_apply_window_if_present_range_bounded(self):
         )
         self.assertEqual(
             result.sql(dialect="bigquery"),
-            "value OVER (ORDER BY `col1` ASC NULLS LAST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)",
+            "value OVER (ORDER BY `col1` ASC RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)",
         )
 
     def test_apply_window_if_present_range_bounded_timedelta(self):
@@ -142,7 +142,7 @@ def test_apply_window_if_present_range_bounded_timedelta(self):
         )
         self.assertEqual(
             result.sql(dialect="bigquery"),
-            "value OVER (ORDER BY `col1` ASC NULLS LAST RANGE BETWEEN 86400000000 PRECEDING AND 43200000000 FOLLOWING)",
+            "value OVER (ORDER BY `col1` ASC RANGE BETWEEN 86400000000 PRECEDING AND 43200000000 FOLLOWING)",
         )
 
     def test_apply_window_if_present_all_params(self):

tests/unit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_binary_remote_function_op/out.sql

@@ -0,0 +1,14 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `float64_col`,
+    `int64_col`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `my_project`.`my_dataset`.`my_routine`(`int64_col`, `float64_col`) AS `bfcol_2`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_2` AS `int64_col`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/test_generic_ops.py

@@ -168,6 +168,43 @@ def test_astype_json_invalid(
         )
 
 
+def test_binary_remote_function_op(scalar_types_df: bpd.DataFrame, snapshot):
+    from google.cloud import bigquery
+
+    from bigframes.functions import udf_def
+
+    bf_df = scalar_types_df[["int64_col", "float64_col"]]
+    op = ops.BinaryRemoteFunctionOp(
+        function_def=udf_def.BigqueryUdf(
+            routine_ref=bigquery.RoutineReference.from_string(
+                "my_project.my_dataset.my_routine"
+            ),
+            signature=udf_def.UdfSignature(
+                input_types=(
+                    udf_def.UdfField(
+                        "x",
+                        bigquery.StandardSqlDataType(
+                            type_kind=bigquery.StandardSqlTypeNames.INT64
+                        ),
+                    ),
+                    udf_def.UdfField(
+                        "y",
+                        bigquery.StandardSqlDataType(
+                            type_kind=bigquery.StandardSqlTypeNames.FLOAT64
+                        ),
+                    ),
+                ),
+                output_bq_type=bigquery.StandardSqlDataType(
+                    type_kind=bigquery.StandardSqlTypeNames.FLOAT64
+                ),
+            ),
+        )
+    )
+    sql = utils._apply_binary_op(bf_df, op, "int64_col", "float64_col")
+
+    snapshot.assert_match(sql, "out.sql")
+
+
 def test_case_when_op(scalar_types_df: bpd.DataFrame, snapshot):
     ops_map = {
         "single_case": ops.case_when_op.as_expr(

tests/unit/core/compile/sqlglot/snapshots/test_compile_window/test_compile_window_w_groupby_rolling/out.sql

@@ -22,10 +22,13 @@ WITH `bfcte_0` AS (
   SELECT
     *,
     CASE
-      WHEN SUM(CAST(NOT `bfcol_7` IS NULL AS INT64)) OVER (
-        PARTITION BY `bfcol_9`
-        ORDER BY `bfcol_9` ASC NULLS LAST, `rowindex` ASC NULLS LAST
-        ROWS BETWEEN 3 PRECEDING AND CURRENT ROW
+      WHEN COALESCE(
+        SUM(CAST(NOT `bfcol_7` IS NULL AS INT64)) OVER (
+          PARTITION BY `bfcol_9`
+          ORDER BY `bfcol_9` ASC NULLS LAST, `rowindex` ASC NULLS LAST
+          ROWS BETWEEN 3 PRECEDING AND CURRENT ROW
+        ),
+        0
       ) < 3
       THEN NULL
       ELSE COALESCE(
@@ -42,10 +45,13 @@ WITH `bfcte_0` AS (
   SELECT
     *,
     CASE
-      WHEN SUM(CAST(NOT `bfcol_8` IS NULL AS INT64)) OVER (
-        PARTITION BY `bfcol_9`
-        ORDER BY `bfcol_9` ASC NULLS LAST, `rowindex` ASC NULLS LAST
-        ROWS BETWEEN 3 PRECEDING AND CURRENT ROW
+      WHEN COALESCE(
+        SUM(CAST(NOT `bfcol_8` IS NULL AS INT64)) OVER (
+          PARTITION BY `bfcol_9`
+          ORDER BY `bfcol_9` ASC NULLS LAST, `rowindex` ASC NULLS LAST
+          ROWS BETWEEN 3 PRECEDING AND CURRENT ROW
+        ),
+        0
       ) < 3
       THEN NULL
       ELSE COALESCE(