Merge branch 'main' into main_chelsealin_aggmean

Author: tswast

bigframes/core/blocks.py

@@ -1375,10 +1375,16 @@ def aggregate(
     ) -> typing.Tuple[Block, typing.Sequence[str]]:
         """
         Apply aggregations to the block.
+
         Arguments:
             by_column_id: column id of the aggregation key, this is preserved through the transform and used as index.
             aggregations: input_column_id, operation tuples
             dropna: whether null keys should be dropped
+
+        Returns:
+            Tuple[Block, Sequence[str]]:
+                The first element is the grouped block. The second is the
+                column IDs corresponding to each applied aggregation.
         """
         if column_labels is None:
             column_labels = pd.Index(range(len(aggregations)))

bigframes/core/compile/sqlglot/__init__.py

@@ -15,7 +15,6 @@
 
 from bigframes.core.compile.sqlglot.compiler import SQLGlotCompiler
 import bigframes.core.compile.sqlglot.expressions.array_ops  # noqa: F401
-import bigframes.core.compile.sqlglot.expressions.binary_compiler  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.blob_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.comparison_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.date_ops  # noqa: F401

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

@@ -21,6 +21,7 @@
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
 
 
 @register_unary_op(ops.obj_fetch_metadata_op)
@@ -31,3 +32,8 @@ def _(expr: TypedExpr) -> sge.Expression:
 @register_unary_op(ops.ObjGetAccessUrl)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("OBJ.GET_ACCESS_URL", expr.expr)
+
+
+@register_binary_op(ops.obj_make_ref_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    return sge.func("OBJ.MAKE_REF", left.expr, right.expr)

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

@@ -19,12 +19,13 @@
 import pandas as pd
 import sqlglot.expressions as sge
 
+from bigframes import dtypes
 from bigframes import operations as ops
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
-import bigframes.dtypes as dtypes
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
 
 
 @register_unary_op(ops.IsInOp, pass_op=True)
@@ -53,7 +54,76 @@ def _(expr: TypedExpr, op: ops.IsInOp) -> sge.Expression:
     )
 
 
+@register_binary_op(ops.eq_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.EQ(this=left_expr, expression=right_expr)
+
+
+@register_binary_op(ops.eq_null_match_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = left.expr
+    if right.dtype != dtypes.BOOL_DTYPE:
+        left_expr = _coerce_bool_to_int(left)
+
+    right_expr = right.expr
+    if left.dtype != dtypes.BOOL_DTYPE:
+        right_expr = _coerce_bool_to_int(right)
+
+    sentinel = sge.convert("$NULL_SENTINEL$")
+    left_coalesce = sge.Coalesce(
+        this=sge.Cast(this=left_expr, to="STRING"), expressions=[sentinel]
+    )
+    right_coalesce = sge.Coalesce(
+        this=sge.Cast(this=right_expr, to="STRING"), expressions=[sentinel]
+    )
+    return sge.EQ(this=left_coalesce, expression=right_coalesce)
+
+
+@register_binary_op(ops.ge_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.GTE(this=left_expr, expression=right_expr)
+
+
+@register_binary_op(ops.gt_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.GT(this=left_expr, expression=right_expr)
+
+
+@register_binary_op(ops.lt_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.LT(this=left_expr, expression=right_expr)
+
+
+@register_binary_op(ops.le_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.LTE(this=left_expr, expression=right_expr)
+
+
+@register_binary_op(ops.ne_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.NEQ(this=left_expr, expression=right_expr)
+
+
 # Helpers
 def _is_null(value) -> bool:
     # float NaN/inf should be treated as distinct from 'true' null values
     return typing.cast(bool, pd.isna(value)) and not isinstance(value, float)
+
+
+def _coerce_bool_to_int(typed_expr: TypedExpr) -> sge.Expression:
+    """Coerce boolean expression to integer."""
+    if typed_expr.dtype == dtypes.BOOL_DTYPE:
+        return sge.Cast(this=typed_expr.expr, to="INT64")
+    return typed_expr.expr

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

@@ -21,6 +21,7 @@
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
+register_binary_op = scalar_compiler.scalar_op_compiler.register_binary_op
 
 
 @register_unary_op(ops.JSONExtract, pass_op=True)
@@ -66,3 +67,8 @@ def _(expr: TypedExpr) -> sge.Expression:
 @register_unary_op(ops.ToJSONString)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("TO_JSON_STRING", expr.expr)
+
+
+@register_binary_op(ops.JSONSet, pass_op=True)
+def _(left: TypedExpr, right: TypedExpr, op) -> sge.Expression:
+    return sge.func("JSON_SET", left.expr, sge.convert(op.json_path), right.expr)