refactor: fix SQLGlot doctests on the window, ai, to_json ops (#2341)

This change can resolve three doctests failures in #2248: 
- Added `to_json` support to the compiler backend.
- Fixed minor bugs on AI ops
- Fixed minor bugs on windows compiler

Fixes internal issue 417774347 🦕

Author: chelsea-lin

bigframes/core/compile/sqlglot/compiler.py

@@ -378,7 +378,7 @@ def compile_window(node: nodes.WindowOpNode, child: ir.SQLGlotIR) -> ir.SQLGlotI
             window_op = sge.Case(ifs=when_expressions, default=window_op)
 
         # TODO: check if we can directly window the expression.
-        result = child.window(
+        result = result.window(
             window_op=window_op,
             output_column_id=cdef.id.sql,
         )

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

@@ -93,6 +93,7 @@ def _construct_prompt(
     for elem in prompt_context:
         if elem is None:
             prompt.append(exprs[column_ref_idx].expr)
+            column_ref_idx += 1
         else:
             prompt.append(sge.Literal.string(elem))
 

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

@@ -69,6 +69,11 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("PARSE_JSON", expr.expr)
 
 
+@register_unary_op(ops.ToJSON)
+def _(expr: TypedExpr) -> sge.Expression:
+    return sge.func("TO_JSON", expr.expr)
+
+
 @register_unary_op(ops.ToJSONString)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("TO_JSON_STRING", expr.expr)

tests/unit/core/compile/sqlglot/expressions/snapshots/test_json_ops/test_to_json/out.sql

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `string_col`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    TO_JSON(`string_col`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `string_col`
+FROM `bfcte_1`

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

@@ -105,6 +105,14 @@ def test_parse_json(scalar_types_df: bpd.DataFrame, snapshot):
     snapshot.assert_match(sql, "out.sql")
 
 
+def test_to_json(scalar_types_df: bpd.DataFrame, snapshot):
+    col_name = "string_col"
+    bf_df = scalar_types_df[[col_name]]
+    sql = utils._apply_ops_to_sql(bf_df, [ops.ToJSON().as_expr(col_name)], [col_name])
+
+    snapshot.assert_match(sql, "out.sql")
+
+
 def test_to_json_string(json_types_df: bpd.DataFrame, snapshot):
     col_name = "json_col"
     bf_df = json_types_df[[col_name]]

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

@@ -12,12 +12,32 @@ WITH `bfcte_0` AS (
     `int64_col` AS `bfcol_8`,
     `bool_col` AS `bfcol_9`
   FROM `bfcte_0`
-), `bfcte_3` AS (
+), `bfcte_2` AS (
   SELECT
     *
   FROM `bfcte_1`
   WHERE
     NOT `bfcol_9` IS NULL
+), `bfcte_3` 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
+      ) < 3
+      THEN NULL
+      ELSE COALESCE(
+        SUM(CAST(`bfcol_7` 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
+      )
+    END AS `bfcol_15`
+  FROM `bfcte_2`
 ), `bfcte_4` AS (
   SELECT
     *,