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..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_1` AS `bfcol_4`, - `bfcol_0` + `bfcol_0` AS `bfcol_5` + `bfcol_0` + `bfcol_0` AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_4` AS `rowindex`, - `bfcol_5` AS `int64_col` + `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 9c4b01a6df..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_1` AS `bfcol_4`, - `bfcol_0` + 1 AS `bfcol_5` + `bfcol_0` + 1 AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_4` AS `rowindex`, - `bfcol_5` AS `int64_col` + `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 7a8ab83df1..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 *, - `bfcol_0` AS `bfcol_4`, - CONCAT(`bfcol_1`, 'a') AS `bfcol_5` + CONCAT(`bfcol_0`, 'a') AS `bfcol_1` FROM `bfcte_0` ) SELECT - `bfcol_4` AS `rowindex`, - `bfcol_5` AS `string_col` + `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 f501dd3b86..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,20 +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` -), `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_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 33a8bded13..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,15 +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 `rowindex`, - `bfcol_4` AS `string_list_col` + `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 34d2225931..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,13 +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 `rowindex`, - `bfcol_4` AS `string_list_col` + `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 d46803ce7c..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,13 +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 `rowindex`, - `bfcol_4` AS `string_list_col` + `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 e0db21f972..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,15 +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 `rowindex`, - `bfcol_4` AS `string_list_col` + `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 2ffb0174a8..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,15 +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 `rowindex`, - `bfcol_4` AS `json_col` + `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 new file mode 100644 index 0000000000..1ddb3999b3 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `json_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_EXTRACT_ARRAY(`bfcol_0`, '$') AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `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 new file mode 100644 index 0000000000..cbc3df74c0 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `json_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_EXTRACT_STRING_ARRAY(`bfcol_0`, '$') AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `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 new file mode 100644 index 0000000000..b5d98b80d2 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `json_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_QUERY(`bfcol_0`, '$') AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `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 new file mode 100644 index 0000000000..1b7a5908eb --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `json_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_QUERY_ARRAY(`bfcol_0`, '$') AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `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 new file mode 100644 index 0000000000..3a84a1a92a --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `json_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + JSON_VALUE(`bfcol_0`, '$') AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `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 d965ea8f1b..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,15 +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 `rowindex`, - `bfcol_4` AS `string_col` + `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 new file mode 100644 index 0000000000..2786973933 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `json_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`json_types` +), `bfcte_1` AS ( + SELECT + *, + TO_JSON_STRING(`bfcol_0`) AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `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 9daff51c9f..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,38 +12,56 @@ # See the License for the specific language governing permissions and # limitations under the License. +import typing + 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, + op: ops.BinaryOp, + 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, col_ids = array_value.compute_values([op_expr]) - bf_df["int64_col"] = bf_df["int64_col"] + bf_df["int64_col"] + # 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]) - snapshot.assert_match(bf_df.sql, "out.sql") + sql = result.session._executor.to_sql(result, enable_cache=False) + return sql -def test_add_numeric_w_scalar(scalar_types_df: bpd.DataFrame, snapshot): +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") - bf_df["int64_col"] = bf_df["int64_col"] + 1 - snapshot.assert_match(bf_df.sql, "out.sql") +def test_add_numeric_w_scalar(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["int64_col"]] + 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") 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..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,74 +14,93 @@ import pytest -import bigframes.bigquery as bbq +from bigframes import operations as ops +from bigframes.operations._op_converters import convert_index, convert_slice import bigframes.pandas as bpd pytest.importorskip("pytest_snapshot") -def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot): - result = bbq.array_to_string(repeated_types_df["string_list_col"], ".") +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, col_ids = array_value.compute_values([op_expr]) - snapshot.assert_match(result.to_frame().sql, "out.sql") + # 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 -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") +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_slice_with_only_start(repeated_types_df: bpd.DataFrame, snapshot): - result = repeated_types_df["string_list_col"].list[1:] +def test_array_index(repeated_types_df: bpd.DataFrame, snapshot): + 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") - snapshot.assert_match(result.to_frame().sql, "out.sql") +def test_array_slice_with_only_start(repeated_types_df: bpd.DataFrame, snapshot): + 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") +def test_array_slice_with_start_and_stop(repeated_types_df: bpd.DataFrame, snapshot): + 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"], "$") - 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") + 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): - 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 +def test_json_extract_array(json_types_df: bpd.DataFrame, snapshot): + 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): - 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 +def test_json_extract_string_array(json_types_df: bpd.DataFrame, snapshot): + 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): - 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 +def test_json_query(json_types_df: bpd.DataFrame, snapshot): + 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): - 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 +def test_json_query_array(json_types_df: bpd.DataFrame, snapshot): + 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): - 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 +def test_json_value(json_types_df: bpd.DataFrame, snapshot): + 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.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")