refactor: rename _apply_unary_ops to _apply_ops_to_sql (#2214)

* refactor: rename _apply_unary_ops to _apply_ops_to_sql

* fix lint

* fix mypy

Author: jialuoo

bigframes/testing/utils.py

@@ -448,12 +448,12 @@ def get_function_name(func, package_requirements=None, is_row_processor=False):
     return f"bigframes_{function_hash}"
 
 
-def _apply_unary_ops(
+def _apply_ops_to_sql(
     obj: bpd.DataFrame,
     ops_list: Sequence[ex.Expression],
     new_names: Sequence[str],
 ) -> str:
-    """Applies a list of unary ops to the given DataFrame and returns the SQL
+    """Applies a list of ops to the given DataFrame and returns the SQL
     representing the resulting DataFrame."""
     array_value = obj._block.expr
     result, old_names = array_value.compute_values(ops_list)
@@ -485,13 +485,6 @@ def _apply_nary_op(
 ) -> 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]: args[0]}).select_columns([args[0]])
-
-    sql = result.session._executor.to_sql(result, enable_cache=False)
+    sql = _apply_ops_to_sql(obj, [op_expr], [args[0]])  # type: ignore
     return sql

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

@@ -39,7 +39,7 @@ def test_ai_generate(scalar_types_df: dataframe.DataFrame, snapshot):
         output_schema=None,
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -58,7 +58,7 @@ def test_ai_generate_with_output_schema(scalar_types_df: dataframe.DataFrame, sn
         output_schema="x INT64, y FLOAT64",
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -82,7 +82,7 @@ def test_ai_generate_with_model_param(scalar_types_df: dataframe.DataFrame, snap
         output_schema=None,
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -100,7 +100,7 @@ def test_ai_generate_bool(scalar_types_df: dataframe.DataFrame, snapshot):
         model_params=None,
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -125,7 +125,7 @@ def test_ai_generate_bool_with_model_param(
         model_params=json.dumps(dict()),
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -144,7 +144,7 @@ def test_ai_generate_int(scalar_types_df: dataframe.DataFrame, snapshot):
         model_params=None,
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -170,7 +170,7 @@ def test_ai_generate_int_with_model_param(
         model_params=json.dumps(dict()),
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -189,7 +189,7 @@ def test_ai_generate_double(scalar_types_df: dataframe.DataFrame, snapshot):
         model_params=None,
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -215,7 +215,7 @@ def test_ai_generate_double_with_model_param(
         model_params=json.dumps(dict()),
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -230,7 +230,7 @@ def test_ai_if(scalar_types_df: dataframe.DataFrame, snapshot):
         connection_id=CONNECTION_ID,
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 
@@ -246,7 +246,7 @@ def test_ai_classify(scalar_types_df: dataframe.DataFrame, snapshot):
         connection_id=CONNECTION_ID,
     )
 
-    sql = utils._apply_unary_ops(scalar_types_df, [op.as_expr(col_name)], ["result"])
+    sql = utils._apply_ops_to_sql(scalar_types_df, [op.as_expr(col_name)], ["result"])
 
     snapshot.assert_match(sql, "out.sql")
 
@@ -259,7 +259,7 @@ def test_ai_score(scalar_types_df: dataframe.DataFrame, snapshot):
         connection_id=CONNECTION_ID,
     )
 
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         scalar_types_df, [op.as_expr(col_name, col_name)], ["result"]
     )
 

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

@@ -25,7 +25,7 @@
 def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot):
     col_name = "string_list_col"
     bf_df = repeated_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.ArrayToStringOp(delimiter=".").as_expr(col_name)], [col_name]
     )
 
@@ -35,7 +35,7 @@ def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot):
 def test_array_index(repeated_types_df: bpd.DataFrame, snapshot):
     col_name = "string_list_col"
     bf_df = repeated_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [convert_index(1).as_expr(col_name)], [col_name]
     )
 
@@ -45,7 +45,7 @@ def test_array_index(repeated_types_df: bpd.DataFrame, snapshot):
 def test_array_slice_with_only_start(repeated_types_df: bpd.DataFrame, snapshot):
     col_name = "string_list_col"
     bf_df = repeated_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [convert_slice(slice(1, None)).as_expr(col_name)], [col_name]
     )
 
@@ -55,7 +55,7 @@ def test_array_slice_with_only_start(repeated_types_df: bpd.DataFrame, snapshot)
 def test_array_slice_with_start_and_stop(repeated_types_df: bpd.DataFrame, snapshot):
     col_name = "string_list_col"
     bf_df = repeated_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [convert_slice(slice(1, 5)).as_expr(col_name)], [col_name]
     )
 

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

@@ -40,7 +40,7 @@ def test_is_in(scalar_types_df: bpd.DataFrame, snapshot):
         "float_in_ints": ops.IsInOp(values=(1, 2, 3, None)).as_expr(float_col),
     }
 
-    sql = utils._apply_unary_ops(bf_df, list(ops_map.values()), list(ops_map.keys()))
+    sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
     snapshot.assert_match(sql, "out.sql")
 
 

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

@@ -25,15 +25,15 @@
 def test_date(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.date_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.date_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_day(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.day_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.day_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
@@ -43,14 +43,14 @@ def test_dayofweek(scalar_types_df: bpd.DataFrame, snapshot):
     bf_df = scalar_types_df[col_names]
     ops_map = {col_name: ops.dayofweek_op.as_expr(col_name) for col_name in col_names}
 
-    sql = utils._apply_unary_ops(bf_df, list(ops_map.values()), list(ops_map.keys()))
+    sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_dayofyear(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.dayofyear_op.as_expr(col_name)], [col_name]
     )
 
@@ -75,7 +75,7 @@ def test_floor_dt(scalar_types_df: bpd.DataFrame, snapshot):
         "datetime_col_us": ops.FloorDtOp("us").as_expr("datetime_col"),
     }
 
-    sql = utils._apply_unary_ops(bf_df, list(ops_map.values()), list(ops_map.keys()))
+    sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
     snapshot.assert_match(sql, "out.sql")
 
 
@@ -85,7 +85,7 @@ def test_floor_dt_op_invalid_freq(scalar_types_df: bpd.DataFrame):
     with pytest.raises(
         NotImplementedError, match="Unsupported freq paramater: invalid"
     ):
-        utils._apply_unary_ops(
+        utils._apply_ops_to_sql(
             bf_df,
             [ops.FloorDtOp(freq="invalid").as_expr(col_name)],  # type:ignore
             [col_name],
@@ -95,31 +95,31 @@ def test_floor_dt_op_invalid_freq(scalar_types_df: bpd.DataFrame):
 def test_hour(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.hour_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.hour_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_minute(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.minute_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.minute_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_month(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.month_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.month_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_normalize(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.normalize_op.as_expr(col_name)], [col_name]
     )
 
@@ -129,23 +129,23 @@ def test_normalize(scalar_types_df: bpd.DataFrame, snapshot):
 def test_quarter(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.quarter_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.quarter_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_second(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.second_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.second_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_strftime(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.StrftimeOp("%Y-%m-%d").as_expr(col_name)], [col_name]
     )
 
@@ -155,15 +155,15 @@ def test_strftime(scalar_types_df: bpd.DataFrame, snapshot):
 def test_time(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.time_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.time_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_to_datetime(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "int64_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.ToDatetimeOp().as_expr(col_name)], [col_name]
     )
 
@@ -173,7 +173,7 @@ def test_to_datetime(scalar_types_df: bpd.DataFrame, snapshot):
 def test_to_timestamp(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "int64_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.ToTimestampOp().as_expr(col_name)], [col_name]
     )
 
@@ -183,7 +183,7 @@ def test_to_timestamp(scalar_types_df: bpd.DataFrame, snapshot):
 def test_unix_micros(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.UnixMicros().as_expr(col_name)], [col_name]
     )
 
@@ -193,7 +193,7 @@ def test_unix_micros(scalar_types_df: bpd.DataFrame, snapshot):
 def test_unix_millis(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.UnixMillis().as_expr(col_name)], [col_name]
     )
 
@@ -203,7 +203,7 @@ def test_unix_millis(scalar_types_df: bpd.DataFrame, snapshot):
 def test_unix_seconds(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(
+    sql = utils._apply_ops_to_sql(
         bf_df, [ops.UnixSeconds().as_expr(col_name)], [col_name]
     )
 
@@ -213,31 +213,35 @@ def test_unix_seconds(scalar_types_df: bpd.DataFrame, snapshot):
 def test_year(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.year_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.year_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_iso_day(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.iso_day_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(bf_df, [ops.iso_day_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_iso_week(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.iso_week_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(
+        bf_df, [ops.iso_week_op.as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_iso_year(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = utils._apply_unary_ops(bf_df, [ops.iso_year_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_ops_to_sql(
+        bf_df, [ops.iso_year_op.as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")