Skip to content

Commit

Permalink
take duplicates into account
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Sep 26, 2023
1 parent c14932d commit 049d793
Showing 1 changed file with 11 additions and 3 deletions.
14 changes: 11 additions & 3 deletions crates/polars-time/src/windows/group_by.rs
Original file line number Diff line number Diff line change
Expand Up @@ -423,14 +423,22 @@ pub(crate) fn group_by_values_iter_lookahead(
};
let mut start = start_offset;
let mut end = start;
let mut previous = 0i64;

time[start_offset..upper_bound].iter().map(move |lower| {
let lower = add(&offset, *lower, tz.as_ref())?;
let upper = add(&period, lower, tz.as_ref())?;

let b = Bounds::new(lower, upper);

let mut duplicate_count = 0;
for &t in &time[start..] {
if t == previous {
duplicate_count += 1;
} else {
duplicate_count = 0;
}
previous = t;
if b.is_member(t, closed_window) {
break;
}
Expand All @@ -447,9 +455,9 @@ pub(crate) fn group_by_values_iter_lookahead(

let len = end - start;
let offset = start as IdxSize;
// -1 for boundary effects
start = start.saturating_sub(1);
end = end.saturating_sub(1);
// Subtract 1 slot for boundary effects and subtract the duplicates.
start = start.saturating_sub(1 + duplicate_count);
end = end.saturating_sub(1 + duplicate_count);

Ok((offset, len as IdxSize))
})
Expand Down

0 comments on commit 049d793

Please sign in to comment.