scalar tests use array_value rather than open APIs

Author: chelsea-lin

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`

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`

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`

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`
+  `bfcol_0` AS `bfuid_col_5`,
+  `bfcol_1` AS `json_col`,
+  `bfcol_4` AS `bfuid_col_6`
+FROM `bfcte_1`

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")