Skip to content

Commit 07bce8e

Browse files
chelsea-linchelsea-lin
authored
chore: migrate 15 scalar operators to SQLGlot (#1925)
This change contains three commits generated by Gemini CLI tool: Migrated cosh_op, tanh_op, arcsinh_op, arccosh_op, and arctanh_op scalar operators to SQLGlot. Migrated abs_op, capitalize_op, ceil_op, date_op, and day_op scalar operators to SQLGlot. Migrated dayofweek_op, dayofyear_op, exp_op, expm1_op, and floor_op scalar operators to SQLGlot. Fixes internal issue 430133370
1 parent c0cefd3 commit 07bce8e

File tree

17 files changed

+429
-0
lines changed

17 files changed

+429
-0
lines changed

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

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,24 @@ def compile(op: ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
3838
return UNARY_OP_REGISTRATION[op](op, expr)
3939

4040

41+
@UNARY_OP_REGISTRATION.register(ops.abs_op)
42+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
43+
return sge.Abs(this=expr.expr)
44+
45+
46+
@UNARY_OP_REGISTRATION.register(ops.arccosh_op)
47+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
48+
return sge.Case(
49+
ifs=[
50+
sge.If(
51+
this=expr.expr < sge.convert(1),
52+
true=_NAN,
53+
)
54+
],
55+
default=sge.func("ACOSH", expr.expr),
56+
)
57+
58+
4159
@UNARY_OP_REGISTRATION.register(ops.arccos_op)
4260
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
4361
return sge.Case(
@@ -64,11 +82,29 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
6482
)
6583

6684

85+
@UNARY_OP_REGISTRATION.register(ops.arcsinh_op)
86+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
87+
return sge.func("ASINH", expr.expr)
88+
89+
6790
@UNARY_OP_REGISTRATION.register(ops.arctan_op)
6891
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
6992
return sge.func("ATAN", expr.expr)
7093

7194

95+
@UNARY_OP_REGISTRATION.register(ops.arctanh_op)
96+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
97+
return sge.Case(
98+
ifs=[
99+
sge.If(
100+
this=sge.func("ABS", expr.expr) > sge.convert(1),
101+
true=_NAN,
102+
)
103+
],
104+
default=sge.func("ATANH", expr.expr),
105+
)
106+
107+
72108
@UNARY_OP_REGISTRATION.register(ops.ArrayToStringOp)
73109
def _(op: ops.ArrayToStringOp, expr: TypedExpr) -> sge.Expression:
74110
return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
@@ -111,11 +147,88 @@ def _(op: ops.ArraySliceOp, expr: TypedExpr) -> sge.Expression:
111147
return sge.array(selected_elements)
112148

113149

150+
@UNARY_OP_REGISTRATION.register(ops.capitalize_op)
151+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
152+
return sge.Initcap(this=expr.expr)
153+
154+
155+
@UNARY_OP_REGISTRATION.register(ops.ceil_op)
156+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
157+
return sge.Ceil(this=expr.expr)
158+
159+
114160
@UNARY_OP_REGISTRATION.register(ops.cos_op)
115161
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
116162
return sge.func("COS", expr.expr)
117163

118164

165+
@UNARY_OP_REGISTRATION.register(ops.cosh_op)
166+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
167+
return sge.Case(
168+
ifs=[
169+
sge.If(
170+
this=sge.func("ABS", expr.expr) > sge.convert(709.78),
171+
true=_INF,
172+
)
173+
],
174+
default=sge.func("COSH", expr.expr),
175+
)
176+
177+
178+
@UNARY_OP_REGISTRATION.register(ops.date_op)
179+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
180+
return sge.Date(this=expr.expr)
181+
182+
183+
@UNARY_OP_REGISTRATION.register(ops.day_op)
184+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
185+
return sge.Extract(this=sge.Identifier(this="DAY"), expression=expr.expr)
186+
187+
188+
@UNARY_OP_REGISTRATION.register(ops.dayofweek_op)
189+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
190+
# Adjust the 1-based day-of-week index (from SQL) to a 0-based index.
191+
return sge.Extract(
192+
this=sge.Identifier(this="DAYOFWEEK"), expression=expr.expr
193+
) - sge.convert(1)
194+
195+
196+
@UNARY_OP_REGISTRATION.register(ops.dayofyear_op)
197+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
198+
return sge.Extract(this=sge.Identifier(this="DAYOFYEAR"), expression=expr.expr)
199+
200+
201+
@UNARY_OP_REGISTRATION.register(ops.exp_op)
202+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
203+
return sge.Case(
204+
ifs=[
205+
sge.If(
206+
this=expr.expr > _FLOAT64_EXP_BOUND,
207+
true=_INF,
208+
)
209+
],
210+
default=sge.func("EXP", expr.expr),
211+
)
212+
213+
214+
@UNARY_OP_REGISTRATION.register(ops.expm1_op)
215+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
216+
return sge.Case(
217+
ifs=[
218+
sge.If(
219+
this=expr.expr > _FLOAT64_EXP_BOUND,
220+
true=_INF,
221+
)
222+
],
223+
default=sge.func("EXP", expr.expr),
224+
) - sge.convert(1)
225+
226+
227+
@UNARY_OP_REGISTRATION.register(ops.floor_op)
228+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
229+
return sge.Floor(this=expr.expr)
230+
231+
119232
@UNARY_OP_REGISTRATION.register(ops.hash_op)
120233
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
121234
return sge.func("FARM_FINGERPRINT", expr.expr)
@@ -154,6 +267,11 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
154267
return sge.func("TAN", expr.expr)
155268

156269

270+
@UNARY_OP_REGISTRATION.register(ops.tanh_op)
271+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
272+
return sge.func("TANH", expr.expr)
273+
274+
157275
# JSON Ops
158276
@UNARY_OP_REGISTRATION.register(ops.JSONExtract)
159277
def _(op: ops.JSONExtract, expr: TypedExpr) -> sge.Expression:

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_abs/out.sql

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+
ABS(`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/snapshots/test_unary_compiler/test_arccosh/out.sql

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 `bfcol_0` < 1 THEN CAST('NaN' AS FLOAT64) ELSE ACOSH(`bfcol_0`) END 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/snapshots/test_unary_compiler/test_arcsinh/out.sql

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+
ASINH(`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/snapshots/test_unary_compiler/test_arctanh/out.sql

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 CAST('NaN' AS FLOAT64) ELSE ATANH(`bfcol_0`) END 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/snapshots/test_unary_compiler/test_capitalize/out.sql

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+
`string_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
INITCAP(`bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `string_col`
13+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_ceil/out.sql

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+
CEIL(`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/snapshots/test_unary_compiler/test_cosh/out.sql

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 CAST('Infinity' AS FLOAT64)
11+
ELSE COSH(`bfcol_0`)
12+
END AS `bfcol_1`
13+
FROM `bfcte_0`
14+
)
15+
SELECT
16+
`bfcol_1` AS `float64_col`
17+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_date/out.sql

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+
`timestamp_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
DATE(`bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `timestamp_col`
13+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_day/out.sql

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+
`timestamp_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
EXTRACT(DAY FROM `bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `timestamp_col`
13+
FROM `bfcte_1`

0 commit comments

Comments
 (0)