From cbd7857628ee9bcb1dbf25369abad8946c11f935 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Sat, 23 Nov 2024 15:50:00 +0100
Subject: [PATCH] perf: Memoize duplicates in rolling-gb-dyn (#19939)

---
 crates/polars-time/src/windows/group_by.rs | 35 ++++++++++++++++++++--
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/crates/polars-time/src/windows/group_by.rs b/crates/polars-time/src/windows/group_by.rs
index af5c9f8ec492..c5b75bafc371 100644
--- a/crates/polars-time/src/windows/group_by.rs
+++ b/crates/polars-time/src/windows/group_by.rs
@@ -264,15 +264,14 @@ pub(crate) fn group_by_values_iter_lookbehind(
         .iter()
         .enumerate()
         .map(move |(mut i, t)| {
-            i += start_offset;
-
             // Fast path for duplicates.
-            if *t == last && i > start_offset {
+            if *t == last && i > 0 {
                 let len = end - start;
                 let offset = start as IdxSize;
                 return Ok((offset, len as IdxSize));
             }
             last = *t;
+            i += start_offset;
 
             let lower = add(&offset, *t, tz.as_ref())?;
             let upper = *t;
@@ -327,7 +326,17 @@ pub(crate) fn group_by_values_iter_window_behind_t(
 
     let mut start = 0;
     let mut end = start;
+    let mut last = time[0];
+    let mut started = false;
     time.iter().map(move |lower| {
+        // Fast path for duplicates.
+        if *lower == last && started {
+            let len = end - start;
+            let offset = start as IdxSize;
+            return Ok((offset, len as IdxSize));
+        }
+        last = *lower;
+        started = true;
         let lower = add(&offset, *lower, tz.as_ref())?;
         let upper = add(&period, lower, tz.as_ref())?;
 
@@ -377,7 +386,16 @@ pub(crate) fn group_by_values_iter_partial_lookbehind(
 
     let mut start = 0;
     let mut end = start;
+    let mut last = time[0];
     time.iter().enumerate().map(move |(i, lower)| {
+        // Fast path for duplicates.
+        if *lower == last && i > 0 {
+            let len = end - start;
+            let offset = start as IdxSize;
+            return Ok((offset, len as IdxSize));
+        }
+        last = *lower;
+
         let lower = add(&offset, *lower, tz.as_ref())?;
         let upper = add(&period, lower, tz.as_ref())?;
 
@@ -429,7 +447,18 @@ pub(crate) fn group_by_values_iter_lookahead(
     let mut start = start_offset;
     let mut end = start;
 
+    let mut last = time[start_offset];
+    let mut started = false;
     time[start_offset..upper_bound].iter().map(move |lower| {
+        // Fast path for duplicates.
+        if *lower == last && started {
+            let len = end - start;
+            let offset = start as IdxSize;
+            return Ok((offset, len as IdxSize));
+        }
+        started = true;
+        last = *lower;
+
         let lower = add(&offset, *lower, tz.as_ref())?;
         let upper = add(&period, lower, tz.as_ref())?;