Merge branch 'main' into tswast-bbq-obj

Author: tswast

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/comparison_ops.py

@@ -16,11 +16,13 @@
 
 import typing
 
+import bigframes_vendored.sqlglot as sg
 import bigframes_vendored.sqlglot.expressions as sge
 import pandas as pd
 
 from bigframes import dtypes
 from bigframes import operations as ops
+from bigframes.core.compile.sqlglot import sqlglot_ir
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
@@ -62,6 +64,10 @@ def _(expr: TypedExpr, op: ops.IsInOp) -> sge.Expression:
 
 @register_binary_op(ops.eq_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if sqlglot_ir._is_null_literal(left.expr):
+        return sge.Is(this=right.expr, expression=sge.Null())
+    if sqlglot_ir._is_null_literal(right.expr):
+        return sge.Is(this=left.expr, expression=sge.Null())
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.EQ(this=left_expr, expression=right_expr)
@@ -139,6 +145,17 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.ne_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    if sqlglot_ir._is_null_literal(left.expr):
+        return sge.Is(
+            this=sge.paren(right.expr, copy=False),
+            expression=sg.not_(sge.Null(), copy=False),
+        )
+    if sqlglot_ir._is_null_literal(right.expr):
+        return sge.Is(
+            this=sge.paren(left.expr, copy=False),
+            expression=sg.not_(sge.Null(), copy=False),
+        )
+
     left_expr = _coerce_bool_to_int(left)
     right_expr = _coerce_bool_to_int(right)
     return sge.NEQ(this=left_expr, expression=right_expr)

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)

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

@@ -19,7 +19,7 @@
 
 from bigframes import dtypes
 from bigframes import operations as ops
-from bigframes.core.compile.sqlglot import sqlglot_types
+from bigframes.core.compile.sqlglot import sqlglot_ir, sqlglot_types
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
@@ -101,11 +101,23 @@ def _(expr: TypedExpr) -> sge.Expression:
 def _(expr: TypedExpr, op: ops.MapOp) -> sge.Expression:
     if len(op.mappings) == 0:
         return expr.expr
+
+    mappings = [
+        (
+            sqlglot_ir._literal(key, dtypes.is_compatible(key, expr.dtype)),
+            sqlglot_ir._literal(value, dtypes.is_compatible(value, expr.dtype)),
+        )
+        for key, value in op.mappings
+    ]
     return sge.Case(
-        this=expr.expr,
         ifs=[
-            sge.If(this=sge.convert(key), true=sge.convert(value))
-            for key, value in op.mappings
+            sge.If(
+                this=sge.EQ(this=expr.expr, expression=key)
+                if not sqlglot_ir._is_null_literal(key)
+                else sge.Is(this=expr.expr, expression=sge.Null()),
+                true=value,
+            )
+            for key, value in mappings
         ],
         default=expr.expr,
     )

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -642,6 +642,15 @@ def _select_to_cte(expr: sge.Select, cte_name: sge.Identifier) -> sge.Select:
     return new_select_expr
 
 
+def _is_null_literal(expr: sge.Expression) -> bool:
+    """Checks if the given expression is a NULL literal."""
+    if isinstance(expr, sge.Null):
+        return True
+    if isinstance(expr, sge.Cast) and isinstance(expr.this, sge.Null):
+        return True
+    return False
+
+
 def _literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
     sqlglot_type = sgt.from_bigframes_dtype(dtype) if dtype else None
     if sqlglot_type is None:

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`