Skip to content

Commit a537ad1

Browse files
jialuoojialuoo
committed
fix
1 parent af06fdd commit a537ad1

File tree

6 files changed

+206
-5
lines changed
  • tests/unit/core/compile/sqlglot/expressions

6 files changed

+206
-5
lines changed

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_integer_label_to_datetime_fixed/out.sql‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`rowindex`,
4+
`timestamp_col`
5+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
6+
), `bfcte_1` AS (
7+
SELECT
8+
*,
9+
CAST(TIMESTAMP_MICROS(
10+
CAST(CAST(`rowindex` AS BIGNUMERIC) * 86400000000 + CAST(UNIX_MICROS(CAST(`timestamp_col` AS TIMESTAMP)) AS BIGNUMERIC) AS INT64)
11+
) AS TIMESTAMP) AS `bfcol_2`
12+
FROM `bfcte_0`
13+
)
14+
SELECT
15+
`bfcol_2` AS `fixed_freq`
16+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_integer_label_to_datetime_month/out.sql‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`rowindex`,
4+
`timestamp_col`
5+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
6+
), `bfcte_1` AS (
7+
SELECT
8+
*,
9+
CAST(TIMESTAMP(
10+
DATETIME(
11+
CASE
12+
WHEN MOD(
13+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 12 + EXTRACT(MONTH FROM `timestamp_col`) - 1,
14+
12
15+
) + 1 = 12
16+
THEN CAST(FLOOR(
17+
IEEE_DIVIDE(
18+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 12 + EXTRACT(MONTH FROM `timestamp_col`) - 1,
19+
12
20+
)
21+
) AS INT64) + 1
22+
ELSE CAST(FLOOR(
23+
IEEE_DIVIDE(
24+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 12 + EXTRACT(MONTH FROM `timestamp_col`) - 1,
25+
12
26+
)
27+
) AS INT64)
28+
END,
29+
CASE
30+
WHEN MOD(
31+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 12 + EXTRACT(MONTH FROM `timestamp_col`) - 1,
32+
12
33+
) + 1 = 12
34+
THEN 1
35+
ELSE MOD(
36+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 12 + EXTRACT(MONTH FROM `timestamp_col`) - 1,
37+
12
38+
) + 1 + 1
39+
END,
40+
1,
41+
0,
42+
0,
43+
0
44+
)
45+
) - INTERVAL 1 DAY AS TIMESTAMP) AS `bfcol_2`
46+
FROM `bfcte_0`
47+
)
48+
SELECT
49+
`bfcol_2` AS `non_fixed_freq_monthly`
50+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_integer_label_to_datetime_quarter/out.sql‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`rowindex`,
4+
`timestamp_col`
5+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
6+
), `bfcte_1` AS (
7+
SELECT
8+
*,
9+
CAST(DATETIME(
10+
CASE
11+
WHEN (
12+
MOD(
13+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 4 + EXTRACT(QUARTER FROM `timestamp_col`) - 1,
14+
4
15+
) + 1
16+
) * 3 = 12
17+
THEN CAST(FLOOR(
18+
IEEE_DIVIDE(
19+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 4 + EXTRACT(QUARTER FROM `timestamp_col`) - 1,
20+
4
21+
)
22+
) AS INT64) + 1
23+
ELSE CAST(FLOOR(
24+
IEEE_DIVIDE(
25+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 4 + EXTRACT(QUARTER FROM `timestamp_col`) - 1,
26+
4
27+
)
28+
) AS INT64)
29+
END,
30+
CASE
31+
WHEN (
32+
MOD(
33+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 4 + EXTRACT(QUARTER FROM `timestamp_col`) - 1,
34+
4
35+
) + 1
36+
) * 3 = 12
37+
THEN 1
38+
ELSE (
39+
MOD(
40+
`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) * 4 + EXTRACT(QUARTER FROM `timestamp_col`) - 1,
41+
4
42+
) + 1
43+
) * 3 + 1
44+
END,
45+
1,
46+
0,
47+
0,
48+
0
49+
) - INTERVAL 1 DAY AS TIMESTAMP) AS `bfcol_2`
50+
FROM `bfcte_0`
51+
)
52+
SELECT
53+
`bfcol_2` AS `non_fixed_freq`
54+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_integer_label_to_datetime_week/out.sql‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`rowindex`,
4+
`timestamp_col`
5+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
6+
), `bfcte_1` AS (
7+
SELECT
8+
*,
9+
CAST(CAST(TIMESTAMP_MICROS(
10+
CAST(CAST(`rowindex` AS BIGNUMERIC) * 604800000000 + CAST(UNIX_MICROS(
11+
TIMESTAMP_TRUNC(CAST(`timestamp_col` AS TIMESTAMP), WEEK(MONDAY)) + INTERVAL 6 DAY
12+
) AS BIGNUMERIC) AS INT64)
13+
) AS TIMESTAMP) AS TIMESTAMP) AS `bfcol_2`
14+
FROM `bfcte_0`
15+
)
16+
SELECT
17+
`bfcol_2` AS `non_fixed_freq_weekly`
18+
FROM `bfcte_1`

‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_integer_label_to_datetime_year/out.sql‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
WITH `bfcte_0` AS (
2+
SELECT
3+
`rowindex`,
4+
`timestamp_col`
5+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
6+
), `bfcte_1` AS (
7+
SELECT
8+
*,
9+
CAST(TIMESTAMP(DATETIME(`rowindex` * 1 + EXTRACT(YEAR FROM `timestamp_col`) + 1, 1, 1, 0, 0, 0)) - INTERVAL 1 DAY AS TIMESTAMP) AS `bfcol_2`
10+
FROM `bfcte_0`
11+
)
12+
SELECT
13+
`bfcol_2` AS `non_fixed_freq_yearly`
14+
FROM `bfcte_1`

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

Lines changed: 54 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -279,17 +279,66 @@ def test_sub_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
279279
snapshot.assert_match(bf_df.sql, "out.sql")
280280

281281

282-
def test_integer_label_to_datetime(scalar_types_df: bpd.DataFrame, snapshot):
282+
def test_integer_label_to_datetime_fixed(scalar_types_df: bpd.DataFrame, snapshot):
283283
col_names = ["rowindex", "timestamp_col"]
284284
bf_df = scalar_types_df[col_names]
285285
ops_map = {
286286
"fixed_freq": ops.IntegerLabelToDatetimeOp(
287-
freq=pd.tseries.offsets.Day(), origin="start", label="left" # type: ignore
287+
freq=pd.tseries.offsets.Day(), origin="start", label="left" # type: ignore
288288
).as_expr("rowindex", "timestamp_col"),
289+
}
290+
291+
sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
292+
snapshot.assert_match(sql, "out.sql")
293+
294+
295+
def test_integer_label_to_datetime_week(scalar_types_df: bpd.DataFrame, snapshot):
296+
col_names = ["rowindex", "timestamp_col"]
297+
bf_df = scalar_types_df[col_names]
298+
ops_map = {
299+
"non_fixed_freq_weekly": ops.IntegerLabelToDatetimeOp(
300+
freq=pd.tseries.offsets.Week(weekday=6), origin="start", label="left"
301+
).as_expr("rowindex", "timestamp_col"),
302+
}
303+
304+
sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
305+
snapshot.assert_match(sql, "out.sql")
306+
307+
308+
def test_integer_label_to_datetime_month(scalar_types_df: bpd.DataFrame, snapshot):
309+
col_names = ["rowindex", "timestamp_col"]
310+
bf_df = scalar_types_df[col_names]
311+
ops_map = {
312+
"non_fixed_freq_monthly": ops.IntegerLabelToDatetimeOp(
313+
freq=pd.tseries.offsets.MonthEnd(), origin="start", label="left"
314+
).as_expr("rowindex", "timestamp_col"),
315+
}
316+
317+
sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
318+
snapshot.assert_match(sql, "out.sql")
319+
320+
321+
def test_integer_label_to_datetime_quarter(scalar_types_df: bpd.DataFrame, snapshot):
322+
col_names = ["rowindex", "timestamp_col"]
323+
bf_df = scalar_types_df[col_names]
324+
ops_map = {
289325
"non_fixed_freq": ops.IntegerLabelToDatetimeOp(
290-
freq=pd.tseries.offsets.QuarterEnd(startingMonth=12), # type: ignore
291-
origin="start",
292-
label="left",
326+
freq=pd.tseries.offsets.QuarterEnd(startingMonth=12), # type: ignore
327+
origin="start",
328+
label="left",
329+
).as_expr("rowindex", "timestamp_col"),
330+
}
331+
332+
sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
333+
snapshot.assert_match(sql, "out.sql")
334+
335+
336+
def test_integer_label_to_datetime_year(scalar_types_df: bpd.DataFrame, snapshot):
337+
col_names = ["rowindex", "timestamp_col"]
338+
bf_df = scalar_types_df[col_names]
339+
ops_map = {
340+
"non_fixed_freq_yearly": ops.IntegerLabelToDatetimeOp(
341+
freq=pd.tseries.offsets.YearEnd(month=12), origin="start", label="left"
293342
).as_expr("rowindex", "timestamp_col"),
294343
}
295344

0 commit comments

Comments
 (0)