diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py b/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py index c02f40a88dc..9e139aae75f 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/datetime.py @@ -152,6 +152,7 @@ def from_polars(cls, obj: polars._expr_nodes.TemporalFunction) -> Self: Name.CastTimeUnit, Name.Truncate, Name.DaysInMonth, + Name.Quarter, *_CENTURY_MILLENNIUM_DIVISOR.keys(), *_TOTAL_COMPONENT_NANOSECONDS.keys(), } @@ -239,6 +240,12 @@ def do_evaluate( plc.datetime.days_in_month(column.obj, stream=df.stream), dtype=DataType(pl.Int16()), ).astype(self.dtype, stream=df.stream) + elif self.name is TemporalFunction.Name.Quarter: + (column,) = columns + return Column( + plc.datetime.extract_quarter(column.obj, stream=df.stream), + dtype=DataType(pl.Int16()), + ).astype(self.dtype, stream=df.stream) elif self.name in self._CENTURY_MILLENNIUM_DIVISOR: (column,) = columns int32 = plc.DataType(plc.TypeId.INT32) diff --git a/python/cudf_polars/tests/expressions/test_datetime_basic.py b/python/cudf_polars/tests/expressions/test_datetime_basic.py index b34d4693b55..daddc07c070 100644 --- a/python/cudf_polars/tests/expressions/test_datetime_basic.py +++ b/python/cudf_polars/tests/expressions/test_datetime_basic.py @@ -529,6 +529,27 @@ def test_datetime_truncate_unsupported(engine: pl.GPUEngine, every: str): 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_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( "dtype", [pl.Date(), pl.Datetime("ms"), pl.Datetime("us"), pl.Datetime("ns")] )