support bool

Author: jialuoo

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

@@ -79,7 +79,9 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_binary_op(ops.arctan2_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
-    return sge.func("ATAN2", left.expr, right.expr)
+    left_expr = _coerce_bool_to_int(left)
+    right_expr = _coerce_bool_to_int(right)
+    return sge.func("ATAN2", left_expr, right_expr)
 
 
 @register_unary_op(ops.arctan_op)

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

@@ -1,14 +1,17 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_0`,
-    `float64_col` AS `bfcol_1`
+    `bool_col` AS `bfcol_0`,
+    `int64_col` AS `bfcol_1`,
+    `float64_col` AS `bfcol_2`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    ATAN2(`bfcol_0`, `bfcol_1`) AS `bfcol_2`
+    ATAN2(`bfcol_1`, `bfcol_2`) AS `bfcol_6`,
+    ATAN2(CAST(`bfcol_0` AS INT64), `bfcol_2`) AS `bfcol_7`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_2` AS `int64_col`
+  `bfcol_6` AS `int64_col`,
+  `bfcol_7` AS `bool_col`
 FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py

@@ -56,9 +56,16 @@ def test_arcsinh(scalar_types_df: bpd.DataFrame, snapshot):
 
 
 def test_arctan2(scalar_types_df: bpd.DataFrame, snapshot):
-    bf_df = scalar_types_df[["int64_col", "float64_col"]]
-    sql = utils._apply_binary_op(bf_df, ops.arctan2_op, "int64_col", "float64_col")
-
+    bf_df = scalar_types_df[["int64_col", "float64_col", "bool_col"]]
+
+    sql = utils._apply_ops_to_sql(
+        bf_df,
+        [
+            ops.arctan2_op.as_expr("int64_col", "float64_col"),
+            ops.arctan2_op.as_expr("bool_col", "float64_col"),
+        ],
+        ["int64_col", "bool_col"],
+    )
     snapshot.assert_match(sql, "out.sql")