Skip to content

Commit c6e277e

Browse files
chelsea-linchelsea-lin
committed
chore: Migrate up to 15 scalar operators to SQLGlot
Migrated the following unary scalar operators to SQLGlot: - StrftimeOp - UnixSeconds - UnixMicros - UnixMillis - FloorDtOp - geo_st_boundary_op - geo_st_geogfromtext_op - geo_st_isclosed_op - GeoStLengthOp - StructFieldOp - AsTypeOp - IsInOp - ToDatetimeOp - ToTimestampOp - ToTimedeltaOp
1 parent 0468a4d commit c6e277e

File tree

15 files changed

+345
-0
lines changed

15 files changed

+345
-0
lines changed

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

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,11 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
105105
)
106106

107107

108+
@UNARY_OP_REGISTRATION.register(ops.AsTypeOp)
109+
def _(op: ops.AsTypeOp, expr: TypedExpr) -> sge.Expression:
110+
return sge.Cast(this=expr.expr, to=op.to_type)
111+
112+
108113
@UNARY_OP_REGISTRATION.register(ops.ArrayToStringOp)
109114
def _(op: ops.ArrayToStringOp, expr: TypedExpr) -> sge.Expression:
110115
return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
@@ -234,6 +239,11 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
234239
) - sge.convert(1)
235240

236241

242+
@UNARY_OP_REGISTRATION.register(ops.FloorDtOp)
243+
def _(op: ops.FloorDtOp, expr: TypedExpr) -> sge.Expression:
244+
return sge.TimestampTrunc(this=expr.expr, unit=sge.Identifier(this=op.freq))
245+
246+
237247
@UNARY_OP_REGISTRATION.register(ops.floor_op)
238248
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
239249
return sge.Floor(this=expr.expr)
@@ -249,6 +259,26 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
249259
return sge.func("ST_ASTEXT", expr.expr)
250260

251261

262+
@UNARY_OP_REGISTRATION.register(ops.geo_st_boundary_op)
263+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
264+
return sge.func("ST_BOUNDARY", expr.expr)
265+
266+
267+
@UNARY_OP_REGISTRATION.register(ops.geo_st_geogfromtext_op)
268+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
269+
return sge.func("SAFE.ST_GEOGFROMTEXT", expr.expr)
270+
271+
272+
@UNARY_OP_REGISTRATION.register(ops.geo_st_isclosed_op)
273+
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
274+
return sge.func("ST_ISCLOSED", expr.expr)
275+
276+
277+
@UNARY_OP_REGISTRATION.register(ops.GeoStLengthOp)
278+
def _(op: ops.GeoStLengthOp, expr: TypedExpr) -> sge.Expression:
279+
return sge.func("ST_LENGTH", expr.expr)
280+
281+
252282
@UNARY_OP_REGISTRATION.register(ops.geo_x_op)
253283
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
254284
return sge.func("SAFE.ST_X", expr.expr)
@@ -274,6 +304,11 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
274304
return sge.BitwiseNot(this=expr.expr)
275305

276306

307+
@UNARY_OP_REGISTRATION.register(ops.IsInOp)
308+
def _(op: ops.IsInOp, expr: TypedExpr) -> sge.Expression:
309+
return sge.In(this=expr.expr, expressions=[sge.convert(v) for v in op.values])
310+
311+
277312
@UNARY_OP_REGISTRATION.register(ops.isalnum_op)
278313
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
279314
return sge.RegexpLike(this=expr.expr, expression=sge.convert(r"^(\p{N}|\p{L})+$"))
@@ -517,6 +552,16 @@ def _(op: ops.StrSliceOp, expr: TypedExpr) -> sge.Expression:
517552
)
518553

519554

