refactor: increase test coverages on sqlglot compiler (#2527)

Author: chelsea-lin

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -663,7 +663,7 @@ def datetime_to_integer_label_non_fixed_frequency(
             .else_((x_int - first - 1) // us + 1)  # type: ignore
             .end()
         )
-    elif rule_code == "ME":  # Monthly
+    elif rule_code in ("M", "ME"):  # Monthly
         x_int = x.year() * 12 + x.month() - 1  # type: ignore
         first = y.year() * 12 + y.month() - 1  # type: ignore
         x_int_label = (
@@ -672,7 +672,7 @@ def datetime_to_integer_label_non_fixed_frequency(
             .else_((x_int - first - 1) // n + 1)  # type: ignore
             .end()
         )
-    elif rule_code == "QE-DEC":  # Quarterly
+    elif rule_code in ("Q-DEC", "QE-DEC"):  # Quarterly
         x_int = x.year() * 4 + x.quarter() - 1  # type: ignore
         first = y.year() * 4 + y.quarter() - 1  # type: ignore
         x_int_label = (
@@ -681,7 +681,7 @@ def datetime_to_integer_label_non_fixed_frequency(
             .else_((x_int - first - 1) // n + 1)  # type: ignore
             .end()
         )
-    elif rule_code == "YE-DEC":  # Yearly
+    elif rule_code in ("A-DEC", "Y-DEC", "YE-DEC"):  # Yearly
         x_int = x.year()  # type: ignore
         first = y.year()  # type: ignore
         x_int_label = (
@@ -749,7 +749,7 @@ def integer_label_to_datetime_op_non_fixed_frequency(
             .cast(ibis_dtypes.Timestamp(timezone="UTC"))
             .cast(y.type())
         )
-    elif rule_code == "ME":  # Monthly
+    elif rule_code in ("M", "ME"):  # Monthly
         one = ibis_types.literal(1)
         twelve = ibis_types.literal(12)
         first = y.year() * twelve + y.month() - one  # type: ignore
@@ -769,7 +769,7 @@ def integer_label_to_datetime_op_non_fixed_frequency(
             0,
         )
         x_label = next_month_date - ibis_api.interval(days=1)
-    elif rule_code == "QE-DEC":  # Quarterly
+    elif rule_code in ("Q-DEC", "QE-DEC"):  # Quarterly
         one = ibis_types.literal(1)
         three = ibis_types.literal(3)
         four = ibis_types.literal(4)
@@ -792,7 +792,7 @@ def integer_label_to_datetime_op_non_fixed_frequency(
         )
 
         x_label = next_month_date - ibis_api.interval(days=1)
-    elif rule_code == "YE-DEC":  # Yearly
+    elif rule_code in ("A-DEC", "Y-DEC", "YE-DEC"):  # Yearly
         one = ibis_types.literal(1)
         first = y.year()  # type: ignore
         x = x * n + first  # type: ignore

bigframes/core/compile/sqlglot/aggregate_compiler.py

@@ -70,7 +70,5 @@ def compile_analytic(
             aggregate.arg.output_type,
         )
         return unary_compiler.compile(aggregate.op, column, window)
-    elif isinstance(aggregate, agg_expressions.BinaryAggregation):
-        raise NotImplementedError("binary analytic operations not yet supported")
     else:
         raise ValueError(f"Unexpected analytic operation: {aggregate}")

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

@@ -105,31 +105,6 @@ def _coerce_bool_to_int(typed_expr: TypedExpr) -> sge.Expression:
     return typed_expr.expr
 
 
-def _string_slice(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:
-    # local name for each element in the array
-    el = sg.to_identifier("el")
-    # local name for the index in the array
-    slice_idx = sg.to_identifier("slice_idx")
-
-    conditions: typing.List[sge.Predicate] = [slice_idx >= op.start]
-    if op.stop is not None:
-        conditions.append(slice_idx < op.stop)
-
-    selected_elements = (
-        sge.select(el)
-        .from_(
-            sge.Unnest(
-                expressions=[expr.expr],
-                alias=sge.TableAlias(columns=[el]),
-                offset=slice_idx,
-            )
-        )
-        .where(*conditions)
-    )
-
-    return sge.array(selected_elements)
-
-
 def _array_slice(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:
     # local name for each element in the array
     el = sg.to_identifier("el")

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

@@ -125,7 +125,7 @@ def _datetime_to_integer_label_non_fixed_frequency(
                 expression=sge.convert(1),
             ),
         )
-    elif rule_code == "ME":  # Monthly
+    elif rule_code in ("M", "ME"):  # Monthly
         x_int = sge.Paren(  # type: ignore
             this=sge.Add(
                 this=sge.Mul(
@@ -182,7 +182,7 @@ def _datetime_to_integer_label_non_fixed_frequency(
                 expression=sge.convert(1),
             ),
         )
-    elif rule_code == "QE-DEC":  # Quarterly
+    elif rule_code in ("Q-DEC", "QE-DEC"):  # Quarterly
         x_int = sge.Paren(  # type: ignore
             this=sge.Add(
                 this=sge.Mul(
@@ -239,7 +239,7 @@ def _datetime_to_integer_label_non_fixed_frequency(
                 expression=sge.convert(1),
             ),
         )
-    elif rule_code == "YE-DEC":  # Yearly
+    elif rule_code in ("A-DEC", "Y-DEC", "YE-DEC"):  # Yearly
         x_int = sge.Extract(this=sge.Identifier(this="YEAR"), expression=x.expr)
         first = sge.Extract(this=sge.Identifier(this="YEAR"), expression=y.expr)
         return sge.Case(

tests/unit/core/compile/sqlglot/aggregations/test_op_registration.py

@@ -42,3 +42,23 @@ def test_func(input: sge.Expression) -> sge.Expression:
         ValueError, match=r".*first parameter must be a window operator.*"
     ):
         test_func(sge.to_identifier("A"))
+
+
+def test_register_already_registered_raise_error():
+    reg = op_registration.OpRegistration()
+
+    @reg.register(agg_ops.SizeOp)
+    def test_func1(op, input):
+        return input
+
+    with pytest.raises(ValueError, match=r".*is already registered.*"):
+
+        @reg.register(agg_ops.SizeOp)
+        def test_func2(op, input):
+            return input
+
+
+def test_getitem_not_registered_raise_error():
+    reg = op_registration.OpRegistration()
+    with pytest.raises(ValueError, match=r".*is not registered.*"):
+        _ = reg[agg_ops.SizeOp()]

tests/unit/core/compile/sqlglot/expressions/snapshots/test_array_ops/test_array_index/out.sql

@@ -1,3 +1,4 @@
 SELECT
-  `string_list_col`[SAFE_OFFSET(1)] AS `string_list_col`
-FROM `bigframes-dev`.`sqlglot_test`.`repeated_types` AS `bft_0`
+  IF(SUBSTRING(`string_col`, 2, 1) <> '', SUBSTRING(`string_col`, 2, 1), NULL) AS `string_index`,
+  [`int64_col`, `int64_too`][SAFE_OFFSET(1)] AS `array_index`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_array_ops/test_array_reduce_op/out.sql

@@ -18,5 +18,10 @@ SELECT
     SELECT
       COALESCE(LOGICAL_OR(bf_arr_reduce_uid), FALSE)
     FROM UNNEST(`bool_list_col`) AS bf_arr_reduce_uid
-  ) AS `any_bool`
+  ) AS `any_bool`,
+  (
+    SELECT
+      ARRAY_AGG(bf_arr_reduce_uid IGNORE NULLS)
+    FROM UNNEST(`string_list_col`) AS bf_arr_reduce_uid
+  ) AS `array_agg_str`
 FROM `bigframes-dev`.`sqlglot_test`.`repeated_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_array_ops/test_array_slice/out.sql

@@ -0,0 +1,17 @@
+SELECT
+  SUBSTRING(`string_col`, 2, 4) AS `string_slice`,
+  ARRAY(
+    SELECT
+      el
+    FROM UNNEST([`int64_col`, `int64_too`]) AS el WITH OFFSET AS slice_idx
+    WHERE
+      slice_idx >= 1
+  ) AS `slice_only_start`,
+  ARRAY(
+    SELECT
+      el
+    FROM UNNEST([`int64_col`, `int64_too`]) AS el WITH OFFSET AS slice_idx
+    WHERE
+      slice_idx >= 1 AND slice_idx < 5
+  ) AS `slice_start_stop`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`