refactor: increase test coverages on sqlglot compiler

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")

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`

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

@@ -5,5 +5,6 @@ SELECT
   CAST(TIMESTAMP_MICROS(CAST(TRUNC(`int64_col` * 1000) AS INT64)) AS TIMESTAMP) AS `int64_col_ms`,
   CAST(TIMESTAMP_MICROS(CAST(TRUNC(`int64_col`) AS INT64)) AS TIMESTAMP) AS `int64_col_us`,
   CAST(TIMESTAMP_MICROS(CAST(TRUNC(`int64_col` * 0.001) AS INT64)) AS TIMESTAMP) AS `int64_col_ns`,
-  TIMESTAMP(`datetime_col`) AS `datetime_col`
+  TIMESTAMP(`datetime_col`) AS `datetime_col`,
+  PARSE_TIMESTAMP('%Y-%m-%d', `string_col`, 'UTC') AS `string_col_fmt`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_json_ops/test_json_value_array/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  JSON_VALUE_ARRAY(`json_col`, '$') AS `json_col`
+FROM `bigframes-dev`.`sqlglot_test`.`json_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_add_numeric/out.sql

@@ -4,6 +4,7 @@ SELECT
   `bool_col`,
   `int64_col` + `int64_col` AS `int_add_int`,
   `int64_col` + 1 AS `int_add_1`,
+  NULL AS `int_add_null`,
   `int64_col` + CAST(`bool_col` AS INT64) AS `int_add_bool`,
   CAST(`bool_col` AS INT64) + `int64_col` AS `bool_add_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_div_numeric/out.sql

@@ -6,6 +6,7 @@ SELECT
   IEEE_DIVIDE(`int64_col`, `int64_col`) AS `int_div_int`,
   IEEE_DIVIDE(`int64_col`, 1) AS `int_div_1`,
   IEEE_DIVIDE(`int64_col`, 0.0) AS `int_div_0`,
+  IEEE_DIVIDE(`int64_col`, NULL) AS `int_div_null`,
   IEEE_DIVIDE(`int64_col`, `float64_col`) AS `int_div_float`,
   IEEE_DIVIDE(`float64_col`, `int64_col`) AS `float_div_int`,
   IEEE_DIVIDE(`float64_col`, 0.0) AS `float_div_0`,

tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_floordiv_numeric/out.sql

@@ -0,0 +1,47 @@
+SELECT
+  `rowindex`,
+  `int64_col`,
+  `bool_col`,
+  `float64_col`,
+  CASE
+    WHEN `int64_col` = CAST(0 AS INT64)
+    THEN CAST(0 AS INT64) * `int64_col`
+    ELSE CAST(FLOOR(IEEE_DIVIDE(`int64_col`, `int64_col`)) AS INT64)
+  END AS `int_div_int`,
+  CASE
+    WHEN 1 = CAST(0 AS INT64)
+    THEN CAST(0 AS INT64) * `int64_col`
+    ELSE CAST(FLOOR(IEEE_DIVIDE(`int64_col`, 1)) AS INT64)
+  END AS `int_div_1`,
+  CASE
+    WHEN 0.0 = CAST(0 AS INT64)
+    THEN CAST('Infinity' AS FLOAT64) * `int64_col`
+    ELSE CAST(FLOOR(IEEE_DIVIDE(`int64_col`, 0.0)) AS INT64)
+  END AS `int_div_0`,
+  NULL AS `int_div_null`,
+  CASE
+    WHEN `float64_col` = CAST(0 AS INT64)
+    THEN CAST('Infinity' AS FLOAT64) * `int64_col`
+    ELSE CAST(FLOOR(IEEE_DIVIDE(`int64_col`, `float64_col`)) AS INT64)
+  END AS `int_div_float`,
+  CASE
+    WHEN `int64_col` = CAST(0 AS INT64)
+    THEN CAST('Infinity' AS FLOAT64) * `float64_col`
+    ELSE CAST(FLOOR(IEEE_DIVIDE(`float64_col`, `int64_col`)) AS INT64)
+  END AS `float_div_int`,
+  CASE
+    WHEN 0.0 = CAST(0 AS INT64)
+    THEN CAST('Infinity' AS FLOAT64) * `float64_col`
+    ELSE CAST(FLOOR(IEEE_DIVIDE(`float64_col`, 0.0)) AS INT64)
+  END AS `float_div_0`,
+  CASE
+    WHEN CAST(`bool_col` AS INT64) = CAST(0 AS INT64)
+    THEN CAST(0 AS INT64) * `int64_col`
+    ELSE CAST(FLOOR(IEEE_DIVIDE(`int64_col`, CAST(`bool_col` AS INT64))) AS INT64)
+  END AS `int_div_bool`,
+  CASE
+    WHEN `int64_col` = CAST(0 AS INT64)
+    THEN CAST(0 AS INT64) * CAST(`bool_col` AS INT64)
+    ELSE CAST(FLOOR(IEEE_DIVIDE(CAST(`bool_col` AS INT64), `int64_col`)) AS INT64)
+  END AS `bool_div_int`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mul_numeric/out.sql

@@ -4,6 +4,7 @@ SELECT
   `bool_col`,
   `int64_col` * `int64_col` AS `int_mul_int`,
   `int64_col` * 1 AS `int_mul_1`,
+  NULL AS `int_mul_null`,
   `int64_col` * CAST(`bool_col` AS INT64) AS `int_mul_bool`,
   CAST(`bool_col` AS INT64) * `int64_col` AS `bool_mul_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_sub_numeric/out.sql

@@ -2,8 +2,9 @@ SELECT
   `rowindex`,
   `int64_col`,
   `bool_col`,
-  `int64_col` - `int64_col` AS `int_add_int`,
-  `int64_col` - 1 AS `int_add_1`,
-  `int64_col` - CAST(`bool_col` AS INT64) AS `int_add_bool`,
-  CAST(`bool_col` AS INT64) - `int64_col` AS `bool_add_int`
+  `int64_col` - `int64_col` AS `int_sub_int`,
+  `int64_col` - 1 AS `int_sub_1`,
+  NULL AS `int_sub_null`,
+  `int64_col` - CAST(`bool_col` AS INT64) AS `int_sub_bool`,
+  CAST(`bool_col` AS INT64) - `int64_col` AS `bool_sub_int`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`