Improve scalar logic for timestamps

dask_planner/src/sql.rs

@@ -244,18 +244,35 @@ impl ContextProvider for DaskSQLContext {
                             DataType::Int64,
                             DataType::Timestamp(TimeUnit::Nanosecond, None),
                         ]),
+                        TypeSignature::Exact(vec![
+                            DataType::Utf8,
+                            DataType::Int64,
+                            DataType::Int64,
+                        ]),
                     ],
                     Volatility::Immutable,
                 );
                 let rtf: ReturnTypeFunction = Arc::new(|_| Ok(Arc::new(DataType::Date64)));
                 return Some(Arc::new(ScalarUDF::new(name, &sig, &rtf, &fun)));
             }
             "timestampdiff" => {
-                let sig = Signature::exact(
+                let sig = Signature::one_of(
                     vec![
-                        DataType::Utf8,
-                        DataType::Timestamp(TimeUnit::Nanosecond, None),
-                        DataType::Timestamp(TimeUnit::Nanosecond, None),
+                        TypeSignature::Exact(vec![
+                            DataType::Utf8,
+                            DataType::Timestamp(TimeUnit::Nanosecond, None),
+                            DataType::Timestamp(TimeUnit::Nanosecond, None),
+                        ]),
+                        TypeSignature::Exact(vec![
+                            DataType::Utf8,
+                            DataType::Date64,
+                            DataType::Date64,
+                        ]),
+                        TypeSignature::Exact(vec![
+                            DataType::Utf8,
+                            DataType::Int64,
+                            DataType::Int64,
+                        ]),
                     ],
                     Volatility::Immutable,
                 );

dask_sql/_compat.py

@@ -18,3 +18,5 @@
 # Parquet predicate-support version checks
 PQ_NOT_IN_SUPPORT = parseVersion(dask.__version__) > parseVersion("2023.5.1")
 PQ_IS_SUPPORT = parseVersion(dask.__version__) >= parseVersion("2023.3.1")
+
+DASK_CUDF_TODATETIME_SUPPORT = _dask_version >= parseVersion("2023.5.1")

dask_sql/physical/rex/core/call.py

@@ -51,6 +51,10 @@ def as_timelike(op):
         raise ValueError(f"Don't know how to make {type(op)} timelike")
 
 
+def is_timestamp_nano(obj):
+    return "int" in str(type(obj)) or "int" in str(getattr(obj, "dtype", ""))
+
+
 class Operation:
     """Helper wrapper around a function, which is used as operator"""
 
@@ -137,6 +141,9 @@ def reduce(self, *operands, **kwargs):
                 )
             ):
                 operands = tuple(map(as_timelike, operands))
+            if is_cudf_type(operands[0]) and isinstance(operands[1], np.timedelta64):
+                operands = (dd.to_datetime(operands[0], unit="s"), operands[1])
+
             return reduce(partial(self.operation, **kwargs), operands)
         else:
             return self.unary_operation(*operands, **kwargs)
@@ -250,6 +257,9 @@ def cast(self, operand, rex=None) -> SeriesOrScalar:
         if output_type == "DECIMAL":
             sql_type_args = rex.getPrecisionScale()
 
+        if output_type == "TIMESTAMP" and is_timestamp_nano(operand):
+            operand = operand * 10**9
+
         if not is_frame(operand):  # pragma: no cover
             return sql_to_python_value(sql_type, operand)
 
@@ -612,17 +622,8 @@ def to_timestamp(self, df, format):
         format = format.replace('"', "")
         format = format.replace("'", "")
 
-        # TODO: format timestamps for GPU tests
-        if is_cudf_type(df):
-            if format != default_format:
-                raise RuntimeError("Non-default timestamp formats not supported on GPU")
-            if df.dtype == "object":
-                return df
-            else:
-                nanoseconds_to_seconds = 10**9
-                return df * nanoseconds_to_seconds
         # String cases
-        elif type(df) == str:
+        if type(df) == str:
             return np.datetime64(datetime.strptime(df, format))
         elif df.dtype == "object":
             return dd.to_datetime(df, format=format)
@@ -634,7 +635,10 @@ def to_timestamp(self, df, format):
         else:
             if format != default_format:
                 raise RuntimeError("Integer input does not accept a format argument")
-            return dd.to_datetime(df, unit="s")
+            if is_cudf_type(df):
+                return df
+            else:
+                return dd.to_datetime(df, unit="s")
 
 
 class YearOperation(Operation):
@@ -655,7 +659,10 @@ def timestampadd(self, unit, interval, df: SeriesOrScalar):
         interval = int(interval)
         if interval < 0:
             raise RuntimeError(f"Negative time interval {interval} is not supported.")
