diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py b/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py index b6b582738188..e95e16a13b70 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py @@ -133,6 +133,11 @@ 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, + } _valid_ops: ClassVar[set[Name]] = { *_COMPONENT_MAP.keys(), Name.IsLeapYear, @@ -145,6 +150,7 @@ def from_polars(cls, obj: polars._expr_nodes.TemporalFunction) -> Self: Name.TimeStamp, Name.CastTimeUnit, Name.Truncate, + Name.Century, *_TOTAL_COMPONENT_NANOSECONDS.keys(), } @@ -225,6 +231,51 @@ 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.CastTimeUnit: (column,) = columns return Column( diff --git a/python/cudf_polars/tests/expressions/test_datetime_basic.py b/python/cudf_polars/tests/expressions/test_datetime_basic.py index 26d4c01a4766..1eb8c0c7c760 100644 --- a/python/cudf_polars/tests/expressions/test_datetime_basic.py +++ b/python/cudf_polars/tests/expressions/test_datetime_basic.py @@ -213,6 +213,45 @@ def test_duration_total_component_extract(engine: pl.GPUEngine, field, dtype): assert_gpu_result_equal(q, engine=engine) +@pytest.mark.parametrize("method", ["century"]) +@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"]) +@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.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")] )