10BC0 refactor: sqlglot scalar unit tests expressions by chelsea-lin · Pull Request #1916 · googleapis/python-bigquery-dataframes · GitHub
[go: up one dir, main page]
Skip to content

refactor: sqlglot scalar unit tests expressions #1916

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chelsea-lin merged 4 commits into from
Jul 18, 2025
Merged

refactor: sqlglot scalar unit tests expressions #1916

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
dbccc47
uniform json golden sql tests
chelsea-lin Jul 16, 2025
cc85630
scalar tests use array_value rather than open APIs
chelsea-lin Jul 17, 2025
a1d3d76
convert unary tests too
chelsea-lin Jul 17, 2025
cbcc887
fix undeterministic sql
chelsea-lin Jul 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 3 additions & 6 deletions .../core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
9 changes: 3 additions & 6 deletions ...pile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
9 changes: 3 additions & 6 deletions ...t/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
15 changes: 4 additions & 11 deletions ...nit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
`bfcol_1` AS `json_col`
FROM `bfcte_1`
8 changes: 3 additions & 5 deletions ...t/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
10 changes: 4 additions & 6 deletions ...qlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,21 +1,19 @@
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(
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`
10 changes: 4 additions & 6 deletions ...ot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,21 +1,19 @@
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(
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`
8 changes: 3 additions & 5 deletions ...re/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
8 changes: 3 additions & 5 deletions .../core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
13 changes: 13 additions & 0 deletions ...compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
13 changes: 13 additions & 0 deletions .../sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
13 changes: 13 additions & 0 deletions ...it/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
13 changes: 13 additions & 0 deletions ...e/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql
View file Delete file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
13 changes: 13 additions & 0 deletions ...it/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
8 changes: 3 additions & 5 deletions ...it/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
13 changes: 13 additions & 0 deletions ...ore/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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`
46 changes: 32 additions & 14 deletions tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Loading
0