-        df = df.astype("datetime64[ns]")
+        if is_timestamp_nano(df):
+            df = df.astype("datetime64[s]")
+        else:
+            df = df.astype("datetime64[ns]")
 
         if is_cudf_type(df):
             from cudf import DateOffset
@@ -699,11 +706,23 @@ def __init__(self):
         super().__init__(self.datetime_sub)
 
     def datetime_sub(self, unit, df1, df2):
+        if is_timestamp_nano(df1):
+            df1 = df1 * 10**9
+        if is_timestamp_nano(df2):
+            df2 = df2 * 10**9
+        if "datetime64[s]" == str(getattr(df1, "dtype", "")):
+            df1 = df1.astype("datetime64[ns]")
+        if "datetime64[s]" == str(getattr(df2, "dtype", "")):
+            df2 = df2.astype("datetime64[ns]")
+
         subtraction_op = ReduceOperation(
             operation=operator.sub, unary_operation=lambda x: -x
         )
         result = subtraction_op(df2, df1)
 
+        if is_cudf_type(df1):
+            result = result.astype("int")
+
         if unit in {"NANOSECOND", "NANOSECONDS"}:
             return result
         elif unit in {"MICROSECOND", "MICROSECONDS"}:

tests/integration/test_rex.py

@@ -5,6 +5,7 @@
 import pandas as pd
 import pytest
 
+from dask_sql._compat import DASK_CUDF_TODATETIME_SUPPORT
 from tests.utils import assert_eq
 
 
@@ -1054,3 +1055,131 @@ def test_totimestamp(c, gpu):
         }
     )
     assert_eq(df, expected_df, check_dtype=False)
