refactor: increase test coverages on sqlglot compiler

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`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_get_access_url_with_duration/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  OBJ.GET_ACCESS_URL(`string_col`, 'READ', INTERVAL 3600 MICROSECOND) AS `string_col`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_make_ref_json/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  OBJ.MAKE_REF(`string_col`) AS `string_col`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_and_op/out.sql

@@ -4,5 +4,6 @@ SELECT
   `int64_col`,
   `int64_col` & `int64_col` AS `int_and_int`,
   `bool_col` AND `bool_col` AS `bool_and_bool`,
-  IF(`bool_col` = FALSE, `bool_col`, NULL) AS `bool_and_null`
+  IF(`bool_col` = FALSE, `bool_col`, NULL) AS `bool_and_null`,
+  IF(`bool_col` = FALSE, `bool_col`, NULL) AS `null_and_bool`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_or_op/out.sql

@@ -4,5 +4,6 @@ SELECT
   `int64_col`,
   `int64_col` | `int64_col` AS `int_and_int`,
   `bool_col` OR `bool_col` AS `bool_and_bool`,
-  IF(`bool_col` = TRUE, `bool_col`, NULL) AS `bool_and_null`
+  IF(`bool_col` = TRUE, `bool_col`, NULL) AS `bool_and_null`,
+  IF(`bool_col` = TRUE, `bool_col`, NULL) AS `null_and_bool`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_bool_ops/test_xor_op/out.sql

@@ -13,5 +13,11 @@ SELECT
   )
   OR (
     NOT `bool_col` AND CAST(NULL AS BOOLEAN)
-  ) AS `bool_and_null`
+  ) AS `bool_and_null`,
+  (
+    `bool_col` AND NOT CAST(NULL AS BOOLEAN)
+  )
+  OR (
+    NOT `bool_col` AND CAST(NULL AS BOOLEAN)
+  ) AS `null_and_bool`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_eq_numeric/out.sql

@@ -5,6 +5,7 @@ SELECT
   `int64_col` = `int64_col` AS `int_eq_int`,
   `int64_col` = 1 AS `int_eq_1`,
   `int64_col` IS NULL AS `int_eq_null`,
+  `int64_col` IS NULL AS `null_eq_int`,
   `int64_col` = CAST(`bool_col` AS INT64) AS `int_eq_bool`,
   CAST(`bool_col` AS INT64) = `int64_col` AS `bool_eq_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_ge_numeric/out.sql

@@ -4,6 +4,7 @@ SELECT
   `bool_col`,
   `int64_col` >= `int64_col` AS `int_ge_int`,
   `int64_col` >= 1 AS `int_ge_1`,
+  NULL AS `null_ge_int`,
   `int64_col` >= CAST(`bool_col` AS INT64) AS `int_ge_bool`,
   CAST(`bool_col` AS INT64) >= `int64_col` AS `bool_ge_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_gt_numeric/out.sql

@@ -4,6 +4,7 @@ SELECT
   `bool_col`,
   `int64_col` > `int64_col` AS `int_gt_int`,
   `int64_col` > 1 AS `int_gt_1`,
+  NULL AS `null_gt_int`,
   `int64_col` > CAST(`bool_col` AS INT64) AS `int_gt_bool`,
   CAST(`bool_col` AS INT64) > `int64_col` AS `bool_gt_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_is_in/out.sql

@@ -10,5 +10,10 @@ SELECT
   COALESCE(`int64_col` IN (123456), FALSE) AS `ints_wo_match_nulls`,
   (
     `float64_col` IS NULL
-  ) OR `float64_col` IN (1, 2, 3) AS `float_in_ints`
+  ) OR `float64_col` IN (1, 2, 3) AS `float_in_ints`,
+  (
+    `int64_col` IS NULL
+  ) OR `int64_col` IN (2) AS `mixed_with_null`,
+  COALESCE(CAST(`bool_col` AS INT64) IN (1, 2.5), FALSE) AS `bool_in_mixed`,
+  `int64_col` IS NULL AS `only_null_match`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_le_numeric/out.sql

@@ -4,6 +4,7 @@ SELECT
   `bool_col`,
   `int64_col` <= `int64_col` AS `int_le_int`,
   `int64_col` <= 1 AS `int_le_1`,
+  NULL AS `null_le_int`,
   `int64_col` <= CAST(`bool_col` AS INT64) AS `int_le_bool`,
   CAST(`bool_col` AS INT64) <= `int64_col` AS `bool_le_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_lt_numeric/out.sql

@@ -4,6 +4,7 @@ SELECT
   `bool_col`,
   `int64_col` < `int64_col` AS `int_lt_int`,
   `int64_col` < 1 AS `int_lt_1`,
+  NULL AS `null_lt_int`,
   `int64_col` < CAST(`bool_col` AS INT64) AS `int_lt_bool`,
   CAST(`bool_col` AS INT64) < `int64_col` AS `bool_lt_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_comparison_ops/test_ne_numeric/out.sql

