fix: Optimize memory for groups iter in NotAggregated state

Author: kdn36

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

@@ -105,7 +105,12 @@ impl ApplyExpr {
         &self,
         mut ac: AggregationContext<'a>,
     ) -> PolarsResult<AggregationContext<'a>> {
+        dbg!("start apply_single_group_aware"); //kdn
+        dbg!(&self.expr);
+        // dbg!(&ac);
+
         let s = ac.get_values();
+        let name = s.name().clone();
 
         #[allow(clippy::nonminimal_bool)]
         {
@@ -116,20 +121,6 @@ impl ApplyExpr {
             );
         }
 
-        let name = s.name().clone();
-        let agg = ac.aggregated();
-        // Collection of empty list leads to a null dtype. See: #3687.
-        if agg.is_empty() {
-            // Create input for the function to determine the output dtype, see #3946.
-            let agg = agg.list().unwrap();
-            let input_dtype = agg.inner_dtype();
-            let input = Column::full_null(name.clone(), 0, input_dtype);
-
-            let output = self.eval_and_flatten(&mut [input])?;
-            let ca = ListChunked::full(name, output.as_materialized_series(), 0);
-            return self.finish_apply_groups(ac, ca);
-        }
-
         let f = |opt_s: Option<Series>| match opt_s {
             None => Ok(None),
             Some(mut s) => {
@@ -144,6 +135,36 @@ impl ApplyExpr {
             },
         };
 
+        // In case of overlapping (rolling) groups, we build groups in a lazy manner to avoid
+        // memory explosion.
+        // kdn TODO: TBD - do we want to follow this path for *all* Slice
+        if matches!(ac.agg_state(), AggState::NotAggregated(_))
+            && let GroupsType::Slice { rolling: 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 = ac.aggregated();
+
+        // Collection of empty list leads to a null dtype. See: #3687.
+        if agg.is_empty() {
+            // Create input for the function to determine the output dtype, see #3946.
+            let agg = agg.list().unwrap();
+            let input_dtype = agg.inner_dtype();
+            let input = Column::full_null(name.clone(), 0, input_dtype);
+
+            let output = self.eval_and_flatten(&mut [input])?;
+            let ca = ListChunked::full(name, output.as_materialized_series(), 0);
+            return self.finish_apply_groups(ac, ca);
+        }
+
         let ca: ListChunked = if self.allow_threading {
             let dtype = if self.output_field.dtype.is_known() && !self.output_field.dtype.is_null()
             {

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

@@ -73,6 +73,41 @@ impl AggregationContext<'_> {
     }
 }
 
+impl AggregationContext<'_> {
+    /// Iterate over groups 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(); //TODO - do we require rechunk?
+                let name = if keep_names {
+                    c.name().clone()
+                } else {
+                    PlSmallStr::EMPTY
+                };
+                let iter = self.groups().iter();
+
+                // Safety:
+                // kdn TODO
+                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 +221,63 @@ 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
+    /// kdn TODO
+    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() {
+            match g {
+                GroupsIndicator::Idx(_) => todo!(), //kdn 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))
+    }
+}