reuse the testing utils.py method

Author: chelsea-lin

bigframes/testing/utils.py

@@ -467,21 +467,20 @@ def _apply_unary_ops(
     return sql
 
 
-def _apply_binary_op(
+def _apply_nary_op(
     obj: bpd.DataFrame,
-    op: ops.BinaryOp,
-    l_arg: str,
-    r_arg: Union[str, ex.Expression],
+    op: Union[ops.BinaryOp, ops.NaryOp],
+    *args: Union[str, ex.Expression],
 ) -> str:
-    """Applies a binary op to the given DataFrame and return the SQL representing
+    """Applies a nary op to the given DataFrame and return the SQL representing
     the resulting DataFrame."""
     array_value = obj._block.expr
-    op_expr = op.as_expr(l_arg, r_arg)
+    op_expr = op.as_expr(*args)
     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])
+    result = result.rename_columns({col_ids[0]: args[0]}).select_columns([args[0]])
 
     sql = result.session._executor.to_sql(result, enable_cache=False)
     return sql

tests/unit/core/compile/sqlglot/expressions/snapshots/test_struct_ops/test_struct_op/out.sql

@@ -17,5 +17,5 @@ WITH `bfcte_0` AS (
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_4` AS `result_col`
+  `bfcol_4` AS `bool_col`
 FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/test_comparison_ops.py

@@ -46,7 +46,7 @@ def test_is_in(scalar_types_df: bpd.DataFrame, snapshot):
 
 def test_eq_null_match(scalar_types_df: bpd.DataFrame, snapshot):
     bf_df = scalar_types_df[["int64_col", "bool_col"]]
-    sql = utils._apply_binary_op(bf_df, ops.eq_null_match_op, "int64_col", "bool_col")
+    sql = utils._apply_nary_op(bf_df, ops.eq_null_match_op, "int64_col", "bool_col")
     snapshot.assert_match(sql, "out.sql")
 
 

tests/unit/core/compile/sqlglot/expressions/test_json_ops.py

@@ -102,7 +102,7 @@ def test_to_json_string(json_types_df: bpd.DataFrame, snapshot):
 
 def test_json_set(json_types_df: bpd.DataFrame, snapshot):
     bf_df = json_types_df[["json_col"]]
-    sql = utils._apply_binary_op(
+    sql = utils._apply_nary_op(
         bf_df, ops.JSONSet(json_path="$.a"), "json_col", ex.const(100)
     )
 

tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py

@@ -221,7 +221,7 @@ def test_add_numeric(scalar_types_df: bpd.DataFrame, snapshot):
 
 def test_add_string(scalar_types_df: bpd.DataFrame, snapshot):
     bf_df = scalar_types_df[["string_col"]]
-    sql = utils._apply_binary_op(bf_df, ops.add_op, "string_col", ex.const("a"))
+    sql = utils._apply_nary_op(bf_df, ops.add_op, "string_col", ex.const("a"))
 
     snapshot.assert_match(sql, "out.sql")
 
@@ -241,10 +241,10 @@ def test_add_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
 
 def test_add_unsupported_raises(scalar_types_df: bpd.DataFrame):
     with pytest.raises(TypeError):
-        utils._apply_binary_op(scalar_types_df, ops.add_op, "timestamp_col", "date_col")
+        utils._apply_nary_op(scalar_types_df, ops.add_op, "timestamp_col", "date_col")
 
     with pytest.raises(TypeError):
-        utils._apply_binary_op(scalar_types_df, ops.add_op, "int64_col", "string_col")
+        utils._apply_nary_op(scalar_types_df, ops.add_op, "int64_col", "string_col")
 
 
 def test_div_numeric(scalar_types_df: bpd.DataFrame, snapshot):
@@ -361,7 +361,7 @@ def test_sub_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
 
 def test_sub_unsupported_raises(scalar_types_df: bpd.DataFrame):
     with pytest.raises(TypeError):
-        utils._apply_binary_op(scalar_types_df, ops.sub_op, "string_col", "string_col")
+        utils._apply_nary_op(scalar_types_df, ops.sub_op, "string_col", "string_col")
 
     with pytest.raises(TypeError):
-        utils._apply_binary_op(scalar_types_df, ops.sub_op, "int64_col", "string_col")
+        utils._apply_nary_op(scalar_types_df, ops.sub_op, "int64_col", "string_col")

tests/unit/core/compile/sqlglot/expressions/test_string_ops.py

@@ -308,13 +308,13 @@ def test_zfill(scalar_types_df: bpd.DataFrame, snapshot):
 
 def test_add_string(scalar_types_df: bpd.DataFrame, snapshot):
     bf_df = scalar_types_df[["string_col"]]
-    sql = utils._apply_binary_op(bf_df, ops.add_op, "string_col", ex.const("a"))
+    sql = utils._apply_nary_op(bf_df, ops.add_op, "string_col", ex.const("a"))
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_strconcat(scalar_types_df: bpd.DataFrame, snapshot):
     bf_df = scalar_types_df[["string_col"]]
-    sql = utils._apply_binary_op(bf_df, ops.strconcat_op, "string_col", ex.const("a"))
+    sql = utils._apply_nary_op(bf_df, ops.strconcat_op, "string_col", ex.const("a"))
 
     snapshot.assert_match(sql, "out.sql")

tests/unit/core/compile/sqlglot/expressions/test_struct_ops.py

@@ -12,39 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import typing
-
 import pytest
 
 from bigframes import operations as ops
-from bigframes.core import expression as ex
 import bigframes.pandas as bpd
 from bigframes.testing import utils
 
 pytest.importorskip("pytest_snapshot")
 
 
-def _apply_nary_op(
-    obj: bpd.DataFrame,
-    op: ops.NaryOp,
-    *args: typing.Union[str, ex.Expression],
-) -> str:
-    """Applies a nary op to the given DataFrame and return the SQL representing
-    the resulting DataFrame."""
-    array_value = obj._block.expr
-    op_expr = op.as_expr(*args)
-    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]: "result_col"}).select_columns(
-        ["result_col"]
-    )
-
-    sql = result.session._executor.to_sql(result, enable_cache=False)
-    return sql
-
-
 def test_struct_field(nested_structs_types_df: bpd.DataFrame, snapshot):
     col_name = "people"
     bf_df = nested_structs_types_df[[col_name]]
@@ -63,6 +39,6 @@ def test_struct_field(nested_structs_types_df: bpd.DataFrame, snapshot):
 def test_struct_op(scalar_types_df: bpd.DataFrame, snapshot):
     bf_df = scalar_types_df[["bool_col", "int64_col", "float64_col", "string_col"]]
     op = ops.StructOp(column_names=tuple(bf_df.columns.tolist()))
-    sql = _apply_nary_op(bf_df, op, *bf_df.columns.tolist())
+    sql = utils._apply_nary_op(bf_df, op, *bf_df.columns.tolist())
 
     snapshot.assert_match(sql, "out.sql")