fix: Optimize memory on rolling groups in ApplyExpr

crates/polars-expr/src/expressions/apply.rs

@@ -108,18 +108,51 @@ impl ApplyExpr {
     fn eval_and_flatten(&self, inputs: &mut [Column]) -> PolarsResult<Column> {
         self.function.call_udf(inputs)
     }
+
     fn apply_single_group_aware<'a>(
         &self,
         mut ac: AggregationContext<'a>,
     ) -> PolarsResult<AggregationContext<'a>> {
         // Fix up groups for AggregatedScalar, so that we can pretend they are just normal groups.
         ac.set_groups_for_undefined_agg_states();
 
+        let name = ac.get_values().name().clone();
+        let f = |opt_s: Option<Series>| match opt_s {
+            None => Ok(None),
+            Some(mut s) => {
+                if self.flags.contains(FunctionFlags::PASS_NAME_TO_APPLY) {
+                    s.rename(name.clone());
+                }
+                Ok(Some(
+                    self.function
+                        .call_udf(&mut [Column::from(s)])?
+                        .take_materialized_series(),
+                ))
+            },
+        };
+
+        // In case of overlapping (rolling) groups, we build groups in a lazy manner to avoid
+        // memory explosion.
+        // TODO: Add parallel iterator path; support Idx GroupsType.
+        if matches!(ac.agg_state(), AggState::NotAggregated(_))
+            && let GroupsType::Slice {
+                overlapping: true, ..
+            } = ac.groups.as_ref().as_ref()
+        {
+            let ca: ChunkedArray<_> = ac
+                .iter_groups_lazy(false)
+                .map(|opt| opt.map(|s| s.as_ref().clone()))
+                .map(f)
+                .collect::<PolarsResult<_>>()?;
+
+            return self.finish_apply_groups(ac, ca.with_name(name));
+        }
+
+        // At this point, calling aggregated() will not lead to memory explosion.
         let agg = match ac.agg_state() {
             AggState::AggregatedScalar(s) => s.as_list().into_column(),
             _ => ac.aggregated(),
         };
-        let name = agg.name().clone();
 
         // Collection of empty list leads to a null dtype. See: #3687.
         if agg.is_empty() {
@@ -133,20 +166,6 @@ impl ApplyExpr {
             return self.finish_apply_groups(ac, ca);
         }
 
-        let f = |opt_s: Option<Series>| match opt_s {
-            None => Ok(None),
-            Some(mut s) => {
-                if self.flags.contains(FunctionFlags::PASS_NAME_TO_APPLY) {
-                    s.rename(name.clone());
-                }
-                Ok(Some(
-                    self.function
-                        .call_udf(&mut [Column::from(s)])?
-                        .take_materialized_series(),
-                ))
-            },
-        };
-
         let ca: ListChunked = if self.allow_threading {
             let lst = agg.list().unwrap();
             let iter = lst.par_iter().map(f);

crates/polars-expr/src/expressions/group_iter.rs

@@ -73,6 +73,40 @@ impl AggregationContext<'_> {
     }
 }
 
+impl AggregationContext<'_> {
+    /// Iterate over groups lazily, i.e., without greedy aggregation into an AggList.
+    pub(super) fn iter_groups_lazy(
+        &mut self,
+        keep_names: bool,
+    ) -> Box<dyn Iterator<Item = Option<AmortSeries>> + '_> {
+        match self.agg_state() {
+            AggState::NotAggregated(_) => {
+                let groups = self.groups();
+                let len = groups.len();
+                let c = self.get_values().rechunk();
+                let name = if keep_names {
+                    c.name().clone()
+                } else {
+                    PlSmallStr::EMPTY
+                };
+                let iter = self.groups().iter();
+
+                // SAFETY: dtype is correct
+                unsafe {
+                    Box::new(NotAggLazyIter::new(
+                        c.as_materialized_series().array_ref(0).clone(),
+                        iter,
+                        len,
+                        c.dtype(),
+                        name,
+                    ))
+                }
+            },
+            _ => self.iter_groups(keep_names),
+        }
+    }
+}
+
 struct LitIter {
     len: usize,
     offset: usize,
@@ -186,3 +220,65 @@ impl Iterator for FlatIter {
         (self.len - self.offset, Some(self.len - self.offset))
     }
 }
+
+struct NotAggLazyIter<'a, I: Iterator<Item = GroupsIndicator<'a>>> {
+    array: ArrayRef,
+    iter: I,
+    groups_idx: usize,
+    len: usize,
+    // AmortSeries referenced that series
+    #[allow(dead_code)]
+    series_container: Rc<Series>,
+    item: AmortSeries,
+}
+
+impl<'a, I: Iterator<Item = GroupsIndicator<'a>>> NotAggLazyIter<'a, I> {
+    /// # Safety
+    /// Caller must ensure the given `logical` dtype belongs to `array`.
+    unsafe fn new(
+        array: ArrayRef,
+        iter: I,
+        len: usize,
+        logical: &DataType,
+        name: PlSmallStr,
+    ) -> Self {
+        let series_container = Rc::new(Series::from_chunks_and_dtype_unchecked(
+            name,
+            vec![array.clone()],
+            logical,
+        ));
+        Self {
+            array,
+            iter,
+            groups_idx: 0,
+            len,
+            series_container: series_container.clone(),
+            item: AmortSeries::new(series_container),
+        }
+    }
+}
+
+impl<'a, I: Iterator<Item = GroupsIndicator<'a>>> Iterator for NotAggLazyIter<'a, I> {
+    type Item = Option<AmortSeries>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(g) = self.iter.next() {
+            self.groups_idx += 1;
+            match g {
+                // TODO: Implement for Idx GroupsType
+                GroupsIndicator::Idx(_) => todo!(),
+                GroupsIndicator::Slice(s) => {
+                    let mut arr =
+                        unsafe { self.array.sliced_unchecked(s[0] as usize, s[1] as usize) };
+                    unsafe { self.item.swap(&mut arr) };
+                    Some(Some(self.item.clone()))
+                },
+            }
+        } else {
+            None
+        }
+    }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.len - self.groups_idx, Some(self.len - self.groups_idx))
+    }
+}