Implement more datetime polars.Expressions

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

@@ -133,6 +133,12 @@ def from_polars(cls, obj: polars._expr_nodes.TemporalFunction) -> Self:
         Name.TotalMicroseconds: 1_000,
         Name.TotalNanoseconds: 1,
     }
+    # Divisor used to derive the century/millennium from the calendar year:
+    # ``(year - 1) // divisor + 1`` (floor division, matching polars).
+    _CENTURY_MILLENNIUM_DIVISOR: ClassVar[dict[Name, int]] = {
+        Name.Century: 100,
+        Name.Millennium: 1_000,
+    }
     _valid_ops: ClassVar[set[Name]] = {
         *_COMPONENT_MAP.keys(),
         Name.IsLeapYear,
@@ -145,6 +151,12 @@ def from_polars(cls, obj: polars._expr_nodes.TemporalFunction) -> Self:
         Name.TimeStamp,
         Name.CastTimeUnit,
         Name.Truncate,
+        Name.Round,
+        Name.Century,
+        Name.Millennium,
+        Name.Date,
+        Name.DaysInMonth,
+        Name.Quarter,
         *_TOTAL_COMPONENT_NANOSECONDS.keys(),
     }
 
@@ -166,12 +178,15 @@ def __init__(
             self.children[0].dtype.plc_type
         ):
             raise NotImplementedError("ToString is not supported on duration types")
-        elif self.name is TemporalFunction.Name.Truncate:
+        elif self.name in {
+            TemporalFunction.Name.Truncate,
+            TemporalFunction.Name.Round,
+        }:
             every = cast("Literal", self.children[1]).value
             match = re.fullmatch(r"(\d+)(ns|us|ms|s|m|h|d)", every)
             if match is None or int(match.group(1)) != 1:
                 # https://github.com/rapidsai/cudf/issues/18654 to support non-1 buckets
-                raise NotImplementedError(f"Unsupported truncate bucket: {every!r}")
+                raise NotImplementedError(f"Unsupported bucket: {every!r}")
             self.options = (self._TRUNCATE_FREQ_MAP[match.group(2)],)
 
     def do_evaluate(
@@ -225,6 +240,89 @@ def do_evaluate(
                 ),
                 dtype=self.dtype,
             )
