pub mod types;
pub mod assigner;
pub mod trigger;
pub mod manager;
pub use types::{Window, WindowBounds, WindowType};
pub use assigner::{
WindowAssigner,
TumblingWindowAssigner,
SlidingWindowAssigner,
SessionWindowAssigner,
};
pub use trigger::{
WindowTrigger,
TriggerResult,
TriggerContext,
OnWatermarkTrigger,
ProcessingTimeTrigger,
CountTrigger,
CompositeTrigger,
CompositeMode,
ImmediateTrigger,
NeverTrigger,
};
pub use manager::{WindowManager, WindowMetadata, WindowStats};
use chrono::DateTime;
use std::collections::HashMap;
#[derive(Debug)]
pub struct WindowMerger {
windows: HashMap<String, Window>,
session_gap: chrono::Duration,
}
impl WindowMerger {
pub fn new(session_gap: chrono::Duration) -> Self {
Self {
windows: HashMap::new(),
session_gap,
}
}
pub fn add_and_merge(&mut self, window: Window) -> (Window, Vec<Window>) {
let mut merged_windows = Vec::new();
let mut result_window = window.clone();
let mut to_remove = Vec::new();
for (id, existing_window) in &self.windows {
if result_window.can_merge(existing_window, self.session_gap) {
result_window = result_window.merge(existing_window);
merged_windows.push(existing_window.clone());
to_remove.push(id.clone());
}
}
for id in to_remove {
self.windows.remove(&id);
}
self.windows.insert(result_window.id.clone(), result_window.clone());
(result_window, merged_windows)
}
pub fn windows(&self) -> Vec<&Window> {
self.windows.values().collect()
}
pub fn remove(&mut self, window_id: &str) -> Option<Window> {
self.windows.remove(window_id)
}
pub fn remove_old_windows(&mut self, watermark: DateTime<chrono::Utc>) -> Vec<Window> {
let mut removed = Vec::new();
self.windows.retain(|_, window| {
if window.bounds.end <= watermark {
removed.push(window.clone());
false
} else {
true
}
});
removed
}
pub fn clear(&mut self) {
self.windows.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::{Duration, TimeZone, Utc};
fn create_timestamp(millis: i64) -> DateTime<Utc> {
Utc.timestamp_millis_opt(millis).unwrap()
}
fn create_window(start_millis: i64, end_millis: i64) -> Window {
Window::new(WindowBounds::new(
create_timestamp(start_millis),
create_timestamp(end_millis),
))
}
#[test]
fn test_window_merger_no_overlap() {
let mut merger = WindowMerger::new(Duration::milliseconds(1000));
let window1 = create_window(0, 1000);
let (result, merged) = merger.add_and_merge(window1.clone());
assert_eq!(result, window1);
assert_eq!(merged.len(), 0);
let window2 = create_window(3000, 4000);
let (result, merged) = merger.add_and_merge(window2.clone());
assert_eq!(result, window2);
assert_eq!(merged.len(), 0);
assert_eq!(merger.windows().len(), 2);
}
#[test]
fn test_window_merger_with_overlap() {
let mut merger = WindowMerger::new(Duration::milliseconds(1000));
let window1 = create_window(0, 1000);
merger.add_and_merge(window1);
let window2 = create_window(500, 1500);
let (result, merged) = merger.add_and_merge(window2);
assert_eq!(result.bounds.start, create_timestamp(0));
assert_eq!(result.bounds.end, create_timestamp(1500));
assert_eq!(merged.len(), 1);
assert_eq!(merger.windows().len(), 1);
}
#[test]
fn test_window_merger_within_gap() {
let mut merger = WindowMerger::new(Duration::milliseconds(1000));
let window1 = create_window(0, 1000);
merger.add_and_merge(window1);
let window2 = create_window(1500, 2500);
let (result, merged) = merger.add_and_merge(window2);
assert_eq!(result.bounds.start, create_timestamp(0));
assert_eq!(result.bounds.end, create_timestamp(2500));
assert_eq!(merged.len(), 1);
}
#[test]
fn test_window_merger_multiple_merges() {
let mut merger = WindowMerger::new(Duration::milliseconds(500));
merger.add_and_merge(create_window(0, 1000));
merger.add_and_merge(create_window(2000, 3000));
merger.add_and_merge(create_window(5000, 6000));
assert_eq!(merger.windows().len(), 3);
let window = create_window(1000, 2000);
let (result, merged) = merger.add_and_merge(window);
assert_eq!(result.bounds.start, create_timestamp(0));
assert_eq!(result.bounds.end, create_timestamp(3000));
assert_eq!(merged.len(), 2);
assert_eq!(merger.windows().len(), 2);
}
#[test]
fn test_window_merger_remove_old() {
let mut merger = WindowMerger::new(Duration::milliseconds(500));
merger.add_and_merge(create_window(0, 1000));
merger.add_and_merge(create_window(2000, 3000));
merger.add_and_merge(create_window(4000, 5000));
let removed = merger.remove_old_windows(create_timestamp(3000));
assert_eq!(removed.len(), 2); assert_eq!(merger.windows().len(), 1); }
#[test]
fn test_tumbling_window_complete_flow() {
use assigner::TumblingWindowAssigner;
use trigger::OnWatermarkTrigger;
let assigner = TumblingWindowAssigner::new(Duration::milliseconds(1000));
let mut trigger = OnWatermarkTrigger::new();
let windows = assigner.assign_windows(create_timestamp(500));
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].bounds.start, create_timestamp(0));
assert_eq!(windows[0].bounds.end, create_timestamp(1000));
let ctx = TriggerContext {
window: windows[0].clone(),
event_time: Some(create_timestamp(500)),
processing_time: create_timestamp(1500),
watermark: Some(create_timestamp(900)),
element_count: 1,
};
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Continue);
let ctx = TriggerContext {
window: windows[0].clone(),
event_time: Some(create_timestamp(500)),
processing_time: create_timestamp(1500),
watermark: Some(create_timestamp(1000)),
element_count: 1,
};
assert_eq!(trigger.on_watermark(&ctx), TriggerResult::Fire);
}
#[test]
fn test_sliding_window_complete_flow() {
use assigner::SlidingWindowAssigner;
let assigner = SlidingWindowAssigner::new(
Duration::milliseconds(1000),
Duration::milliseconds(500),
);
let windows = assigner.assign_windows(create_timestamp(700));
assert_eq!(windows.len(), 2);
assert_eq!(windows[0].bounds.start, create_timestamp(0));
assert_eq!(windows[0].bounds.end, create_timestamp(1000));
assert_eq!(windows[1].bounds.start, create_timestamp(500));
assert_eq!(windows[1].bounds.end, create_timestamp(1500));
}
#[test]
fn test_session_window_complete_flow() {
use assigner::SessionWindowAssigner;
let assigner = SessionWindowAssigner::new(Duration::milliseconds(1000));
let mut merger = WindowMerger::new(Duration::milliseconds(1000));
let windows = assigner.assign_windows(create_timestamp(0));
let (_window1, _) = merger.add_and_merge(windows[0].clone());
let windows = assigner.assign_windows(create_timestamp(500));
let (window2, merged) = merger.add_and_merge(windows[0].clone());
assert_eq!(merged.len(), 1);
assert_eq!(window2.bounds.start, create_timestamp(0));
assert_eq!(window2.bounds.end, create_timestamp(1500));
let windows = assigner.assign_windows(create_timestamp(3000));
let (window3, merged) = merger.add_and_merge(windows[0].clone());
assert_eq!(merged.len(), 0);
assert_eq!(window3.bounds.start, create_timestamp(3000));
assert_eq!(window3.bounds.end, create_timestamp(4000));
}
}