Skip to main content

polars_time/windows/
group_by.rs

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                // Skip windows that start before the range. `get_stride` never overshoots,
38                // so windows close to the range are still visited one by one.
39                stride = iter.get_stride(placement.start_range.start());
40                continue 'bounds;
41            }
42        }
43
44        let mut has_member = false;
45        // find starting point of window
46        for &t in &time[start..time.len().saturating_sub(1)] {
47            // the window is behind the time values.
48            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            // element drops out of the window
57            start += 1;
58        }
59
60        // update stride so we can fast-forward in case of sparse data
61        stride = if has_member {
62            0
63        } else {
64            debug_assert!(start < time.len());
65            iter.get_stride(time[start])
66        };
67
68        // find members of this window
69        let mut end = start;
70
71        // last value isn't always added
72        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/// Window boundaries are created based on the given `Window`, which is defined by:
104/// - every
105/// - period
106/// - offset
107///
108/// And every window boundary we search for the values that fit that window by the given
109/// `ClosedWindow`. The groups are return as `GroupTuples` together with the lower bound and upper
110/// bound timestamps. These timestamps indicate the start (lower) and end (upper) of the window of
111/// that group.
112///
113/// If `include_boundaries` is `false` those `lower` and `upper` vectors will be empty.
114#[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    // the boundary we define here is not yet correct. It doesn't take 'period' into account
128    // and it doesn't have the proper starting point. This boundary is used as a proxy to find
129    // the proper 'boundary' in  'window.get_overlapping_bounds_iter'.
130    let boundary = if time.len() > 1 {
131        // +1 because left or closed boundary could match the next window if it is on the boundary
132        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// t is right at the end of the window
198// ------t---
199// [------]
200#[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    // Use binary search to find the initial start as that is behind.
217    let mut start = if let Some(&t) = time.get(start_offset) {
218        let lower = offset.add(tu, t, tz.as_ref())?;
219        // We have `period == -offset`, so `t + offset + period` is equal to `t`,
220        // and `upper` is trivially equal to `t` itself. Using the trivial calculation,
221        // instead of `upper = lower + period`, avoids issues around
222        // `t - 1mo + 1mo` not round-tripping.
223        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            // Fast path for duplicates.
237            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            // faster path, check if `i` is member.
258            if b.is_member_exit(*t, closed_window) {
259                end = i;
260            } else {
261                end = std::cmp::max(end, start);
262            }
263            // we still must loop to consume duplicates
264            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// this one is correct for all lookbehind/lookaheads, but is slower
279// window is completely behind t and t itself is not a member
280// ---------------t---
281//  [---]
282#[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        // Fast path for duplicates.
300        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// window is with -1 periods of t
338// ----t---
339//  [---]
340#[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        // Fast path for duplicates.
358        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
393/// The first row of `time[..start_offset]` that can be a member of the window of row
394/// `start_offset`, so a window search need not walk the rows before it.
395fn 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)]
411// window is completely ahead of t and t itself is not a member
412// --t-----------
413//        [---]
414pub(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        // Fast path for duplicates.
433        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    // t is at the right endpoint of the window
480    group_by_values_iter_lookbehind(period, offset, time, closed_window, tu, tz, 0, None)
481}
482
483/// Checks if the boundary elements don't split on duplicates.
484/// If they do we remove them
485fn 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            // Only push left block
502            new.push(window[0])
503        }
504    }
505    // Check last block
506    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    // We pruned invalid blocks, now we must correct the lengths.
513    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        // Call again to check.
525        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
580/// Different from `group_by_windows`, where define window buckets and search which values fit that
581/// pre-defined bucket.
582///
583/// This function defines every window based on the:
584///     - timestamp (lower bound)
585///     - timestamp + period (upper bound)
586/// where timestamps are the individual values in the array `time`
587///
588/// Only the rows in `rows` get a window; the windows still index into the whole of `time`.
589/// `rows` must not start or end inside a run of equal values.
590pub 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    // The threads split the whole of `time` and each takes the part of its split inside `rows`,
607    // so the windows do not depend on `rows`.
608    let mut thread_offsets = _split_offsets(time.len(), RAYON.current_num_threads());
609    // there are duplicates in the splits, so we opt for a single partition
610    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    // If we start from within parallel work we will do this single threaded.
619    let run_parallel = !RAYON.current_thread_has_pending_tasks().unwrap_or(false);
620
621    // we have a (partial) lookbehind window
622    if offset.negative && !offset.is_zero() {
623        // lookbehind
624        if offset.duration_ns() == period.duration_ns() {
625            // t is right at the end of the window
626            // ------t---
627            // [------]
628            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            // window is completely behind t and t itself is not a member
668            // ---------------t---
669            //  [---]
670            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        // partial lookbehind
684        // this one is still single threaded
685        // can make it parallel later, its a bit more complicated because the boundaries are unknown
686        // window is with -1 periods of t
687        // ----t---
688        //  [---]
689        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        // window is completely ahead of t and t itself is not a member
708        // --t-----------
709        //        [---]
710
711        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        // Offset is 0 and window is closed on the left:
762        // it must be that the window starts at t and t is a member
763        // --t-----------
764        //  [---]
765        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    /// Insert new values into the windower.
865    ///
866    /// This should be given all the old values that were not processed yet.
867    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); // skip over empty lists
874        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            // For datetime arithmetic, it does *NOT* hold 0 + a - a == 0. Therefore, we make sure
882            // that if `offset` and `period` are inverses we keep the `t`.
883            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    /// Process all remaining items and signal that no more items are coming.
920    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    /// No window can start in the placement's range any more.
987    past_range: bool,
988
989    num_seen: IdxSize,
990    /// Window `k` starts at `origin + every * k`.
991    origin: i64,
992    /// Index and start of the next window that may open.
993    next_k: i64,
994    next_lower_bound: i64,
995    active: VecDeque<ActiveDynWindow>,
996
997    /// Upper bound of the most recently opened window.
998    prev_upper_bound: Option<i64>,
999    // Kept track of because of clamping/DST
1000    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    /// Find the first window from window `k` on whose upper bound is not behind `target`.
1049    /// Returns its index and bounds if it contains `target`, otherwise its index and start.
1050    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            // Close every window that `t` has moved past.
1112            // This sets flags for monotonicity.
1113            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    /// Whether no further input can produce a window: every window that could start in the
1217    /// placement's range has been opened and closed.
1218    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        //                     |--|------------|----|---------|
1280        //                     0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
1281        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}