add cumcount expr

polars/polars-lazy/src/dsl.rs

@@ -22,7 +22,7 @@ use polars_arrow::array::default_arrays::FromData;
 use polars_core::frame::select::Selection;
 #[cfg(feature = "diff")]
 use polars_core::series::ops::NullBehavior;
-use polars_core::utils::get_supertype;
+use polars_core::utils::{get_supertype, NoNull};
 
 /// A wrapper trait for any closure `Fn(Vec<Series>) -> Result<Series>`
 pub trait SeriesUdf: Send + Sync {
@@ -1757,6 +1757,26 @@ impl Expr {
         };
         self.apply(move |s| s.reshape(&dims), output_type)
     }
+
+    /// Cumulatively count values from 0 to len.
+    pub fn cumcount(self, reverse: bool) -> Self {
+        self.apply(
+            move |s| {
+                if reverse {
+                    let ca: NoNull<UInt32Chunked> = (0u32..s.len() as u32).rev().collect();
+                    let mut ca = ca.into_inner();
+                    ca.rename(s.name());
+                    Ok(ca.into_series())
+                } else {
+                    let ca: NoNull<UInt32Chunked> = (0u32..s.len() as u32).collect();
+                    let mut ca = ca.into_inner();
+                    ca.rename(s.name());
+                    Ok(ca.into_series())
+                }
+            },
+            GetOutput::from_type(DataType::UInt32),
+        )
+    }
 }
 
 /// Create a Column Expression based on a column name.

py-polars/docs/source/reference/expression.rst

@@ -116,6 +116,7 @@ Computations
     Expr.cummin
     Expr.cummax
     Expr.cumprod
+    Expr.cumcount
     Expr.dot
     Expr.mode
     Expr.n_unique

py-polars/polars/internals/expr.py

@@ -614,6 +614,18 @@ def cummax(self, reverse: bool = False) -> "Expr":
         """
         return wrap_expr(self._pyexpr.cummax(reverse))
 
+    def cumcount(self, reverse: bool = False) -> "Expr":
+        """
+        Get an array with the cumulative count computed at every element.
+        Counting from 0 to len
+
+        Parameters
+        ----------
+        reverse
+            Reverse the operation.
+        """
+        return wrap_expr(self._pyexpr.cumcount(reverse))
+
     def floor(self) -> "Expr":
         """
         Floor underlying floating point array to the lowest integers smaller or equal to the float value.

py-polars/src/lazy/dsl.rs

@@ -957,6 +957,10 @@ impl PyExpr {
     pub fn reshape(&self, dims: Vec<i64>) -> Self {
         self.inner.clone().reshape(&dims).into()
     }
+
+    pub fn cumcount(&self, reverse: bool) -> Self {
+        self.inner.clone().cumcount(reverse).into()
+    }
 }
 
 impl From<dsl::Expr> for PyExpr {

py-polars/tests/test_exprs.py

@@ -20,3 +20,14 @@ def test_prefix(fruits_cars: pl.DataFrame) -> None:
     df = fruits_cars
     out = df.select([pl.all().prefix("reverse_")])
     assert out.columns == ["reverse_A", "reverse_fruits", "reverse_B", "reverse_cars"]
+
+
+def test_cumcount() -> None:
+    df = pl.DataFrame([["a"], ["a"], ["a"], ["b"], ["b"], ["a"]], columns=["A"])
+
+    out = df.groupby("A", maintain_order=True).agg(
+        [pl.col("A").cumcount(reverse=False).alias("foo")]
+    )
+
+    assert out["foo"][0].to_list() == [0, 1, 2, 3]
+    assert out["foo"][1].to_list() == [0, 1]