From f9fbf5e38ec3050fef957e6b171894a2e1376c84 Mon Sep 17 00:00:00 2001 From: Sarah Yurick <53962159+sarahyurick@users.noreply.github.com> Date: Tue, 18 Oct 2022 10:10:29 -0700 Subject: [PATCH] Check for `np.timedelta64` in `as_timelike` (#860) * check np.timedelta64 case * add test * Update dask_sql/physical/rex/core/call.py Co-authored-by: Ayush Dattagupta Co-authored-by: Ayush Dattagupta --- dask_sql/physical/rex/core/call.py | 2 +- tests/integration/test_rex.py | 19 ++++++++++++++++++- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/dask_sql/physical/rex/core/call.py b/dask_sql/physical/rex/core/call.py index c0b364d89..1903c8fd9 100644 --- a/dask_sql/physical/rex/core/call.py +++ b/dask_sql/physical/rex/core/call.py @@ -44,7 +44,7 @@ def as_timelike(op): return np.timedelta64(op, "D") elif isinstance(op, str): return np.datetime64(op) - elif pd.api.types.is_datetime64_dtype(op): + elif pd.api.types.is_datetime64_dtype(op) or isinstance(op, np.timedelta64): return op else: raise ValueError(f"Don't know how to make {type(op)} timelike") diff --git a/tests/integration/test_rex.py b/tests/integration/test_rex.py index 00153fc43..655ff69de 100644 --- a/tests/integration/test_rex.py +++ b/tests/integration/test_rex.py @@ -59,7 +59,6 @@ def test_intervals(c): """SELECT INTERVAL '3' DAY as "IN" """ ) - expected_df = pd.DataFrame( { "IN": [pd.to_timedelta("3d")], @@ -67,6 +66,24 @@ def test_intervals(c): ) assert_eq(df, expected_df) + date1 = datetime(2021, 10, 3, 15, 53, 42, 47) + date2 = datetime(2021, 2, 28, 15, 53, 42, 47) + dates = dd.from_pandas(pd.DataFrame({"d": [date1, date2]}), npartitions=1) + c.register_dask_table(dates, "dates") + df = c.sql( + """SELECT d + INTERVAL '5 days' AS "Plus_5_days" FROM dates + """ + ) + expected_df = pd.DataFrame( + { + "Plus_5_days": [ + datetime(2021, 10, 8, 15, 53, 42, 47), + datetime(2021, 3, 5, 15, 53, 42, 47), + ] + } + ) + assert_eq(df, expected_df) + def test_literals(c): df = c.sql(