Skip to content

Commit e3f8ca7

Browse files
chelsea-linchelsea-lin
authored
chore: Migrate up to 15 scalar operators to SQLGlot (#1936)
Migrated the following unary scalar operators to SQLGlot: - pos_op - neg_op - sqrt_op - log10_op - ln_op - log1p_op - len_op - reverse_op - lower_op - upper_op - StrLstripOp - StrRstripOp - StrStripOp - StrContainsOp - StrContainsRegexOp
1 parent 31b17b0 commit e3f8ca7

File tree

17 files changed

+407
-0
lines changed

17 files changed

+407
-0
lines changed

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

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,16 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
175175
)
176176

177177

178+
@UNARY_OP_REGISTRATION.register(ops.StrContainsRegexOp)
179+
def _(op: ops.StrContainsRegexOp, expr: TypedExpr) -> sge.Expression:
180+
return sge.RegexpLike(this=expr.expr, expression=sge.convert(op.pat))
181+
182+
183+
@UNARY_OP_REGISTRATION.register(ops.StrContainsOp)
184+
def _(op: ops.StrContainsOp, expr: TypedExpr) -> sge.Expression:
185+
return sge.Like(this=expr.expr, expression=sge.convert(f"%{op.pat}%"))
186+
187+
178188
@UNARY_OP_REGISTRATION.register(ops.date_op)
179189
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
180190
return sge.Date(this=expr.expr)
@@ -302,6 +312,98 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
302312
)
303313

304314

315+
@UNARY_OP_REGISTRATION.register(ops.len_op)
316+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
317+
return sge.Length(this=expr.expr)
318+
319+
320+
@UNARY_OP_REGISTRATION.register(ops.ln_op)
321+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
322+
return sge.Case(
323+
ifs=[
324+
sge.If(
325+
this=expr.expr < sge.convert(0),
326+
true=_NAN,
327+
)
328+
],
329+
default=sge.Ln(this=expr.expr),
330+
)
331+
332+
333+
@UNARY_OP_REGISTRATION.register(ops.log10_op)
334+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
335+
return sge.Case(
336+
ifs=[
337+
sge.If(
338+
this=expr.expr < sge.convert(0),
339+
true=_NAN,
340+
)
341+
],
342+
default=sge.Log(this=expr.expr, expression=sge.convert(10)),
343+
)
344+
345+
346+
@UNARY_OP_REGISTRATION.register(ops.log1p_op)
347+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
348+
return sge.Case(
349+
ifs=[
350+
sge.If(
351+
this=expr.expr < sge.convert(-1),
352+
true=_NAN,
353+
)
354+
],
355+
default=sge.Ln(this=sge.convert(1) + expr.expr),
356+
)
357+
358+
359+
@UNARY_OP_REGISTRATION.register(ops.lower_op)
360+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
361+
return sge.Lower(this=expr.expr)
362+
363+
364+
@UNARY_OP_REGISTRATION.register(ops.StrLstripOp)
365+
def _(op: ops.StrLstripOp, expr: TypedExpr) -> sge.Expression:
366+
return sge.Trim(this=expr.expr, expression=sge.convert(op.to_strip), side="LEFT")
367+
368+
369+
@UNARY_OP_REGISTRATION.register(ops.neg_op)
370+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
371+
return sge.Neg(this=expr.expr)
372+
373+
374+
@UNARY_OP_REGISTRATION.register(ops.pos_op)
375+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
376+
return expr.expr
377+
378+
379+
@UNARY_OP_REGISTRATION.register(ops.reverse_op)
380+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
381+
return sge.func("REVERSE", expr.expr)
382+
383+
384+
@UNARY_OP_REGISTRATION.register(ops.StrRstripOp)
385+
def _(op: ops.StrRstripOp, expr: TypedExpr) -> sge.Expression:
386+
return sge.Trim(this=expr.expr, expression=sge.convert(op.to_strip), side="RIGHT")
387+
388+
389+
@UNARY_OP_REGISTRATION.register(ops.sqrt_op)
390+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
391+
return sge.Case(
392+
ifs=[
393+
sge.If(
394+
this=expr.expr < sge.convert(0),
395+
true=_NAN,
396+
)
397+
],
398+
default=sge.Sqrt(this=expr.expr),
399+
)
400+
401+
402+
@UNARY_OP_REGISTRATION.register(ops.StrStripOp)
403+
def _(op: ops.StrStripOp, expr: TypedExpr) -> sge.Expression:
404+
return sge.Trim(this=sge.convert(op.to_strip), expression=expr.expr)
405+
406+
305407
@UNARY_OP_REGISTRATION.register(ops.iso_day_op)
306408
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
307409
return sge.Extract(this=sge.Identifier(this="DAYOFWEEK"), expression=expr.expr)
@@ -394,3 +496,8 @@ def _(op: ops.ParseJSON, expr: TypedExpr) -> sge.Expression:
394496
@UNARY_OP_REGISTRATION.register(ops.ToJSONString)
395497
def _(op: ops.ToJSONString, expr: TypedExpr) -> sge.Expression:
396498
return sge.func("TO_JSON_STRING", expr.expr)
499+
500+
501+
@UNARY_OP_REGISTRATION.register(ops.upper_op)
502+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
503+
return sge.Upper(this=expr.expr)

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_len/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+
LENGTH(`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_ln/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` < 0 THEN CAST('NaN' AS FLOAT64) ELSE LN(`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_log10/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` < 0 THEN CAST('NaN' AS FLOAT64) ELSE LOG(10, `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_log1p/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 LN(1 + `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_lower/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+
LOWER(`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_lstrip/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+
TRIM(`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_neg/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+
-`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_pos/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+
`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_reverse/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+
REVERSE(`bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `string_col`
13+
FROM `bfcte_1`

0 commit comments

Comments
 (0)