refactor: fix some aggregation ops in the sqlglot compiler (#2382)

This change fixes several aggregation-related test failures in #2248 by
including the following operators:
- `agg_ops.DiffOp`: Fixes test_date_series_diff_agg
- `agg_ops.AllOp/AnyOp`: Fixes test_list_apply_callable
- `agg_ops.QuantileOp`: Fixes test_dataframe_aggregates_median
- `agg_ops.ProductOp`: Fixes test_dataframe_groupby_analytic

Fixes internal issue 417774347 🦕

Author: chelsea-lin

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -23,6 +23,7 @@
 from bigframes.core import window_spec
 import bigframes.core.compile.sqlglot.aggregations.op_registration as reg
 from bigframes.core.compile.sqlglot.aggregations.windows import apply_window_if_present
+from bigframes.core.compile.sqlglot.expressions import constants
 import bigframes.core.compile.sqlglot.expressions.typed_expr as typed_expr
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 from bigframes.operations import aggregations as agg_ops
@@ -44,9 +45,13 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    # BQ will return null for empty column, result would be false in pandas.
-    result = apply_window_if_present(sge.func("LOGICAL_AND", column.expr), window)
-    return sge.func("IFNULL", result, sge.true())
+    expr = column.expr
+    if column.dtype != dtypes.BOOL_DTYPE:
+        expr = sge.NEQ(this=expr, expression=sge.convert(0))
+    expr = apply_window_if_present(sge.func("LOGICAL_AND", expr), window)
+
+    # BQ will return null for empty column, result would be true in pandas.
+    return sge.func("COALESCE", expr, sge.convert(True))
 
 
 @UNARY_OP_REGISTRATION.register(agg_ops.AnyOp)
@@ -56,6 +61,8 @@ def _(
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
     expr = column.expr
+    if column.dtype != dtypes.BOOL_DTYPE:
+        expr = sge.NEQ(this=expr, expression=sge.convert(0))
     expr = apply_window_if_present(sge.func("LOGICAL_OR", expr), window)
 
     # BQ will return null for empty column, result would be false in pandas.
@@ -326,6 +333,15 @@ def _(
             unit=sge.Identifier(this="MICROSECOND"),
         )
 
+    if column.dtype == dtypes.DATE_DTYPE:
+        date_diff = sge.DateDiff(
+            this=column.expr, expression=shifted, unit=sge.Identifier(this="DAY")
+        )
+        return sge.Cast(
+            this=sge.Floor(this=date_diff * constants._DAY_TO_MICROSECONDS),
+            to="INT64",
+        )
+
     raise TypeError(f"Cannot perform diff on type {column.dtype}")
 
 
@@ -410,24 +426,28 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=expr, to="INT64")
+
     # Need to short-circuit as log with zeroes is illegal sql
-    is_zero = sge.EQ(this=column.expr, expression=sge.convert(0))
+    is_zero = sge.EQ(this=expr, expression=sge.convert(0))
 
     # There is no product sql aggregate function, so must implement as a sum of logs, and then
     # apply power after. Note, log and power base must be equal! This impl uses natural log.
-    logs = (
-        sge.Case()
-        .when(is_zero, sge.convert(0))
-        .else_(sge.func("LN", sge.func("ABS", column.expr)))
+    logs = sge.If(
+        this=is_zero,
+        true=sge.convert(0),
+        false=sge.func("LOG", sge.convert(2), sge.func("ABS", expr)),
     )
     logs_sum = apply_window_if_present(sge.func("SUM", logs), window)
-    magnitude = sge.func("EXP", logs_sum)
+    magnitude = sge.func("POWER", sge.convert(2), logs_sum)
 
     # Can't determine sign from logs, so have to determine parity of count of negative inputs
     is_negative = (
         sge.Case()
         .when(
-            sge.LT(this=sge.func("SIGN", column.expr), expression=sge.convert(0)),
+            sge.EQ(this=sge.func("SIGN", expr), expression=sge.convert(-1)),
             sge.convert(1),
         )
         .else_(sge.convert(0))
@@ -445,11 +465,7 @@ def _(
         .else_(
             sge.Mul(
                 this=magnitude,
-                expression=sge.If(
-                    this=sge.EQ(this=negative_count_parity, expression=sge.convert(1)),
-                    true=sge.convert(-1),
-                    false=sge.convert(1),
-                ),
+                expression=sge.func("POWER", sge.convert(-1), negative_count_parity),
             )
         )
     )
@@ -499,14 +515,18 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    # TODO: Support interpolation argument
-    # TODO: Support percentile_disc
-    result: sge.Expression = sge.func("PERCENTILE_CONT", column.expr, sge.convert(op.q))
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=expr, to="INT64")
+
+    result: sge.Expression = sge.func("PERCENTILE_CONT", expr, sge.convert(op.q))
     if window is None:
-        # PERCENTILE_CONT is a navigation function, not an aggregate function, so it always needs an OVER clause.
+        # PERCENTILE_CONT is a navigation function, not an aggregate function,
+        # so it always needs an OVER clause.
         result = sge.Window(this=result)
     else:
         result = apply_window_if_present(result, window)
+
     if op.should_floor_result:
         result = sge.Cast(this=sge.func("FLOOR", result), to="INT64")
     return result

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

@@ -20,6 +20,7 @@
 _NAN = sge.Cast(this=sge.convert("NaN"), to="FLOAT64")
 _INF = sge.Cast(this=sge.convert("Infinity"), to="FLOAT64")
 _NEG_INF = sge.Cast(this=sge.convert("-Infinity"), to="FLOAT64")
+_DAY_TO_MICROSECONDS = sge.convert(86400000000)
 
 # Approx Highest number you can pass in to EXP function and get a valid FLOAT64 result
 # FLOAT64 has 11 exponent bits, so max values is about 2**(2**10)

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all/out.sql

@@ -1,12 +1,15 @@
 WITH `bfcte_0` AS (
   SELECT
-    `bool_col`
+    `bool_col`,
+    `int64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    COALESCE(LOGICAL_AND(`bool_col`), TRUE) AS `bfcol_1`
+    COALESCE(LOGICAL_AND(`bool_col`), TRUE) AS `bfcol_2`,
+    COALESCE(LOGICAL_AND(`int64_col` <> 0), TRUE) AS `bfcol_3`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `bool_col`
+  `bfcol_2` AS `bool_col`,
+  `bfcol_3` AS `int64_col`
 FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all/window_partition_out.sql

@@ -1,12 +1,15 @@
 WITH `bfcte_0` AS (
   SELECT
-    `bool_col`
+    `bool_col`,
+    `int64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    COALESCE(LOGICAL_OR(`bool_col`), FALSE) AS `bfcol_1`
+    COALESCE(LOGICAL_OR(`bool_col`), FALSE) AS `bfcol_2`,
+    COALESCE(LOGICAL_OR(`int64_col` <> 0), FALSE) AS `bfcol_3`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `bool_col`
+  `bfcol_2` AS `bool_col`,
+  `bfcol_3` AS `int64_col`
 FROM `bfcte_1`

…t_unary_compiler/test_all/window_out.sql …unary_compiler/test_all_w_window/out.sqltests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all/window_out.sql renamed to tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all_w_window/out.sql tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all/window_out.sql renamed to tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all_w_window/out.sql

@@ -0,0 +1,15 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `date_col`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    CAST(FLOOR(
+      DATE_DIFF(`date_col`, LAG(`date_col`, 1) OVER (ORDER BY `date_col` ASC NULLS LAST), DAY) * 86400000000
+    ) AS INT64) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `diff_date`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_any/out.sql

@@ -7,7 +7,7 @@ WITH `bfcte_0` AS (
     CASE
       WHEN LOGICAL_OR(`int64_col` = 0)
       THEN 0
-      ELSE EXP(SUM(CASE WHEN `int64_col` = 0 THEN 0 ELSE LN(ABS(`int64_col`)) END)) * IF(MOD(SUM(CASE WHEN SIGN(`int64_col`) < 0 THEN 1 ELSE 0 END), 2) = 1, -1, 1)
+      ELSE POWER(2, SUM(IF(`int64_col` = 0, 0, LOG(ABS(`int64_col`), 2)))) * POWER(-1, MOD(SUM(CASE WHEN SIGN(`int64_col`) = -1 THEN 1 ELSE 0 END), 2))
     END AS `bfcol_1`
   FROM `bfcte_0`
 )

…t_unary_compiler/test_any/window_out.sql …unary_compiler/test_any_w_window/out.sqltests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_any/window_out.sql renamed to tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_any_w_window/out.sql tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_any/window_out.sql renamed to tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_any_w_window/out.sql

@@ -9,15 +9,15 @@ WITH `bfcte_0` AS (
     CASE
       WHEN LOGICAL_OR(`int64_col` = 0) OVER (PARTITION BY `string_col`)
       THEN 0
-      ELSE EXP(
-        SUM(CASE WHEN `int64_col` = 0 THEN 0 ELSE LN(ABS(`int64_col`)) END) OVER (PARTITION BY `string_col`)
-      ) * IF(
+      ELSE POWER(
+        2,
+        SUM(IF(`int64_col` = 0, 0, LOG(ABS(`int64_col`), 2))) OVER (PARTITION BY `string_col`)
+      ) * POWER(
+        -1,
         MOD(
-          SUM(CASE WHEN SIGN(`int64_col`) < 0 THEN 1 ELSE 0 END) OVER (PARTITION BY `string_col`),
+          SUM(CASE WHEN SIGN(`int64_col`) = -1 THEN 1 ELSE 0 END) OVER (PARTITION BY `string_col`),
           2
-        ) = 1,
-        -1,
-        1
+        )
       )
     END AS `bfcol_2`
   FROM `bfcte_0`