From 4ad65b67527a7ce967a2ff7b3efff49c1da4675a Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Tue, 16 Dec 2025 09:02:36 +0000 Subject: [PATCH 1/3] refactor: fix groupby.GroupBy.first doctest for sqlglot --- bigframes/core/compile/sqlglot/compiler.py | 2 +- .../out.sql | 22 ++++++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/bigframes/core/compile/sqlglot/compiler.py b/bigframes/core/compile/sqlglot/compiler.py index 7ecc15f6a2..501243fe8e 100644 --- a/bigframes/core/compile/sqlglot/compiler.py +++ b/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, ) diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_window/test_compile_window_w_groupby_rolling/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_window/test_compile_window_w_groupby_rolling/out.sql index b1d498bc76..e8fabd1129 100644 --- a/tests/unit/core/compile/sqlglot/snapshots/test_compile_window/test_compile_window_w_groupby_rolling/out.sql +++ b/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 *, From ed77af9d26e45fb647c03e8ad0b3ec31475d5873 Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Wed, 17 Dec 2025 08:40:37 +0000 Subject: [PATCH 2/3] refactor: bbq.ai.generate_bool --- bigframes/core/compile/sqlglot/expressions/ai_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bigframes/core/compile/sqlglot/expressions/ai_ops.py b/bigframes/core/compile/sqlglot/expressions/ai_ops.py index 680e35c511..a8a36cb6c0 100644 --- a/bigframes/core/compile/sqlglot/expressions/ai_ops.py +++ b/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)) From 7c8e3462f97546245d253b714491fdae62940170 Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Wed, 17 Dec 2025 12:54:07 +0000 Subject: [PATCH 3/3] refactor: add to_json compiler --- .../core/compile/sqlglot/expressions/json_ops.py | 5 +++++ .../snapshots/test_json_ops/test_to_json/out.sql | 13 +++++++++++++ .../compile/sqlglot/expressions/test_json_ops.py | 8 ++++++++ 3 files changed, 26 insertions(+) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_json_ops/test_to_json/out.sql diff --git a/bigframes/core/compile/sqlglot/expressions/json_ops.py b/bigframes/core/compile/sqlglot/expressions/json_ops.py index ef55f6edac..0a38e8e138 100644 --- a/bigframes/core/compile/sqlglot/expressions/json_ops.py +++ b/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) diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_json_ops/test_to_json/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_json_ops/test_to_json/out.sql new file mode 100644 index 0000000000..ebca0c51c5 --- /dev/null +++ b/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` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_json_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_json_ops.py index 4ae3eb3fcc..1c5894fc96 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_json_ops.py +++ b/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]]