perf(rust, python): impove performance reducing window functions with…

… numeric output `~-14%` (#5841)
pola-rs · Dec 18, 2022 · eb3c6c1 · eb3c6c1
1 parent 33d2a11
commit eb3c6c1
Show file tree

Hide file tree

Showing 3 changed files with 370 additions and 181 deletions.
diff --git a/polars/polars-core/src/frame/groupby/mod.rs b/polars/polars-core/src/frame/groupby/mod.rs
@@ -318,6 +318,10 @@ impl<'df> GroupBy<'df> {
         self.groups
     }
 
+    pub fn take_groups_mut(&mut self) -> GroupsProxy {
+        std::mem::take(&mut self.groups)
+    }
+
     pub fn keys_sliced(&self, slice: Option<(i64, usize)>) -> Vec<Series> {
         #[allow(unused_assignments)]
         // needed to keep the lifetimes valid for this scope

diff --git a/polars/polars-lazy/src/physical_plan/expressions/take.rs b/polars/polars-lazy/src/physical_plan/expressions/take.rs
@@ -151,7 +151,7 @@ impl PhysicalExpr for TakeExpr {
                                 };
                                 let taken = ac.flat_naive().take(&idx.into_inner())?;
                                 ac.with_series(taken, true);
-                                ac.with_update_groups(UpdateGroups::WithSeriesLen);
+                                ac.with_update_groups(UpdateGroups::WithGroupsLen);
                                 Ok(ac)
                             }
                         }