address comments

Author: chelsea-lin

bigframes/core/compile/sqlglot/expressions/generic_ops.py

@@ -83,7 +83,6 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_nary_op(ops.SqlScalarOp, pass_op=True)
 def _(*operands: TypedExpr, op: ops.SqlScalarOp) -> sge.Expression:
-    # TODO: can we include a string in the sqlglot expression without parsing?
     return sg.parse_one(
         op.sql_template.format(
             *[operand.expr.sql(dialect="bigquery") for operand in operands]

bigframes/testing/utils.py

@@ -475,13 +475,23 @@ def _apply_binary_op(
 ) -> str:
     """Applies a binary op to the given DataFrame and return the SQL representing
     the resulting DataFrame."""
+    return _apply_nary_op(obj, op, l_arg, r_arg)
+
+
+def _apply_nary_op(
+    obj: bpd.DataFrame,
+    op: Union[ops.BinaryOp, ops.NaryOp],
+    *args: 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(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