rskit_stream/operators/
windowing.rs

1//! Leading-edge windowing and batching operators.
2//!
3//! These operators partition a stream into `Vec<T>` windows whose timers run **leading-edge**,
4//! from the first buffered item (`tumbling_window`, `batch`), or by item count (`sliding_window`).
5//! Contrast the trailing-edge rate operators in [`crate::operators::rate`],
6//! whose timers are driven by quiet gaps between arrivals.
7
8use std::collections::VecDeque;
9use std::time::Duration;
10
11use futures::{Stream, StreamExt as _};
12
13use super::buffer::{take_window, window_capacity};
14
15/// Collect items into non-overlapping time windows.
16///
17/// A window is emitted after `duration` has elapsed since the first item.
18pub fn tumbling_window<S, T>(
19    stream: S,
20    duration: Duration,
21    max_items: usize,
22) -> impl Stream<Item = Vec<T>> + Send + 'static
23where
24    S: Stream<Item = T> + Send + 'static,
25    T: Send + 'static,
26{
27    async_stream::stream! {
28        tokio::pin!(stream);
29        let (max_items, initial_capacity) = window_capacity(max_items);
30        let mut buf: Vec<T> = Vec::with_capacity(initial_capacity);
31        let mut deadline = tokio::time::Instant::now() + duration;
32        loop {
33            tokio::select! {
34                item = stream.next() => {
35                    if let Some(v) = item {
36                        buf.push(v);
37                        if buf.len() >= max_items {
38                            yield take_window(&mut buf, initial_capacity);
39                            deadline = tokio::time::Instant::now() + duration;
40                        }
41                    } else {
42                        if !buf.is_empty() {
43                            yield take_window(&mut buf, initial_capacity);
44                        }
45                        break;
46                    }
47                }
48                () = tokio::time::sleep_until(deadline) => {
49                    if !buf.is_empty() {
50                        yield take_window(&mut buf, initial_capacity);
51                    }
52                    deadline = tokio::time::Instant::now() + duration;
53                }
54            }
55        }
56    }
57}
58
59/// Collect up to `size` items into a batch.
60///
61/// A batch is emitted either when `size` items arrive
62/// or when `timeout` elapses since the first item in the batch.
63pub fn batch<S, T>(
64    stream: S,
65    size: usize,
66    timeout: Duration,
67) -> impl Stream<Item = Vec<T>> + Send + 'static
68where
69    S: Stream<Item = T> + Send + 'static,
70    T: Send + 'static,
71{
72    async_stream::stream! {
73        tokio::pin!(stream);
74        let size = size.max(1);
75        let mut buf: Vec<T> = Vec::with_capacity(size);
76        let mut deadline: Option<tokio::time::Instant> = None;
77
78        loop {
79            tokio::select! {
80                item = stream.next() => {
81                    if let Some(v) = item {
82                        if buf.is_empty() {
83                            deadline = Some(tokio::time::Instant::now() + timeout);
84                        }
85                        buf.push(v);
86                        if buf.len() >= size {
87                            deadline = None;
88                            yield std::mem::take(&mut buf);
89                        }
90                    } else {
91                        if !buf.is_empty() {
92                            yield std::mem::take(&mut buf);
93                        }
94                        break;
95                    }
96                }
97                () = async {
98                    match deadline {
99                        Some(d) => tokio::time::sleep_until(d).await,
100                        None => std::future::pending::<()>().await,
101                    }
102                } => {
103                    deadline = None;
104                    if !buf.is_empty() {
105                        yield std::mem::take(&mut buf);
106                    }
107                }
108            }
109        }
110    }
111}
112
113/// Emit sliding windows of `size` items, advancing by `step` items each time.
114pub fn sliding_window<S, T>(
115    stream: S,
116    size: usize,
117    step: usize,
118) -> impl Stream<Item = Vec<T>> + Send + 'static
119where
120    S: Stream<Item = T> + Send + 'static,
121    T: Clone + Send + 'static,
122{
123    async_stream::stream! {
124        tokio::pin!(stream);
125        let mut window = VecDeque::with_capacity(size.max(1));
126        let effective_step = step.max(1);
127
128        while let Some(item) = stream.next().await {
129            window.push_back(item);
130            while window.len() > size {
131                window.pop_front();
132            }
133            if size > 0 && window.len() == size {
134                yield window.iter().cloned().collect();
135                for _ in 0..effective_step.min(window.len()) {
136                    window.pop_front();
137                }
138            }
139        }
140    }
141}
rskit_stream/operators/windowing.rs

rskit_stream/operators/
windowing.rs