+
+
+@pytest.mark.parametrize(
+    "gpu",
+    [
+        False,
+        pytest.param(
+            True,
+            marks=(
+                pytest.mark.gpu,
+                pytest.mark.skipif(
+                    not DASK_CUDF_TODATETIME_SUPPORT,
+                    reason="Requires https://github.com/dask/dask/pull/9881",
+                ),
+            ),
+        ),
+    ],
+)
+def test_scalar_timestamps(c, gpu):
+    df = pd.DataFrame({"d": [1203073300, 1503073700]})
+    c.create_table("df", df, gpu=gpu)
+
+    expected_df = pd.DataFrame(
+        {
+            "dt": [datetime(2008, 2, 20, 11, 1, 40), datetime(2017, 8, 23, 16, 28, 20)],
+        }
+    )
+
+    df1 = c.sql("SELECT to_timestamp(d) + INTERVAL '5 days' AS dt FROM df")
+    assert_eq(df1, expected_df)
+    df2 = c.sql("SELECT CAST(d AS TIMESTAMP) + INTERVAL '5 days' AS dt FROM df")
+    assert_eq(df2, expected_df)
+
+    df1 = c.sql("SELECT TIMESTAMPADD(DAY, 5, to_timestamp(d)) AS dt FROM df")
+    assert_eq(df1, expected_df)
+    df2 = c.sql("SELECT TIMESTAMPADD(DAY, 5, d) AS dt FROM df")
+    assert_eq(df2, expected_df)
+    df3 = c.sql("SELECT TIMESTAMPADD(DAY, 5, CAST(d AS TIMESTAMP)) AS dt FROM df")
+    assert_eq(df3, expected_df)
+
+    expected_df = pd.DataFrame({"day": [15, 18]})
+    df1 = c.sql("SELECT EXTRACT(DAY FROM to_timestamp(d)) AS day FROM df")
+    assert_eq(df1, expected_df, check_dtype=False)
+    df2 = c.sql("SELECT EXTRACT(DAY FROM CAST(d AS TIMESTAMP)) AS day FROM df")
+    assert_eq(df2, expected_df, check_dtype=False)
+
+    expected_df = pd.DataFrame(
+        {
+            "ceil_to_day": [datetime(2008, 2, 16), datetime(2017, 8, 19)],
+        }
+    )
+    df1 = c.sql("SELECT CEIL(to_timestamp(d) TO DAY) AS ceil_to_day FROM df")
+    assert_eq(df1, expected_df)
+    df2 = c.sql("SELECT CEIL(CAST(d AS TIMESTAMP) TO DAY) AS ceil_to_day FROM df")
+    assert_eq(df2, expected_df)
+
+    expected_df = pd.DataFrame(
+        {
+            "floor_to_day": [datetime(2008, 2, 15), datetime(2017, 8, 18)],
+        }
+    )
+    df1 = c.sql("SELECT FLOOR(to_timestamp(d) TO DAY) AS floor_to_day FROM df")
+    assert_eq(df1, expected_df)
+    df2 = c.sql("SELECT FLOOR(CAST(d AS TIMESTAMP) TO DAY) AS floor_to_day FROM df")
+    assert_eq(df2, expected_df)
+
+    df = pd.DataFrame({"d1": [1203073300], "d2": [1503073700]})
+    c.create_table("df", df, gpu=gpu)
+    expected_df = pd.DataFrame({"dt": [3472]})
+    df1 = c.sql(
+        "SELECT TIMESTAMPDIFF(DAY, to_timestamp(d1), to_timestamp(d2)) AS dt FROM df"
+    )
+    assert_eq(df1, expected_df)
+    df2 = c.sql("SELECT TIMESTAMPDIFF(DAY, d1, d2) AS dt FROM df")
+    assert_eq(df2, expected_df, check_dtype=False)
+    df3 = c.sql(
+        "SELECT TIMESTAMPDIFF(DAY, CAST(d1 AS TIMESTAMP), CAST(d2 AS TIMESTAMP)) AS dt FROM df"
+    )
+    assert_eq(df3, expected_df)
+
+    scalar1 = 1203073300
+    scalar2 = 1503073700
+
+    expected_df = pd.DataFrame({"dt": [datetime(2008, 2, 20, 11, 1, 40)]})
+
+    df1 = c.sql(f"SELECT to_timestamp({scalar1}) + INTERVAL '5 days' AS dt")
+    assert_eq(df1, expected_df)
+    # TODO: Fix seconds/nanoseconds conversion
+    # df2 = c.sql(f"SELECT CAST({scalar1} AS TIMESTAMP) + INTERVAL '5 days' AS dt")
+    # assert_eq(df2, expected_df)
+
+    df1 = c.sql(f"SELECT TIMESTAMPADD(DAY, 5, to_timestamp({scalar1})) AS dt")
+    assert_eq(df1, expected_df)
+    df2 = c.sql(f"SELECT TIMESTAMPADD(DAY, 5, {scalar1}) AS dt")
+    assert_eq(df2, expected_df)
+    df3 = c.sql(f"SELECT TIMESTAMPADD(DAY, 5, CAST({scalar1} AS TIMESTAMP)) AS dt")
+    assert_eq(df3, expected_df)
+
+    expected_df = pd.DataFrame({"day": [15]})
+    df1 = c.sql(f"SELECT EXTRACT(DAY FROM to_timestamp({scalar1})) AS day")
+    assert_eq(df1, expected_df, check_dtype=False)
+    # TODO: Fix seconds/nanoseconds conversion
+    # df2 = c.sql(f"SELECT EXTRACT(DAY FROM CAST({scalar1} AS TIMESTAMP)) AS day")
+    # assert_eq(df2, expected_df, check_dtype=False)
+
+    expected_df = pd.DataFrame({"ceil_to_day": [datetime(2008, 2, 16)]})
+    df1 = c.sql(f"SELECT CEIL(to_timestamp({scalar1}) TO DAY) AS ceil_to_day")
+    assert_eq(df1, expected_df)
+    df2 = c.sql(f"SELECT CEIL(CAST({scalar1} AS TIMESTAMP) TO DAY) AS ceil_to_day")
+    assert_eq(df2, expected_df)
+
+    expected_df = pd.DataFrame({"floor_to_day": [datetime(2008, 2, 15)]})
+    df1 = c.sql(f"SELECT FLOOR(to_timestamp({scalar1}) TO DAY) AS floor_to_day")
+    assert_eq(df1, expected_df)
+    df2 = c.sql(f"SELECT FLOOR(CAST({scalar1} AS TIMESTAMP) TO DAY) AS floor_to_day")
+    assert_eq(df2, expected_df)
+
+    expected_df = pd.DataFrame({"dt": [3472]})
+    df1 = c.sql(
+        f"SELECT TIMESTAMPDIFF(DAY, to_timestamp({scalar1}), to_timestamp({scalar2})) AS dt"
+    )
+    assert_eq(df1, expected_df)
+    df2 = c.sql(f"SELECT TIMESTAMPDIFF(DAY, {scalar1}, {scalar2}) AS dt")
+    assert_eq(df2, expected_df, check_dtype=False)
+    df3 = c.sql(
+        f"SELECT TIMESTAMPDIFF(DAY, CAST({scalar1} AS TIMESTAMP), CAST({scalar2} AS TIMESTAMP)) AS dt"
+    )
+    assert_eq(df3, expected_df)