Support cudf-polars epoch

python/cudf_polars/cudf_polars/dsl/expressions/datetime.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # TODO: Document TemporalFunction to remove noqa
 # ruff: noqa: D101
@@ -10,17 +10,19 @@
 from enum import IntEnum, auto
 from typing import TYPE_CHECKING, Any, ClassVar, cast
 
+import polars as pl
+
 import pylibcudf as plc
 
-from cudf_polars.containers import Column
+from cudf_polars.containers import Column, DataType
 from cudf_polars.dsl.expressions.base import ExecutionContext, Expr
 
 if TYPE_CHECKING:
     from typing import Self
 
     from polars import polars  # type: ignore[attr-defined]
 
-    from cudf_polars.containers import DataFrame, DataType
+    from cudf_polars.containers import DataFrame
     from cudf_polars.dsl.expressions.literal import Literal
 
 __all__ = ["TemporalFunction"]
@@ -121,6 +123,7 @@ def from_polars(cls, obj: polars._expr_nodes.TemporalFunction) -> Self:
         Name.IsoYear,
         Name.MonthStart,
         Name.MonthEnd,
+        Name.TimeStamp,
         Name.CastTimeUnit,
         Name.Truncate,
     }
@@ -150,13 +153,27 @@ def __init__(
                 # https://github.com/rapidsai/cudf/issues/18654 to support non-1 buckets
                 raise NotImplementedError(f"Unsupported truncate bucket: {every!r}")
             self.options = (self._TRUNCATE_FREQ_MAP[match.group(2)],)
+        elif self.name is TemporalFunction.Name.TimeStamp:
+            (time_unit,) = self.options
+            if time_unit not in {"ms", "us", "ns"}:
+                raise NotImplementedError(
+                    f"Unsupported epoch/timestamp time unit: {time_unit!r}"
+                )
 
     def do_evaluate(
         self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME
     ) -> Column:
         """Evaluate this expression given a dataframe for context."""
         columns = [child.evaluate(df, context=context) for child in self.children]
-        if self.name is TemporalFunction.Name.Truncate:
+        if self.name is TemporalFunction.Name.TimeStamp:
+            (column,) = columns
+            (time_unit,) = self.options
+            # Rescale the timestamp to the requested resolution
+            df_stream = df.stream
+            return column.astype(
+                DataType(pl.Datetime(time_unit)), stream=df_stream
+            ).astype(self.dtype, stream=df_stream)
+        elif self.name is TemporalFunction.Name.Truncate:
             (column, _) = columns
             return Column(
                 plc.datetime.floor_datetimes(

python/cudf_polars/tests/expressions/test_datetime_basic.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
@@ -301,6 +301,31 @@ def test_isoyear(engine: pl.GPUEngine):
     assert_gpu_result_equal(q, engine=engine)
 
 
+@pytest.mark.parametrize(
+    "dtype",
+    [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")],
+    ids=repr,
+)
+@pytest.mark.parametrize("time_unit", ["ms", "us", "ns", "s", "d"])
+def test_epoch(engine: pl.GPUEngine, dtype, time_unit):
+    ldf = pl.LazyFrame(
+        {
+            "datetimes": pl.Series(
+                [
+                    datetime.datetime(2001, 1, 1),
+                    datetime.datetime(2001, 1, 2, 12, 30, 15),
+                    datetime.datetime(2020, 2, 29, 23, 59, 59),
+                    datetime.datetime(2024, 12, 31, 23, 59, 59),
+                ],
+                dtype=dtype,
+            )
+        }
+    )
+
+    q = ldf.select(pl.col("datetimes").dt.epoch(time_unit))
+    assert_gpu_result_equal(q, engine=engine)
+
+
 @pytest.mark.parametrize(
     "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")]
 )