Skip to content

Commit c8bae76

Browse files
committed
chore: Migrate up to 5 scalar operators to SQLGlot
Migrated operators: - abs_op - capitalize_op - ceil_op - date_op - day_op
1 parent 4d91581 commit c8bae76

File tree

7 files changed

+120
-0
lines changed

7 files changed

+120
-0
lines changed

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,11 @@ 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+
4146
@UNARY_OP_REGISTRATION.register(ops.arccosh_op)
4247
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
4348
return sge.Case(
@@ -142,6 +147,16 @@ def _(op: ops.ArraySliceOp, expr: TypedExpr) -> sge.Expression:
142147
return sge.array(selected_elements)
143148

144149

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+
145160
@UNARY_OP_REGISTRATION.register(ops.cos_op)
146161
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
147162
return sge.func("COS", expr.expr)
@@ -160,6 +175,16 @@ def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
160175
)
161176

162177

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+
163188
@UNARY_OP_REGISTRATION.register(ops.hash_op)
164189
def _(op: ops.base_ops.UnaryOp, expr: TypedExpr) -> sge.Expression:
165190
return sge.func("FARM_FINGERPRINT", expr.expr)
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`
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`
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`
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`
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`

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,36 @@ def test_arctanh(scalar_types_df: bpd.DataFrame, snapshot):
7373
snapshot.assert_match(sql, "out.sql")
7474

7575

76+
def test_abs(scalar_types_df: bpd.DataFrame, snapshot):
77+
bf_df = scalar_types_df[["float64_col"]]
78+
sql = _apply_unary_op(bf_df, ops.abs_op, "float64_col")
79+
snapshot.assert_match(sql, "out.sql")
80+
81+
82+
def test_capitalize(scalar_types_df: bpd.DataFrame, snapshot):
83+
bf_df = scalar_types_df[["string_col"]]
84+
sql = _apply_unary_op(bf_df, ops.capitalize_op, "string_col")
85+
snapshot.assert_match(sql, "out.sql")
86+
87+
88+
def test_ceil(scalar_types_df: bpd.DataFrame, snapshot):
89+
bf_df = scalar_types_df[["float64_col"]]
90+
sql = _apply_unary_op(bf_df, ops.ceil_op, "float64_col")
91+
snapshot.assert_match(sql, "out.sql")
92+
93+
94+
def test_date(scalar_types_df: bpd.DataFrame, snapshot):
95+
bf_df = scalar_types_df[["timestamp_col"]]
96+
sql = _apply_unary_op(bf_df, ops.date_op, "timestamp_col")
97+
snapshot.assert_match(sql, "out.sql")
98+
99+
100+
def test_day(scalar_types_df: bpd.DataFrame, snapshot):
101+
bf_df = scalar_types_df[["timestamp_col"]]
102+
sql = _apply_unary_op(bf_df, ops.day_op, "timestamp_col")
103+
snapshot.assert_match(sql, "out.sql")
104+
105+
76106
def test_array_to_string(repeated_types_df: bpd.DataFrame, snapshot):
77107
bf_df = repeated_types_df[["string_list_col"]]
78108
sql = _apply_unary_op(bf_df, ops.ArrayToStringOp(delimiter="."), "string_list_col")

0 commit comments

Comments
 (0)