polars_time/windows/

group_by.rs

use std::collections::VecDeque;

use polars_arrow::legacy::time_zone::Tz;
use polars_arrow::trusted_len::TrustedLen;
use polars_core::prelude::*;
use polars_core::runtime::RAYON;
use polars_core::utils::_split_offsets;
use polars_core::utils::flatten::flatten_par;
use polars_defs::time::duration::Duration;
use polars_defs::time::group_by::{ClosedWindow, DynamicWindowPlacement, StartBy};
use rayon::prelude::*;

use crate::prelude::*;

#[allow(clippy::too_many_arguments)]
fn update_groups_and_bounds(
    bounds_iter: BoundsIter<'_>,
    mut start: usize,
    time: &[i64],
    closed_window: ClosedWindow,
    include_lower_bound: bool,
    include_upper_bound: bool,
    placement: Option<&DynamicWindowPlacement>,
    lower_bound: &mut Vec<i64>,
    upper_bound: &mut Vec<i64>,
    groups: &mut Vec<[IdxSize; 2]>,
) {
    let mut iter = bounds_iter.into_iter();
    let mut stride = 0;

    'bounds: while let Some(bi) = iter.nth(stride) {
        if let Some(placement) = placement {
            if placement.start_range.is_past(bi.start) {
                break;
            }
            if placement.start_range.is_before(bi.start) {
                // Skip windows that start before the range. `get_stride` never overshoots,
                // so windows close to the range are still visited one by one.
                stride = iter.get_stride(placement.start_range.start());
                continue 'bounds;
            }
        }

        let mut has_member = false;
        // find starting point of window
        for &t in &time[start..time.len().saturating_sub(1)] {
            // the window is behind the time values.
            if bi.is_future(t, closed_window) {
                stride = iter.get_stride(t);
                continue 'bounds;
            }
            if bi.is_member_entry(t, closed_window) {
                has_member = true;
                break;
            }
            // element drops out of the window
            start += 1;
        }

        // update stride so we can fast-forward in case of sparse data
        stride = if has_member {
            0
        } else {
            debug_assert!(start < time.len());
            iter.get_stride(time[start])
        };

        // find members of this window
        let mut end = start;

        // last value isn't always added
        if end == time.len() - 1 {
            let t = time[end];
            if bi.is_member(t, closed_window) {
                if include_lower_bound {
                    lower_bound.push(bi.start);
                }
                if include_upper_bound {
                    upper_bound.push(bi.stop);
                }
                groups.push([end as IdxSize, 1])
            }
            continue;
        }
        for &t in &time[end..] {
            if !bi.is_member_exit(t, closed_window) {
                break;
            }
            end += 1;
        }
        let len = end - start;

        if include_lower_bound {
            lower_bound.push(bi.start);
        }
        if include_upper_bound {
            upper_bound.push(bi.stop);
        }
        groups.push([start as IdxSize, len as IdxSize])
    }
}

/// Window boundaries are created based on the given `Window`, which is defined by:
/// - every
/// - period
/// - offset
///
/// And every window boundary we search for the values that fit that window by the given
/// `ClosedWindow`. The groups are returned as `GroupTuples` together with the lower bound and upper
/// bound timestamps. These timestamps indicate the start (lower) and end (upper) of the window of
/// that group.
///
/// If `include_boundaries` is `false` those `lower` and `upper` vectors will be empty.
#[allow(clippy::too_many_arguments)]
pub fn group_by_windows(
    window: Window,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: &Option<TimeZone>,
    include_lower_bound: bool,
    include_upper_bound: bool,
    start_by: StartBy,
    placement: Option<DynamicWindowPlacement>,
) -> PolarsResult<(GroupsSlice, Vec<i64>, Vec<i64>)> {
    let start = time[0];
    // the boundary we define here is not yet correct. It doesn't take 'period' into account
    // and it doesn't have the proper starting point. This boundary is used as a proxy to find
    // the proper 'boundary' in  'window.get_overlapping_bounds_iter'.
    let boundary = if time.len() > 1 {
        // +1 because left or closed boundary could match the next window if it is on the boundary
        let stop = time[time.len() - 1] + 1;
        Bounds::new_checked(start, stop)
    } else {
        let stop = start + 1;
        Bounds::new_checked(start, stop)
    };

    let size = window.estimate_overlapping_bounds(tu, boundary);
    let size_lower = if include_lower_bound { size } else { 0 };
    let size_upper = if include_upper_bound { size } else { 0 };
    let mut lower_bound = Vec::with_capacity(size_lower);
    let mut upper_bound = Vec::with_capacity(size_upper);

    let mut groups = Vec::with_capacity(size);
    let start_offset = 0;

    match tz {
        #[cfg(feature = "timezones")]
        Some(tz) => {
            update_groups_and_bounds(
                window.get_overlapping_bounds_iter(
                    boundary,
                    closed_window,
                    tu,
                    tz.parse::<Tz>().ok().as_ref(),
                    start_by,
                    placement.map(|p| p.origin),
                )?,
                start_offset,
                time,
                closed_window,
                include_lower_bound,
                include_upper_bound,
                placement.as_ref(),
                &mut lower_bound,
                &mut upper_bound,
                &mut groups,
            );
        },
        _ => {
            update_groups_and_bounds(
                window.get_overlapping_bounds_iter(
                    boundary,
                    closed_window,
                    tu,
                    None,
                    start_by,
                    placement.map(|p| p.origin),
                )?,
                start_offset,
                time,
                closed_window,
                include_lower_bound,
                include_upper_bound,
                placement.as_ref(),
                &mut lower_bound,
                &mut upper_bound,
                &mut groups,
            );
        },
    };

    Ok((groups, lower_bound, upper_bound))
}

