Skip to content

Commit b1e9700

Browse files
chelsea-linchelsea-lin
committed
refactor: add some numeric tests in the golden sql
1 parent 5ec3cc0 commit b1e9700

File tree

4 files changed

+244
-0
lines changed

4 files changed

+244
-0
lines changed

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_add_timedelta/out.sql‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`date_col` AS `bfcol_0`,
4+
`rowindex` AS `bfcol_1`,
5+
`timestamp_col` AS `bfcol_2`
6+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
7+
), `bfcte_1` AS (
8+
SELECT
9+
*,
10+
`bfcol_1` AS `bfcol_6`,
11+
`bfcol_2` AS `bfcol_7`,
12+
`bfcol_0` AS `bfcol_8`,
13+
TIMESTAMP_ADD(CAST(`bfcol_0` AS DATETIME), INTERVAL 86400000000 MICROSECOND) AS `bfcol_9`
14+
FROM `bfcte_0`
15+
), `bfcte_2` AS (
16+
SELECT
17+
*,
18+
`bfcol_6` AS `bfcol_14`,
19+
`bfcol_7` AS `bfcol_15`,
20+
`bfcol_8` AS `bfcol_16`,
21+
`bfcol_9` AS `bfcol_17`,
22+
TIMESTAMP_ADD(`bfcol_7`, INTERVAL 86400000000 MICROSECOND) AS `bfcol_18`
23+
FROM `bfcte_1`
24+
), `bfcte_3` AS (
25+
SELECT
26+
*,
27+
`bfcol_14` AS `bfcol_24`,
28+
`bfcol_15` AS `bfcol_25`,
29+
`bfcol_16` AS `bfcol_26`,
30+
`bfcol_17` AS `bfcol_27`,
31+
`bfcol_18` AS `bfcol_28`,
32+
TIMESTAMP_ADD(CAST(`bfcol_16` AS DATETIME), INTERVAL 86400000000 MICROSECOND) AS `bfcol_29`
33+
FROM `bfcte_2`
34+
), `bfcte_4` AS (
35+
SELECT
36+
*,
37+
`bfcol_24` AS `bfcol_36`,
38+
`bfcol_25` AS `bfcol_37`,
39+
`bfcol_26` AS `bfcol_38`,
40+
`bfcol_27` AS `bfcol_39`,
41+
`bfcol_28` AS `bfcol_40`,
42+
`bfcol_29` AS `bfcol_41`,
43+
TIMESTAMP_ADD(`bfcol_25`, INTERVAL 86400000000 MICROSECOND) AS `bfcol_42`
44+
FROM `bfcte_3`
45+
), `bfcte_5` AS (
46+
SELECT
47+
*,
48+
172800000000 AS `bfcol_50`
49+
FROM `bfcte_4`
50+
)
51+
SELECT
52+
`bfcol_36` AS `rowindex`,
53+
`bfcol_37` AS `timestamp_col`,
54+
`bfcol_38` AS `date_col`,
55+
`bfcol_39` AS `date_add_timedelta`,
56+
`bfcol_40` AS `timestamp_add_timedelta`,
57+
`bfcol_41` AS `timedelta_add_date`,
58+
`bfcol_42` AS `timedelta_add_timestamp`,
59+
`bfcol_50` AS `timedelta_add_timedelta`
60+
FROM `bfcte_5`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mul_timedelta/out.sql‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`int64_col` AS `bfcol_0`,
4+
`rowindex` AS `bfcol_1`,
5+
`timestamp_col` AS `bfcol_2`,
6+
`duration_col` AS `bfcol_3`
7+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
8+
), `bfcte_1` AS (
9+
SELECT
10+
*,
11+
`bfcol_1` AS `bfcol_8`,
12+
`bfcol_2` AS `bfcol_9`,
13+
`bfcol_0` AS `bfcol_10`,
14+
`bfcol_3` AS `bfcol_11`
15+
FROM `bfcte_0`
16+
), `bfcte_2` AS (
17+
SELECT
18+
*,
19+
`bfcol_8` AS `bfcol_16`,
20+
`bfcol_9` AS `bfcol_17`,
21+
`bfcol_10` AS `bfcol_18`,
22+
`bfcol_11` AS `bfcol_19`,
23+
CAST(FLOOR(`bfcol_11` * `bfcol_10`) AS INT64) AS `bfcol_20`
24+
FROM `bfcte_1`
25+
), `bfcte_3` AS (
26+
SELECT
27+
*,
28+
`bfcol_16` AS `bfcol_26`,
29+
`bfcol_17` AS `bfcol_27`,
30+
`bfcol_18` AS `bfcol_28`,
31+
`bfcol_19` AS `bfcol_29`,
32+
`bfcol_20` AS `bfcol_30`,
33+
CAST(FLOOR(`bfcol_18` * `bfcol_19`) AS INT64) AS `bfcol_31`
34+
FROM `bfcte_2`
35+
)
36+
SELECT
37+
`bfcol_26` AS `rowindex`,
38+
`bfcol_27` AS `timestamp_col`,
39+
`bfcol_28` AS `int64_col`,
40+
`bfcol_29` AS `duration_col`,
41+
`bfcol_30` AS `timedelta_mul_numeric`,
42+
`bfcol_31` AS `numeric_mul_timedelta`
43+
FROM `bfcte_3`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_sub_timedelta/out.sql‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`date_col` AS `bfcol_0`,
4+
`rowindex` AS `bfcol_1`,
5+
`timestamp_col` AS `bfcol_2`,
6+
`duration_col` AS `bfcol_3`
7+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
8+
), `bfcte_1` AS (
9+
SELECT
10+
*,
11+
`bfcol_1` AS `bfcol_8`,
12+
`bfcol_2` AS `bfcol_9`,
13+
`bfcol_0` AS `bfcol_10`,
14+
`bfcol_3` AS `bfcol_11`
15+
FROM `bfcte_0`
16+
), `bfcte_2` AS (
17+
SELECT
18+
*,
19+
`bfcol_8` AS `bfcol_16`,
20+
`bfcol_9` AS `bfcol_17`,
21+
`bfcol_11` AS `bfcol_18`,
22+
`bfcol_10` AS `bfcol_19`,
23+
TIMESTAMP_SUB(CAST(`bfcol_10` AS DATETIME), INTERVAL `bfcol_11` MICROSECOND) AS `bfcol_20`
24+
FROM `bfcte_1`
25+
), `bfcte_3` AS (
26+
SELECT
27+
*,
28+
`bfcol_16` AS `bfcol_26`,
29+
`bfcol_17` AS `bfcol_27`,
30+
`bfcol_18` AS `bfcol_28`,
31+
`bfcol_19` AS `bfcol_29`,
32+
`bfcol_20` AS `bfcol_30`,
33+
TIMESTAMP_SUB(`bfcol_17`, INTERVAL `bfcol_18` MICROSECOND) AS `bfcol_31`
34+
FROM `bfcte_2`
35+
), `bfcte_4` AS (
36+
SELECT
37+
*,
38+
`bfcol_26` AS `bfcol_38`,
39+
`bfcol_27` AS `bfcol_39`,
40+
`bfcol_28` AS `bfcol_40`,
41+
`bfcol_29` AS `bfcol_41`,
42+
`bfcol_30` AS `bfcol_42`,
43+
`bfcol_31` AS `bfcol_43`,
44+
TIMESTAMP_DIFF(CAST(`bfcol_29` AS DATETIME), CAST(`bfcol_29` AS DATETIME), MICROSECOND) AS `bfcol_44`
45+
FROM `bfcte_3`
46+
), `bfcte_5` AS (
47+
SELECT
48+
*,
49+
`bfcol_38` AS `bfcol_52`,
50+
`bfcol_39` AS `bfcol_53`,
51+
`bfcol_40` AS `bfcol_54`,
52+
`bfcol_41` AS `bfcol_55`,
53+
`bfcol_42` AS `bfcol_56`,
54+
`bfcol_43` AS `bfcol_57`,
55+
`bfcol_44` AS `bfcol_58`,
56+
TIMESTAMP_DIFF(`bfcol_39`, `bfcol_39`, MICROSECOND) AS `bfcol_59`
57+
FROM `bfcte_4`
58+
), `bfcte_6` AS (
59+
SELECT
60+
*,
61+
`bfcol_52` AS `bfcol_68`,
62+
`bfcol_53` AS `bfcol_69`,
63+
`bfcol_54` AS `bfcol_70`,
64+
`bfcol_55` AS `bfcol_71`,
65+
`bfcol_56` AS `bfcol_72`,
66+
`bfcol_57` AS `bfcol_73`,
67+
`bfcol_58` AS `bfcol_74`,
68+
`bfcol_59` AS `bfcol_75`,
69+
`bfcol_54` - `bfcol_54` AS `bfcol_76`
70+
FROM `bfcte_5`
71+
)
72+
SELECT
73+
`bfcol_68` AS `rowindex`,
74+
`bfcol_69` AS `timestamp_col`,
75+
`bfcol_70` AS `duration_col`,
76+
`bfcol_71` AS `date_col`,
77+
`bfcol_72` AS `date_sub_timedelta`,
78+
`bfcol_73` AS `timestamp_sub_timedelta`,
79+
`bfcol_74` AS `timestamp_sub_date`,
80+
`bfcol_75` AS `date_sub_timestamp`,
81+
`bfcol_76` AS `timedelta_sub_timedelta`
82+
FROM `bfcte_6`

