Support pl.Expr.cum_count in cudf_polars

Author: mroeschke

python/cudf_polars/cudf_polars/dsl/expressions/unary.py

@@ -121,6 +121,7 @@ class UnaryFunction(Expr):
     )
     _supported_cum_aggs = frozenset(
         {
+            "cum_count",
             "cum_min",
             "cum_max",
             "cum_prod",
@@ -552,6 +553,22 @@ def do_evaluate(
             )
         elif self.name in UnaryFunction._supported_cum_aggs:
             column = self.children[0].evaluate(df, context=context)
+            if self.name == "cum_count":
+                # cum_count is the cumulative count of non-null values.
+                counts = plc.unary.cast(
+                    plc.unary.is_valid(column.obj, stream=df.stream),
+                    self.dtype.plc_type,
+                    stream=df.stream,
+                )
+                return Column(
+                    plc.reduce.scan(
+                        counts,
+                        plc.aggregation.sum(),
+                        plc.reduce.ScanType.INCLUSIVE,
+                        stream=df.stream,
+                    ),
+                    dtype=self.dtype,
+                )
             plc_col = column.obj
             col_type = column.dtype.plc_type
             # cum_sum casts

python/cudf_polars/tests/expressions/test_agg.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
 
@@ -34,6 +34,7 @@
         "std",
         "var",
         # scan aggs from UnaryFunction
+        "cum_count",
         "cum_min",
         "cum_max",
         "cum_prod",
@@ -112,6 +113,21 @@ def test_bool_agg(engine: pl.GPUEngine, agg):
     assert_gpu_result_equal(q, engine=engine, check_exact=False)
 
 
+@pytest.mark.parametrize(
+    "data",
+    [
+        [1, None, 3, None, 5],
+        [None, None, None],
+        [1, 2, 3],
+        [],
+    ],
+)
+def test_cum_count(engine: pl.GPUEngine, data):
+    df = pl.LazyFrame({"a": pl.Series(data, dtype=pl.Int64())})
+    q = df.select(pl.col("a").cum_count())
+    assert_gpu_result_equal(q, engine=engine)
+
+
 @pytest.mark.parametrize("cum_agg", sorted(expr.UnaryFunction._supported_cum_aggs))
 def test_cum_agg_reverse_unsupported(engine: pl.GPUEngine, cum_agg):
     df = pl.LazyFrame({"a": [1, 2, 3]})