Support cudf-polars total_xxx datetime extraction methods

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):

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

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(

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")]
 )