‎tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,34 @@ def test_add_numeric(scalar_types_df: bpd.DataFrame, snapshot):
218218
snapshot.assert_match(bf_df.sql, "out.sql")
219219

220220

221+
def test_add_string(scalar_types_df: bpd.DataFrame, snapshot):
222+
bf_df = scalar_types_df[["string_col"]]
223+
sql = utils._apply_binary_op(bf_df, ops.add_op, "string_col", ex.const("a"))
224+
225+
snapshot.assert_match(sql, "out.sql")
226+
227+
228+
def test_add_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
229+
bf_df = scalar_types_df[["timestamp_col", "date_col"]]
230+
timedelta = pd.Timedelta(1, unit="d")
231+
232+
bf_df["date_add_timedelta"] = bf_df["date_col"] + timedelta
233+
bf_df["timestamp_add_timedelta"] = bf_df["timestamp_col"] + timedelta
234+
bf_df["timedelta_add_date"] = timedelta + bf_df["date_col"]
235+
bf_df["timedelta_add_timestamp"] = timedelta + bf_df["timestamp_col"]
236+
bf_df["timedelta_add_timedelta"] = timedelta + timedelta
237+
238+
snapshot.assert_match(bf_df.sql, "out.sql")
239+
240+
241+
def test_add_unsupported_raises(scalar_types_df: bpd.DataFrame):
242+
with pytest.raises(TypeError):
243+
utils._apply_binary_op(scalar_types_df, ops.add_op, "timestamp_col", "date_col")
244+
245+
with pytest.raises(TypeError):
246+
utils._apply_binary_op(scalar_types_df, ops.add_op, "int64_col", "string_col")
247+
248+
221249
def test_div_numeric(scalar_types_df: bpd.DataFrame, snapshot):
222250
bf_df = scalar_types_df[["int64_col", "bool_col", "float64_col"]]
223251