+        elif self.name in self._CENTURY_MILLENNIUM_DIVISOR:
+            (column,) = columns
+            int32 = plc.DataType(plc.TypeId.INT32)
+            # YEAR extraction yields INT16; cast up so the arithmetic (and the
+            # INT32 output polars produces) does not overflow or need promotion.
+            year = plc.unary.cast(
+                plc.datetime.extract_datetime_component(
+                    column.obj,
+                    plc.datetime.DatetimeComponent.YEAR,
+                    stream=df.stream,
+                ),
+                int32,
+                stream=df.stream,
+            )
+            # polars computes ``(year - 1) // divisor + 1`` using floor division;
+            # evaluate the whole arithmetic as a single fused libcudf AST
+            # expression rather than three separate binaryop kernels.
+            one = plc.expressions.Literal(
+                plc.Scalar.from_py(1, int32, stream=df.stream)
+            )
+            predicate = plc.expressions.Operation(
+                plc.expressions.ASTOperator.ADD,
+                plc.expressions.Operation(
+                    plc.expressions.ASTOperator.FLOOR_DIV,
+                    plc.expressions.Operation(
+                        plc.expressions.ASTOperator.SUB,
+                        plc.expressions.ColumnReference(0),
+                        one,
+                    ),
+                    plc.expressions.Literal(
+                        plc.Scalar.from_py(
+                            self._CENTURY_MILLENNIUM_DIVISOR[self.name],
+                            int32,
+                            stream=df.stream,
+                        )
+                    ),
+                ),
+                one,
+            )
+            return Column(
+                plc.transform.compute_column(
+                    plc.Table([year]), predicate, stream=df.stream
+                ),
+                dtype=self.dtype,
+            )
+        elif self.name is TemporalFunction.Name.Round:
+            (column, _) = columns
+            return Column(
+                plc.datetime.round_datetimes(
+                    column.obj,
+                    self.options[0],
+                    stream=df.stream,
+                ),
+                dtype=self.dtype,
+            )
+        elif self.name is TemporalFunction.Name.DaysInMonth:
+            (column,) = columns
+            return Column(
+                plc.unary.cast(
+                    plc.datetime.days_in_month(column.obj, stream=df.stream),
+                    self.dtype.plc_type,
+                    stream=df.stream,
+                ),
+                dtype=self.dtype,
+            )
+        elif self.name is TemporalFunction.Name.Quarter:
+            (column,) = columns
+            return Column(
+                plc.unary.cast(
+                    plc.datetime.extract_quarter(column.obj, stream=df.stream),
+                    self.dtype.plc_type,
+                    stream=df.stream,
+                ),
+                dtype=self.dtype,
+            )
+        elif self.name is TemporalFunction.Name.Date:
+            (column,) = columns
+            # Casting the timestamp to TIMESTAMP_DAYS (the storage of ``pl.Date``)
+            # drops the sub-day component.
+            return Column(
+                plc.unary.cast(column.obj, self.dtype.plc_type, stream=df.stream),
+                dtype=self.dtype,
+            )
         elif self.name is TemporalFunction.Name.CastTimeUnit:
             (column,) = columns
             return Column(

python/cudf_polars/tests/expressions/test_datetime_basic.py

@@ -213,6 +213,96 @@ def test_duration_total_component_extract(engine: pl.GPUEngine, field, dtype):
     assert_gpu_result_equal(q, engine=engine)
 
 
+@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_fractional(
+    engine: pl.GPUEngine, field, dtype
+):
+    ldf = pl.LazyFrame(
+        {
+            "durations": pl.Series(
+                [
+                    0,
+                    1,
+                    15,
+                    -1500,
+                    1000,
+                    1111,
+                    1500,
+                    11111,
+                    -134234534,
+                    134234534,
+                    5857593848682946,
+                    -5857593848682946,
+                ],
+                dtype=dtype,
+            ),
+        }
+    )
+    q = ldf.select(getattr(pl.col("durations").dt, field)(fractional=True))
+    assert_gpu_result_equal(q, engine=engine, check_exact=False)
+
+
+@pytest.mark.parametrize("method", ["century", "millennium"])
+@pytest.mark.parametrize(
+    "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")]
+)
+def test_century_millennium(engine: pl.GPUEngine, method, dtype):
+    data = pl.Series(
+        [
+            datetime.date(1897, 5, 7),
+            datetime.date(1900, 12, 31),
+            datetime.date(1901, 1, 1),
+            datetime.date(2000, 1, 1),
+            datetime.date(2001, 7, 5),
+            None,
+        ],
+        dtype=pl.Date(),
+    ).cast(dtype)
+    ldf = pl.LazyFrame({"dates": data})
+    q = ldf.select(getattr(pl.col("dates").dt, method)())
+    assert_gpu_result_equal(q, engine=engine)
+
+
+@pytest.mark.parametrize("method", ["century", "millennium"])
+@pytest.mark.parametrize(
+    "days",
+    [
+        # ``Date`` supports a much wider year range than ``Datetime``; include
+        # pre-year-1 offsets to exercise the floor-division branch of the
+        # century/millennium formula.
+        [-1_000_000, -800_000, -365, 0],
+        [364_000, 376_000],  # years ~2966 and ~3000
+    ],
+)
+def test_century_millennium_date_extreme_years(engine: pl.GPUEngine, method, days):
+    dates = pl.Series(days, dtype=pl.Int32).cast(pl.Date())
+    ldf = pl.LazyFrame({"dates": dates})
+    q = ldf.select(getattr(pl.col("dates").dt, method)())
+    assert_gpu_result_equal(q, engine=engine)
+
+
+@pytest.mark.parametrize(
+    "dtype", [pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")]
+)
+def test_datetime_date(engine: pl.GPUEngine, dtype):
+    data = pl.Series(
+        [
+            datetime.datetime(1978, 1, 1, 1, 1, 1),
+            datetime.datetime(1969, 12, 31, 23, 59, 59),  # pre-epoch (floors down)
+            datetime.datetime(2024, 10, 13, 5, 30, 14, 500_000),
+            datetime.datetime(2065, 1, 1, 10, 20, 30, 60_000),
+            None,
+        ],
+        dtype=dtype,
+    )
+    ldf = pl.LazyFrame({"datetimes": data})
+    q = ldf.select(pl.col("datetimes").dt.date())
+    assert_gpu_result_equal(q, engine=engine)
+
+
 @pytest.mark.parametrize(
     "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")]
 )