// t is right at the end of the window
// ------t---
// [------]
#[inline]
#[allow(clippy::too_many_arguments)]
pub(crate) fn group_by_values_iter_lookbehind(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
    start_offset: usize,
    upper_bound: Option<usize>,
) -> PolarsResult<impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_> {
    debug_assert_eq!(offset, -period);

    let upper_bound = upper_bound.unwrap_or(time.len());
    // Use binary search to find the initial start as that is behind.
    let mut start = if let Some(&t) = time.get(start_offset) {
        let lower = offset.add(tu, t, tz.as_ref())?;
        // We have `period == -offset`, so `t + offset + period` is equal to `t`,
        // and `upper` is trivially equal to `t` itself. Using the trivial calculation,
        // instead of `upper = lower + period`, avoids issues around
        // `t - 1mo + 1mo` not round-tripping.
        let upper = t;
        let b = Bounds::new(lower, upper);
        let slice = &time[..start_offset];
        slice.partition_point(|v| !b.is_member(*v, closed_window))
    } else {
        0
    };
    let mut end = start;
    let mut last = time[start_offset];
    Ok(time[start_offset..upper_bound]
        .iter()
        .enumerate()
        .map(move |(mut i, t)| {
            // Fast path for duplicates.
            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 = offset.add(tu, *t, tz.as_ref())?;
            let upper = *t;

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

            for &t in unsafe { time.get_unchecked(start..i) } {
                if b.is_member_entry(t, closed_window) {
                    break;
                }
                start += 1;
            }

            // faster path, check if `i` is member.
            if b.is_member_exit(*t, closed_window) {
                end = i;
            } else {
                end = std::cmp::max(end, start);
            }
            // we still must loop to consume duplicates
            for &t in unsafe { time.get_unchecked(end..) } {
                if !b.is_member_exit(t, closed_window) {
                    break;
                }
                end += 1;
            }

            let len = end - start;
            let offset = start as IdxSize;

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

// this one is correct for all lookbehind/lookaheads, but is slower
// window is completely behind t and t itself is not a member
// ---------------t---
//  [---]
#[allow(clippy::too_many_arguments)]
pub(crate) fn group_by_values_iter_window_behind_t(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
    start_offset: usize,
    upper_bound: Option<usize>,
) -> PolarsResult<impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_> {
    let upper_bound = upper_bound.unwrap_or(time.len());
    let mut start = first_member_row(period, offset, time, closed_window, tu, tz, start_offset)?;
    let mut end = start;
    let mut last = time[start_offset];
    let mut started = false;
    Ok(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));
        }
        last = *lower;
        started = true;
        let lower = offset.add(tu, *lower, tz.as_ref())?;
        let upper = period.add(tu, lower, tz.as_ref())?;

        let b = Bounds::new(lower, upper);
        if b.is_future(time[0], closed_window) {
            Ok((0, 0))
        } else {
            for &t in &time[start..] {
                if b.is_member_entry(t, closed_window) {
                    break;
                }
                start += 1;
            }

            end = std::cmp::max(start, end);
            for &t in &time[end..] {
                if !b.is_member_exit(t, closed_window) {
                    break;
                }
                end += 1;
            }

            let len = end - start;
            let offset = start as IdxSize;

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

// window is within -1 periods of t
// ----t---
//  [---]
#[allow(clippy::too_many_arguments)]
pub(crate) fn group_by_values_iter_partial_lookbehind(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
    start_offset: usize,
    upper_bound: Option<usize>,
) -> PolarsResult<impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_> {
    let upper_bound = upper_bound.unwrap_or(time.len());
    let mut start = first_member_row(period, offset, time, closed_window, tu, tz, start_offset)?;
    let mut end = start;
    let mut last = time[start_offset];
    let rows = &time[start_offset..upper_bound];
    Ok(rows.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 i = i + start_offset;

        let lower = offset.add(tu, *lower, tz.as_ref())?;
        let upper = period.add(tu, lower, tz.as_ref())?;

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

        for &t in &time[start..] {
            if b.is_member_entry(t, closed_window) || start == i {
                break;
            }
            start += 1;
        }

        end = std::cmp::max(start, end);
        for &t in &time[end..] {
            if !b.is_member_exit(t, closed_window) {
                break;
            }
            end += 1;
        }

        let len = end - start;
        let offset = start as IdxSize;

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

/// The first row of `time[..start_offset]` that can be a member of the window of row
/// `start_offset`, so a window search need not walk the rows before it.
fn first_member_row(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
    start_offset: usize,
) -> PolarsResult<usize> {
    let lower = offset.add(tu, time[start_offset], tz.as_ref())?;
    let upper = period.add(tu, lower, tz.as_ref())?;
    let b = Bounds::new(lower, upper);
    Ok(time[..start_offset].partition_point(|&t| !b.is_member_entry(t, closed_window)))
}

#[allow(clippy::too_many_arguments)]
// window is completely ahead of t and t itself is not a member
// --t-----------
//        [---]
pub(crate) fn group_by_values_iter_lookahead(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
    start_offset: usize,
    upper_bound: Option<usize>,
) -> impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_ {
    let upper_bound = upper_bound.unwrap_or(time.len());

    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 = offset.add(tu, *lower, tz.as_ref())?;
        let upper = period.add(tu, lower, tz.as_ref())?;

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

        for &t in &time[start..] {
            if b.is_member_entry(t, closed_window) {
                break;
            }
            start += 1;
        }

        end = std::cmp::max(start, end);
        for &t in &time[end..] {
            if !b.is_member_exit(t, closed_window) {
                break;
            }
            end += 1;
        }

        let len = end - start;
        let offset = start as IdxSize;

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

#[cfg(feature = "rolling_window_by")]
#[inline]
pub(crate) fn group_by_values_iter(
    period: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
) -> PolarsResult<impl TrustedLen<Item = PolarsResult<(IdxSize, IdxSize)>> + '_> {
    let mut offset = period;
    offset.negative = true;
    // t is at the right endpoint of the window
    group_by_values_iter_lookbehind(period, offset, time, closed_window, tu, tz, 0, None)
}

/// Checks if the boundary elements don't split on duplicates.
/// If they do we remove them
fn prune_splits_on_duplicates(time: &[i64], thread_offsets: &mut Vec<(usize, usize)>) {
    let is_valid = |window: &[(usize, usize)]| -> bool {
        debug_assert_eq!(window.len(), 2);
        let left_block_end = window[0].0 + window[0].1.saturating_sub(1);
        let right_block_start = window[1].0;
        time[left_block_end] != time[right_block_start]
    };

    if time.is_empty() || thread_offsets.len() <= 1 || thread_offsets.windows(2).all(is_valid) {
        return;
    }

    let mut new = vec![];
    for window in thread_offsets.windows(2) {
        let this_block_is_valid = is_valid(window);
        if this_block_is_valid {
            // Only push left block
            new.push(window[0])
        }
    }
    // Check last block
    if thread_offsets.len().is_multiple_of(2) {
        let window = &thread_offsets[thread_offsets.len() - 2..];
        if is_valid(window) {
            new.push(thread_offsets[thread_offsets.len() - 1])
        }
    }
    // We pruned invalid blocks, now we must correct the lengths.
    if new.len() <= 1 {
        new = vec![(0, time.len())];
    } else {
        let mut previous_start = time.len();
        for window in new.iter_mut().rev() {
            window.1 = previous_start - window.0;
            previous_start = window.0;
        }
        new[0].0 = 0;
        new[0].1 = new[1].0;
        debug_assert_eq!(new.iter().map(|w| w.1).sum::<usize>(), time.len());
        // Call again to check.
        prune_splits_on_duplicates(time, &mut new)
    }
    std::mem::swap(thread_offsets, &mut new);
}

#[allow(clippy::too_many_arguments)]
fn group_by_values_iter_lookbehind_collected(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
    start_offset: usize,
    upper_bound: Option<usize>,
) -> PolarsResult<Vec<[IdxSize; 2]>> {
    let iter = group_by_values_iter_lookbehind(
        period,
        offset,
        time,
        closed_window,
        tu,
        tz,
        start_offset,
        upper_bound,
    )?;
    iter.map(|result| result.map(|(offset, len)| [offset, len]))
        .collect::<PolarsResult<Vec<_>>>()
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn group_by_values_iter_lookahead_collected(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
    start_offset: usize,
    upper_bound: Option<usize>,
) -> PolarsResult<Vec<[IdxSize; 2]>> {
    let iter = group_by_values_iter_lookahead(
        period,
        offset,
        time,
        closed_window,
        tu,
        tz,
        start_offset,
        upper_bound,
    );
    iter.map(|result| result.map(|(offset, len)| [offset as IdxSize, len]))
        .collect::<PolarsResult<Vec<_>>>()
}

/// Different from `group_by_windows`, where define window buckets and search which values fit that
/// pre-defined bucket.
///
/// This function defines every window based on the:
///     - timestamp (lower bound)
///     - timestamp + period (upper bound)
/// where timestamps are the individual values in the array `time`
///
/// Only the rows in `rows` get a window; the windows still index into the whole of `time`.
/// `rows` must not start or end inside a run of equal values.
pub fn group_by_values(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    tz: Option<Tz>,
    rows: std::ops::Range<usize>,
) -> PolarsResult<GroupsSlice> {
    if rows.is_empty() {
        return Ok(GroupsSlice::from(vec![]));
    }
    assert!(rows.start == 0 || time[rows.start - 1] < time[rows.start]);
    assert!(rows.end == time.len() || time[rows.end - 1] < time[rows.end]);
    let (start_offset, upper_bound) = (rows.start, rows.end);

    // The threads split the whole of `time` and each takes the part of its split inside `rows`,
    // so the windows do not depend on `rows`.
    let mut thread_offsets = _split_offsets(time.len(), RAYON.current_num_threads());
    // there are duplicates in the splits, so we opt for a single partition
    prune_splits_on_duplicates(time, &mut thread_offsets);
    thread_offsets.retain_mut(|(offset, len)| {
        let start = (*offset).max(start_offset);
        let end = (*offset + *len).min(upper_bound);
        (*offset, *len) = (start, end.saturating_sub(start));
        *len > 0
    });

    // If we start from within parallel work we will do this single threaded.
    let run_parallel = !RAYON.current_thread_has_pending_tasks().unwrap_or(false);

    // we have a (partial) lookbehind window
    if offset.negative && !offset.is_zero() {
        // lookbehind
        if offset == -period {
            // t is right at the end of the window
            // ------t---
            // [------]
            if !run_parallel {
                let vecs = group_by_values_iter_lookbehind_collected(
                    period,
                    offset,
                    time,
                    closed_window,
                    tu,
                    tz,
                    start_offset,
                    Some(upper_bound),
                )?;
                return Ok(GroupsSlice::from(vecs));
            }

            RAYON.install(|| {
                let vals = thread_offsets
                    .par_iter()
                    .copied()
                    .map(|(base_offset, len)| {
                        let upper_bound = base_offset + len;
                        group_by_values_iter_lookbehind_collected(
                            period,
                            offset,
                            time,
                            closed_window,
                            tu,
                            tz,
                            base_offset,
                            Some(upper_bound),
                        )
                    })
                    .collect::<PolarsResult<Vec<_>>>()?;
                Ok(flatten_par(&vals))
            })
        } else if ((offset.duration_ns() >= period.duration_ns())
            && matches!(closed_window, ClosedWindow::Left | ClosedWindow::None))
            || ((offset.duration_ns() > period.duration_ns())
                && matches!(closed_window, ClosedWindow::Right | ClosedWindow::Both))
        {
            // The window is likely completely behind t. `duration_ns` is only an estimate for
            // calendar units, but both this and the partial lookbehind below are correct for any
            // negative offset, so it only picks the faster one.
            // ---------------t---
            //  [---]
            let iter = group_by_values_iter_window_behind_t(
                period,
                offset,
                time,
                closed_window,
                tu,
                tz,
                start_offset,
                Some(upper_bound),
            )?;
            iter.map(|result| result.map(|(offset, len)| [offset, len]))
                .collect::<PolarsResult<_>>()
        }
        // partial lookbehind
        // this one is still single threaded
        // can make it parallel later, it's a bit more complicated because the boundaries are unknown
        // window is within -1 periods of t
        // ----t---
        //  [---]
        else {
            let iter = group_by_values_iter_partial_lookbehind(
                period,
                offset,
                time,
                closed_window,
                tu,
                tz,
                start_offset,
                Some(upper_bound),
            )?;
            iter.map(|result| result.map(|(offset, len)| [offset, len]))
                .collect::<PolarsResult<_>>()
        }
    } else if !offset.is_zero()
        || closed_window == ClosedWindow::Right
        || closed_window == ClosedWindow::None
    {
        // window is completely ahead of t and t itself is not a member
        // --t-----------
        //        [---]

        if !run_parallel {
            let vecs = group_by_values_iter_lookahead_collected(
                period,
                offset,
                time,
                closed_window,
                tu,
                tz,
                start_offset,
                Some(upper_bound),
            )?;
            return Ok(GroupsSlice::from(vecs));
        }

        RAYON.install(|| {
            let vals = thread_offsets
                .par_iter()
                .copied()
                .map(|(base_offset, len)| {
                    let lower_bound = base_offset;
                    let upper_bound = base_offset + len;
                    group_by_values_iter_lookahead_collected(
                        period,
                        offset,
                        time,
                        closed_window,
                        tu,
                        tz,
                        lower_bound,
                        Some(upper_bound),
                    )
                })
                .collect::<PolarsResult<Vec<_>>>()?;
            Ok(flatten_par(&vals))
        })
    } else {
        if !run_parallel {
            let vecs = group_by_values_iter_lookahead_collected(
                period,
                offset,
                time,
                closed_window,
                tu,
                tz,
                start_offset,
                Some(upper_bound),
            )?;
            return Ok(GroupsSlice::from(vecs));
        }

        // Offset is 0 and window is closed on the left:
        // it must be that the window starts at t and t is a member
        // --t-----------
        //  [---]
        RAYON.install(|| {
            let vals = thread_offsets
                .par_iter()
                .copied()
                .map(|(base_offset, len)| {
                    let lower_bound = base_offset;
                    let upper_bound = base_offset + len;
                    group_by_values_iter_lookahead_collected(
                        period,
                        offset,
                        time,
                        closed_window,
                        tu,
                        tz,
                        lower_bound,
                        Some(upper_bound),
                    )
                })
                .collect::<PolarsResult<Vec<_>>>()?;
            Ok(flatten_par(&vals))
        })
    }
}

pub struct RollingWindower {
    period: Duration,
    offset: Duration,
    closed: ClosedWindow,

    tu: TimeUnit,
    tz: Option<Tz>,

    start: IdxSize,
    end: IdxSize,
    length: IdxSize,

    active: VecDeque<ActiveWindow>,
}

struct ActiveWindow {
    start: i64,
    end: i64,
}

impl ActiveWindow {
    #[inline(always)]
    fn above_lower_bound(&self, t: i64, closed: ClosedWindow) -> bool {
        (t > self.start)
            | (matches!(closed, ClosedWindow::Left | ClosedWindow::Both) & (t == self.start))
    }

    #[inline(always)]
    fn below_upper_bound(&self, t: i64, closed: ClosedWindow) -> bool {
        (t < self.end)
            | (matches!(closed, ClosedWindow::Right | ClosedWindow::Both) & (t == self.end))
    }
}

fn skip_in_2d_list(l: &[&[i64]], mut n: usize) -> (usize, usize) {
    let mut y = 0;
    while y < l.len() && (n >= l[y].len() || l[y].is_empty()) {
        n -= l[y].len();
        y += 1;
    }
    assert!(n == 0 || y < l.len());
    (n, y)
}
fn increment_2d(x: &mut usize, y: &mut usize, l: &[&[i64]]) {
    *x += 1;
    while *y < l.len() && *x == l[*y].len() {
        *y += 1;
        *x = 0;
    }
}

impl RollingWindower {
    pub fn new(
        period: Duration,
        offset: Duration,
        closed: ClosedWindow,
        tu: TimeUnit,
        tz: Option<Tz>,
    ) -> Self {
        Self {
            period,
            offset,
            closed,

            tu,
            tz,

            start: 0,
            end: 0,
            length: 0,

            active: Default::default(),
        }
    }

    /// Insert new values into the windower.
    ///
    /// This should be given all the old values that were not processed yet.
    pub fn insert(
        &mut self,
        time: &[&[i64]],
        windows: &mut Vec<[IdxSize; 2]>,
    ) -> PolarsResult<IdxSize> {
        let (mut i_x, mut i_y) = skip_in_2d_list(time, (self.length - self.start) as usize);
        let (mut s_x, mut s_y) = skip_in_2d_list(time, 0); // skip over empty lists
        let (mut e_x, mut e_y) = skip_in_2d_list(time, (self.end - self.start) as usize);

        let time_start = self.start;
        let mut i = self.length;
        while i_y < time.len() {
            let t = time[i_y][i_x];
            let window_start = self.offset.add(self.tu, t, self.tz.as_ref())?;
            // For datetime arithmetic, it does *NOT* hold 0 + a - a == 0. Therefore, we make sure
            // that if `offset` and `period` are inverses we keep the `t`.
            let window_end = if self.offset == -self.period {
                t
            } else {
                self.period.add(self.tu, window_start, self.tz.as_ref())?
            };

            self.active.push_back(ActiveWindow {
                start: window_start,
                end: window_end,
            });

            while let Some(w) = self.active.front() {
                if w.below_upper_bound(t, self.closed) {
                    break;
                }

                let w = self.active.pop_front().unwrap();
                while self.start < i && !w.above_lower_bound(time[s_y][s_x], self.closed) {
                    increment_2d(&mut s_x, &mut s_y, time);
                    self.start += 1;
                }
                while self.end < i && w.below_upper_bound(time[e_y][e_x], self.closed) {
                    increment_2d(&mut e_x, &mut e_y, time);
                    self.end += 1;
                }
                windows.push([self.start, self.end - self.start]);
            }

            increment_2d(&mut i_x, &mut i_y, time);
            i += 1;
        }

        self.length = i;
        Ok(self.start - time_start)
    }

    /// Process all remaining items and signal that no more items are coming.
    pub fn finalize(&mut self, time: &[&[i64]], windows: &mut Vec<[IdxSize; 2]>) {
        assert_eq!(
            time.iter().map(|t| t.len()).sum::<usize>() as IdxSize,
            self.length - self.start
        );

        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);

        windows.extend(self.active.drain(..).map(|w| {
            while self.start < self.length && !w.above_lower_bound(time[s_y][s_x], self.closed) {
                increment_2d(&mut s_x, &mut s_y, time);
                self.start += 1;
            }
            while self.end < self.length && w.below_upper_bound(time[e_y][e_x], self.closed) {
                increment_2d(&mut e_x, &mut e_y, time);
                self.end += 1;
            }
            [self.start, self.end - self.start]
        }));

        self.start = 0;
        self.end = 0;
        self.length = 0;
    }

    pub fn reset(&mut self) {
        self.active.clear();
        self.start = 0;
        self.end = 0;
        self.length = 0;
    }
}

#[derive(Debug)]
struct ActiveDynWindow {
    start: IdxSize,
    end: Option<IdxSize>,
    lower_bound: i64,
    upper_bound: i64,
}

#[inline(always)]
fn is_above_lower_bound(t: i64, lb: i64, closed: ClosedWindow) -> bool {
    (t > lb) | (matches!(closed, ClosedWindow::Left | ClosedWindow::Both) & (t == lb))
}
#[inline(always)]
fn is_below_upper_bound(t: i64, ub: i64, closed: ClosedWindow) -> bool {
    (t < ub) | (matches!(closed, ClosedWindow::Right | ClosedWindow::Both) & (t == ub))
}

pub struct GroupByDynamicWindower {
    period: Duration,
    offset: Duration,
    every: Duration,
    closed: ClosedWindow,

    start_by: StartBy,

    tu: TimeUnit,
    tz: Option<Tz>,

    include_lower_bound: bool,
    include_upper_bound: bool,

    placement: Option<DynamicWindowPlacement>,
    /// No window can start in the placement's range any more.
    past_range: bool,

    num_seen: IdxSize,
    /// Window `k` starts at `origin + every * k`.
    origin: i64,
    /// Index and start of the next window that may open.
    next_k: i64,
    next_lower_bound: i64,
    active: VecDeque<ActiveDynWindow>,

    /// Upper bound of the most recently opened window.
    prev_upper_bound: Option<i64>,
    // Kept track of because of clamping/DST
    non_monotonic_upper_bounds: bool,
}

impl GroupByDynamicWindower {
    #[expect(clippy::too_many_arguments)]
    pub fn new(
        period: Duration,
        offset: Duration,
        every: Duration,
        start_by: StartBy,
        closed: ClosedWindow,
        tu: TimeUnit,
        tz: Option<Tz>,
        include_lower_bound: bool,
        include_upper_bound: bool,
        placement: Option<DynamicWindowPlacement>,
    ) -> Self {
        Self {
            period,
            offset,
            every,
            closed,

            start_by,

            tu,
            tz,

            include_lower_bound,
            include_upper_bound,

            placement,
            past_range: false,

            num_seen: 0,
            origin: 0,
            next_k: 0,
            next_lower_bound: 0,
            active: Default::default(),
            prev_upper_bound: None,
            non_monotonic_upper_bounds: false,
        }
    }

    fn window_start(&self, k: i64) -> PolarsResult<i64> {
        (self.every * k).add(self.tu, self.origin, self.tz.as_ref())
    }

    /// Find the first window from window `k` on whose upper bound is not behind `target`.
    /// Returns its index and bounds if it contains `target`, otherwise its index and start.
    fn find_first_window_around(&self, mut k: i64, target: i64) -> PolarsResult<FirstWindowAround> {
        let mut lower_bound = self.window_start(k)?;
        let mut upper_bound = self.period.add(self.tu, lower_bound, self.tz.as_ref())?;
        while !is_below_upper_bound(target, upper_bound, self.closed) {
            let gap = target - lower_bound;
            let every = self.every.nte_duration(self.tu);
            let period = self.period.nte_duration(self.tu);
            let nth = if gap > every + period {
                (gap - period) / every
            } else {
                1
            };

            k += nth;
            lower_bound = self.window_start(k)?;
            upper_bound = self.period.add(self.tu, lower_bound, self.tz.as_ref())?;
        }

        if is_above_lower_bound(target, lower_bound, self.closed) {
            Ok(FirstWindowAround::ContainingTarget {
                k,
                lower_bound,
                upper_bound,
            })
        } else {
            Ok(FirstWindowAround::NotContainingTarget { k, lower_bound })
        }
    }

    fn start_lower_bound(&self, first: i64) -> PolarsResult<i64> {
        match self.placement {
            Some(placement) => Ok(placement.origin),
            None => Window::new(self.every, self.period, self.offset).first_window_start(
                first,
                self.closed,
                self.tu,
                self.tz.as_ref(),
                self.start_by,
            ),
        }
    }

    pub fn insert(
        &mut self,
        time: &[i64],
        windows: &mut Vec<[IdxSize; 2]>,
        lower_bound: &mut Vec<i64>,
        upper_bound: &mut Vec<i64>,
    ) -> PolarsResult<()> {
        if time.is_empty() {
            return Ok(());
        }

        if self.num_seen == 0 {
            debug_assert!(self.active.is_empty());
            self.origin = self.start_lower_bound(time[0])?;
            self.next_k = 0;
            self.next_lower_bound = self.origin;
        }

        for &t in time {
            // Close every window that `t` has moved past.
            // This sets flags for monotonicity.
            if self.non_monotonic_upper_bounds {
                for w in self.active.iter_mut() {
                    if w.end.is_none() && !is_below_upper_bound(t, w.upper_bound, self.closed) {
                        w.end = Some(self.num_seen);
                    }
                }

                while let Some(w) = self.active.front()
                    && w.end.is_some()
                {
                    let w = self.active.pop_front().unwrap();
                    Self::emit(
                        &w,
                        self.num_seen,
                        self.include_lower_bound,
                        self.include_upper_bound,
                        windows,
                        lower_bound,
                        upper_bound,
                    );
                }
            } else {
                while let Some(w) = self.active.front()
                    && !is_below_upper_bound(t, w.upper_bound, self.closed)
                {
                    let w = self.active.pop_front().unwrap();
                    Self::emit(
                        &w,
                        self.num_seen,
                        self.include_lower_bound,
                        self.include_upper_bound,
                        windows,
                        lower_bound,
                        upper_bound,
                    );
                }
            }

            while !self.past_range && is_above_lower_bound(t, self.next_lower_bound, self.closed) {
                match self.find_first_window_around(self.next_k, t)? {
                    FirstWindowAround::ContainingTarget {
                        k,
                        lower_bound,
                        upper_bound,
                    } => {
                        self.next_k = k + 1;
                        self.next_lower_bound = self.window_start(self.next_k)?;
                        if let Some(placement) = &self.placement {
                            if placement.start_range.is_past(lower_bound) {
                                self.past_range = true;
                                break;
                            }
                            if placement.start_range.is_before(lower_bound) {
                                continue;
                            }
                        }
                        self.non_monotonic_upper_bounds |=
                            self.prev_upper_bound.is_some_and(|prev| upper_bound < prev);
                        self.prev_upper_bound = Some(upper_bound);
                        self.active.push_back(ActiveDynWindow {
                            start: self.num_seen,
                            end: None,
                            lower_bound,
                            upper_bound,
                        });
                    },
                    FirstWindowAround::NotContainingTarget { k, lower_bound } => {
                        self.next_k = k;
                        self.next_lower_bound = lower_bound;
                        break;
                    },
                }
            }

            self.num_seen += 1
        }

        Ok(())
    }

    fn emit(
        w: &ActiveDynWindow,
        num_seen: IdxSize,
        include_lower_bound: bool,
        include_upper_bound: bool,
        windows: &mut Vec<[IdxSize; 2]>,
        lower_bound: &mut Vec<i64>,
        upper_bound: &mut Vec<i64>,
    ) {
        let end = w.end.unwrap_or(num_seen);
        windows.push([w.start, end - w.start]);
        if include_lower_bound {
            lower_bound.push(w.lower_bound);
        }
        if include_upper_bound {
            upper_bound.push(w.upper_bound);
        }
    }

    pub fn lowest_needed_index(&self) -> IdxSize {
        self.active.front().map_or(self.num_seen, |w| w.start)
    }

    /// Whether no further input can produce a window: every window that could start in the
    /// placement's range has been opened and closed.
    pub fn is_done(&self) -> bool {
        self.past_range && self.active.is_empty()
    }

    pub fn finalize(
        &mut self,
        windows: &mut Vec<[IdxSize; 2]>,
        lower_bound: &mut Vec<i64>,
        upper_bound: &mut Vec<i64>,
    ) {
        let num_seen = self.num_seen;
        let (include_lower_bound, include_upper_bound) =
            (self.include_lower_bound, self.include_upper_bound);
        for w in self.active.drain(..) {
            Self::emit(
                &w,
                num_seen,
                include_lower_bound,
                include_upper_bound,
                windows,
                lower_bound,
                upper_bound,
            );
        }

        self.origin = 0;
        self.next_k = 0;
        self.next_lower_bound = 0;
        self.num_seen = 0;
        self.prev_upper_bound = None;
        self.non_monotonic_upper_bounds = false;
        self.past_range = false;
    }

    pub fn num_seen(&self) -> IdxSize {
        self.num_seen
    }

    pub fn time_unit(&self) -> TimeUnit {
        self.tu
    }
}

enum FirstWindowAround {
    ContainingTarget {
        k: i64,
        lower_bound: i64,
        upper_bound: i64,
    },
    NotContainingTarget {
        k: i64,
        lower_bound: i64,
    },
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_prune_duplicates() {
        //                     |--|------------|----|---------|
        //                     0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
        let time = &[0, 1, 1, 2, 2, 2, 3, 4, 5, 6, 5];
        let mut splits = vec![(0, 2), (2, 4), (6, 2), (8, 3)];
        prune_splits_on_duplicates(time, &mut splits);
        assert_eq!(splits, &[(0, 6), (6, 2), (8, 3)]);
    }
}