@@ -279,6 +307,16 @@ def test_mul_numeric(scalar_types_df: bpd.DataFrame, snapshot):
279307
snapshot.assert_match(bf_df.sql, "out.sql")
280308

281309

310+
def test_mul_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
311+
bf_df = scalar_types_df[["timestamp_col", "int64_col", "duration_col"]]
312+
bf_df["duration_col"] = bpd.to_timedelta(bf_df["duration_col"], unit="us")
313+
314+
bf_df["timedelta_mul_numeric"] = bf_df["duration_col"] * bf_df["int64_col"]
315+
bf_df["numeric_mul_timedelta"] = bf_df["int64_col"] * bf_df["duration_col"]
316+
317+
snapshot.assert_match(bf_df.sql, "out.sql")
318+
319+
282320
def test_mod_numeric(scalar_types_df: bpd.DataFrame, snapshot):
283321
bf_df = scalar_types_df[["int64_col", "float64_col"]]
284322

@@ -305,3 +343,24 @@ def test_sub_numeric(scalar_types_df: bpd.DataFrame, snapshot):
305343
bf_df["bool_add_int"] = bf_df["bool_col"] - bf_df["int64_col"]
306344

307345
snapshot.assert_match(bf_df.sql, "out.sql")
346+
347+
348+
def test_sub_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
349+
bf_df = scalar_types_df[["timestamp_col", "duration_col", "date_col"]]
350+
bf_df["duration_col"] = bpd.to_timedelta(bf_df["duration_col"], unit="us")
351+
352+
bf_df["date_sub_timedelta"] = bf_df["date_col"] - bf_df["duration_col"]
353+
bf_df["timestamp_sub_timedelta"] = bf_df["timestamp_col"] - bf_df["duration_col"]
354+
bf_df["timestamp_sub_date"] = bf_df["date_col"] - bf_df["date_col"]
355+
bf_df["date_sub_timestamp"] = bf_df["timestamp_col"] - bf_df["timestamp_col"]
356+
bf_df["timedelta_sub_timedelta"] = bf_df["duration_col"] - bf_df["duration_col"]
357+
358+
snapshot.assert_match(bf_df.sql, "out.sql")
359+
360+
361+
def test_sub_unsupported_raises(scalar_types_df: bpd.DataFrame):
362+
with pytest.raises(TypeError):
363+
utils._apply_binary_op(scalar_types_df, ops.sub_op, "string_col", "string_col")
364+
365+
with pytest.raises(TypeError):
366+
utils._apply_binary_op(scalar_types_df, ops.sub_op, "int64_col", "string_col")

0 commit comments

Comments
 (0)