chore: Migrate up to 15 scalar operators to SQLGlot (#1941)

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

Author: chelsea-lin

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

@@ -16,6 +16,8 @@
 
 import typing
 
+import pandas as pd
+import pyarrow as pa
 import sqlglot
 import sqlglot.expressions as sge
 
@@ -105,6 +107,12 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     )
 
 
+@UNARY_OP_REGISTRATION.register(ops.AsTypeOp)
+def _(op: ops.AsTypeOp, expr: TypedExpr) -> sge.Expression:
+    # TODO: Support more types for casting, such as JSON, etc.
+    return sge.Cast(this=expr.expr, to=op.to_type)
+
+
 @UNARY_OP_REGISTRATION.register(ops.ArrayToStringOp)
 def _(op: ops.ArrayToStringOp, expr: TypedExpr) -> sge.Expression:
     return sge.ArrayToString(this=expr.expr, expression=f"'{op.delimiter}'")
@@ -234,6 +242,12 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     ) - sge.convert(1)
 
 
+@UNARY_OP_REGISTRATION.register(ops.FloorDtOp)
+def _(op: ops.FloorDtOp, expr: TypedExpr) -> sge.Expression:
+    # TODO: Remove this method when it is covered by ops.FloorOp
+    return sge.TimestampTrunc(this=expr.expr, unit=sge.Identifier(this=op.freq))
+
+
 @UNARY_OP_REGISTRATION.register(ops.floor_op)
 def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.Floor(this=expr.expr)
@@ -249,6 +263,26 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.func("ST_ASTEXT", expr.expr)
 
 
+@UNARY_OP_REGISTRATION.register(ops.geo_st_boundary_op)
+def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("ST_BOUNDARY", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.geo_st_geogfromtext_op)
+def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("SAFE.ST_GEOGFROMTEXT", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.geo_st_isclosed_op)
+def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("ST_ISCLOSED", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.GeoStLengthOp)
+def _(op: ops.GeoStLengthOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("ST_LENGTH", expr.expr)
+
+
 @UNARY_OP_REGISTRATION.register(ops.geo_x_op)
 def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.func("SAFE.ST_X", expr.expr)
@@ -274,6 +308,11 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.BitwiseNot(this=expr.expr)
 
 
+@UNARY_OP_REGISTRATION.register(ops.IsInOp)
+def _(op: ops.IsInOp, expr: TypedExpr) -> sge.Expression:
+    return sge.In(this=expr.expr, expressions=[sge.convert(v) for v in op.values])
+
+
 @UNARY_OP_REGISTRATION.register(ops.isalnum_op)
 def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.RegexpLike(this=expr.expr, expression=sge.convert(r"^(\p{N}|\p{L})+$"))
@@ -517,6 +556,26 @@ def _(op: ops.StrSliceOp, expr: TypedExpr) -> sge.Expression:
     )
 
 
+@UNARY_OP_REGISTRATION.register(ops.StrftimeOp)
+def _(op: ops.StrftimeOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("FORMAT_TIMESTAMP", sge.convert(op.date_format), expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.StructFieldOp)
+def _(op: ops.StructFieldOp, expr: TypedExpr) -> sge.Expression:
+    if isinstance(op.name_or_index, str):
+        name = op.name_or_index
+    else:
+        pa_type = typing.cast(pd.ArrowDtype, expr.dtype)
+        pa_struct_type = typing.cast(pa.StructType, pa_type.pyarrow_dtype)
+        name = pa_struct_type.field(op.name_or_index).name
+
+    return sge.Column(
+        this=sge.to_identifier(name, quoted=True),
+        catalog=expr.expr,
+    )
+
+
 @UNARY_OP_REGISTRATION.register(ops.tan_op)
 def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.func("TAN", expr.expr)
@@ -537,6 +596,36 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
     return sge.Floor(this=expr.expr)
 
 
+@UNARY_OP_REGISTRATION.register(ops.ToDatetimeOp)
+def _(op: ops.ToDatetimeOp, expr: TypedExpr) -> sge.Expression:
+    return sge.Cast(this=sge.func("TIMESTAMP_SECONDS", expr.expr), to="DATETIME")
+
+
+@UNARY_OP_REGISTRATION.register(ops.ToTimestampOp)
+def _(op: ops.ToTimestampOp, expr: TypedExpr) -> sge.Expression:
+    return sge.func("TIMESTAMP_SECONDS", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.ToTimedeltaOp)
+def _(op: ops.ToTimedeltaOp, expr: TypedExpr) -> sge.Expression:
+    return sge.Interval(this=expr.expr, unit=sge.Identifier(this="SECOND"))
+
+
+@UNARY_OP_REGISTRATION.register(ops.UnixMicros)
+def _(op: ops.UnixMicros, expr: TypedExpr) -> sge.Expression:
+    return sge.func("UNIX_MICROS", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.UnixMillis)
+def _(op: ops.UnixMillis, expr: TypedExpr) -> sge.Expression:
+    return sge.func("UNIX_MILLIS", expr.expr)
+
+
+@UNARY_OP_REGISTRATION.register(ops.UnixSeconds)
+def _(op: ops.UnixSeconds, expr: TypedExpr) -> sge.Expression:
+    return sge.func("UNIX_SECONDS", expr.expr)
+
+
 # JSON Ops
 @UNARY_OP_REGISTRATION.register(ops.JSONExtract)
 def _(op: ops.JSONExtract, expr: TypedExpr) -> sge.Expression:

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `timestamp_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    TIMESTAMP_TRUNC(`bfcol_0`, DAY) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `timestamp_col`
+FROM `bfcte_1`

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `geography_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    ST_BOUNDARY(`bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `geography_col`
+FROM `bfcte_1`

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `string_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    SAFE.ST_GEOGFROMTEXT(`bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `string_col`
+FROM `bfcte_1`

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `geography_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    ST_ISCLOSED(`bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `geography_col`
+FROM `bfcte_1`

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `geography_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    ST_LENGTH(`bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `geography_col`
+FROM `bfcte_1`

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_0` IN (1, 2, 3) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `int64_col`
+FROM `bfcte_1`

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `timestamp_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    FORMAT_TIMESTAMP('%Y-%m-%d', `bfcol_0`) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `timestamp_col`
+FROM `bfcte_1`

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `people` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`nested_structs_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_0`.`name` AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `people`
+FROM `bfcte_1`

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

@@ -0,0 +1,13 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    CAST(TIMESTAMP_SECONDS(`bfcol_0`) AS DATETIME) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `int64_col`
+FROM `bfcte_1`