fix import errors

Author: chelsea-lin

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

@@ -17,15 +17,16 @@
 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 . 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 +36,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 +56,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,8 @@
 
 from bigframes import operations as ops
 import bigframes.pandas as bpd
-from tests.unit.core.compile.sqlglot.expressions.utils import _apply_unary_ops
+
+from . import utils
 
 pytest.importorskip("pytest_snapshot")
 
@@ -40,5 +41,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,112 @@
 
 from bigframes import operations as ops
 import bigframes.pandas as bpd
-from tests.unit.core.compile.sqlglot.expressions.utils import _apply_unary_ops
+
+from . 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 +131,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,40 @@
 
 from bigframes import operations as ops
 import bigframes.pandas as bpd
-from tests.unit.core.compile.sqlglot.expressions.utils import _apply_unary_ops
+
+from . 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],