From dbccc4715af0b3796f5ed85e3675d85039469309 Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Wed, 16 Jul 2025 22:54:25 +0000 Subject: [PATCH 1/4] uniform json golden sql tests --- .../test_json_extract_array/out.sql | 15 +++++++++++ .../test_json_extract_string_array/out.sql | 15 +++++++++++ .../test_json_query/out.sql | 15 +++++++++++ .../test_json_query_array/out.sql | 15 +++++++++++ .../test_json_value/out.sql | 15 +++++++++++ .../expressions/test_unary_compiler.py | 27 +++++++------------ 6 files changed, 85 insertions(+), 17 deletions(-) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql new file mode 100644 index 0000000000..db3172de62 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql @@ -0,0 +1,15 @@ +WITH `bfcte_0` AS ( + SELECT + `rowindex` AS `bfcol_0`, + `json_col` AS `bfcol_1` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_EXTRACT_ARRAY(`bfcol_1`, '$') AS `bfcol_4` + FROM `bfcte_0` +) +SELECT + `bfcol_0` AS `rowindex`, + `bfcol_4` AS `json_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql new file mode 100644 index 0000000000..21db7bc5b5 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql @@ -0,0 +1,15 @@ +WITH `bfcte_0` AS ( + SELECT + `rowindex` AS `bfcol_0`, + `json_col` AS `bfcol_1` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_EXTRACT_STRING_ARRAY(`bfcol_1`, '$') AS `bfcol_4` + FROM `bfcte_0` +) +SELECT + `bfcol_0` AS `rowindex`, + `bfcol_4` AS `json_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql new file mode 100644 index 0000000000..a5288a9dd9 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql @@ -0,0 +1,15 @@ +WITH `bfcte_0` AS ( + SELECT + `rowindex` AS `bfcol_0`, + `json_col` AS `bfcol_1` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_QUERY(`bfcol_1`, '$') AS `bfcol_4` + FROM `bfcte_0` +) +SELECT + `bfcol_0` AS `rowindex`, + `bfcol_4` AS `json_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql new file mode 100644 index 0000000000..9de4780c94 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql @@ -0,0 +1,15 @@ +WITH `bfcte_0` AS ( + SELECT + `rowindex` AS `bfcol_0`, + `json_col` AS `bfcol_1` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_QUERY_ARRAY(`bfcol_1`, '$') AS `bfcol_4` + FROM `bfcte_0` +) +SELECT + `bfcol_0` AS `rowindex`, + `bfcol_4` AS `json_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql new file mode 100644 index 0000000000..ce895204ba --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql @@ -0,0 +1,15 @@ +WITH `bfcte_0` AS ( + SELECT + `rowindex` AS `bfcol_0`, + `json_col` AS `bfcol_1` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_VALUE(`bfcol_1`, '$') AS `bfcol_4` + FROM `bfcte_0` +) +SELECT + `bfcol_0` AS `rowindex`, + `bfcol_4` AS `json_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index 6d9101aff0..f74d42c404 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py @@ -47,39 +47,32 @@ def test_array_slice_with_start_and_stop(repeated_types_df: bpd.DataFrame, snaps # JSON Ops def test_json_extract(json_types_df: bpd.DataFrame, snapshot): result = bbq.json_extract(json_types_df["json_col"], "$") - expected_sql = "JSON_EXTRACT(`bfcol_1`, '$') AS `bfcol_4`" - assert expected_sql in result.to_frame().sql snapshot.assert_match(result.to_frame().sql, "out.sql") -def test_json_extract_array(json_types_df: bpd.DataFrame): +def test_json_extract_array(json_types_df: bpd.DataFrame, snapshot): result = bbq.json_extract_array(json_types_df["json_col"], "$") - expected_sql = "JSON_EXTRACT_ARRAY(`bfcol_1`, '$') AS `bfcol_4`" - assert expected_sql in result.to_frame().sql + snapshot.assert_match(result.to_frame().sql, "out.sql") -def test_json_extract_string_array(json_types_df: bpd.DataFrame): +def test_json_extract_string_array(json_types_df: bpd.DataFrame, snapshot): result = bbq.json_extract_string_array(json_types_df["json_col"], "$") - expected_sql = "JSON_EXTRACT_STRING_ARRAY(`bfcol_1`, '$') AS `bfcol_4`" - assert expected_sql in result.to_frame().sql + snapshot.assert_match(result.to_frame().sql, "out.sql") -def test_json_query(json_types_df: bpd.DataFrame): +def test_json_query(json_types_df: bpd.DataFrame, snapshot): result = bbq.json_query(json_types_df["json_col"], "$") - expected_sql = "JSON_QUERY(`bfcol_1`, '$') AS `bfcol_4`" - assert expected_sql in result.to_frame().sql + snapshot.assert_match(result.to_frame().sql, "out.sql") -def test_json_query_array(json_types_df: bpd.DataFrame): +def test_json_query_array(json_types_df: bpd.DataFrame, snapshot): result = bbq.json_query_array(json_types_df["json_col"], "$") - expected_sql = "JSON_QUERY_ARRAY(`bfcol_1`, '$') AS `bfcol_4`" - assert expected_sql in result.to_frame().sql + snapshot.assert_match(result.to_frame().sql, "out.sql") -def test_json_value(json_types_df: bpd.DataFrame): +def test_json_value(json_types_df: bpd.DataFrame, snapshot): result = bbq.json_value(json_types_df["json_col"], "$") - expected_sql = "JSON_VALUE(`bfcol_1`, '$') AS `bfcol_4`" - assert expected_sql in result.to_frame().sql + snapshot.assert_match(result.to_frame().sql, "out.sql") def test_parse_json(scalar_types_df: bpd.DataFrame, snapshot): From cc85630a0a580321fe09fd61e08f023e06055419 Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Thu, 17 Jul 2025 18:31:45 +0000 Subject: [PATCH 2/4] scalar tests use array_value rather than open APIs --- .../test_add_numeric/out.sql | 8 ++-- .../test_add_numeric_w_scalar/out.sql | 8 ++-- .../test_add_string/out.sql | 8 ++-- .../test_json_set/out.sql | 12 ++---- .../expressions/test_binary_compiler.py | 39 ++++++++++++------- 5 files changed, 40 insertions(+), 35 deletions(-) diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql index 1496f89f28..18e1be495d 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql @@ -6,11 +6,11 @@ WITH `bfcte_0` AS ( ), `bfcte_1` AS ( SELECT *, - `bfcol_1` AS `bfcol_4`, - `bfcol_0` + `bfcol_0` AS `bfcol_5` + `bfcol_0` + `bfcol_0` AS `bfcol_4` FROM `bfcte_0` ) SELECT - `bfcol_4` AS `rowindex`, - `bfcol_5` AS `int64_col` + `bfcol_1` AS `bfuid_col_1`, + `bfcol_0` AS `int64_col`, + `bfcol_4` AS `bfuid_col_2` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql index 9c4b01a6df..2162f71631 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql @@ -6,11 +6,11 @@ WITH `bfcte_0` AS ( ), `bfcte_1` AS ( SELECT *, - `bfcol_1` AS `bfcol_4`, - `bfcol_0` + 1 AS `bfcol_5` + `bfcol_0` + 1 AS `bfcol_4` FROM `bfcte_0` ) SELECT - `bfcol_4` AS `rowindex`, - `bfcol_5` AS `int64_col` + `bfcol_1` AS `bfuid_col_1`, + `bfcol_0` AS `int64_col`, + `bfcol_4` AS `bfuid_col_3` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql index 7a8ab83df1..c30e5d4cc6 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql @@ -6,11 +6,11 @@ WITH `bfcte_0` AS ( ), `bfcte_1` AS ( SELECT *, - `bfcol_0` AS `bfcol_4`, - CONCAT(`bfcol_1`, 'a') AS `bfcol_5` + CONCAT(`bfcol_1`, 'a') AS `bfcol_4` FROM `bfcte_0` ) SELECT - `bfcol_4` AS `rowindex`, - `bfcol_5` AS `string_col` + `bfcol_0` AS `bfuid_col_1`, + `bfcol_1` AS `string_col`, + `bfcol_4` AS `bfuid_col_4` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql index f501dd3b86..e5096b9412 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql @@ -8,13 +8,9 @@ WITH `bfcte_0` AS ( *, JSON_SET(`bfcol_1`, '$.a', 100) AS `bfcol_4` FROM `bfcte_0` -), `bfcte_2` AS ( - SELECT - *, - JSON_SET(`bfcol_4`, '$.b', 'hi') AS `bfcol_7` - FROM `bfcte_1` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_7` AS `json_col` -FROM `bfcte_2` \ No newline at end of file + `bfcol_0` AS `bfuid_col_5`, + `bfcol_1` AS `json_col`, + `bfcol_4` AS `bfuid_col_6` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py index 9daff51c9f..57c0a4b831 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py @@ -14,36 +14,45 @@ import pytest -import bigframes.bigquery as bbq +from bigframes import operations as ops +import bigframes.core.expression as ex import bigframes.pandas as bpd pytest.importorskip("pytest_snapshot") -def test_add_numeric(scalar_types_df: bpd.DataFrame, snapshot): - bf_df = scalar_types_df[["int64_col"]] +def _apply_binary_op( + obj: bpd.DataFrame | bpd.Series, + op: ops.BinaryOp, + l_arg: str | ex.Expression, + r_arg: str | ex.Expression, +) -> str: + array_value = obj._block.expr + op_expr = op.as_expr(l_arg, r_arg) + result, _ = array_value.compute_values([op_expr]) + sql = result.session._executor.to_sql(result, enable_cache=False) + return sql - bf_df["int64_col"] = bf_df["int64_col"] + bf_df["int64_col"] - snapshot.assert_match(bf_df.sql, "out.sql") +def test_add_numeric(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["int64_col"]] + sql = _apply_binary_op(bf_df, ops.add_op, "int64_col", "int64_col") + snapshot.assert_match(sql, "out.sql") def test_add_numeric_w_scalar(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["int64_col"]] - - bf_df["int64_col"] = bf_df["int64_col"] + 1 - - snapshot.assert_match(bf_df.sql, "out.sql") + sql = _apply_binary_op(bf_df, ops.add_op, "int64_col", ex.const(1)) + snapshot.assert_match(sql, "out.sql") def test_add_string(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["string_col"]] - - bf_df["string_col"] = bf_df["string_col"] + "a" - - snapshot.assert_match(bf_df.sql, "out.sql") + sql = _apply_binary_op(bf_df, ops.add_op, "string_col", ex.const("a")) + snapshot.assert_match(sql, "out.sql") def test_json_set(json_types_df: bpd.DataFrame, snapshot): - result = bbq.json_set(json_types_df["json_col"], [("$.a", 100), ("$.b", "hi")]) - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = json_types_df[["json_col"]] + sql = _apply_binary_op(bf_df, ops.JSONSet(json_path="$.a"), "json_col", ex.const(100)) + snapshot.assert_match(sql, "out.sql") From a1d3d76f5883e48a2f662cc4a49e680cac980961 Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Thu, 17 Jul 2025 18:48:35 +0000 Subject: [PATCH 3/4] convert unary tests too --- .../test_array_index/out.sql | 5 +- .../test_array_slice_with_only_start/out.sql | 5 +- .../out.sql | 5 +- .../test_array_to_string/out.sql | 5 +- .../test_json_extract/out.sql | 5 +- .../test_json_extract_array/out.sql | 5 +- .../test_json_extract_string_array/out.sql | 5 +- .../test_json_query/out.sql | 5 +- .../test_json_query_array/out.sql | 5 +- .../test_json_value/out.sql | 5 +- .../test_parse_json/out.sql | 5 +- .../test_to_json_string/out.sql | 16 ++++ .../expressions/test_unary_compiler.py | 82 +++++++++++++------ 13 files changed, 104 insertions(+), 49 deletions(-) create mode 100644 tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql index 33a8bded13..b874459b14 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `string_list_col` + `bfcol_0` AS `bfuid_col_1`, + `bfcol_1` AS `string_list_col`, + `bfcol_4` AS `bfuid_col_3` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql index 34d2225931..38eb5f2891 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql @@ -16,6 +16,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `string_list_col` + `bfcol_0` AS `bfuid_col_1`, + `bfcol_1` AS `string_list_col`, + `bfcol_4` AS `bfuid_col_4` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql index d46803ce7c..4947b5af12 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql @@ -16,6 +16,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `string_list_col` + `bfcol_0` AS `bfuid_col_1`, + `bfcol_1` AS `string_list_col`, + `bfcol_4` AS `bfuid_col_5` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql index e0db21f972..aa327a7fe4 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `string_list_col` + `bfcol_0` AS `bfuid_col_1`, + `bfcol_1` AS `string_list_col`, + `bfcol_4` AS `bfuid_col_2` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql index 2ffb0174a8..9781f2ed8d 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `json_col` + `bfcol_0` AS `bfuid_col_6`, + `bfcol_1` AS `json_col`, + `bfcol_4` AS `bfuid_col_7` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql index db3172de62..fd59f5dbd8 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `json_col` + `bfcol_0` AS `bfuid_col_6`, + `bfcol_1` AS `json_col`, + `bfcol_4` AS `bfuid_col_8` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql index 21db7bc5b5..024f214eab 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `json_col` + `bfcol_0` AS `bfuid_col_6`, + `bfcol_1` AS `json_col`, + `bfcol_4` AS `bfuid_col_9` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql index a5288a9dd9..fc61ca2402 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `json_col` + `bfcol_0` AS `bfuid_col_6`, + `bfcol_1` AS `json_col`, + `bfcol_4` AS `bfuid_col_10` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql index 9de4780c94..4313601883 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `json_col` + `bfcol_0` AS `bfuid_col_6`, + `bfcol_1` AS `json_col`, + `bfcol_4` AS `bfuid_col_11` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql index ce895204ba..7be7f6e0b2 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `json_col` + `bfcol_0` AS `bfuid_col_6`, + `bfcol_1` AS `json_col`, + `bfcol_4` AS `bfuid_col_12` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql index d965ea8f1b..ee47a61a39 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql @@ -10,6 +10,7 @@ WITH `bfcte_0` AS ( FROM `bfcte_0` ) SELECT - `bfcol_0` AS `rowindex`, - `bfcol_4` AS `string_col` + `bfcol_0` AS `bfuid_col_13`, + `bfcol_1` AS `string_col`, + `bfcol_4` AS `bfuid_col_14` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql new file mode 100644 index 0000000000..62abb07b94 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql @@ -0,0 +1,16 @@ +WITH `bfcte_0` AS ( + SELECT + `rowindex` AS `bfcol_0`, + `json_col` AS `bfcol_1` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + TO_JSON_STRING(`bfcol_1`) AS `bfcol_4` + FROM `bfcte_0` +) +SELECT + `bfcol_0` AS `bfuid_col_6`, + `bfcol_1` AS `json_col`, + `bfcol_4` AS `bfuid_col_15` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index f74d42c404..f51801c9df 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py @@ -17,64 +17,92 @@ import bigframes.bigquery as bbq import bigframes.pandas as bpd +from bigframes import operations as ops +import bigframes.core.expression as ex +from bigframes.operations._op_converters import convert_index, convert_slice +import bigframes.pandas as bpd + pytest.importorskip("pytest_snapshot") +def _apply_unary_op( + obj: bpd.DataFrame | bpd.Series, + op: ops.UnaryOp, + arg: str | ex.Expression, +) -> str: + array_value = obj._block.expr + op_expr = op.as_expr(arg) + result, _ = array_value.compute_values([op_expr]) + sql = result.session._executor.to_sql(result, enable_cache=False) + return sql -def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot): - result = bbq.array_to_string(repeated_types_df["string_list_col"], ".") - snapshot.assert_match(result.to_frame().sql, "out.sql") +def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot): + bf_df = repeated_types_df[["string_list_col"]] + sql = _apply_unary_op(bf_df, ops.ArrayToStringOp(delimiter="."), "string_list_col") + snapshot.assert_match(sql, "out.sql") def test_array_index(repeated_types_df: bpd.DataFrame, snapshot): - result = repeated_types_df["string_list_col"].list[1] - - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = repeated_types_df[["string_list_col"]] + sql = _apply_unary_op(bf_df, convert_index(1), "string_list_col") + snapshot.assert_match(sql, "out.sql") def test_array_slice_with_only_start(repeated_types_df: bpd.DataFrame, snapshot): - result = repeated_types_df["string_list_col"].list[1:] - - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = repeated_types_df[["string_list_col"]] + sql = _apply_unary_op(bf_df, convert_slice(slice(1, None)), "string_list_col") + snapshot.assert_match(sql, "out.sql") def test_array_slice_with_start_and_stop(repeated_types_df: bpd.DataFrame, snapshot): - result = repeated_types_df["string_list_col"].list[1:5] - - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = repeated_types_df[["string_list_col"]] + sql = _apply_unary_op(bf_df, convert_slice(slice(1, 5)), "string_list_col") + snapshot.assert_match(sql, "out.sql") -# JSON Ops def test_json_extract(json_types_df: bpd.DataFrame, snapshot): - result = bbq.json_extract(json_types_df["json_col"], "$") - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = json_types_df[["json_col"]] + sql = _apply_unary_op(bf_df, ops.JSONExtract(json_path="$"), "json_col") + snapshot.assert_match(sql, "out.sql") def test_json_extract_array(json_types_df: bpd.DataFrame, snapshot): - result = bbq.json_extract_array(json_types_df["json_col"], "$") - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = json_types_df[["json_col"]] + sql = _apply_unary_op(bf_df, ops.JSONExtractArray(json_path="$"), "json_col") + snapshot.assert_match(sql, "out.sql") def test_json_extract_string_array(json_types_df: bpd.DataFrame, snapshot): - result = bbq.json_extract_string_array(json_types_df["json_col"], "$") - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = json_types_df[["json_col"]] + sql = _apply_unary_op(bf_df, ops.JSONExtractStringArray(json_path="$"), "json_col") + snapshot.assert_match(sql, "out.sql") def test_json_query(json_types_df: bpd.DataFrame, snapshot): - result = bbq.json_query(json_types_df["json_col"], "$") - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = json_types_df[["json_col"]] + sql = _apply_unary_op(bf_df, ops.JSONQuery(json_path="$"), "json_col") + snapshot.assert_match(sql, "out.sql") def test_json_query_array(json_types_df: bpd.DataFrame, snapshot): - result = bbq.json_query_array(json_types_df["json_col"], "$") - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = json_types_df[["json_col"]] + sql = _apply_unary_op(bf_df, ops.JSONQueryArray(json_path="$"), "json_col") + snapshot.assert_match(sql, "out.sql") def test_json_value(json_types_df: bpd.DataFrame, snapshot): - result = bbq.json_value(json_types_df["json_col"], "$") - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = json_types_df[["json_col"]] + sql = _apply_unary_op(bf_df, ops.JSONValue(json_path="$"), "json_col") + snapshot.assert_match(sql, "out.sql") def test_parse_json(scalar_types_df: bpd.DataFrame, snapshot): - result = bbq.json_value(scalar_types_df["string_col"], "$") - snapshot.assert_match(result.to_frame().sql, "out.sql") + bf_df = scalar_types_df[["string_col"]] + sql = _apply_unary_op(bf_df, ops.JSONValue(json_path="$"), "string_col") + snapshot.assert_match(sql, "out.sql") + + +def test_to_json_string(json_types_df: bpd.DataFrame, snapshot): + bf_df = json_types_df[["json_col"]] + sql = _apply_unary_op(bf_df, ops.ToJSONString(), "json_col") + snapshot.assert_match(sql, "out.sql") \ No newline at end of file From cbcc887899669c94f513c8f846d8900edfe553a2 Mon Sep 17 00:00:00 2001 From: Chelsea Lin Date: Thu, 17 Jul 2025 20:14:57 +0000 Subject: [PATCH 4/4] fix undeterministic sql --- .../test_add_numeric/out.sql | 9 +++----- .../test_add_numeric_w_scalar/out.sql | 9 +++----- .../test_add_string/out.sql | 9 +++----- .../test_json_set/out.sql | 9 +++----- .../test_array_index/out.sql | 9 +++----- .../test_array_slice_with_only_start/out.sql | 11 ++++------ .../out.sql | 11 ++++------ .../test_array_to_string/out.sql | 9 +++----- .../test_json_extract/out.sql | 9 +++----- .../test_json_extract_array/out.sql | 9 +++----- .../test_json_extract_string_array/out.sql | 9 +++----- .../test_json_query/out.sql | 9 +++----- .../test_json_query_array/out.sql | 9 +++----- .../test_json_value/out.sql | 9 +++----- .../test_parse_json/out.sql | 9 +++----- .../test_to_json_string/out.sql | 9 +++----- .../expressions/test_binary_compiler.py | 19 +++++++++++----- .../expressions/test_unary_compiler.py | 22 +++++++++---------- 18 files changed, 74 insertions(+), 115 deletions(-) diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql index 18e1be495d..e8dc2edb80 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `int64_col` AS `bfcol_0`, - `rowindex` AS `bfcol_1` + `int64_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` ), `bfcte_1` AS ( SELECT *, - `bfcol_0` + `bfcol_0` AS `bfcol_4` + `bfcol_0` + `bfcol_0` AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_1` AS `bfuid_col_1`, - `bfcol_0` AS `int64_col`, - `bfcol_4` AS `bfuid_col_2` + `bfcol_1` AS `int64_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql index 2162f71631..7c4cc2c770 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `int64_col` AS `bfcol_0`, - `rowindex` AS `bfcol_1` + `int64_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` ), `bfcte_1` AS ( SELECT *, - `bfcol_0` + 1 AS `bfcol_4` + `bfcol_0` + 1 AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_1` AS `bfuid_col_1`, - `bfcol_0` AS `int64_col`, - `bfcol_4` AS `bfuid_col_3` + `bfcol_1` AS `int64_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql index c30e5d4cc6..de5129a6a3 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `string_col` AS `bfcol_1` + `string_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` ), `bfcte_1` AS ( SELECT *, - CONCAT(`bfcol_1`, 'a') AS `bfcol_4` + CONCAT(`bfcol_0`, 'a') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_1`, - `bfcol_1` AS `string_col`, - `bfcol_4` AS `bfuid_col_4` + `bfcol_1` AS `string_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql index e5096b9412..b226066b16 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `json_col` AS `bfcol_1` + `json_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`json_types` ), `bfcte_1` AS ( SELECT *, - JSON_SET(`bfcol_1`, '$.a', 100) AS `bfcol_4` + JSON_SET(`bfcol_0`, '$.a', 100) AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_5`, - `bfcol_1` AS `json_col`, - `bfcol_4` AS `bfuid_col_6` + `bfcol_1` AS `json_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql index b874459b14..4398084227 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `string_list_col` AS `bfcol_1` + `string_list_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`repeated_types` ), `bfcte_1` AS ( SELECT *, - `bfcol_1`[SAFE_OFFSET(1)] AS `bfcol_4` + `bfcol_0`[SAFE_OFFSET(1)] AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_1`, - `bfcol_1` AS `string_list_col`, - `bfcol_4` AS `bfuid_col_3` + `bfcol_1` AS `string_list_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql index 38eb5f2891..1ffc3ee8f9 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql @@ -1,7 +1,6 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `string_list_col` AS `bfcol_1` + `string_list_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`repeated_types` ), `bfcte_1` AS ( SELECT @@ -9,14 +8,12 @@ WITH `bfcte_0` AS ( ARRAY( SELECT el - FROM UNNEST(`bfcol_1`) AS el WITH OFFSET AS slice_idx + FROM UNNEST(`bfcol_0`) AS el WITH OFFSET AS slice_idx WHERE slice_idx >= 1 - ) AS `bfcol_4` + ) AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_1`, - `bfcol_1` AS `string_list_col`, - `bfcol_4` AS `bfuid_col_4` + `bfcol_1` AS `string_list_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql index 4947b5af12..878b60e5e2 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql @@ -1,7 +1,6 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `string_list_col` AS `bfcol_1` + `string_list_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`repeated_types` ), `bfcte_1` AS ( SELECT @@ -9,14 +8,12 @@ WITH `bfcte_0` AS ( ARRAY( SELECT el - FROM UNNEST(`bfcol_1`) AS el WITH OFFSET AS slice_idx + FROM UNNEST(`bfcol_0`) AS el WITH OFFSET AS slice_idx WHERE slice_idx >= 1 AND slice_idx < 5 - ) AS `bfcol_4` + ) AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_1`, - `bfcol_1` AS `string_list_col`, - `bfcol_4` AS `bfuid_col_5` + `bfcol_1` AS `string_list_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql index aa327a7fe4..4dbd602bea 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `string_list_col` AS `bfcol_1` + `string_list_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`repeated_types` ), `bfcte_1` AS ( SELECT *, - ARRAY_TO_STRING(`bfcol_1`, '.') AS `bfcol_4` + ARRAY_TO_STRING(`bfcol_0`, '.') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_1`, - `bfcol_1` AS `string_list_col`, - `bfcol_4` AS `bfuid_col_2` + `bfcol_1` AS `string_list_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql index 9781f2ed8d..3d23bd1e3e 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `json_col` AS `bfcol_1` + `json_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`json_types` ), `bfcte_1` AS ( SELECT *, - JSON_EXTRACT(`bfcol_1`, '$') AS `bfcol_4` + JSON_EXTRACT(`bfcol_0`, '$') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_6`, - `bfcol_1` AS `json_col`, - `bfcol_4` AS `bfuid_col_7` + `bfcol_1` AS `json_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql index fd59f5dbd8..1ddb3999b3 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `json_col` AS `bfcol_1` + `json_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`json_types` ), `bfcte_1` AS ( SELECT *, - JSON_EXTRACT_ARRAY(`bfcol_1`, '$') AS `bfcol_4` + JSON_EXTRACT_ARRAY(`bfcol_0`, '$') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_6`, - `bfcol_1` AS `json_col`, - `bfcol_4` AS `bfuid_col_8` + `bfcol_1` AS `json_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql index 024f214eab..cbc3df74c0 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `json_col` AS `bfcol_1` + `json_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`json_types` ), `bfcte_1` AS ( SELECT *, - JSON_EXTRACT_STRING_ARRAY(`bfcol_1`, '$') AS `bfcol_4` + JSON_EXTRACT_STRING_ARRAY(`bfcol_0`, '$') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_6`, - `bfcol_1` AS `json_col`, - `bfcol_4` AS `bfuid_col_9` + `bfcol_1` AS `json_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql index fc61ca2402..b5d98b80d2 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `json_col` AS `bfcol_1` + `json_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`json_types` ), `bfcte_1` AS ( SELECT *, - JSON_QUERY(`bfcol_1`, '$') AS `bfcol_4` + JSON_QUERY(`bfcol_0`, '$') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_6`, - `bfcol_1` AS `json_col`, - `bfcol_4` AS `bfuid_col_10` + `bfcol_1` AS `json_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql index 4313601883..1b7a5908eb 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `json_col` AS `bfcol_1` + `json_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`json_types` ), `bfcte_1` AS ( SELECT *, - JSON_QUERY_ARRAY(`bfcol_1`, '$') AS `bfcol_4` + JSON_QUERY_ARRAY(`bfcol_0`, '$') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_6`, - `bfcol_1` AS `json_col`, - `bfcol_4` AS `bfuid_col_11` + `bfcol_1` AS `json_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql index 7be7f6e0b2..3a84a1a92a 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `json_col` AS `bfcol_1` + `json_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`json_types` ), `bfcte_1` AS ( SELECT *, - JSON_VALUE(`bfcol_1`, '$') AS `bfcol_4` + JSON_VALUE(`bfcol_0`, '$') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_6`, - `bfcol_1` AS `json_col`, - `bfcol_4` AS `bfuid_col_12` + `bfcol_1` AS `json_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql index ee47a61a39..cdb091ae39 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `string_col` AS `bfcol_1` + `string_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` ), `bfcte_1` AS ( SELECT *, - JSON_VALUE(`bfcol_1`, '$') AS `bfcol_4` + PARSE_JSON(`bfcol_0`) AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_13`, - `bfcol_1` AS `string_col`, - `bfcol_4` AS `bfuid_col_14` + `bfcol_1` AS `string_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql index 62abb07b94..2786973933 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql @@ -1,16 +1,13 @@ WITH `bfcte_0` AS ( SELECT - `rowindex` AS `bfcol_0`, - `json_col` AS `bfcol_1` + `json_col` AS `bfcol_0` FROM `bigframes-dev`.`sqlglot_test`.`json_types` ), `bfcte_1` AS ( SELECT *, - TO_JSON_STRING(`bfcol_1`) AS `bfcol_4` + TO_JSON_STRING(`bfcol_0`) AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_0` AS `bfuid_col_6`, - `bfcol_1` AS `json_col`, - `bfcol_4` AS `bfuid_col_15` + `bfcol_1` AS `json_col` FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py index 57c0a4b831..c046d389f6 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import typing + import pytest from bigframes import operations as ops @@ -22,14 +24,19 @@ def _apply_binary_op( - obj: bpd.DataFrame | bpd.Series, + obj: bpd.DataFrame, op: ops.BinaryOp, - l_arg: str | ex.Expression, - r_arg: str | ex.Expression, + l_arg: str, + r_arg: typing.Union[str, ex.Expression], ) -> str: array_value = obj._block.expr op_expr = op.as_expr(l_arg, r_arg) - result, _ = array_value.compute_values([op_expr]) + result, col_ids = array_value.compute_values([op_expr]) + + # Rename columns for deterministic golden SQL results. + assert len(col_ids) == 1 + result = result.rename_columns({col_ids[0]: l_arg}).select_columns([l_arg]) + sql = result.session._executor.to_sql(result, enable_cache=False) return sql @@ -54,5 +61,7 @@ def test_add_string(scalar_types_df: bpd.DataFrame, snapshot): def test_json_set(json_types_df: bpd.DataFrame, snapshot): bf_df = json_types_df[["json_col"]] - sql = _apply_binary_op(bf_df, ops.JSONSet(json_path="$.a"), "json_col", ex.const(100)) + sql = _apply_binary_op( + bf_df, ops.JSONSet(json_path="$.a"), "json_col", ex.const(100) + ) snapshot.assert_match(sql, "out.sql") diff --git a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index f51801c9df..d74e9c9d48 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py @@ -14,24 +14,22 @@ import pytest -import bigframes.bigquery as bbq -import bigframes.pandas as bpd - from bigframes import operations as ops -import bigframes.core.expression as ex from bigframes.operations._op_converters import convert_index, convert_slice import bigframes.pandas as bpd pytest.importorskip("pytest_snapshot") -def _apply_unary_op( - obj: bpd.DataFrame | bpd.Series, - op: ops.UnaryOp, - arg: str | ex.Expression, -) -> str: + +def _apply_unary_op(obj: bpd.DataFrame, op: ops.UnaryOp, arg: str) -> str: array_value = obj._block.expr op_expr = op.as_expr(arg) - result, _ = array_value.compute_values([op_expr]) + result, col_ids = array_value.compute_values([op_expr]) + + # Rename columns for deterministic golden SQL results. + assert len(col_ids) == 1 + result = result.rename_columns({col_ids[0]: arg}).select_columns([arg]) + sql = result.session._executor.to_sql(result, enable_cache=False) return sql @@ -98,11 +96,11 @@ def test_json_value(json_types_df: bpd.DataFrame, snapshot): def test_parse_json(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["string_col"]] - sql = _apply_unary_op(bf_df, ops.JSONValue(json_path="$"), "string_col") + sql = _apply_unary_op(bf_df, ops.ParseJSON(), "string_col") snapshot.assert_match(sql, "out.sql") def test_to_json_string(json_types_df: bpd.DataFrame, snapshot): bf_df = json_types_df[["json_col"]] sql = _apply_unary_op(bf_df, ops.ToJSONString(), "json_col") - snapshot.assert_match(sql, "out.sql") \ No newline at end of file + snapshot.assert_match(sql, "out.sql")