diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py b/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py index 3489bf360ade..4b6036e5c713 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py +++ b/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,9 +10,11 @@ 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: @@ -20,7 +22,7 @@ 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, } @@ -156,7 +159,15 @@ def do_evaluate( ) -> 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( diff --git a/python/cudf_polars/tests/expressions/test_datetime_basic.py b/python/cudf_polars/tests/expressions/test_datetime_basic.py index c5fcbbde7b25..01c876b57dde 100644 --- a/python/cudf_polars/tests/expressions/test_datetime_basic.py +++ b/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")] )