555+
@UNARY_OP_REGISTRATION.register(ops.StrftimeOp)
556+
def _(op: ops.StrftimeOp, expr: TypedExpr) -> sge.Expression:
557+
return sge.func("FORMAT_TIMESTAMP", sge.convert(op.date_format), expr.expr)
558+
559+
560+
@UNARY_OP_REGISTRATION.register(ops.StructFieldOp)
561+
def _(op: ops.StructFieldOp, expr: TypedExpr) -> sge.Expression:
562+
return sge.StructExtract(this=expr.expr, expression=sge.convert(op.name_or_index))
563+
564+
520565
@UNARY_OP_REGISTRATION.register(ops.tan_op)
521566
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
522567
return sge.func("TAN", expr.expr)
@@ -537,6 +582,36 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
537582
return sge.Floor(this=expr.expr)
538583

539584

585+
@UNARY_OP_REGISTRATION.register(ops.ToDatetimeOp)
586+
def _(op: ops.ToDatetimeOp, expr: TypedExpr) -> sge.Expression:
587+
return sge.Cast(this=sge.func("TIMESTAMP_SECONDS", expr.expr), to="DATETIME")
588+
589+
590+
@UNARY_OP_REGISTRATION.register(ops.ToTimestampOp)
591+
def _(op: ops.ToTimestampOp, expr: TypedExpr) -> sge.Expression:
592+
return sge.func("TIMESTAMP_SECONDS", expr.expr)
593+
594+
595+
@UNARY_OP_REGISTRATION.register(ops.ToTimedeltaOp)
596+
def _(op: ops.ToTimedeltaOp, expr: TypedExpr) -> sge.Expression:
597+
return sge.Interval(this=expr.expr, unit=sge.Identifier(this="SECOND"))
598+
599+
600+
@UNARY_OP_REGISTRATION.register(ops.UnixMicros)
601+
def _(op: ops.UnixMicros, expr: TypedExpr) -> sge.Expression:
602+
return sge.func("UNIX_MICROS", expr.expr)
603+
604+
605+
@UNARY_OP_REGISTRATION.register(ops.UnixMillis)
606+
def _(op: ops.UnixMillis, expr: TypedExpr) -> sge.Expression:
607+
return sge.func("UNIX_MILLIS", expr.expr)
608+
609+
610+
@UNARY_OP_REGISTRATION.register(ops.UnixSeconds)
611+
def _(op: ops.UnixSeconds, expr: TypedExpr) -> sge.Expression:
612+
return sge.func("UNIX_SECONDS", expr.expr)
613+
614+
540615
# JSON Ops
541616
@UNARY_OP_REGISTRATION.register(ops.JSONExtract)
542617
def _(op: ops.JSONExtract, expr: TypedExpr) -> sge.Expression:

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_floor_dt/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+
TIMESTAMP_TRUNC(`bfcol_0`, DAY) 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_geo_st_boundary/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+
`geography_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
ST_BOUNDARY(`bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `geography_col`
13+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_geogfromtext/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+
SAFE.ST_GEOGFROMTEXT(`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_geo_st_isclosed/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+
`geography_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
ST_ISCLOSED(`bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `geography_col`
13+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_geo_st_length/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+
`geography_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
ST_LENGTH(`bfcol_0`) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `geography_col`
13+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_is_in/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+
`int64_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
`bfcol_0` IN (1, 2, 3) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `int64_col`
13+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_strftime/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+
FORMAT_TIMESTAMP('%Y-%m-%d', `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_to_datetime/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+
`int64_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
CAST(TIMESTAMP_SECONDS(`bfcol_0`) AS DATETIME) AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `int64_col`
13+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_timedelta/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+
`int64_col` AS `bfcol_0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
5+
), `bfcte_1` AS (
6+
SELECT
7+
*,
8+
INTERVAL `bfcol_0` SECOND AS `bfcol_1`
9+
FROM `bfcte_0`
10+
)
11+
SELECT
12+
`bfcol_1` AS `int64_col`
13+
FROM `bfcte_1`

0 commit comments

Comments
 (0)