@@ -7,6 +7,9 @@ SELECT
   (
     `int64_col`
   ) IS NOT NULL AS `int_ne_null`,
+  (
+    `int64_col`
+  ) IS NOT NULL AS `null_ne_int`,
   `int64_col` <> CAST(`bool_col` AS INT64) AS `int_ne_bool`,
   CAST(`bool_col` AS INT64) <> `int64_col` AS `bool_ne_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_datetime_to_integer_label/out.sql

@@ -5,6 +5,13 @@ SELECT
       86400000000
     )
   ) AS INT64) AS `fixed_freq`,
+  CAST(FLOOR(IEEE_DIVIDE(UNIX_MICROS(CAST(`datetime_col` AS TIMESTAMP)) - 0, 86400000000)) AS INT64) AS `origin_epoch`,
+  CAST(FLOOR(
+    IEEE_DIVIDE(
+      UNIX_MICROS(CAST(`datetime_col` AS TIMESTAMP)) - UNIX_MICROS(CAST(CAST(`timestamp_col` AS DATE) AS TIMESTAMP)),
+      86400000000
+    )
+  ) AS INT64) AS `origin_start_day`,
   CASE
     WHEN UNIX_MICROS(
       CAST(TIMESTAMP_TRUNC(`datetime_col`, WEEK(MONDAY)) + INTERVAL 6 DAY AS TIMESTAMP)
@@ -22,5 +29,48 @@ SELECT
         604800000000
       )
     ) AS INT64) + 1
-  END AS `non_fixed_freq_weekly`
+  END AS `non_fixed_freq_weekly`,
+  CASE
+    WHEN (
+      EXTRACT(YEAR FROM `datetime_col`) * 12 + EXTRACT(MONTH FROM `datetime_col`) - 1
+    ) = (
+      EXTRACT(YEAR FROM `timestamp_col`) * 12 + EXTRACT(MONTH FROM `timestamp_col`) - 1
+    )
+    THEN 0
+    ELSE CAST(FLOOR(
+      IEEE_DIVIDE(
+        (
+          EXTRACT(YEAR FROM `datetime_col`) * 12 + EXTRACT(MONTH FROM `datetime_col`) - 1
+        ) - (
+          EXTRACT(YEAR FROM `timestamp_col`) * 12 + EXTRACT(MONTH FROM `timestamp_col`) - 1
+        ) - 1,
+        1
+      )
+    ) AS INT64) + 1
+  END AS `non_fixed_freq_monthly`,
+  CASE
+    WHEN (
+      EXTRACT(YEAR FROM `datetime_col`) * 4 + EXTRACT(QUARTER FROM `datetime_col`) - 1
+    ) = (
+      EXTRACT(YEAR FROM `timestamp_col`) * 4 + EXTRACT(QUARTER FROM `timestamp_col`) - 1
+    )
+    THEN 0
+    ELSE CAST(FLOOR(
+      IEEE_DIVIDE(
+        (
+          EXTRACT(YEAR FROM `datetime_col`) * 4 + EXTRACT(QUARTER FROM `datetime_col`) - 1
+        ) - (
+          EXTRACT(YEAR FROM `timestamp_col`) * 4 + EXTRACT(QUARTER FROM `timestamp_col`) - 1
+        ) - 1,
+        1
+      )
+    ) AS INT64) + 1
+  END AS `non_fixed_freq_quarterly`,
+  CASE
+    WHEN EXTRACT(YEAR FROM `datetime_col`) = EXTRACT(YEAR FROM `timestamp_col`)
+    THEN 0
+    ELSE CAST(FLOOR(
+      IEEE_DIVIDE(EXTRACT(YEAR FROM `datetime_col`) - EXTRACT(YEAR FROM `timestamp_col`) - 1, 1)
+    ) AS INT64) + 1
+  END AS `non_fixed_freq_yearly`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_to_datetime/out.sql

@@ -2,5 +2,6 @@ SELECT
   CAST(TIMESTAMP_MICROS(CAST(TRUNC(`int64_col` * 0.001) AS INT64)) AS DATETIME) AS `int64_col`,
   SAFE_CAST(`string_col` AS DATETIME),
   CAST(TIMESTAMP_MICROS(CAST(TRUNC(`float64_col` * 0.001) AS INT64)) AS DATETIME) AS `float64_col`,
-  SAFE_CAST(`timestamp_col` AS DATETIME)
+  SAFE_CAST(`timestamp_col` AS DATETIME),
+  CAST(PARSE_TIMESTAMP('%Y-%m-%d', `string_col`, 'UTC') AS DATETIME) AS `string_col_fmt`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`