I Built a Streaming Window Aggregator in Pure Python and Finally Understood How Flink Handles Late Data

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional

class WindowType(Enum):
    TUMBLING = "tumbling"
    SLIDING = "sliding"
    SESSION = "session"

@dataclass
class WindowSpec:
    window_type: WindowType
    size_ms: int               # window duration in milliseconds
    slide_ms: Optional[int] = None    # sliding only: step between windows
    gap_ms: Optional[int] = None      # session only: inactivity gap

    def validate(self):
        if self.window_type == WindowType.SLIDING and not self.slide_ms:
            raise ValueError("Sliding windows need slide_ms")
        if self.window_type == WindowType.SESSION and not self.gap_ms:
            raise ValueError("Session windows need gap_ms")

@dataclass
class Event:
    key: str
    event_time_ms: int   # time the event actually happened (from the source)
    value: float

class WatermarkTracker:
    def __init__(self, allowed_lateness_ms: int = 5000):
        self.allowed_lateness_ms = allowed_lateness_ms
        self.current_watermark_ms: int = 0

    def advance(self, event_time_ms: int) -> int:
        """Advance watermark based on new event. Returns new watermark."""
        candidate = event_time_ms - self.allowed_lateness_ms
        if candidate > self.current_watermark_ms:
            self.current_watermark_ms = candidate
        return self.current_watermark_ms

    def is_late(self, event_time_ms: int) -> bool:
        """True if this event arrived after its window already closed."""
        return event_time_ms < self.current_watermark_ms

from collections import defaultdict
from typing import Dict, List, Tuple

@dataclass
class WindowState:
    window_start_ms: int
    window_end_ms: int
    key: str
    count: int = 0
    total: float = 0.0

    def add(self, value: float):
        self.count += 1
        self.total += value

    def result(self) -> dict:
        return {
            "key": self.key,
            "window_start_ms": self.window_start_ms,
            "window_end_ms": self.window_end_ms,
            "count": self.count,
            "sum": self.total,
            "avg": self.total / self.count if self.count > 0 else 0,
        }

class StreamWindowAggregator:
    def __init__(self, spec: WindowSpec, allowed_lateness_ms: int = 5000):
        self.spec = spec
        self.spec.validate()
        self.watermark = WatermarkTracker(allowed_lateness_ms)
        # key: (key, window_start_ms) -> WindowState
        self.windows: Dict[Tuple[str, int], WindowState] = {}
        self.emitted: List[dict] = []

    def _get_tumbling_windows(self, event_time_ms: int) -> List[Tuple[int, int]]:
        start = (event_time_ms // self.spec.size_ms) * self.spec.size_ms
        return [(start, start + self.spec.size_ms)]

    def _get_sliding_windows(self, event_time_ms: int) -> List[Tuple[int, int]]:
        slide = self.spec.slide_ms
        size = self.spec.size_ms
        # find all sliding windows that contain this event
        latest_start = (event_time_ms // slide) * slide
        windows = []
        start = latest_start
        while start > event_time_ms - size:
            end = start + size
            if start <= event_time_ms < end:
                windows.append((start, end))
            start -= slide
        return windows

    def _get_session_windows(self, key: str, event_time_ms: int) -> List[Tuple[int, int]]:
        gap = self.spec.gap_ms
        # find existing session for this key that this event extends
        for (k, w_start), state in list(self.windows.items()):
            if k != key:
                continue
            # extend session if within gap
            if state.window_end_ms >= event_time_ms - gap:
                new_end = max(state.window_end_ms, event_time_ms + gap)
                state.window_end_ms = new_end
                return []  # event gets merged into existing session
        # new session
        return [(event_time_ms, event_time_ms + gap)]

    def process(self, event: Event) -> List[dict]:
        """Process one event. Returns any windows that are now ready to emit."""
        if self.watermark.is_late(event.event_time_ms):
            # late event: drop it (you could also route to a side output)
            return []

        self.watermark.advance(event.event_time_ms)

        # assign event to window(s)
        if self.spec.window_type == WindowType.TUMBLING:
            assigned = self._get_tumbling_windows(event.event_time_ms)
        elif self.spec.window_type == WindowType.SLIDING:
            assigned = self._get_sliding_windows(event.event_time_ms)
        else:
            assigned = self._get_session_windows(event.key, event.event_time_ms)

        for (w_start, w_end) in assigned:
            wkey = (event.key, w_start)
            if wkey not in self.windows:
                self.windows[wkey] = WindowState(w_start, w_end, event.key)
            self.windows[wkey].add(event.value)

        # check if any windows are past the watermark (ready to emit)
        return self._flush_ready_windows()

    def _flush_ready_windows(self) -> List[dict]:
        ready = []
        current_wm = self.watermark.current_watermark_ms
        to_remove = []
        for wkey, state in self.windows.items():
            if state.window_end_ms <= current_wm:
                ready.append(state.result())
                to_remove.append(wkey)
        for wkey in to_remove:
            del self.windows[wkey]
        self.emitted.extend(ready)
        return ready

    def flush_all(self) -> List[dict]:
        """Force-emit all open windows (call at end of stream)."""
        remaining = [state.result() for state in self.windows.values()]
        self.windows.clear()
        self.emitted.extend(remaining)
        return remaining

import time

def demo_tumbling():
    spec = WindowSpec(WindowType.TUMBLING, size_ms=10_000)
    agg = StreamWindowAggregator(spec, allowed_lateness_ms=3_000)

    # Simulate events arriving out of order
    events = [
        Event("user_A", 1_000,  10.0),   # t=1s
        Event("user_A", 5_000,  20.0),   # t=5s
        Event("user_B", 3_000,  15.0),   # t=3s
        Event("user_A", 12_000, 30.0),   # t=12s  (pushes watermark to 9s)
        Event("user_A", 6_000,  5.0),    # t=6s   (late? watermark=9s, t=6s: NOT late yet)
        Event("user_B", 8_000,  25.0),   # t=8s
        Event("user_A", 25_000, 40.0),   # t=25s  (watermark=22s, closes window [0,10))
    ]

    for event in events:
        results = agg.process(event)
        for r in results:
            print(f"EMIT: {r}")

    final = agg.flush_all()
    for r in final:
        print(f"FINAL: {r}")

demo_tumbling()

EMIT: {'key': 'user_A', 'window_start_ms': 0, 'window_end_ms': 10000, 'count': 3, 'sum': 35.0, 'avg': 11.67}
EMIT: {'key': 'user_B', 'window_start_ms': 0, 'window_end_ms': 10000, 'count': 2, 'sum': 40.0, 'avg': 20.0}
FINAL: {'key': 'user_A', 'window_start_ms': 10000, 'window_end_ms': 20000, 'count': 1, 'sum': 30.0, 'avg': 30.0}
FINAL: {'key': 'user_A', 'window_start_ms': 20000, 'window_end_ms': 30000, 'count': 1, 'sum': 40.0, 'avg': 40.0}

session_A: [100ms, 5100ms]   (gap=5000ms)
session_B: [8000ms, 13000ms]
event arrives at t=4500ms     # within gap of session_A, so merge into A
new session_A: [100ms, 9500ms]