1use std::collections::VecDeque;
2
3use polars_arrow::legacy::time_zone::Tz;
4use polars_arrow::trusted_len::TrustedLen;
5use polars_core::prelude::*;
6use polars_core::runtime::RAYON;
7use polars_core::utils::_split_offsets;
8use polars_core::utils::flatten::flatten_par;
9use polars_defs::time::duration::Duration;
10use polars_defs::time::group_by::{ClosedWindow, DynamicWindowPlacement, StartBy};
11use rayon::prelude::*;
12
13use crate::prelude::*;
14
15#[allow(clippy::too_many_arguments)]
16fn update_groups_and_bounds(
17 bounds_iter: BoundsIter<'_>,
18 mut start: usize,
19 time: &[i64],
20 closed_window: ClosedWindow,
21 include_lower_bound: bool,
22 include_upper_bound: bool,
23 placement: Option<&DynamicWindowPlacement>,
24 lower_bound: &mut Vec<i64>,
25 upper_bound: &mut Vec<i64>,
26 groups: &mut Vec<[IdxSize; 2]>,
27) {
28 let mut iter = bounds_iter.into_iter();
29 let mut stride = 0;
30
31 'bounds: while let Some(bi) = iter.nth(stride) {
32 if let Some(placement) = placement {
33 if placement.start_range.is_past(bi.start) {
34 break;
35 }
36 if placement.start_range.is_before(bi.start) {
37 stride = iter.get_stride(placement.start_range.start());
40 continue 'bounds;
41 }
42 }
43
44 let mut has_member = false;
45 for &t in &time[start..time.len().saturating_sub(1)] {
47 if bi.is_future(t, closed_window) {
49 stride = iter.get_stride(t);
50 continue 'bounds;
51 }
52 if bi.is_member_entry(t, closed_window) {
53 has_member = true;
54 break;
55 }
56 start += 1;
58 }
59
60 stride = if has_member {
62 0
63 } else {
64 debug_assert!(start < time.len());
65 iter.get_stride(time[start])
66 };
67
68 let mut end = start;
70
71 if end == time.len() - 1 {
73 let t = time[end];
74 if bi.is_member(t, closed_window) {
75 if include_lower_bound {
76 lower_bound.push(bi.start);
77 }
78 if include_upper_bound {
79 upper_bound.push(bi.stop);
80 }
81 groups.push([end as IdxSize, 1])
82 }
83 continue;
84 }
85 for &t in &time[end..] {
86 if !bi.is_member_exit(t, closed_window) {
87 break;
88 }
89 end += 1;
90 }
91 let len = end - start;
92
93 if include_lower_bound {
94 lower_bound.push(bi.start);
95 }
96 if include_upper_bound {
97 upper_bound.push(bi.stop);
98 }
99 groups.push([start as IdxSize, len as IdxSize])
100 }
101}
102
103#[allow(clippy::too_many_arguments)]
115pub fn group_by_windows(
116 window: Window,
117 time: &[i64],
118 closed_window: ClosedWindow,
119 tu: TimeUnit,
120 tz: &Option<TimeZone>,
121 include_lower_bound: bool,
122 include_upper_bound: bool,
123 start_by: StartBy,
124 placement: Option<DynamicWindowPlacement>,
125) -> PolarsResult<(GroupsSlice, Vec<i64>, Vec<i64>)> {
126 let start = time[0];
127 let boundary = if time.len() > 1 {
131 let stop = time[time.len() - 1] + 1;
133 Bounds::new_checked(start, stop)
134 } else {
135 let stop = start + 1;
136 Bounds::new_checked(start, stop)
137 };
138
139 let size = window.estimate_overlapping_bounds(tu, boundary);
140 let size_lower = if include_lower_bound { size } else { 0 };
141 let size_upper = if include_upper_bound { size } else { 0 };
142 let mut lower_bound = Vec::with_capacity(size_lower);
143 let mut upper_bound = Vec::with_capacity(size_upper);
144
145 let mut groups = Vec::with_capacity(size);
146 let start_offset = 0;
147
148 match tz {
149 #[cfg(feature = "timezones")]
150 Some(tz) => {
151 update_groups_and_bounds(
152 window.get_overlapping_bounds_iter(
153 boundary,
154 closed_window,
155 tu,
156 tz.parse::<Tz>().ok().as_ref(),
157 start_by,
158 placement.map(|p| p.origin),
159 )?,
160 start_offset,
161 time,
162 closed_window,
163 include_lower_bound,
164 include_upper_bound,
165 placement.as_ref(),
166 &mut lower_bound,
167 &mut upper_bound,
168 &mut groups,
169 );
170 },
171 _ => {
172 update_groups_and_bounds(
173 window.get_overlapping_bounds_iter(
174 boundary,
175 closed_window,
176 tu,
177 None,
178 start_by,
179 placement.map(|p| p.origin),
180 )?,
181 start_offset,
182 time,
183 closed_window,
184 include_lower_bound,
185 include_upper_bound,
186 placement.as_ref(),
187 &mut lower_bound,
188 &mut upper_bound,
189 &mut groups,
190 );
191 },
192 };
193
194 Ok((groups, lower_bound, upper_bound))
195}
196
197#[inline]
201#[allow(clippy::too_many_arguments)]
202pub(crate) fn group_by_values_iter_lookbehind(
203 period: Duration,
204 offset: Duration,
205 time: &[i64],
206 closed_window: ClosedWindow,
207 tu: TimeUnit,
208 tz: Option<Tz>,
209 start_offset: usize,
210 upper_bound: Option<usize>,
211) -> PolarsResult<impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_> {
212 debug_assert!(offset.duration_ns() == period.duration_ns());
213 debug_assert!(offset.negative);
214
215 let upper_bound = upper_bound.unwrap_or(time.len());
216 let mut start = if let Some(&t) = time.get(start_offset) {
218 let lower = offset.add(tu, t, tz.as_ref())?;
219 let upper = t;
224 let b = Bounds::new(lower, upper);
225 let slice = &time[..start_offset];
226 slice.partition_point(|v| !b.is_member(*v, closed_window))
227 } else {
228 0
229 };
230 let mut end = start;
231 let mut last = time[start_offset];
232 Ok(time[start_offset..upper_bound]
233 .iter()
234 .enumerate()
235 .map(move |(mut i, t)| {
236 if *t == last && i > 0 {
238 let len = end - start;
239 let offset = start as IdxSize;
240 return Ok((offset, len as IdxSize));
241 }
242 last = *t;
243 i += start_offset;
244
245 let lower = offset.add(tu, *t, tz.as_ref())?;
246 let upper = *t;
247
248 let b = Bounds::new(lower, upper);
249
250 for &t in unsafe { time.get_unchecked(start..i) } {
251 if b.is_member_entry(t, closed_window) {
252 break;
253 }
254 start += 1;
255 }
256
257 if b.is_member_exit(*t, closed_window) {
259 end = i;
260 } else {
261 end = std::cmp::max(end, start);
262 }
263 for &t in unsafe { time.get_unchecked(end..) } {
265 if !b.is_member_exit(t, closed_window) {
266 break;
267 }
268 end += 1;
269 }
270
271 let len = end - start;
272 let offset = start as IdxSize;
273
274 Ok((offset, len as IdxSize))
275 }))
276}
277
278#[allow(clippy::too_many_arguments)]
283pub(crate) fn group_by_values_iter_window_behind_t(
284 period: Duration,
285 offset: Duration,
286 time: &[i64],
287 closed_window: ClosedWindow,
288 tu: TimeUnit,
289 tz: Option<Tz>,
290 start_offset: usize,
291 upper_bound: Option<usize>,
292) -> PolarsResult<impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_> {
293 let upper_bound = upper_bound.unwrap_or(time.len());
294 let mut start = first_member_row(period, offset, time, closed_window, tu, tz, start_offset)?;
295 let mut end = start;
296 let mut last = time[start_offset];
297 let mut started = false;
298 Ok(time[start_offset..upper_bound].iter().map(move |lower| {
299 if *lower == last && started {
301 let len = end - start;
302 let offset = start as IdxSize;
303 return Ok((offset, len as IdxSize));
304 }
305 last = *lower;
306 started = true;
307 let lower = offset.add(tu, *lower, tz.as_ref())?;
308 let upper = period.add(tu, lower, tz.as_ref())?;
309
310 let b = Bounds::new(lower, upper);
311 if b.is_future(time[0], closed_window) {
312 Ok((0, 0))
313 } else {
314 for &t in &time[start..] {
315 if b.is_member_entry(t, closed_window) {
316 break;
317 }
318 start += 1;
319 }
320
321 end = std::cmp::max(start, end);
322 for &t in &time[end..] {
323 if !b.is_member_exit(t, closed_window) {
324 break;
325 }
326 end += 1;
327 }
328
329 let len = end - start;
330 let offset = start as IdxSize;
331
332 Ok((offset, len as IdxSize))
333 }
334 }))
335}
336
337#[allow(clippy::too_many_arguments)]
341pub(crate) fn group_by_values_iter_partial_lookbehind(
342 period: Duration,
343 offset: Duration,
344 time: &[i64],
345 closed_window: ClosedWindow,
346 tu: TimeUnit,
347 tz: Option<Tz>,
348 start_offset: usize,
349 upper_bound: Option<usize>,
350) -> PolarsResult<impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_> {
351 let upper_bound = upper_bound.unwrap_or(time.len());
352 let mut start = first_member_row(period, offset, time, closed_window, tu, tz, start_offset)?;
353 let mut end = start;
354 let mut last = time[start_offset];
355 let rows = &time[start_offset..upper_bound];
356 Ok(rows.iter().enumerate().map(move |(i, lower)| {
357 if *lower == last && i > 0 {
359 let len = end - start;
360 let offset = start as IdxSize;
361 return Ok((offset, len as IdxSize));
362 }
363 last = *lower;
364 let i = i + start_offset;
365
366 let lower = offset.add(tu, *lower, tz.as_ref())?;
367 let upper = period.add(tu, lower, tz.as_ref())?;
368
369 let b = Bounds::new(lower, upper);
370
371 for &t in &time[start..] {
372 if b.is_member_entry(t, closed_window) || start == i {
373 break;
374 }
375 start += 1;
376 }
377
378 end = std::cmp::max(start, end);
379 for &t in &time[end..] {
380 if !b.is_member_exit(t, closed_window) {
381 break;
382 }
383 end += 1;
384 }
385
386 let len = end - start;
387 let offset = start as IdxSize;
388
389 Ok((offset, len as IdxSize))
390 }))
391}
392
393fn first_member_row(
396 period: Duration,
397 offset: Duration,
398 time: &[i64],
399 closed_window: ClosedWindow,
400 tu: TimeUnit,
401 tz: Option<Tz>,
402 start_offset: usize,
403) -> PolarsResult<usize> {
404 let lower = offset.add(tu, time[start_offset], tz.as_ref())?;
405 let upper = period.add(tu, lower, tz.as_ref())?;
406 let b = Bounds::new(lower, upper);
407 Ok(time[..start_offset].partition_point(|&t| !b.is_member_entry(t, closed_window)))
408}
409
410#[allow(clippy::too_many_arguments)]
411pub(crate) fn group_by_values_iter_lookahead(
415 period: Duration,
416 offset: Duration,
417 time: &[i64],
418 closed_window: ClosedWindow,
419 tu: TimeUnit,
420 tz: Option<Tz>,
421 start_offset: usize,
422 upper_bound: Option<usize>,
423) -> impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_ {
424 let upper_bound = upper_bound.unwrap_or(time.len());
425
426 let mut start = start_offset;
427 let mut end = start;
428
429 let mut last = time[start_offset];
430 let mut started = false;
431 time[start_offset..upper_bound].iter().map(move |lower| {
432 if *lower == last && started {
434 let len = end - start;
435 let offset = start as IdxSize;
436 return Ok((offset, len as IdxSize));
437 }
438 started = true;
439 last = *lower;
440
441 let lower = offset.add(tu, *lower, tz.as_ref())?;
442 let upper = period.add(tu, lower, tz.as_ref())?;
443
444 let b = Bounds::new(lower, upper);
445
446 for &t in &time[start..] {
447 if b.is_member_entry(t, closed_window) {
448 break;
449 }
450 start += 1;
451 }
452
453 end = std::cmp::max(start, end);
454 for &t in &time[end..] {
455 if !b.is_member_exit(t, closed_window) {
456 break;
457 }
458 end += 1;
459 }
460
461 let len = end - start;
462 let offset = start as IdxSize;
463
464 Ok((offset, len as IdxSize))
465 })
466}
467
468#[cfg(feature = "rolling_window_by")]
469#[inline]
470pub(crate) fn group_by_values_iter(
471 period: Duration,
472 time: &[i64],
473 closed_window: ClosedWindow,
474 tu: TimeUnit,
475 tz: Option<Tz>,
476) -> PolarsResult<impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_> {
477 let mut offset = period;
478 offset.negative = true;
479 group_by_values_iter_lookbehind(period, offset, time, closed_window, tu, tz, 0, None)
481}
482
483fn prune_splits_on_duplicates(time: &[i64], thread_offsets: &mut Vec<(usize, usize)>) {
486 let is_valid = |window: &[(usize, usize)]| -> bool {
487 debug_assert_eq!(window.len(), 2);
488 let left_block_end = window[0].0 + window[0].1.saturating_sub(1);
489 let right_block_start = window[1].0;
490 time[left_block_end] != time[right_block_start]
491 };
492
493 if time.is_empty() || thread_offsets.len() <= 1 || thread_offsets.windows(2).all(is_valid) {
494 return;
495 }
496
497 let mut new = vec![];
498 for window in thread_offsets.windows(2) {
499 let this_block_is_valid = is_valid(window);
500 if this_block_is_valid {
501 new.push(window[0])
503 }
504 }
505 if thread_offsets.len().is_multiple_of(2) {
507 let window = &thread_offsets[thread_offsets.len() - 2..];
508 if is_valid(window) {
509 new.push(thread_offsets[thread_offsets.len() - 1])
510 }
511 }
512 if new.len() <= 1 {
514 new = vec![(0, time.len())];
515 } else {
516 let mut previous_start = time.len();
517 for window in new.iter_mut().rev() {
518 window.1 = previous_start - window.0;
519 previous_start = window.0;
520 }
521 new[0].0 = 0;
522 new[0].1 = new[1].0;
523 debug_assert_eq!(new.iter().map(|w| w.1).sum::<usize>(), time.len());
524 prune_splits_on_duplicates(time, &mut new)
526 }
527 std::mem::swap(thread_offsets, &mut new);
528}
529
530#[allow(clippy::too_many_arguments)]
531fn group_by_values_iter_lookbehind_collected(
532 period: Duration,
533 offset: Duration,
534 time: &[i64],
535 closed_window: ClosedWindow,
536 tu: TimeUnit,
537 tz: Option<Tz>,
538 start_offset: usize,
539 upper_bound: Option<usize>,
540) -> PolarsResult<Vec<[IdxSize; 2]>> {
541 let iter = group_by_values_iter_lookbehind(
542 period,
543 offset,
544 time,
545 closed_window,
546 tu,
547 tz,
548 start_offset,
549 upper_bound,
550 )?;
551 iter.map(|result| result.map(|(offset, len)| [offset, len]))
552 .collect::<PolarsResult<Vec<_>>>()
553}
554
555#[allow(clippy::too_many_arguments)]
556pub(crate) fn group_by_values_iter_lookahead_collected(
557 period: Duration,
558 offset: Duration,
559 time: &[i64],
560 closed_window: ClosedWindow,
561 tu: TimeUnit,
562 tz: Option<Tz>,
563 start_offset: usize,
564 upper_bound: Option<usize>,
565) -> PolarsResult<Vec<[IdxSize; 2]>> {
566 let iter = group_by_values_iter_lookahead(
567 period,
568 offset,
569 time,
570 closed_window,
571 tu,
572 tz,
573 start_offset,
574 upper_bound,
575 );
576 iter.map(|result| result.map(|(offset, len)| [offset as IdxSize, len]))
577 .collect::<PolarsResult<Vec<_>>>()
578}
579
580pub fn group_by_values(
591 period: Duration,
592 offset: Duration,
593 time: &[i64],
594 closed_window: ClosedWindow,
595 tu: TimeUnit,
596 tz: Option<Tz>,
597 rows: std::ops::Range<usize>,
598) -> PolarsResult<GroupsSlice> {
599 if rows.is_empty() {
600 return Ok(GroupsSlice::from(vec![]));
601 }
602 assert!(rows.start == 0 || time[rows.start - 1] < time[rows.start]);
603 assert!(rows.end == time.len() || time[rows.end - 1] < time[rows.end]);
604 let (start_offset, upper_bound) = (rows.start, rows.end);
605
606 let mut thread_offsets = _split_offsets(time.len(), RAYON.current_num_threads());
609 prune_splits_on_duplicates(time, &mut thread_offsets);
611 thread_offsets.retain_mut(|(offset, len)| {
612 let start = (*offset).max(start_offset);
613 let end = (*offset + *len).min(upper_bound);
614 (*offset, *len) = (start, end.saturating_sub(start));
615 *len > 0
616 });
617
618 let run_parallel = !RAYON.current_thread_has_pending_tasks().unwrap_or(false);
620
621 if offset.negative && !offset.is_zero() {
623 if offset.duration_ns() == period.duration_ns() {
625 if !run_parallel {
629 let vecs = group_by_values_iter_lookbehind_collected(
630 period,
631 offset,
632 time,
633 closed_window,
634 tu,
635 tz,
636 start_offset,
637 Some(upper_bound),
638 )?;
639 return Ok(GroupsSlice::from(vecs));
640 }
641
642 RAYON.install(|| {
643 let vals = thread_offsets
644 .par_iter()
645 .copied()
646 .map(|(base_offset, len)| {
647 let upper_bound = base_offset + len;
648 group_by_values_iter_lookbehind_collected(
649 period,
650 offset,
651 time,
652 closed_window,
653 tu,
654 tz,
655 base_offset,
656 Some(upper_bound),
657 )
658 })
659 .collect::<PolarsResult<Vec<_>>>()?;
660 Ok(flatten_par(&vals))
661 })
662 } else if ((offset.duration_ns() >= period.duration_ns())
663 && matches!(closed_window, ClosedWindow::Left | ClosedWindow::None))
664 || ((offset.duration_ns() > period.duration_ns())
665 && matches!(closed_window, ClosedWindow::Right | ClosedWindow::Both))
666 {
667 let iter = group_by_values_iter_window_behind_t(
671 period,
672 offset,
673 time,
674 closed_window,
675 tu,
676 tz,
677 start_offset,
678 Some(upper_bound),
679 )?;
680 iter.map(|result| result.map(|(offset, len)| [offset, len]))
681 .collect::<PolarsResult<_>>()
682 }
683 else {
690 let iter = group_by_values_iter_partial_lookbehind(
691 period,
692 offset,
693 time,
694 closed_window,
695 tu,
696 tz,
697 start_offset,
698 Some(upper_bound),
699 )?;
700 iter.map(|result| result.map(|(offset, len)| [offset, len]))
701 .collect::<PolarsResult<_>>()
702 }
703 } else if !offset.is_zero()
704 || closed_window == ClosedWindow::Right
705 || closed_window == ClosedWindow::None
706 {
707 if !run_parallel {
712 let vecs = group_by_values_iter_lookahead_collected(
713 period,
714 offset,
715 time,
716 closed_window,
717 tu,
718 tz,
719 start_offset,
720 Some(upper_bound),
721 )?;
722 return Ok(GroupsSlice::from(vecs));
723 }
724
725 RAYON.install(|| {
726 let vals = thread_offsets
727 .par_iter()
728 .copied()
729 .map(|(base_offset, len)| {
730 let lower_bound = base_offset;
731 let upper_bound = base_offset + len;
732 group_by_values_iter_lookahead_collected(
733 period,
734 offset,
735 time,
736 closed_window,
737 tu,
738 tz,
739 lower_bound,
740 Some(upper_bound),
741 )
742 })
743 .collect::<PolarsResult<Vec<_>>>()?;
744 Ok(flatten_par(&vals))
745 })
746 } else {
747 if !run_parallel {
748 let vecs = group_by_values_iter_lookahead_collected(
749 period,
750 offset,
751 time,
752 closed_window,
753 tu,
754 tz,
755 start_offset,
756 Some(upper_bound),
757 )?;
758 return Ok(GroupsSlice::from(vecs));
759 }
760
761 RAYON.install(|| {
766 let vals = thread_offsets
767 .par_iter()
768 .copied()
769 .map(|(base_offset, len)| {
770 let lower_bound = base_offset;
771 let upper_bound = base_offset + len;
772 group_by_values_iter_lookahead_collected(
773 period,
774 offset,
775 time,
776 closed_window,
777 tu,
778 tz,
779 lower_bound,
780 Some(upper_bound),
781 )
782 })
783 .collect::<PolarsResult<Vec<_>>>()?;
784 Ok(flatten_par(&vals))
785 })
786 }
787}
788
789pub struct RollingWindower {
790 period: Duration,
791 offset: Duration,
792 closed: ClosedWindow,
793
794 tu: TimeUnit,
795 tz: Option<Tz>,
796
797 start: IdxSize,
798 end: IdxSize,
799 length: IdxSize,
800
801 active: VecDeque<ActiveWindow>,
802}
803
804struct ActiveWindow {
805 start: i64,
806 end: i64,
807}
808
809impl ActiveWindow {
810 #[inline(always)]
811 fn above_lower_bound(&self, t: i64, closed: ClosedWindow) -> bool {
812 (t > self.start)
813 | (matches!(closed, ClosedWindow::Left | ClosedWindow::Both) & (t == self.start))
814 }
815
816 #[inline(always)]
817 fn below_upper_bound(&self, t: i64, closed: ClosedWindow) -> bool {
818 (t < self.end)
819 | (matches!(closed, ClosedWindow::Right | ClosedWindow::Both) & (t == self.end))
820 }
821}
822
823fn skip_in_2d_list(l: &[&[i64]], mut n: usize) -> (usize, usize) {
824 let mut y = 0;
825 while y < l.len() && (n >= l[y].len() || l[y].is_empty()) {
826 n -= l[y].len();
827 y += 1;
828 }
829 assert!(n == 0 || y < l.len());
830 (n, y)
831}
832fn increment_2d(x: &mut usize, y: &mut usize, l: &[&[i64]]) {
833 *x += 1;
834 while *y < l.len() && *x == l[*y].len() {
835 *y += 1;
836 *x = 0;
837 }
838}
839
840impl RollingWindower {
841 pub fn new(
842 period: Duration,
843 offset: Duration,
844 closed: ClosedWindow,
845 tu: TimeUnit,
846 tz: Option<Tz>,
847 ) -> Self {
848 Self {
849 period,
850 offset,
851 closed,
852
853 tu,
854 tz,
855
856 start: 0,
857 end: 0,
858 length: 0,
859
860 active: Default::default(),
861 }
862 }
863
864 pub fn insert(
868 &mut self,
869 time: &[&[i64]],
870 windows: &mut Vec<[IdxSize; 2]>,
871 ) -> PolarsResult<IdxSize> {
872 let (mut i_x, mut i_y) = skip_in_2d_list(time, (self.length - self.start) as usize);
873 let (mut s_x, mut s_y) = skip_in_2d_list(time, 0); let (mut e_x, mut e_y) = skip_in_2d_list(time, (self.end - self.start) as usize);
875
876 let time_start = self.start;
877 let mut i = self.length;
878 while i_y < time.len() {
879 let t = time[i_y][i_x];
880 let window_start = self.offset.add(self.tu, t, self.tz.as_ref())?;
881 let window_end = if self.offset == -self.period {
884 t
885 } else {
886 self.period.add(self.tu, window_start, self.tz.as_ref())?
887 };
888
889 self.active.push_back(ActiveWindow {
890 start: window_start,
891 end: window_end,
892 });
893
894 while let Some(w) = self.active.front() {
895 if w.below_upper_bound(t, self.closed) {
896 break;
897 }
898
899 let w = self.active.pop_front().unwrap();
900 while self.start < i && !w.above_lower_bound(time[s_y][s_x], self.closed) {
901 increment_2d(&mut s_x, &mut s_y, time);
902 self.start += 1;
903 }
904 while self.end < i && w.below_upper_bound(time[e_y][e_x], self.closed) {
905 increment_2d(&mut e_x, &mut e_y, time);
906 self.end += 1;
907 }
908 windows.push([self.start, self.end - self.start]);
909 }
910
911 increment_2d(&mut i_x, &mut i_y, time);
912 i += 1;
913 }
914
915 self.length = i;
916 Ok(self.start - time_start)
917 }
918
919 pub fn finalize(&mut self, time: &[&[i64]], windows: &mut Vec<[IdxSize; 2]>) {
921 assert_eq!(
922 time.iter().map(|t| t.len()).sum::<usize>() as IdxSize,
923 self.length - self.start
924 );
925
926 let (mut s_x, mut s_y) = skip_in_2d_list(time, 0);
927 let (mut e_x, mut e_y) = skip_in_2d_list(time, (self.end - self.start) as usize);
928
929 windows.extend(self.active.drain(..).map(|w| {
930 while self.start < self.length && !w.above_lower_bound(time[s_y][s_x], self.closed) {
931 increment_2d(&mut s_x, &mut s_y, time);
932 self.start += 1;
933 }
934 while self.end < self.length && w.below_upper_bound(time[e_y][e_x], self.closed) {
935 increment_2d(&mut e_x, &mut e_y, time);
936 self.end += 1;
937 }
938 [self.start, self.end - self.start]
939 }));
940
941 self.start = 0;
942 self.end = 0;
943 self.length = 0;
944 }
945
946 pub fn reset(&mut self) {
947 self.active.clear();
948 self.start = 0;
949 self.end = 0;
950 self.length = 0;
951 }
952}
953
954#[derive(Debug)]
955struct ActiveDynWindow {
956 start: IdxSize,
957 end: Option<IdxSize>,
958 lower_bound: i64,
959 upper_bound: i64,
960}
961
962#[inline(always)]
963fn is_above_lower_bound(t: i64, lb: i64, closed: ClosedWindow) -> bool {
964 (t > lb) | (matches!(closed, ClosedWindow::Left | ClosedWindow::Both) & (t == lb))
965}
966#[inline(always)]
967fn is_below_upper_bound(t: i64, ub: i64, closed: ClosedWindow) -> bool {
968 (t < ub) | (matches!(closed, ClosedWindow::Right | ClosedWindow::Both) & (t == ub))
969}
970
971pub struct GroupByDynamicWindower {
972 period: Duration,
973 offset: Duration,
974 every: Duration,
975 closed: ClosedWindow,
976
977 start_by: StartBy,
978
979 tu: TimeUnit,
980 tz: Option<Tz>,
981
982 include_lower_bound: bool,
983 include_upper_bound: bool,
984
985 placement: Option<DynamicWindowPlacement>,
986 past_range: bool,
988
989 num_seen: IdxSize,
990 origin: i64,
992 next_k: i64,
994 next_lower_bound: i64,
995 active: VecDeque<ActiveDynWindow>,
996
997 prev_upper_bound: Option<i64>,
999 non_monotonic_upper_bounds: bool,
1001}
1002
1003impl GroupByDynamicWindower {
1004 #[expect(clippy::too_many_arguments)]
1005 pub fn new(
1006 period: Duration,
1007 offset: Duration,
1008 every: Duration,
1009 start_by: StartBy,
1010 closed: ClosedWindow,
1011 tu: TimeUnit,
1012 tz: Option<Tz>,
1013 include_lower_bound: bool,
1014 include_upper_bound: bool,
1015 placement: Option<DynamicWindowPlacement>,
1016 ) -> Self {
1017 Self {
1018 period,
1019 offset,
1020 every,
1021 closed,
1022
1023 start_by,
1024
1025 tu,
1026 tz,
1027
1028 include_lower_bound,
1029 include_upper_bound,
1030
1031 placement,
1032 past_range: false,
1033
1034 num_seen: 0,
1035 origin: 0,
1036 next_k: 0,
1037 next_lower_bound: 0,
1038 active: Default::default(),
1039 prev_upper_bound: None,
1040 non_monotonic_upper_bounds: false,
1041 }
1042 }
1043
1044 fn window_start(&self, k: i64) -> PolarsResult<i64> {
1045 (self.every * k).add(self.tu, self.origin, self.tz.as_ref())
1046 }
1047
1048 fn find_first_window_around(&self, mut k: i64, target: i64) -> PolarsResult<FirstWindowAround> {
1051 let mut lower_bound = self.window_start(k)?;
1052 let mut upper_bound = self.period.add(self.tu, lower_bound, self.tz.as_ref())?;
1053 while !is_below_upper_bound(target, upper_bound, self.closed) {
1054 let gap = target - lower_bound;
1055 let every = self.every.nte_duration(self.tu);
1056 let period = self.period.nte_duration(self.tu);
1057 let nth = if gap > every + period {
1058 (gap - period) / every
1059 } else {
1060 1
1061 };
1062
1063 k += nth;
1064 lower_bound = self.window_start(k)?;
1065 upper_bound = self.period.add(self.tu, lower_bound, self.tz.as_ref())?;
1066 }
1067
1068 if is_above_lower_bound(target, lower_bound, self.closed) {
1069 Ok(FirstWindowAround::ContainingTarget {
1070 k,
1071 lower_bound,
1072 upper_bound,
1073 })
1074 } else {
1075 Ok(FirstWindowAround::NotContainingTarget { k, lower_bound })
1076 }
1077 }
1078
1079 fn start_lower_bound(&self, first: i64) -> PolarsResult<i64> {
1080 match self.placement {
1081 Some(placement) => Ok(placement.origin),
1082 None => Window::new(self.every, self.period, self.offset).first_window_start(
1083 first,
1084 self.closed,
1085 self.tu,
1086 self.tz.as_ref(),
1087 self.start_by,
1088 ),
1089 }
1090 }
1091
1092 pub fn insert(
1093 &mut self,
1094 time: &[i64],
1095 windows: &mut Vec<[IdxSize; 2]>,
1096 lower_bound: &mut Vec<i64>,
1097 upper_bound: &mut Vec<i64>,
1098 ) -> PolarsResult<()> {
1099 if time.is_empty() {
1100 return Ok(());
1101 }
1102
1103 if self.num_seen == 0 {
1104 debug_assert!(self.active.is_empty());
1105 self.origin = self.start_lower_bound(time[0])?;
1106 self.next_k = 0;
1107 self.next_lower_bound = self.origin;
1108 }
1109
1110 for &t in time {
1111 if self.non_monotonic_upper_bounds {
1114 for w in self.active.iter_mut() {
1115 if w.end.is_none() && !is_below_upper_bound(t, w.upper_bound, self.closed) {
1116 w.end = Some(self.num_seen);
1117 }
1118 }
1119
1120 while let Some(w) = self.active.front()
1121 && w.end.is_some()
1122 {
1123 let w = self.active.pop_front().unwrap();
1124 Self::emit(
1125 &w,
1126 self.num_seen,
1127 self.include_lower_bound,
1128 self.include_upper_bound,
1129 windows,
1130 lower_bound,
1131 upper_bound,
1132 );
1133 }
1134 } else {
1135 while let Some(w) = self.active.front()
1136 && !is_below_upper_bound(t, w.upper_bound, self.closed)
1137 {
1138 let w = self.active.pop_front().unwrap();
1139 Self::emit(
1140 &w,
1141 self.num_seen,
1142 self.include_lower_bound,
1143 self.include_upper_bound,
1144 windows,
1145 lower_bound,
1146 upper_bound,
1147 );
1148 }
1149 }
1150
1151 while !self.past_range && is_above_lower_bound(t, self.next_lower_bound, self.closed) {
1152 match self.find_first_window_around(self.next_k, t)? {
1153 FirstWindowAround::ContainingTarget {
1154 k,
1155 lower_bound,
1156 upper_bound,
1157 } => {
1158 self.next_k = k + 1;
1159 self.next_lower_bound = self.window_start(self.next_k)?;
1160 if let Some(placement) = &self.placement {
1161 if placement.start_range.is_past(lower_bound) {
1162 self.past_range = true;
1163 break;
1164 }
1165 if placement.start_range.is_before(lower_bound) {
1166 continue;
1167 }
1168 }
1169 self.non_monotonic_upper_bounds |=
1170 self.prev_upper_bound.is_some_and(|prev| upper_bound < prev);
1171 self.prev_upper_bound = Some(upper_bound);
1172 self.active.push_back(ActiveDynWindow {
1173 start: self.num_seen,
1174 end: None,
1175 lower_bound,
1176 upper_bound,
1177 });
1178 },
1179 FirstWindowAround::NotContainingTarget { k, lower_bound } => {
1180 self.next_k = k;
1181 self.next_lower_bound = lower_bound;
1182 break;
1183 },
1184 }
1185 }
1186
1187 self.num_seen += 1
1188 }
1189
1190 Ok(())
1191 }
1192
1193 fn emit(
1194 w: &ActiveDynWindow,
1195 num_seen: IdxSize,
1196 include_lower_bound: bool,
1197 include_upper_bound: bool,
1198 windows: &mut Vec<[IdxSize; 2]>,
1199 lower_bound: &mut Vec<i64>,
1200 upper_bound: &mut Vec<i64>,
1201 ) {
1202 let end = w.end.unwrap_or(num_seen);
1203 windows.push([w.start, end - w.start]);
1204 if include_lower_bound {
1205 lower_bound.push(w.lower_bound);
1206 }
1207 if include_upper_bound {
1208 upper_bound.push(w.upper_bound);
1209 }
1210 }
1211
1212 pub fn lowest_needed_index(&self) -> IdxSize {
1213 self.active.front().map_or(self.num_seen, |w| w.start)
1214 }
1215
1216 pub fn is_done(&self) -> bool {
1219 self.past_range && self.active.is_empty()
1220 }
1221
1222 pub fn finalize(
1223 &mut self,
1224 windows: &mut Vec<[IdxSize; 2]>,
1225 lower_bound: &mut Vec<i64>,
1226 upper_bound: &mut Vec<i64>,
1227 ) {
1228 let num_seen = self.num_seen;
1229 let (include_lower_bound, include_upper_bound) =
1230 (self.include_lower_bound, self.include_upper_bound);
1231 for w in self.active.drain(..) {
1232 Self::emit(
1233 &w,
1234 num_seen,
1235 include_lower_bound,
1236 include_upper_bound,
1237 windows,
1238 lower_bound,
1239 upper_bound,
1240 );
1241 }
1242
1243 self.origin = 0;
1244 self.next_k = 0;
1245 self.next_lower_bound = 0;
1246 self.num_seen = 0;
1247 self.prev_upper_bound = None;
1248 self.non_monotonic_upper_bounds = false;
1249 self.past_range = false;
1250 }
1251
1252 pub fn num_seen(&self) -> IdxSize {
1253 self.num_seen
1254 }
1255
1256 pub fn time_unit(&self) -> TimeUnit {
1257 self.tu
1258 }
1259}
1260
1261enum FirstWindowAround {
1262 ContainingTarget {
1263 k: i64,
1264 lower_bound: i64,
1265 upper_bound: i64,
1266 },
1267 NotContainingTarget {
1268 k: i64,
1269 lower_bound: i64,
1270 },
1271}
1272
1273#[cfg(test)]
1274mod test {
1275 use super::*;
1276
1277 #[test]
1278 fn test_prune_duplicates() {
1279 let time = &[0, 1, 1, 2, 2, 2, 3, 4, 5, 6, 5];
1282 let mut splits = vec![(0, 2), (2, 4), (6, 2), (8, 3)];
1283 prune_splits_on_duplicates(time, &mut splits);
1284 assert_eq!(splits, &[(0, 6), (6, 2), (8, 3)]);
1285 }
1286}