@@ -490,6 +580,89 @@ def test_datetime_truncate_unsupported(engine: pl.GPUEngine, every: str):
     assert_ir_translation_raises(q, engine, NotImplementedError)
 
 
+@pytest.mark.parametrize(
+    "dtype", [pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")]
+)
+# ``every`` finer than a column's storage unit panics in CPU polars, so only
+# use frequencies that are coarser-or-equal to every tested resolution.
+@pytest.mark.parametrize("every", ["1ms", "1s", "1m", "1h", "1d"])
+def test_datetime_round(engine: pl.GPUEngine, dtype, every):
+    # Use an irregular step so no timestamp lands exactly on a half-way point:
+    # libcudf rounds half-to-even while polars rounds half-away, so only the
+    # exact-tie behaviour differs.
+    ldf = pl.LazyFrame(
+        {
+            "datetimes": pl.datetime_range(
+                datetime.datetime(2020, 1, 1),
+                datetime.datetime(2020, 1, 2),
+                "3h14m15s11ms33us999ns",
+                eager=True,
+            ).cast(dtype)
+        }
+    )
+
+    q = ldf.select(pl.col("datetimes").dt.round(every))
+    assert_gpu_result_equal(q, engine=engine)
+
+
+@pytest.mark.parametrize("every", ["30m", "1mo"])
+def test_datetime_round_unsupported(engine: pl.GPUEngine, every: str):
+    ldf = pl.LazyFrame(
+        {
+            "datetimes": pl.datetime_range(
+                datetime.datetime(2020, 1, 1),
+                datetime.datetime(2020, 1, 2),
+                "30m",
+                eager=True,
+            )
+        }
+    )
+
+    q = ldf.select(pl.col("datetimes").dt.round(every))
+    assert_ir_translation_raises(q, engine, NotImplementedError)
+
+
+@pytest.mark.parametrize(
+    "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")]
+)
+def test_datetime_days_in_month(engine: pl.GPUEngine, dtype):
+    data = pl.Series(
+        [
+            datetime.date(2001, 1, 15),
+            datetime.date(2001, 2, 15),  # non-leap February
+            datetime.date(2000, 2, 15),  # leap February
+            datetime.date(2001, 4, 15),
+            datetime.date(2001, 12, 31),
+            None,
+        ],
+        dtype=pl.Date(),
+    ).cast(dtype)
+    ldf = pl.LazyFrame({"dates": data})
+    q = ldf.select(pl.col("dates").dt.days_in_month())
+    assert_gpu_result_equal(q, engine=engine)
+
+
+@pytest.mark.parametrize(
+    "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")]
+)
+def test_datetime_quarter(engine: pl.GPUEngine, dtype):
+    data = pl.Series(
+        [
+            datetime.date(2001, 1, 1),
+            datetime.date(2001, 3, 31),
+            datetime.date(2001, 4, 1),
+            datetime.date(2001, 6, 30),
+            datetime.date(2001, 9, 15),
+            datetime.date(2001, 12, 27),
+            None,
+        ],
+        dtype=pl.Date(),
+    ).cast(dtype)
+    ldf = pl.LazyFrame({"dates": data})
+    q = ldf.select(pl.col("dates").dt.quarter())
+    assert_gpu_result_equal(q, engine=engine)
+
+
 @pytest.mark.parametrize(
     "datetime_dtype",
     [