refactor: sqlglot scalar unit tests expressions (#1916)

chelsea-lin · web-flow · commit ce1aa48b9cde · 2025-07-18T10:44:40.000-07:00
The sqlglot scalar unit tests are being converted to use ArrayValue for operations, since some scalar functions do not have a matching pandas API.
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_numeric_w_scalar/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_add_string/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_json_set/out.sql b/packages/bigframes/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`
+  `bfcol_1` AS `json_col`
+FROM `bfcte_1`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_index/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_only_start/out.sql
@@ -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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_slice_with_start_and_stop/out.sql
@@ -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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_array_to_string/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_array/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_extract_string_array/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_query_array/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_json_value/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_parse_json/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_json_string/out.sql b/packages/bigframes/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`
diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py b/packages/bigframes/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/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py
