add cumprod

polars/polars-core/src/chunked_array/ops/cum_agg.rs

@@ -2,7 +2,7 @@ use crate::prelude::*;
 use crate::utils::CustomIterTools;
 use itertools::__std_iter::FromIterator;
 use num::Bounded;
-use std::ops::{Add, AddAssign};
+use std::ops::{Add, AddAssign, Mul};
 
 fn det_max<T>(state: &mut T, v: Option<T>) -> Option<Option<T>>
 where
@@ -51,6 +51,23 @@ where
     }
 }
 
+fn det_prod<T>(state: &mut Option<T>, v: Option<T>) -> Option<Option<T>>
+where
+    T: Copy + PartialOrd + Mul<Output = T>,
+{
+    match (*state, v) {
+        (Some(state_inner), Some(v)) => {
+            *state = Some(state_inner * v);
+            Some(*state)
+        }
+        (None, Some(v)) => {
+            *state = Some(v);
+            Some(*state)
+        }
+        (_, None) => Some(None),
+    }
+}
+
 impl<T> ChunkCumAgg<T> for ChunkedArray<T>
 where
     T: PolarsNumericType,
@@ -115,6 +132,26 @@ where
         ca.rename(self.name());
         ca
     }
+
+    fn cumprod(&self, reverse: bool) -> ChunkedArray<T> {
+        let init = None;
+        let mut ca: Self = match reverse {
+            false => self
+                .into_iter()
+                .scan(init, det_prod)
+                .trust_my_length(self.len())
+                .collect_trusted(),
+            true => self
+                .into_iter()
+                .rev()
+                .scan(init, det_prod)
+                .trust_my_length(self.len())
+                .collect_reversed(),
+        };
+
+        ca.rename(self.name());
+        ca
+    }
 }
 
 #[cfg(test)]

polars/polars-core/src/chunked_array/ops/mod.rs

@@ -105,6 +105,10 @@ pub trait ChunkCumAgg<T> {
     fn cumsum(&self, _reverse: bool) -> ChunkedArray<T> {
         panic!("operation cumsum not supported for this dtype")
     }
+    /// Get an array with the cumulative product computed at every element
+    fn cumprod(&self, _reverse: bool) -> ChunkedArray<T> {
+        panic!("operation cumprod not supported for this dtype")
+    }
 }
 
 /// Traverse and collect every nth element

polars/polars-core/src/series/implementations/mod.rs

@@ -127,6 +127,11 @@ macro_rules! impl_dyn_series {
                 self.0.cumsum(reverse).into_series()
             }
 
+            #[cfg(feature = "cum_agg")]
+            fn _cumprod(&self, reverse: bool) -> Series {
+                self.0.cumprod(reverse).into_series()
+            }
+
             #[cfg(feature = "asof_join")]
             fn join_asof(&self, other: &Series) -> Result<Vec<Option<u32>>> {
                 self.0.join_asof(other)

polars/polars-core/src/series/mod.rs

@@ -504,6 +504,22 @@ impl Series {
         }
     }
 
+    /// Get an array with the cumulative product computed at every element
+    #[cfg_attr(docsrs, doc(cfg(feature = "cum_agg")))]
+    pub fn cumprod(&self, _reverse: bool) -> Series {
+        #[cfg(feature = "cum_agg")]
+        {
+            match self.dtype() {
+                DataType::Boolean => self.cast(&DataType::UInt32).unwrap()._cumprod(_reverse),
+                _ => self._cumprod(_reverse),
+            }
+        }
+        #[cfg(not(feature = "cum_agg"))]
+        {
+            panic!("activate 'cum_agg' feature")
+        }
+    }
+
     /// Apply a rolling variance to a Series. See:
     #[cfg_attr(docsrs, doc(cfg(feature = "rolling_window")))]
     pub fn rolling_var(&self, _options: RollingOptions) -> Result<Series> {

polars/polars-core/src/series/series_trait.rs

@@ -118,6 +118,12 @@ pub(crate) mod private {
             panic!("operation cumsum not supported for this dtype")
         }
 
+        /// Get an array with the cumulative sum computed at every element
+        #[cfg(feature = "cum_agg")]
+        fn _cumprod(&self, _reverse: bool) -> Series {
+            panic!("operation cumprod not supported for this dtype")
+        }
+
         #[cfg(feature = "asof_join")]
         fn join_asof(&self, _other: &Series) -> Result<Vec<Option<u32>>> {
             unimplemented!()

polars/polars-lazy/src/dsl.rs

@@ -1004,6 +1004,15 @@ impl Expr {
         )
     }
 
+    /// Get an array with the cumulative product computed at every element
+    #[cfg_attr(docsrs, doc(cfg(feature = "cum_agg")))]
+    pub fn cumprod(self, reverse: bool) -> Self {
+        self.apply(
+            move |s: Series| Ok(s.cumprod(reverse)),
+            GetOutput::same_type(),
+        )
+    }
+
     /// Get an array with the cumulative min computed at every element
     #[cfg_attr(docsrs, doc(cfg(feature = "cum_agg")))]
     pub fn cummin(self, reverse: bool) -> Self {

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

@@ -110,6 +110,7 @@ Computations
     Expr.cumsum
     Expr.cummin
     Expr.cummax
+    Expr.cumprod
     Expr.dot
     Expr.mode
     Expr.n_unique

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

@@ -88,6 +88,7 @@ Computations
     Series.cumsum
     Series.cummin
     Series.cummax
+    Series.cumprod
     Series.arg_sort
     Series.arg_true
     Series.arg_unique

py-polars/polars/eager/series.py

@@ -980,6 +980,17 @@ def cummax(self, reverse: bool = False) -> "Series":
         """
         return wrap_s(self._s.cummax(reverse))
 
+    def cumprod(self, reverse: bool = False) -> "Series":
+        """
+        Get an array with the cumulative product computed at every element.
+
+        Parameters
+        ----------
+        reverse
+            reverse the operation.
+        """
+        return wrap_s(self._s.cumprod(reverse))
+
     def limit(self, num_elements: int = 10) -> "Series":
         """
         Take n elements from this Series.

py-polars/polars/lazy/expr.py

@@ -590,6 +590,17 @@ def cumsum(self, reverse: bool = False) -> "Expr":
         """
         return wrap_expr(self._pyexpr.cumsum(reverse))
 
+    def cumprod(self, reverse: bool = False) -> "Expr":
+        """
+        Get an array with the cumulative product computed at every element.
+
+        Parameters
+        ----------
+        reverse
+            Reverse the operation.
+        """
+        return wrap_expr(self._pyexpr.cumprod(reverse))
+
     def cummin(self, reverse: bool = False) -> "Expr":
         """
         Get an array with the cumulative min computed at every element.

py-polars/src/lazy/dsl.rs

@@ -315,6 +315,10 @@ impl PyExpr {
         self.clone().inner.cummin(reverse).into()
     }
 
+    pub fn cumprod(&self, reverse: bool) -> PyExpr {
+        self.clone().inner.cumprod(reverse).into()
+    }
+
     pub fn str_parse_date(&self, fmt: Option<String>) -> PyExpr {
         let function = move |s: Series| {
             let ca = s.utf8()?;

py-polars/src/series.rs

@@ -334,6 +334,10 @@ impl PySeries {
         self.series.cummin(reverse).into()
     }
 
+    pub fn cumprod(&self, reverse: bool) -> Self {
+        self.series.cumprod(reverse).into()
+    }
+
     pub fn chunk_lengths(&self) -> Vec<usize> {
         self.series.chunk_lengths().collect()
     }