Improve scalar logic for timestamps

dask_sql/physical/rex/core/call.py

@@ -52,6 +52,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"""
 
@@ -246,6 +250,9 @@ def cast(self, operand, rex=None) -> SeriesOrScalar:
         output_type = str(rex.getType())
         sql_type = SqlTypeName.fromString(output_type.upper())
 
+        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)
 
@@ -607,11 +614,7 @@ def to_timestamp(self, df, format):
         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
+            return df
         # String cases
         elif type(df) == str:
             return np.datetime64(datetime.strptime(df, format))
@@ -646,7 +649,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
@@ -690,6 +696,15 @@ 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
         )

tests/integration/test_rex.py

@@ -1015,3 +1015,114 @@ 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))])
+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")
+    df2 = c.sql("SELECT CAST(d AS TIMESTAMP) + INTERVAL '5 days' AS dt FROM df")
+    assert_eq(df1, expected_df)
+    assert_eq(df2, expected_df)
+
+    df1 = c.sql("SELECT TIMESTAMPADD(DAY, 5, to_timestamp(d)) AS dt FROM df")
+    df2 = c.sql("SELECT TIMESTAMPADD(DAY, 5, d) AS dt FROM df")
+    df3 = c.sql("SELECT TIMESTAMPADD(DAY, 5, CAST(d AS TIMESTAMP)) AS dt FROM df")
+    assert_eq(df1, expected_df)
+    assert_eq(df2, expected_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")
+    df2 = c.sql("SELECT EXTRACT(DAY FROM CAST(d AS TIMESTAMP)) AS day FROM df")
+    assert_eq(df1, expected_df)
+    assert_eq(df2, expected_df)
+
+    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")
+    df2 = c.sql("SELECT CEIL(CAST(d AS TIMESTAMP) TO DAY) AS ceil_to_day FROM df")
+    assert_eq(df1, expected_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")
+    df2 = c.sql("SELECT FLOOR(CAST(d AS TIMESTAMP) TO DAY) AS floor_to_day FROM df")
+    assert_eq(df1, expected_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"
+    )
+    df2 = c.sql("SELECT TIMESTAMPDIFF(DAY, d1, d2) AS dt FROM df")
+    df3 = c.sql(
+        "SELECT TIMESTAMPDIFF(DAY, CAST(d1 AS TIMESTAMP), CAST(d2 AS TIMESTAMP)) AS dt FROM df"
+    )
+    assert_eq(df1, expected_df)
+    assert_eq(df2, expected_df, check_dtype=False)
+    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")
+    # df2 = c.sql(f"SELECT CAST({scalar1} AS TIMESTAMP) + INTERVAL '5 days' AS dt")
+    assert_eq(df1, expected_df)
+    # assert_eq(df2, expected_df)  # TODO: Incorrect
+
+    df1 = c.sql(f"SELECT TIMESTAMPADD(DAY, 5, to_timestamp({scalar1})) AS dt")
+    df2 = c.sql(f"SELECT TIMESTAMPADD(DAY, 5, {scalar1}) AS dt")
+    df3 = c.sql(f"SELECT TIMESTAMPADD(DAY, 5, CAST({scalar1} AS TIMESTAMP)) AS dt")
+    assert_eq(df1, expected_df)
+    assert_eq(df2, expected_df)
+    assert_eq(df3, expected_df)
+
+    expected_df = pd.DataFrame({"day": [15]})
+    df1 = c.sql(f"SELECT EXTRACT(DAY FROM to_timestamp({scalar1})) AS day")
+    # df2 = c.sql(f"SELECT EXTRACT(DAY FROM CAST({scalar1} AS TIMESTAMP)) AS day")
+    assert_eq(df1, expected_df)
+    # assert_eq(df2, expected_df)  # TODO: Incorrect
+
+    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")
+    df2 = c.sql(f"SELECT CEIL(CAST({scalar1} AS TIMESTAMP) TO DAY) AS ceil_to_day")
+    assert_eq(df1, expected_df)
+    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")
+    df2 = c.sql(f"SELECT FLOOR(CAST({scalar1} AS TIMESTAMP) TO DAY) AS floor_to_day")
+    assert_eq(df1, expected_df)
+    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"
+    )
+    df2 = c.sql(f"SELECT TIMESTAMPDIFF(DAY, {scalar1}, {scalar2}) AS dt")
+    df3 = c.sql(
+        f"SELECT TIMESTAMPDIFF(DAY, CAST({scalar1} AS TIMESTAMP), CAST({scalar2} AS TIMESTAMP)) AS dt"
+    )
+    assert_eq(df1, expected_df)
+    assert_eq(df2, expected_df, check_dtype=False)
+    assert_eq(df3, expected_df)