Skip to content

Commit c83b879

Browse files
committed
chore: Migrate up to 5 scalar operators to SQLGlot
Migrated tan_op, arcsin_op, arccos_op, arctan_op, and sinh_op scalar operators to SQLGlot.
1 parent c03b84e commit c83b879

File tree

7 files changed

+149
-0
lines changed

7 files changed

+149
-0
lines changed

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,37 @@ def compile(op: ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
3030
return UNARY_OP_REGISTRATION[op](op, expr)
3131

3232

33+
@UNARY_OP_REGISTRATION.register(ops.arccos_op)
34+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
35+
return sge.Case(
36+
ifs=[
37+
sge.If(
38+
this=sge.func("ABS", expr.expr) > sge.convert(1),
39+
true=sge.func("IEEE_DIVIDE", sge.convert(0), sge.convert(0)),
40+
)
41+
],
42+
default=sge.func("ACOS", expr.expr),
43+
)
44+
45+
46+
@UNARY_OP_REGISTRATION.register(ops.arcsin_op)
47+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
48+
return sge.Case(
49+
ifs=[
50+
sge.If(
51+
this=sge.func("ABS", expr.expr) > sge.convert(1),
52+
true=sge.func("IEEE_DIVIDE", sge.convert(0), sge.convert(0)),
53+
)
54+
],
55+
default=sge.func("ASIN", expr.expr),
56+
)
57+
58+
59+
@UNARY_OP_REGISTRATION.register(ops.arctan_op)
60+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
61+
return sge.func("ATAN", expr.expr)
62+
63+
3364
@UNARY_OP_REGISTRATION.register(ops.ArrayToStringOp)
3465
def _(op: ops.ArrayToStringOp, expr: TypedExpr) -> sge.Expression:
3566
return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
@@ -97,6 +128,25 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
97128
return sge.func("sin", expr.expr)
98129

99130

131+
@UNARY_OP_REGISTRATION.register(ops.sinh_op)
132+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
133+
return sge.Case(
134+
ifs=[
135+
sge.If(
136+
this=sge.func("ABS", expr.expr) > sge.convert(709.78),
137+
true=sge.func("SIGN", expr.expr)
138+
* sge.func("IEEE_DIVIDE", sge.convert(1), sge.convert(0)),
139+
)
140+
],
141+
default=sge.func("SINH", expr.expr),
142+
)
143+
144+
145+
@UNARY_OP_REGISTRATION.register(ops.tan_op)
146+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
147+
return sge.func("TAN", expr.expr)
148+
149+
100150
# JSON Ops
101151
@UNARY_OP_REGISTRATION.register(ops.JSONExtract)
102152
def _(op: ops.JSONExtract, expr: TypedExpr) -> sge.Expression:
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`float64_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
CASE WHEN ABS(`bfcol_0`) > 1 THEN IEEE_DIVIDE(0, 0) ELSE ACOS(`bfcol_0`) END AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `float64_col`
13+
FROM `bfcte_1`
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`float64_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
CASE WHEN ABS(`bfcol_0`) > 1 THEN IEEE_DIVIDE(0, 0) ELSE ASIN(`bfcol_0`) END AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `float64_col`
13+
FROM `bfcte_1`
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`float64_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
ATAN(`bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `float64_col`
13+
FROM `bfcte_1`
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`float64_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
CASE
9+
WHEN ABS(`bfcol_0`) > 709.78
10+
THEN SIGN(`bfcol_0`) * IEEE_DIVIDE(1, 0)
11+
ELSE SINH(`bfcol_0`)
12+
END AS `bfcol_1`
13+
FROM `bfcte_0`
14+
)
15+
SELECT
16+
`bfcol_1` AS `float64_col`
17+
FROM `bfcte_1`
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`float64_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
TAN(`bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `float64_col`
13+
FROM `bfcte_1`

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,24 @@ def _apply_unary_op(obj: bpd.DataFrame, op: ops.UnaryOp, arg: str) -> str:
3434
return sql
3535

3636

37+
def test_arccos(scalar_types_df: bpd.DataFrame, snapshot):
38+
bf_df = scalar_types_df[["float64_col"]]
39+
sql = _apply_unary_op(bf_df, ops.arccos_op, "float64_col")
40+
snapshot.assert_match(sql, "out.sql")
41+
42+
43+
def test_arcsin(scalar_types_df: bpd.DataFrame, snapshot):
44+
bf_df = scalar_types_df[["float64_col"]]
45+
sql = _apply_unary_op(bf_df, ops.arcsin_op, "float64_col")
46+
snapshot.assert_match(sql, "out.sql")
47+
48+
49+
def test_arctan(scalar_types_df: bpd.DataFrame, snapshot):
50+
bf_df = scalar_types_df[["float64_col"]]
51+
sql = _apply_unary_op(bf_df, ops.arctan_op, "float64_col")
52+
snapshot.assert_match(sql, "out.sql")
53+
54+
3755
def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot):
3856
bf_df = repeated_types_df[["string_list_col"]]
3957
sql = _apply_unary_op(bf_df, ops.ArrayToStringOp(delimiter="."), "string_list_col")
@@ -88,6 +106,18 @@ def test_sin(scalar_types_df: bpd.DataFrame, snapshot):
88106
snapshot.assert_match(sql, "out.sql")
89107

90108

109+
def test_sinh(scalar_types_df: bpd.DataFrame, snapshot):
110+
bf_df = scalar_types_df[["float64_col"]]
111+
sql = _apply_unary_op(bf_df, ops.sinh_op, "float64_col")
112+
snapshot.assert_match(sql, "out.sql")
113+
114+
115+
def test_tan(scalar_types_df: bpd.DataFrame, snapshot):
116+
bf_df = scalar_types_df[["float64_col"]]
117+
sql = _apply_unary_op(bf_df, ops.tan_op, "float64_col")
118+
snapshot.assert_match(sql, "out.sql")
119+
120+
91121
def test_json_extract(json_types_df: bpd.DataFrame, snapshot):
92122
bf_df = json_types_df[["json_col"]]
93123
sql = _apply_unary_op(bf_df, ops.JSONExtract(json_path="$"), "json_col")

0 commit comments

Comments
 (0)