diff --git a/python/cudf_polars/cudf_polars/containers/column.py b/python/cudf_polars/cudf_polars/containers/column.py index 9f8514dca77f..bca694352f6c 100644 --- a/python/cudf_polars/cudf_polars/containers/column.py +++ b/python/cudf_polars/cudf_polars/containers/column.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 """A column, with some properties.""" @@ -353,6 +353,29 @@ def astype(self, dtype: DataType, stream: Stream, *, strict: bool = True) -> Col dtype=dtype, name=self.name, ).sorted_like(self) + elif plc.traits.is_integral_not_bool(plc_dtype) and plc.traits.is_duration( + self.obj.type() + ): + # A duration is stored as an integer tick count, so casting to that + # integer type is a no-op reinterpret of the same bytes. Relabel the + # column instead of launching a cast kernel. + rep = plc.DataType( + plc.TypeId.INT32 + if self.obj.type().id() == plc.TypeId.DURATION_DAYS + else plc.TypeId.INT64 + ) + plc_col = plc.column.Column( + rep, + self.obj.size(), + self.obj.data(), + self.obj.null_mask(), + self.obj.null_count(), + self.obj.offset(), + self.obj.children(), + ) + if rep.id() != plc_dtype.id(): + plc_col = plc.unary.cast(plc_col, plc_dtype, stream=stream) + return Column(plc_col, dtype=dtype, name=self.name).sorted_like(self) elif plc.traits.is_floating_point( self.obj.type() ) and plc.traits.is_fixed_point(plc_dtype): diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py b/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py index 4b6036e5c713..b6b582738188 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py @@ -28,6 +28,15 @@ __all__ = ["TemporalFunction"] +_unit_to_nanoseconds_conversion = { + plc.TypeId.DURATION_NANOSECONDS: 1, + plc.TypeId.DURATION_MICROSECONDS: 1_000, + plc.TypeId.DURATION_MILLISECONDS: 1_000_000, + plc.TypeId.DURATION_SECONDS: 1_000_000_000, + plc.TypeId.DURATION_DAYS: 86_400_000_000_000, +} + + class TemporalFunction(Expr): class Name(IntEnum): """Internal and picklable representation of polars' `TemporalFunction`.""" @@ -114,6 +123,16 @@ def from_polars(cls, obj: polars._expr_nodes.TemporalFunction) -> Self: "ns": plc.datetime.RoundingFrequency.NANOSECOND, } + # Number of nanoseconds represented by one unit of each ``total_*`` component. + _TOTAL_COMPONENT_NANOSECONDS: ClassVar[dict[Name, int]] = { + Name.TotalDays: 86_400_000_000_000, + Name.TotalHours: 3_600_000_000_000, + Name.TotalMinutes: 60_000_000_000, + Name.TotalSeconds: 1_000_000_000, + Name.TotalMilliseconds: 1_000_000, + Name.TotalMicroseconds: 1_000, + Name.TotalNanoseconds: 1, + } _valid_ops: ClassVar[set[Name]] = { *_COMPONENT_MAP.keys(), Name.IsLeapYear, @@ -126,6 +145,7 @@ def from_polars(cls, obj: polars._expr_nodes.TemporalFunction) -> Self: Name.TimeStamp, Name.CastTimeUnit, Name.Truncate, + *_TOTAL_COMPONENT_NANOSECONDS.keys(), } def __init__( @@ -159,6 +179,34 @@ 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 in self._TOTAL_COMPONENT_NANOSECONDS: + (column,) = columns + source_ns = _unit_to_nanoseconds_conversion[column.obj.type().id()] + target_ns = self._TOTAL_COMPONENT_NANOSECONDS[self.name] + # Reinterpret the duration's integer tick count as int64. + casted = column.astype(self.dtype, stream=df.stream) + if source_ns >= target_ns: + # Coarser (or equal) storage unit: exact integer multiply. + op = plc.binaryop.BinaryOperator.MUL + factor = source_ns // target_ns + else: + # Finer storage unit: integer divide. libcudf (like polars) + # truncates toward zero for signed integer division. + op = plc.binaryop.BinaryOperator.DIV + factor = target_ns // source_ns + if factor == 1: + # Storage unit already matches the requested unit. + return casted + result = plc.binaryop.binary_operation( + casted.obj, + plc.Scalar.from_py( + factor, plc.DataType(plc.TypeId.INT64), stream=df.stream + ), + op, + self.dtype.plc_type, + stream=df.stream, + ) + return Column(result, dtype=self.dtype) if self.name is TemporalFunction.Name.TimeStamp: (column,) = columns (time_unit,) = self.options @@ -257,7 +305,6 @@ def do_evaluate( self.dtype.plc_type, stream=df.stream, ) - return Column(result, dtype=self.dtype) elif self.name is TemporalFunction.Name.MonthEnd: (column,) = columns diff --git a/python/cudf_polars/tests/containers/test_column.py b/python/cudf_polars/tests/containers/test_column.py index 674ae10edbab..19b249e2687e 100644 --- a/python/cudf_polars/tests/containers/test_column.py +++ b/python/cudf_polars/tests/containers/test_column.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 @@ -315,6 +315,24 @@ def test_astype_to_string(val, plc_tid, pl_type): assert result.dtype == target_dtype +def test_astype_duration_to_narrower_integer(): + stream = get_cuda_stream() + col = Column( + plc.unary.cast( + plc.Column.from_iterable_of_py( + [1, 2, -3], plc.DataType(plc.TypeId.INT64), stream=stream + ), + plc.DataType(plc.TypeId.DURATION_MICROSECONDS), + stream=stream, + ), + dtype=DataType(pl.Duration(time_unit="us")), + ) + target_dtype = DataType(pl.Int32()) + result = col.astype(target_dtype, stream=stream) + assert result.dtype == target_dtype + assert result.obj.type().id() == plc.TypeId.INT32 + + def test_astype_from_string_unsupported(): stream = get_cuda_stream() col = Column( diff --git a/python/cudf_polars/tests/expressions/test_datetime_basic.py b/python/cudf_polars/tests/expressions/test_datetime_basic.py index 01c876b57dde..26d4c01a4766 100644 --- a/python/cudf_polars/tests/expressions/test_datetime_basic.py +++ b/python/cudf_polars/tests/expressions/test_datetime_basic.py @@ -57,6 +57,16 @@ def test_datetime_dataframe_scan(engine: pl.GPUEngine, dtype): "nanosecond", ] +duration_extract_fields = [ + "total_seconds", + "total_milliseconds", + "total_microseconds", + "total_nanoseconds", + "total_days", + "total_hours", + "total_minutes", +] + @pytest.fixture( ids=datetime_extract_fields, @@ -171,6 +181,38 @@ def test_strftime_duration(engine: pl.GPUEngine, format): assert_ir_translation_raises(q, engine, NotImplementedError) +@pytest.mark.parametrize("field", duration_extract_fields) +@pytest.mark.parametrize( + "dtype", [pl.Duration("ms"), pl.Duration("us"), pl.Duration("ns")] +) +def test_duration_total_component_extract(engine: pl.GPUEngine, field, dtype): + ldf = pl.LazyFrame( + { + "durations": pl.Series( + [ + 0, + 1, + 15, + -1500, + 1000, + 1111, + 1500, + 11111, + -134234534, + 134234534, + # values beyond float64's exact-integer range to guard + # against precision loss in the unit conversion + 5857593848682946, + -5857593848682946, + ], + dtype=dtype, + ), + } + ) + q = ldf.select(getattr(pl.col("durations").dt, field)()) + assert_gpu_result_equal(q, engine=engine) + + @pytest.mark.parametrize( "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")] )