put util to bigframes.testing to avoid the importing error because of sqlglot module names matching with folder name

Author: chelsea-lin

bigframes/testing/utils.py

@@ -14,7 +14,7 @@
 
 import base64
 import decimal
-from typing import Iterable, Optional, Set, Union
+from typing import Iterable, Optional, Sequence, Set, Union
 
 import geopandas as gpd  # type: ignore
 import google.api_core.operation
@@ -25,6 +25,7 @@
 import pyarrow as pa  # type: ignore
 import pytest
 
+from bigframes.core import expression as expr
 import bigframes.functions._utils as bff_utils
 import bigframes.pandas
 
@@ -448,3 +449,22 @@ def get_function_name(func, package_requirements=None, is_row_processor=False):
     function_hash = bff_utils.get_hash(func, package_requirements)
 
     return f"bigframes_{function_hash}"
+
+
+def _apply_unary_ops(
+    obj: bigframes.pandas.DataFrame,
+    ops_list: Sequence[expr.Expression],
+    new_names: Sequence[str],
+) -> str:
+    """Applies a list of unary ops to the given DataFrame and returns the SQL
+    representing the resulting DataFrames."""
+    array_value = obj._block.expr
+    result, old_names = array_value.compute_values(ops_list)
+
+    # Rename columns for deterministic golden SQL results.
+    assert len(old_names) == len(new_names)
+    col_ids = {old_name: new_name for old_name, new_name in zip(old_names, new_names)}
+    result = result.rename_columns(col_ids).select_columns(new_names)
+
+    sql = result.session._executor.to_sql(result, enable_cache=False)
+    return sql

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

@@ -17,15 +17,15 @@
 from bigframes import operations as ops
 from bigframes.operations._op_converters import convert_index, convert_slice
 import bigframes.pandas as bpd
-from tests.unit.core.compile.sqlglot.expressions.utils import _apply_unary_ops
+from bigframes.testing import utils
 
 pytest.importorskip("pytest_snapshot")
 
 
 def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot):
     col_name = "string_list_col"
     bf_df = repeated_types_df[[col_name]]
-    sql = _apply_unary_ops(
+    sql = utils._apply_unary_ops(
         bf_df, [ops.ArrayToStringOp(delimiter=".").as_expr(col_name)], [col_name]
     )
 
@@ -35,15 +35,17 @@ 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 = _apply_unary_ops(bf_df, [convert_index(1).as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [convert_index(1).as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 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 = _apply_unary_ops(
+    sql = utils._apply_unary_ops(
         bf_df, [convert_slice(slice(1, None)).as_expr(col_name)], [col_name]
     )
 
@@ -53,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 = _apply_unary_ops(
+    sql = utils._apply_unary_ops(
         bf_df, [convert_slice(slice(1, 5)).as_expr(col_name)], [col_name]
     )
 

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

@@ -16,7 +16,7 @@
 
 from bigframes import operations as ops
 import bigframes.pandas as bpd
-from tests.unit.core.compile.sqlglot.expressions.utils import _apply_unary_ops
+from bigframes.testing import utils
 
 pytest.importorskip("pytest_snapshot")
 
@@ -40,5 +40,5 @@ 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 = _apply_unary_ops(bf_df, list(ops_map.values()), list(ops_map.keys()))
+    sql = utils._apply_unary_ops(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

@@ -16,103 +16,111 @@
 
 from bigframes import operations as ops
 import bigframes.pandas as bpd
-from tests.unit.core.compile.sqlglot.expressions.utils import _apply_unary_ops
+from bigframes.testing import utils
 
 pytest.importorskip("pytest_snapshot")
 
 
 def test_date(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.date_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.day_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(bf_df, [ops.day_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_dayofweek(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.dayofweek_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.dayofweek_op.as_expr(col_name)], [col_name]
+    )
 
     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 = _apply_unary_ops(bf_df, [ops.dayofyear_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.dayofyear_op.as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_floor_dt(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.FloorDtOp("D").as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.FloorDtOp("D").as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_hour(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.hour_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.minute_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.month_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.normalize_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.normalize_op.as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_quarter(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.quarter_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.second_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(
+    sql = utils._apply_unary_ops(
         bf_df, [ops.StrftimeOp("%Y-%m-%d").as_expr(col_name)], [col_name]
     )
 
@@ -122,78 +130,88 @@ 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 = _apply_unary_ops(bf_df, [ops.time_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.ToDatetimeOp().as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.ToDatetimeOp().as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_to_timestamp(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "int64_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.ToTimestampOp().as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.ToTimestampOp().as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_unix_micros(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.UnixMicros().as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.UnixMicros().as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_unix_millis(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.UnixMillis().as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.UnixMillis().as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_unix_seconds(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.UnixSeconds().as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(
+        bf_df, [ops.UnixSeconds().as_expr(col_name)], [col_name]
+    )
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_year(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "timestamp_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.year_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.iso_day_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.iso_week_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(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 = _apply_unary_ops(bf_df, [ops.iso_year_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(bf_df, [ops.iso_year_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")

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

@@ -16,39 +16,39 @@
 
 from bigframes import operations as ops
 import bigframes.pandas as bpd
-from tests.unit.core.compile.sqlglot.expressions.utils import _apply_unary_ops
+from bigframes.testing import utils
 
 pytest.importorskip("pytest_snapshot")
 
 
 def test_hash(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "string_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.hash_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(bf_df, [ops.hash_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_isnull(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "float64_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.isnull_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(bf_df, [ops.isnull_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_notnull(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "float64_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(bf_df, [ops.notnull_op.as_expr(col_name)], [col_name])
+    sql = utils._apply_unary_ops(bf_df, [ops.notnull_op.as_expr(col_name)], [col_name])
 
     snapshot.assert_match(sql, "out.sql")
 
 
 def test_map(scalar_types_df: bpd.DataFrame, snapshot):
     col_name = "string_col"
     bf_df = scalar_types_df[[col_name]]
-    sql = _apply_unary_ops(
+    sql = utils._apply_unary_ops(
         bf_df,
         [ops.MapOp(mappings=(("value1", "mapped1"),)).as_expr(col_name)],
         [col_name],