From df22b1b75c9580d30093503818c3dba41860c32a Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Thu, 26 May 2022 16:17:15 +0200 Subject: [PATCH] fix accidental quadratic behavior in rolling_groupby (#3510) --- polars/polars-time/src/windows/groupby.rs | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/polars/polars-time/src/windows/groupby.rs b/polars/polars-time/src/windows/groupby.rs index ece3631f6a45..1a608f819388 100644 --- a/polars/polars-time/src/windows/groupby.rs +++ b/polars/polars-time/src/windows/groupby.rs @@ -127,13 +127,8 @@ pub fn groupby_windows( (groups, lower_bound, upper_bound) } -fn find_offset(time: &[i64], b: Bounds, closed: ClosedWindow) -> Option { - time.iter() - .enumerate() - .find_map(|(i, t)| match b.is_member(*t, closed) { - true => None, - false => Some(i), - }) +fn find_offset(time: &[i64], b: Bounds, closed: ClosedWindow) -> usize { + time.partition_point(|v| b.is_member(*v, closed)) } /// Different from `groupby_windows`, where define window buckets and search which values fit that @@ -172,8 +167,10 @@ pub fn groupby_values( lagging_offset += 1; } - let slice = &time[lagging_offset..]; - let len = find_offset(slice, b, closed_window).unwrap_or(slice.len()); + // Safety + // we just iterated over value i. + let slice = unsafe { time.get_unchecked(lagging_offset..) }; + let len = find_offset(slice, b, closed_window); [lagging_offset as IdxSize, len as IdxSize] })