laminar_core/operator/

session_window.rs

//! # Session Window Operators
//!
//! Implementation of session windows for stream processing.
//!
//! Session windows are dynamic windows that group events by activity periods
//! separated by gaps. Unlike tumbling and sliding windows which have fixed
//! boundaries, session windows grow with activity and close after inactivity.
//!
//! ## Key Characteristics
//!
//! - **Dynamic boundaries**: Sessions start with the first event and extend
//!   with each new event within the gap period
//! - **Per-key tracking**: Each key maintains independent session state
//! - **Gap-based closure**: Sessions close when no events arrive within the gap
//! - **Session merging**: Late data can merge previously separate sessions
//!
//! ## Example
//!
//! ```text
//! Gap: 30 seconds
//!
//! Events: [t=0] [t=10] [t=20]  ...gap...  [t=100] [t=110]
//!         |<---- Session 1 ---->|          |<- Session 2 ->|
//!         [0, 50)                          [100, 140)
//! ```
//!
//! ## Usage
//!
//! ```rust,no_run
//! use laminar_core::operator::session_window::SessionWindowOperator;
//! use laminar_core::operator::window::CountAggregator;
//! use std::time::Duration;
//!
//! // Create a session window with 30-second gap
//! let operator = SessionWindowOperator::new(
//!     Duration::from_secs(30),    // gap timeout
//!     CountAggregator::new(),
//!     Duration::from_secs(60),    // allowed lateness
//! );
//! ```

use super::window::{
    Accumulator, Aggregator, ChangelogRecord, EmitStrategy, LateDataConfig, LateDataMetrics,
    ResultToArrow, WindowCloseMetrics, WindowId,
};
use super::{
    Event, Operator, OperatorContext, OperatorError, OperatorState, Output, OutputVec,
    SideOutputData, Timer,
};
use crate::state::{StateStore, StateStoreExt};
use arrow_array::{Array, Int64Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use rkyv::{
    api::high::{HighDeserializer, HighSerializer, HighValidator},
    bytecheck::CheckBytes,
    rancor::Error as RkyvError,
    ser::allocator::ArenaHandle,
    util::AlignedVec,
    Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize,
};
use rustc_hash::FxHashMap;
use smallvec::SmallVec;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

/// State key prefix for session index (4 bytes)
const SESSION_INDEX_PREFIX: &[u8; 4] = b"six:";

/// State key prefix for session accumulator (4 bytes)
const SESSION_ACC_PREFIX: &[u8; 4] = b"sac:";

/// Timer key prefix for session closure (1 byte)
const SESSION_TIMER_PREFIX: u8 = 0x01;

/// Default maximum number of session indices to cache in-memory.
/// Beyond this, the cache is cleared to prevent unbounded memory growth.
const DEFAULT_MAX_CACHED_INDICES: usize = 16_384;

/// Static counter for generating unique operator IDs.
static SESSION_OPERATOR_COUNTER: AtomicU64 = AtomicU64::new(0);

/// Unique identifier for a session instance.
///
/// Generated by XOR-ing a hash of the operator ID with a monotonic counter,
/// ensuring uniqueness within and across operators.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Archive, RkyvSerialize, RkyvDeserialize)]
pub struct SessionId(u64);

impl SessionId {
    /// Generates a new unique session ID.
    ///
    /// Combines a hash of the operator ID with a monotonic counter to produce
    /// IDs that are unique across operators and invocations.
    pub fn generate(operator_id: &str, counter: &AtomicU64) -> Self {
        let op_hash = {
            use std::hash::{Hash, Hasher};
            let mut hasher = rustc_hash::FxHasher::default();
            operator_id.as_bytes().hash(&mut hasher);
            hasher.finish()
        };
        let seq = counter.fetch_add(1, Ordering::Relaxed);
        Self(op_hash ^ seq)
    }

    /// Serializes this session ID to a big-endian byte array.
    #[must_use]
    pub fn to_bytes(self) -> [u8; 8] {
        self.0.to_be_bytes()
    }

    /// Deserializes a session ID from a big-endian byte slice.
    ///
    /// Returns `None` if the slice is not exactly 8 bytes.
    #[must_use]
    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
        let arr: [u8; 8] = bytes.try_into().ok()?;
        Some(Self(u64::from_be_bytes(arr)))
    }

    /// Returns the raw `u64` value.
    #[must_use]
    pub fn as_u64(self) -> u64 {
        self.0
    }
}

/// Metadata for a single session within a key's session index.
///
/// Tracks the temporal bounds and emission state of a session. The accumulator
/// state is stored separately under `sac:<session_id>`.
#[derive(Debug, Clone, Archive, RkyvSerialize, RkyvDeserialize)]
pub struct SessionMetadata {
    /// Unique session identifier.
    pub id: SessionId,
    /// Session start timestamp (inclusive).
    pub start: i64,
    /// Session end timestamp (exclusive, = last event time + gap).
    pub end: i64,
    /// Whether this session has been emitted (used for retraction tracking).
    pub emitted: bool,
}

impl SessionMetadata {
    /// Creates a new session from a first event timestamp.
    ///
    /// The session spans `[timestamp, timestamp + gap_ms)`.
    fn new(id: SessionId, timestamp: i64, gap_ms: i64) -> Self {
        Self {
            id,
            start: timestamp,
            end: timestamp + gap_ms,
            emitted: false,
        }
    }

    /// Returns the window ID for this session.
    #[must_use]
    pub fn window_id(&self) -> WindowId {
        WindowId::new(self.start, self.end)
    }

    /// Checks whether an event at `timestamp` would overlap this session.
    ///
    /// An event overlaps if its potential session `[timestamp, timestamp + gap_ms)`
    /// intersects with this session's bounds `[start, end)`.
    fn overlaps(&self, timestamp: i64, gap_ms: i64) -> bool {
        let event_start = timestamp;
        let event_end = timestamp + gap_ms;
        // Two intervals [a, b) and [c, d) overlap when a < d && c < b
        event_start < self.end && self.start < event_end
    }

    /// Extends this session to include a new event timestamp.
    ///
    /// Widens the bounds: start becomes `min(start, timestamp)`,
    /// end becomes `max(end, timestamp + gap_ms)`.
    fn extend(&mut self, timestamp: i64, gap_ms: i64) {
        self.start = self.start.min(timestamp);
        self.end = self.end.max(timestamp + gap_ms);
    }

    /// Merges another session into this one.
    ///
    /// Takes the union of the two sessions' bounds. Used for
    /// session merging when late data bridges two previously separate sessions.
    fn merge(&mut self, other: &SessionMetadata) {
        self.start = self.start.min(other.start);
        self.end = self.end.max(other.end);
    }
}

/// Index of all active sessions for a single key.
///
/// Sessions are kept sorted by start time for efficient overlap queries.
/// Stored in the state store under `six:<key_hash>`.
#[derive(Debug, Clone, Default, Archive, RkyvSerialize, RkyvDeserialize)]
pub struct SessionIndex {
    /// Active sessions sorted by start time.
    pub sessions: Vec<SessionMetadata>,
}

impl SessionIndex {
    /// Inserts a session and maintains sorted order by start time.
    fn insert(&mut self, session: SessionMetadata) {
        let pos = self
            .sessions
            .binary_search_by_key(&session.start, |s| s.start)
            .unwrap_or_else(|pos| pos);
        self.sessions.insert(pos, session);
    }

    /// Removes a session by ID, returning it if found.
    fn remove(&mut self, id: SessionId) -> Option<SessionMetadata> {
        if let Some(pos) = self.sessions.iter().position(|s| s.id == id) {
            Some(self.sessions.remove(pos))
        } else {
            None
        }
    }

    /// Finds all sessions that overlap with an event at `timestamp`.
    fn find_overlapping(&self, timestamp: i64, gap_ms: i64) -> SmallVec<[SessionId; 4]> {
        self.sessions
            .iter()
            .filter(|s| s.overlaps(timestamp, gap_ms))
            .map(|s| s.id)
            .collect()
    }

    /// Returns a reference to the session with the given ID.
    fn get(&self, id: SessionId) -> Option<&SessionMetadata> {
        self.sessions.iter().find(|s| s.id == id)
    }

    /// Returns a mutable reference to the session with the given ID.
    fn get_mut(&mut self, id: SessionId) -> Option<&mut SessionMetadata> {
        self.sessions.iter_mut().find(|s| s.id == id)
    }

    /// Returns the number of active sessions.
    #[must_use]
    fn len(&self) -> usize {
        self.sessions.len()
    }

    /// Returns `true` if there are no active sessions.
    #[must_use]
    fn is_empty(&self) -> bool {
        self.sessions.is_empty()
    }
}

/// Creates the standard session output schema.
/// Session window operator.
///
/// Groups events by activity periods separated by gaps. Each unique key
/// maintains its own session state independently, supporting **multiple
/// concurrent sessions per key**.
///
/// # Session Lifecycle
///
/// 1. **Start**: First event for a key creates a new session
/// 2. **Extend**: Events within gap period extend the session
/// 3. **Close**: Timer fires when gap expires, emitting results
/// 4. **Merge**: Late data may merge previously separate sessions
///
/// # State Management
///
/// Session state is stored using prefixed keys:
/// - `six:<key_hash>` - [`SessionIndex`] (all sessions for a key)
/// - `sac:<session_id>` - Per-session accumulator state
///
/// # Emit Strategies
///
/// - `OnWatermark`: Emit when watermark passes session end
/// - `OnUpdate`: Emit after every state update
/// - `OnWindowClose`: Only emit on final closure
/// - `Changelog`: Emit CDC records with Z-set weights
/// - `Final`: Suppress all intermediate, drop late data
pub struct SessionWindowOperator<A: Aggregator> {
    /// Gap timeout in milliseconds
    gap_ms: i64,
    /// Aggregator function
    aggregator: A,
    /// Allowed lateness for late data
    allowed_lateness_ms: i64,
    /// Monotonic counter for generating unique session IDs
    session_id_counter: AtomicU64,
    /// Per-key session indices (in-memory cache, backed by state store).
    ///
    /// Bounded by `max_cached_indices` to prevent unbounded memory growth
    /// in high-cardinality key spaces. Evicted entries are reloaded from
    /// the state store on demand.
    session_indices: FxHashMap<u64, SessionIndex>,
    /// Maximum number of session indices to keep in-memory.
    max_cached_indices: usize,
    /// Pending timers: `session_id` -> (`timer_time`, `key_hash`)
    pending_timers: FxHashMap<u64, (i64, u64)>,
    /// Emit strategy
    emit_strategy: EmitStrategy,
    /// Late data configuration
    late_data_config: LateDataConfig,
    /// Late data metrics
    late_data_metrics: LateDataMetrics,
    /// Window close metrics
    window_close_metrics: WindowCloseMetrics,
    /// Operator ID for checkpointing
    operator_id: String,
    /// Cached output schema
    output_schema: SchemaRef,
    /// Key column index for partitioning (None = global session)
    key_column: Option<usize>,
    /// Whether timers need re-registration with `TimerService` after restore
    needs_timer_reregistration: bool,
    /// Phantom data for accumulator type
    _phantom: PhantomData<A::Acc>,
}

impl<A: Aggregator> SessionWindowOperator<A>
where
    A::Acc: Archive + for<'a> RkyvSerialize<HighSerializer<AlignedVec, ArenaHandle<'a>, RkyvError>>,
    <A::Acc as Archive>::Archived: for<'a> CheckBytes<HighValidator<'a, RkyvError>>
        + RkyvDeserialize<A::Acc, HighDeserializer<RkyvError>>,
{
    /// Creates a new session window operator.
    ///
    /// # Arguments
    ///
    /// * `gap` - The inactivity gap that closes a session
    /// * `aggregator` - Aggregation function to apply within sessions
    /// * `allowed_lateness` - Grace period for late data after session close
    ///
    /// # Panics
    ///
    /// Panics if gap or allowed lateness does not fit in i64.
    pub fn new(gap: Duration, aggregator: A, allowed_lateness: Duration) -> Self {
        let operator_num = SESSION_OPERATOR_COUNTER.fetch_add(1, Ordering::Relaxed);
        let output_schema = Arc::new(Schema::new(vec![
            Field::new("window_start", DataType::Int64, false),
            Field::new("window_end", DataType::Int64, false),
            Field::new(
                "result",
                aggregator.output_data_type(),
                aggregator.output_nullable(),
            ),
        ]));
        Self {
            gap_ms: i64::try_from(gap.as_millis()).expect("Gap must fit in i64"),
            aggregator,
            allowed_lateness_ms: i64::try_from(allowed_lateness.as_millis())
                .expect("Allowed lateness must fit in i64"),
            session_id_counter: AtomicU64::new(0),
            session_indices: FxHashMap::default(),
            max_cached_indices: DEFAULT_MAX_CACHED_INDICES,
            pending_timers: FxHashMap::default(),
            emit_strategy: EmitStrategy::default(),
            late_data_config: LateDataConfig::default(),
            late_data_metrics: LateDataMetrics::new(),
            window_close_metrics: WindowCloseMetrics::new(),
            operator_id: format!("session_window_{operator_num}"),
            output_schema,
            key_column: None,
            needs_timer_reregistration: false,
            _phantom: PhantomData,
        }
    }

    /// Creates a new session window operator with a custom operator ID.
    ///
    /// # Panics
    ///
    /// Panics if gap or allowed lateness does not fit in i64.
    pub fn with_id(
        gap: Duration,
        aggregator: A,
        allowed_lateness: Duration,
        operator_id: String,
    ) -> Self {
        let output_schema = Arc::new(Schema::new(vec![
            Field::new("window_start", DataType::Int64, false),
            Field::new("window_end", DataType::Int64, false),
            Field::new(
                "result",
                aggregator.output_data_type(),
                aggregator.output_nullable(),
            ),
        ]));
        Self {
            gap_ms: i64::try_from(gap.as_millis()).expect("Gap must fit in i64"),
            aggregator,
            allowed_lateness_ms: i64::try_from(allowed_lateness.as_millis())
                .expect("Allowed lateness must fit in i64"),
            session_id_counter: AtomicU64::new(0),
            session_indices: FxHashMap::default(),
            max_cached_indices: DEFAULT_MAX_CACHED_INDICES,
            pending_timers: FxHashMap::default(),
            emit_strategy: EmitStrategy::default(),
            late_data_config: LateDataConfig::default(),
            late_data_metrics: LateDataMetrics::new(),
            window_close_metrics: WindowCloseMetrics::new(),
            operator_id,
            output_schema,
            key_column: None,
            needs_timer_reregistration: false,
            _phantom: PhantomData,
        }
    }

    /// Sets the maximum number of session indices to cache in-memory.
    ///
    /// When the cache exceeds this limit, it is cleared and entries are
    /// reloaded from the persistent state store on demand. Default: 16,384.
    pub fn set_max_cached_indices(&mut self, max: usize) {
        self.max_cached_indices = max;
    }

    /// Sets the key column for per-key session tracking.
    ///
    /// If not set, a single global session is maintained.
    pub fn set_key_column(&mut self, column_index: usize) {
        self.key_column = Some(column_index);
    }

    /// Returns the key column index if set.
    #[must_use]
    pub fn key_column(&self) -> Option<usize> {
        self.key_column
    }

    /// Sets the emit strategy for this operator.
    pub fn set_emit_strategy(&mut self, strategy: EmitStrategy) {
        self.emit_strategy = strategy;
    }

    /// Returns the current emit strategy.
    #[must_use]
    pub fn emit_strategy(&self) -> &EmitStrategy {
        &self.emit_strategy
    }

    /// Sets the late data handling configuration.
    pub fn set_late_data_config(&mut self, config: LateDataConfig) {
        self.late_data_config = config;
    }

    /// Returns the current late data configuration.
    #[must_use]
    pub fn late_data_config(&self) -> &LateDataConfig {
        &self.late_data_config
    }

    /// Returns the late data metrics.
    #[must_use]
    pub fn late_data_metrics(&self) -> &LateDataMetrics {
        &self.late_data_metrics
    }

    /// Resets the late data metrics counters.
    pub fn reset_late_data_metrics(&mut self) {
        self.late_data_metrics.reset();
    }

    /// Returns the window close metrics.
    #[must_use]
    pub fn window_close_metrics(&self) -> &WindowCloseMetrics {
        &self.window_close_metrics
    }

    /// Resets the window close metrics counters.
    pub fn reset_window_close_metrics(&mut self) {
        self.window_close_metrics.reset();
    }

    /// Returns the gap timeout in milliseconds.
    #[must_use]
    pub fn gap_ms(&self) -> i64 {
        self.gap_ms
    }

    /// Returns the allowed lateness in milliseconds.
    #[must_use]
    pub fn allowed_lateness_ms(&self) -> i64 {
        self.allowed_lateness_ms
    }

    /// Returns the number of active sessions across all keys.
    #[must_use]
    pub fn active_session_count(&self) -> usize {
        self.session_indices.values().map(SessionIndex::len).sum()
    }

    /// Extracts the key from an event.
    fn extract_key(&self, event: &Event) -> SmallVec<[u8; 16]> {
        use arrow_array::cast::AsArray;
        use arrow_array::types::Int64Type;

        if let Some(col_idx) = self.key_column {
            if col_idx < event.data.num_columns() {
                let column = event.data.column(col_idx);
                if let Some(array) = column.as_primitive_opt::<Int64Type>() {
                    if !array.is_empty() && !array.is_null(0) {
                        return SmallVec::from_slice(&array.value(0).to_be_bytes());
                    }
                }
                // Try string column
                if let Some(array) = column.as_string_opt::<i32>() {
                    if !array.is_empty() && !array.is_null(0) {
                        return SmallVec::from_slice(array.value(0).as_bytes());
                    }
                }
            }
        }
        // Default: global session (empty key)
        SmallVec::new()
    }

    /// Computes a hash for the key.
    fn key_hash(key: &[u8]) -> u64 {
        use std::hash::{Hash, Hasher};
        let mut hasher = rustc_hash::FxHasher::default();
        key.hash(&mut hasher);
        hasher.finish()
    }

    /// Generates the state key for a session index (per key hash).
    fn session_index_key(key_hash: u64) -> [u8; 12] {
        let mut key = [0u8; 12];
        key[..4].copy_from_slice(SESSION_INDEX_PREFIX);
        key[4..12].copy_from_slice(&key_hash.to_be_bytes());
        key
    }

    /// Generates the state key for a session accumulator (per session ID).
    fn session_acc_key(session_id: SessionId) -> [u8; 12] {
        let mut key = [0u8; 12];
        key[..4].copy_from_slice(SESSION_ACC_PREFIX);
        key[4..12].copy_from_slice(&session_id.to_bytes());
        key
    }

    /// Generates the timer key for session closure (keyed by session ID).
    fn timer_key(session_id: SessionId) -> super::TimerKey {
        let mut key = super::TimerKey::new();
        key.push(SESSION_TIMER_PREFIX);
        key.extend_from_slice(&session_id.to_bytes());
        key
    }

    /// Parses a [`SessionId`] from a timer key.
    fn session_id_from_timer(timer_key: &[u8]) -> Option<SessionId> {
        if timer_key.len() != 9 || timer_key[0] != SESSION_TIMER_PREFIX {
            return None;
        }
        SessionId::from_bytes(&timer_key[1..9])
    }

    /// Loads the session index for a key, checking the in-memory cache first
    /// and falling back to the persistent state store.
    ///
    /// If the cache exceeds `max_cached_indices`, it is cleared before
    /// inserting the new entry. This prevents unbounded memory growth
    /// in high-cardinality key spaces while keeping hot entries fast.
    fn load_session_index(&mut self, key_hash: u64, state: &dyn StateStore) -> SessionIndex {
        // Check in-memory cache first
        if let Some(index) = self.session_indices.get(&key_hash) {
            return index.clone();
        }

        // Evict the cache if it exceeds the limit. All entries are backed
        // by the persistent state store, so they can be reloaded on demand.
        if self.session_indices.len() >= self.max_cached_indices {
            self.session_indices.clear();
        }

        // Check persistent state
        let state_key = Self::session_index_key(key_hash);
        if let Ok(Some(index)) = state.get_typed::<SessionIndex>(&state_key) {
            self.session_indices.insert(key_hash, index.clone());
            return index;
        }

        SessionIndex::default()
    }

    /// Persists the session index for a key. Removes the state entry if empty.
    fn store_session_index(
        &mut self,
        key_hash: u64,
        index: &SessionIndex,
        state: &mut dyn StateStore,
    ) -> Result<(), OperatorError> {
        let state_key = Self::session_index_key(key_hash);
        if index.is_empty() {
            state
                .delete(&state_key)
                .map_err(|e| OperatorError::StateAccessFailed(e.to_string()))?;
            self.session_indices.remove(&key_hash);
        } else {
            state
                .put_typed(&state_key, index)
                .map_err(|e| OperatorError::StateAccessFailed(e.to_string()))?;
            self.session_indices.insert(key_hash, index.clone());
        }
        Ok(())
    }

    /// Loads the accumulator for a session, creating a fresh one if absent.
    fn load_accumulator(&self, session_id: SessionId, state: &dyn StateStore) -> A::Acc {
        let acc_key = Self::session_acc_key(session_id);
        state
            .get_typed::<A::Acc>(&acc_key)
            .ok()
            .flatten()
            .unwrap_or_else(|| self.aggregator.create_accumulator())
    }

    /// Persists the accumulator for a session.
    fn store_accumulator(
        session_id: SessionId,
        acc: &A::Acc,
        state: &mut dyn StateStore,
    ) -> Result<(), OperatorError> {
        let acc_key = Self::session_acc_key(session_id);
        state
            .put_typed(&acc_key, acc)
            .map_err(|e| OperatorError::StateAccessFailed(e.to_string()))
    }

    /// Deletes the accumulator for a session.
    fn delete_accumulator(
        session_id: SessionId,
        state: &mut dyn StateStore,
    ) -> Result<(), OperatorError> {
        let acc_key = Self::session_acc_key(session_id);
        state
            .delete(&acc_key)
            .map_err(|e| OperatorError::StateAccessFailed(e.to_string()))
    }

    /// Registers or updates a timer for session closure.
    fn register_timer(
        &mut self,
        session_id: SessionId,
        key_hash: u64,
        session: &SessionMetadata,
        ctx: &mut OperatorContext,
    ) {
        let trigger_time = session.end + self.allowed_lateness_ms;

        // Cancel previous timer if different
        if let Some(&(old_time, _)) = self.pending_timers.get(&session_id.as_u64()) {
            if old_time == trigger_time {
                return; // Timer already set for correct time
            }
            // Note: We can't cancel timers, but the handler will check staleness
        }

        let timer_key = Self::timer_key(session_id);
        ctx.timers
            .register_timer(trigger_time, Some(timer_key), Some(ctx.operator_index));
        self.pending_timers
            .insert(session_id.as_u64(), (trigger_time, key_hash));
    }

    /// Checks if an event is late for its potential session.
    fn is_late(&self, timestamp: i64, watermark: i64) -> bool {
        // An event is late if its session would have already closed
        // Session end = timestamp + gap, cleanup = session end + allowed lateness
        let potential_cleanup = timestamp + self.gap_ms + self.allowed_lateness_ms;
        watermark >= potential_cleanup
    }

    /// Creates an output event from a session.
    fn create_output(&self, session: &SessionMetadata, acc: &A::Acc) -> Option<Event> {
        if acc.is_empty() {
            return None;
        }

        let result = acc.result();
        let result_array = result.to_arrow_array();

        let batch = RecordBatch::try_new(
            Arc::clone(&self.output_schema),
            vec![
                Arc::new(Int64Array::from(vec![session.start])),
                Arc::new(Int64Array::from(vec![session.end])),
                result_array,
            ],
        )
        .ok()?;

        Some(Event::new(session.end, batch))
    }

    /// Merges multiple overlapping sessions into one.
    ///
    /// The first session in `overlapping` becomes the "winner" that absorbs all
    /// others. Each "loser" session has its accumulator merged into the winner,
    /// its timer removed, its state cleaned up, and (for `Changelog` strategy)
    /// a retraction emitted if it was previously emitted.
    ///
    /// Returns the winner's [`SessionId`]. The caller is responsible for
    /// extending the winner to include the new event's timestamp and adding
    /// the current event's value to the accumulator.
    fn merge_sessions(
        &mut self,
        index: &mut SessionIndex,
        overlapping: &[SessionId],
        ctx: &mut OperatorContext,
        output: &mut OutputVec,
    ) -> SessionId {
        let winner_id = overlapping[0];

        // Load the winner's accumulator so we can merge losers into it
        let mut winner_acc = self.load_accumulator(winner_id, ctx.state);

        // If the winner was previously emitted and strategy is Changelog,
        // retract the winner's old state (since the merged result will differ)
        if matches!(self.emit_strategy, EmitStrategy::Changelog) {
            if let Some(winner_meta) = index.get(winner_id) {
                if winner_meta.emitted {
                    if let Some(old_evt) = self.create_output(winner_meta, &winner_acc) {
                        output.push(Output::Changelog(ChangelogRecord::delete(
                            old_evt,
                            ctx.processing_time,
                        )));
                    }
                }
            }
        }

        // Process each loser: merge accumulator, emit retraction, clean up
        for &loser_id in &overlapping[1..] {
            let loser_acc = self.load_accumulator(loser_id, ctx.state);

            // Emit retraction for the loser if it was previously emitted
            if matches!(self.emit_strategy, EmitStrategy::Changelog) {
                if let Some(loser_meta) = index.get(loser_id) {
                    if loser_meta.emitted {
                        if let Some(old_evt) = self.create_output(loser_meta, &loser_acc) {
                            output.push(Output::Changelog(ChangelogRecord::delete(
                                old_evt,
                                ctx.processing_time,
                            )));
                        }
                    }
                }
            }

            // Merge loser's accumulator into winner
            winner_acc.merge(&loser_acc);

            // Merge loser's metadata bounds into winner
            if let Some(loser_meta) = index.get(loser_id).cloned() {
                if let Some(winner_meta) = index.get_mut(winner_id) {
                    winner_meta.merge(&loser_meta);
                }
            }

            // Clean up loser: delete accumulator, remove timer, remove from index
            let _ = Self::delete_accumulator(loser_id, ctx.state);
            self.pending_timers.remove(&loser_id.as_u64());
            index.remove(loser_id);
        }

        // Store the merged accumulator for the winner
        let _ = Self::store_accumulator(winner_id, &winner_acc, ctx.state);

        winner_id
    }
}

impl<A: Aggregator> Operator for SessionWindowOperator<A>
where
    A::Acc: 'static
        + Archive
        + for<'a> RkyvSerialize<HighSerializer<AlignedVec, ArenaHandle<'a>, RkyvError>>,
    <A::Acc as Archive>::Archived: for<'a> CheckBytes<HighValidator<'a, RkyvError>>
        + RkyvDeserialize<A::Acc, HighDeserializer<RkyvError>>,
{
    fn process(&mut self, event: &Event, ctx: &mut OperatorContext) -> OutputVec {
        // Lazy timer re-registration after restore (restore() has no OperatorContext)
        if self.needs_timer_reregistration {
            self.needs_timer_reregistration = false;
            for (&sid_raw, &(trigger_time, _key_hash)) in &self.pending_timers {
                let timer_key = Self::timer_key(SessionId(sid_raw));
                ctx.timers
                    .register_timer(trigger_time, Some(timer_key), Some(ctx.operator_index));
            }
        }

        let event_time = event.timestamp;
        let mut output = OutputVec::new();

        // Update watermark
        let emitted_watermark = ctx.watermark_generator.on_event(event_time);

        // Check if event is late
        let current_wm = ctx.watermark_generator.current_watermark();
        if current_wm > i64::MIN && self.is_late(event_time, current_wm) {
            // EMIT FINAL drops late data entirely
            if self.emit_strategy.drops_late_data() {
                self.late_data_metrics.record_dropped();
                return output;
            }

            if let Some(side_output_name) = self.late_data_config.side_output() {
                self.late_data_metrics.record_side_output();
                output.push(Output::SideOutput(Box::new(SideOutputData {
                    name: Arc::from(side_output_name),
                    event: event.clone(),
                })));
            } else {
                self.late_data_metrics.record_dropped();
                output.push(Output::LateEvent(event.clone()));
            }
            return output;
        }

        // Extract key and compute hash
        let key = self.extract_key(event);
        let key_hash = Self::key_hash(&key);

        // Load session index for this key
        let mut index = self.load_session_index(key_hash, ctx.state);
        let overlapping = index.find_overlapping(event_time, self.gap_ms);

        let session_id;
        match overlapping.len() {
            0 => {
                // No overlapping session — create a new one
                let id = SessionId::generate(&self.operator_id, &self.session_id_counter);
                let session = SessionMetadata::new(id, event_time, self.gap_ms);
                index.insert(session);
                session_id = id;
            }
            1 => {
                // Exactly one overlap — extend it
                session_id = overlapping[0];
                if let Some(session) = index.get_mut(session_id) {
                    session.extend(event_time, self.gap_ms);
                }
            }
            _ => {
                // Multiple overlaps — merge all into one session
                session_id = self.merge_sessions(&mut index, &overlapping, ctx, &mut output);
                // Extend the merged winner to include this event
                if let Some(session) = index.get_mut(session_id) {
                    session.extend(event_time, self.gap_ms);
                }
            }
        }

        // Load and update accumulator (handles multi-row batches)
        let mut acc = self.load_accumulator(session_id, ctx.state);
        let values = self.aggregator.extract_batch(event);
        for value in values {
            acc.add(value);
        }

        // Persist accumulator
        let _ = Self::store_accumulator(session_id, &acc, ctx.state);

        // Register timer for the affected session
        if let Some(session) = index.get(session_id) {
            self.register_timer(session_id, key_hash, session, ctx);

            // Handle emit strategy
            match &self.emit_strategy {
                EmitStrategy::OnUpdate => {
                    if let Some(evt) = self.create_output(session, &acc) {
                        output.push(Output::Event(evt));
                    }
                    if let Some(s) = index.get_mut(session_id) {
                        s.emitted = true;
                    }
                }
                EmitStrategy::Changelog => {
                    if let Some(evt) = self.create_output(session, &acc) {
                        let record = ChangelogRecord::insert(evt, ctx.processing_time);
                        output.push(Output::Changelog(record));
                    }
                    if let Some(s) = index.get_mut(session_id) {
                        s.emitted = true;
                    }
                }
                // Other strategies: no intermediate emission
                EmitStrategy::OnWatermark
                | EmitStrategy::Periodic(_)
                | EmitStrategy::OnWindowClose
                | EmitStrategy::Final => {}
            }
        }

        // Persist session index after emit (so emitted flag is saved)
        let _ = self.store_session_index(key_hash, &index, ctx.state);

        // Emit watermark if generated
        if let Some(wm) = emitted_watermark {
            output.push(Output::Watermark(wm.timestamp()));
        }

        output
    }

    fn on_timer(&mut self, timer: Timer, ctx: &mut OperatorContext) -> OutputVec {
        let mut output = OutputVec::new();

        // Parse session ID from timer key
        let Some(session_id) = Self::session_id_from_timer(&timer.key) else {
            return output;
        };

        // Look up the pending timer entry: (expected_time, key_hash)
        let Some(&(expected_time, key_hash)) = self.pending_timers.get(&session_id.as_u64()) else {
            return output; // Timer was cancelled or session no longer exists
        };

        if expected_time != timer.timestamp {
            return output; // Stale timer, session was extended
        }

        // Load session index for the owning key
        let mut index = self.load_session_index(key_hash, ctx.state);
        let Some(session) = index.get(session_id).cloned() else {
            self.pending_timers.remove(&session_id.as_u64());
            return output;
        };

        // Load accumulator and create output
        let acc = self.load_accumulator(session_id, ctx.state);
        if let Some(event) = self.create_output(&session, &acc) {
            // Record window close metrics
            self.window_close_metrics
                .record_close(session.end, ctx.processing_time);

            match &self.emit_strategy {
                EmitStrategy::Changelog => {
                    let record = ChangelogRecord::insert(event, ctx.processing_time);
                    output.push(Output::Changelog(record));
                }
                _ => {
                    output.push(Output::Event(event));
                }
            }
        }

        // Clean up: remove session from index, delete accumulator, update state
        index.remove(session_id);
        let _ = Self::delete_accumulator(session_id, ctx.state);
        let _ = self.store_session_index(key_hash, &index, ctx.state);
        self.pending_timers.remove(&session_id.as_u64());

        output
    }

    fn checkpoint(&self) -> OperatorState {
        // Serialize pending timers: (session_id, timer_time, key_hash)
        let timer_entries: Vec<(u64, i64, u64)> = self
            .pending_timers
            .iter()
            .map(|(&sid, &(time, kh))| (sid, time, kh))
            .collect();
        let counter_val = self.session_id_counter.load(Ordering::Relaxed);

        // Checkpoint format: (timer_entries, session_id_counter)
        let checkpoint_data: (Vec<(u64, i64, u64)>, u64) = (timer_entries, counter_val);

        let data = rkyv::to_bytes::<RkyvError>(&checkpoint_data)
            .map(|v| v.to_vec())
            .unwrap_or_default();

        OperatorState {
            operator_id: self.operator_id.clone(),
            data,
        }
    }

    fn restore(&mut self, state: OperatorState) -> Result<(), OperatorError> {
        if state.operator_id != self.operator_id {
            return Err(OperatorError::StateAccessFailed(format!(
                "Operator ID mismatch: expected {}, got {}",
                self.operator_id, state.operator_id
            )));
        }

        if state.data.is_empty() {
            return Ok(());
        }

        let archived =
            rkyv::access::<rkyv::Archived<(Vec<(u64, i64, u64)>, u64)>, RkyvError>(&state.data)
                .map_err(|e| OperatorError::SerializationFailed(e.to_string()))?;
        let (timers, counter_val) =
            rkyv::deserialize::<(Vec<(u64, i64, u64)>, u64), RkyvError>(archived)
                .map_err(|e| OperatorError::SerializationFailed(e.to_string()))?;

        self.pending_timers = timers
            .into_iter()
            .map(|(sid, time, kh)| (sid, (time, kh)))
            .collect();
        self.session_id_counter
            .store(counter_val, Ordering::Relaxed);
        self.needs_timer_reregistration = !self.pending_timers.is_empty();
        // Note: session_indices will be populated lazily from state store

        Ok(())
    }
}

/// Session metrics for monitoring.
#[derive(Debug, Clone, Default)]
pub struct SessionMetrics {
    /// Total sessions created
    pub sessions_created: u64,
    /// Total sessions closed
    pub sessions_closed: u64,
    /// Total sessions merged
    pub sessions_merged: u64,
    /// Current active sessions
    pub active_sessions: u64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::operator::window::{CountAggregator, SumAggregator};
    use crate::state::InMemoryStore;
    use crate::time::{BoundedOutOfOrdernessGenerator, TimerService};
    use arrow_array::{Int64Array, RecordBatch};
    use arrow_schema::{DataType, Field, Schema};

    fn create_test_event(timestamp: i64, value: i64) -> Event {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            false,
        )]));
        let batch =
            RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![value]))]).unwrap();
        Event::new(timestamp, batch)
    }

    fn create_keyed_event(timestamp: i64, key: i64, value: i64) -> Event {
        let schema = Arc::new(Schema::new(vec![
            Field::new("key", DataType::Int64, false),
            Field::new("value", DataType::Int64, false),
        ]));
        let batch = RecordBatch::try_new(
            schema,
            vec![
                Arc::new(Int64Array::from(vec![key])),
                Arc::new(Int64Array::from(vec![value])),
            ],
        )
        .unwrap();
        Event::new(timestamp, batch)
    }

    fn create_test_context<'a>(
        timers: &'a mut TimerService,
        state: &'a mut dyn StateStore,
        watermark_gen: &'a mut dyn crate::time::WatermarkGenerator,
    ) -> OperatorContext<'a> {
        OperatorContext {
            event_time: 0,
            processing_time: 0,
            timers,
            state,
            watermark_generator: watermark_gen,
            operator_index: 0,
        }
    }

    /// Helper: get the first (`session_id`, `timer_time`, `key_hash`) from `pending_timers`.
    fn first_pending_timer(op: &SessionWindowOperator<impl Aggregator>) -> (SessionId, i64, u64) {
        let (&sid_raw, &(time, kh)) = op.pending_timers.iter().next().expect("no pending timer");
        (SessionId(sid_raw), time, kh)
    }

    // ---------------------------------------------------------------
    //  SessionId tests
    // ---------------------------------------------------------------

    #[test]
    fn test_session_id_generation() {
        let counter = AtomicU64::new(0);
        let id1 = SessionId::generate("op_a", &counter);
        let id2 = SessionId::generate("op_a", &counter);
        let id3 = SessionId::generate("op_b", &counter);

        // Same operator, different counter → different IDs
        assert_ne!(id1, id2);
        // Different operator → different IDs (overwhelmingly likely)
        assert_ne!(id1, id3);
    }

    #[test]
    fn test_session_id_bytes_roundtrip() {
        let original = SessionId(0x1234_5678_9ABC_DEF0);
        let bytes = original.to_bytes();
        let restored = SessionId::from_bytes(&bytes);

        assert_eq!(restored, Some(original));
        assert_eq!(SessionId::from_bytes(&[0u8; 4]), None); // wrong length
    }

    // ---------------------------------------------------------------
    //  SessionMetadata tests
    // ---------------------------------------------------------------

    #[test]
    fn test_session_metadata_creation() {
        let id = SessionId(42);
        let meta = SessionMetadata::new(id, 1000, 5000);

        assert_eq!(meta.id, id);
        assert_eq!(meta.start, 1000);
        assert_eq!(meta.end, 6000); // start + gap
        assert!(!meta.emitted);
    }

    #[test]
    fn test_session_metadata_overlaps() {
        let id = SessionId(1);
        let meta = SessionMetadata::new(id, 1000, 5000); // [1000, 6000)

        // Event whose session [ts, ts+gap) overlaps [1000, 6000)
        assert!(meta.overlaps(1000, 5000)); // [1000, 6000) exact match
        assert!(meta.overlaps(3000, 5000)); // [3000, 8000) overlaps
        assert!(meta.overlaps(5999, 5000)); // [5999, 10999)
        assert!(meta.overlaps(500, 5000)); // [500, 5500) overlaps

        // Event at -1000 with gap 5000 → [-1000, 4000) overlaps [1000, 6000)
        assert!(meta.overlaps(-1000, 5000));

        // Truly before: [-5001, -1) does not reach 1000
        assert!(!meta.overlaps(-5001, 5000));

        // After: [6000, 11000) starts at session end → no overlap (half-open)
        assert!(!meta.overlaps(6000, 5000));
    }

    #[test]
    fn test_session_metadata_extend() {
        let id = SessionId(1);
        let mut meta = SessionMetadata::new(id, 1000, 5000); // [1000, 6000)

        // Extend forward
        meta.extend(8000, 5000);
        assert_eq!(meta.start, 1000);
        assert_eq!(meta.end, 13000); // 8000 + 5000

        // Extend backward
        meta.extend(500, 5000);
        assert_eq!(meta.start, 500);
        assert_eq!(meta.end, 13000); // Unchanged (max)
    }

    #[test]
    fn test_session_metadata_merge() {
        let mut s1 = SessionMetadata::new(SessionId(1), 1000, 5000); // [1000, 6000)
        let s2 = SessionMetadata::new(SessionId(2), 8000, 5000); // [8000, 13000)

        s1.merge(&s2);
        assert_eq!(s1.start, 1000);
        assert_eq!(s1.end, 13000);
    }

    #[test]
    fn test_session_metadata_window_id() {
        let meta = SessionMetadata::new(SessionId(1), 1000, 5000);
        let wid = meta.window_id();
        assert_eq!(wid.start, 1000);
        assert_eq!(wid.end, 6000);
    }

    // ---------------------------------------------------------------
    //  SessionIndex tests
    // ---------------------------------------------------------------

    #[test]
    fn test_session_index_sorted_insert() {
        let mut idx = SessionIndex::default();

        let s3 = SessionMetadata::new(SessionId(3), 3000, 1000);
        let s1 = SessionMetadata::new(SessionId(1), 1000, 1000);
        let s2 = SessionMetadata::new(SessionId(2), 2000, 1000);

        idx.insert(s3);
        idx.insert(s1);
        idx.insert(s2);

        assert_eq!(idx.len(), 3);
        assert_eq!(idx.sessions[0].start, 1000);
        assert_eq!(idx.sessions[1].start, 2000);
        assert_eq!(idx.sessions[2].start, 3000);
    }

    #[test]
    fn test_session_index_find_overlapping() {
        let mut idx = SessionIndex::default();

        // Two non-overlapping sessions: [100, 600) and [2000, 2500)
        idx.insert(SessionMetadata::new(SessionId(1), 100, 500));
        idx.insert(SessionMetadata::new(SessionId(2), 2000, 500));

        // Event at 300 with gap=500 → [300, 800) overlaps [100, 600)
        let hits = idx.find_overlapping(300, 500);
        assert_eq!(hits.as_slice(), &[SessionId(1)]);

        // Event at 1500 with gap=500 → [1500, 2000) doesn't overlap either
        let hits = idx.find_overlapping(1500, 500);
        assert!(hits.is_empty());

        // Event at 1800 with gap=500 → [1800, 2300) overlaps [2000, 2500)
        let hits = idx.find_overlapping(1800, 500);
        assert_eq!(hits.as_slice(), &[SessionId(2)]);
    }

    #[test]
    fn test_session_index_remove() {
        let mut idx = SessionIndex::default();
        idx.insert(SessionMetadata::new(SessionId(1), 100, 500));
        idx.insert(SessionMetadata::new(SessionId(2), 2000, 500));

        assert_eq!(idx.len(), 2);

        let removed = idx.remove(SessionId(1));
        assert!(removed.is_some());
        assert_eq!(removed.unwrap().start, 100);
        assert_eq!(idx.len(), 1);

        // Removing non-existent ID
        assert!(idx.remove(SessionId(99)).is_none());
    }

    #[test]
    fn test_session_index_get_and_get_mut() {
        let mut idx = SessionIndex::default();
        idx.insert(SessionMetadata::new(SessionId(1), 100, 500));

        assert!(idx.get(SessionId(1)).is_some());
        assert!(idx.get(SessionId(99)).is_none());

        idx.get_mut(SessionId(1)).unwrap().emitted = true;
        assert!(idx.get(SessionId(1)).unwrap().emitted);
    }

    // ---------------------------------------------------------------
    //  Operator creation
    // ---------------------------------------------------------------

    #[test]
    fn test_session_operator_creation() {
        let aggregator = CountAggregator::new();
        let operator = SessionWindowOperator::new(
            Duration::from_secs(30),
            aggregator,
            Duration::from_secs(60),
        );

        assert_eq!(operator.gap_ms(), 30_000);
        assert_eq!(operator.allowed_lateness_ms(), 60_000);
        assert_eq!(operator.active_session_count(), 0);
        assert_eq!(*operator.emit_strategy(), EmitStrategy::OnWatermark);
    }

    #[test]
    fn test_session_operator_with_id() {
        let aggregator = CountAggregator::new();
        let operator = SessionWindowOperator::with_id(
            Duration::from_secs(30),
            aggregator,
            Duration::from_secs(0),
            "test_session".to_string(),
        );

        assert_eq!(operator.operator_id, "test_session");
    }

    // ---------------------------------------------------------------
    //  Timer key roundtrip (now uses SessionId)
    // ---------------------------------------------------------------

    #[test]
    fn test_timer_key_roundtrip() {
        let sid = SessionId(0x1234_5678_9ABC_DEF0);
        let timer_key = SessionWindowOperator::<CountAggregator>::timer_key(sid);
        let parsed = SessionWindowOperator::<CountAggregator>::session_id_from_timer(&timer_key);
        assert_eq!(parsed, Some(sid));
    }

    #[test]
    fn test_timer_key_invalid() {
        // Wrong prefix
        let invalid1 = vec![0x02, 0, 0, 0, 0, 0, 0, 0, 0];
        assert!(
            SessionWindowOperator::<CountAggregator>::session_id_from_timer(&invalid1).is_none()
        );

        // Wrong length
        let invalid2 = vec![SESSION_TIMER_PREFIX, 0, 0, 0];
        assert!(
            SessionWindowOperator::<CountAggregator>::session_id_from_timer(&invalid2).is_none()
        );
    }

    // ---------------------------------------------------------------
    //  Single event
    // ---------------------------------------------------------------

    #[test]
    fn test_session_single_event() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        let event = create_test_event(500, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }

        assert_eq!(operator.active_session_count(), 1);
        assert_eq!(operator.pending_timers.len(), 1);
    }

    // ---------------------------------------------------------------
    //  Multiple events in the same session
    // ---------------------------------------------------------------

    #[test]
    fn test_session_multiple_events_same_session() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Events within gap (1000ms)
        for ts in [100, 500, 900, 1500] {
            let event = create_test_event(ts, 1);
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }

        // All events should be in the same session
        assert_eq!(operator.active_session_count(), 1);

        // Verify accumulator via the session index
        let key_hash = SessionWindowOperator::<CountAggregator>::key_hash(&[]);
        let index = operator.load_session_index(key_hash, &state);
        assert_eq!(index.len(), 1);
        let sid = index.sessions[0].id;
        let acc = operator.load_accumulator(sid, &state);
        assert_eq!(acc.result(), 4);
    }

    // ---------------------------------------------------------------
    //  Gap creates a new session (multi-session per key)
    // ---------------------------------------------------------------

    #[test]
    fn test_session_gap_creates_new_session() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::OnUpdate);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // First session
        let event1 = create_test_event(100, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }

        // Gap > 1000ms ⇒ creates a *second* session in the same index
        let event2 = create_test_event(3000, 1);
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx)
        };

        // Two sessions now exist for the same key
        assert_eq!(operator.active_session_count(), 2);

        // OnUpdate emits for the new session
        let event_count = outputs
            .iter()
            .filter(|o| matches!(o, Output::Event(_)))
            .count();
        assert_eq!(event_count, 1);
    }

    // ---------------------------------------------------------------
    //  Timer triggers emission
    // ---------------------------------------------------------------

    #[test]
    fn test_session_timer_triggers_emission() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create session
        let event = create_test_event(500, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }

        // Get the pending timer (session_id, timer_time, _key_hash)
        let (sid, timer_time, _kh) = first_pending_timer(&operator);

        // Fire timer
        let timer = Timer {
            key: SessionWindowOperator::<CountAggregator>::timer_key(sid),
            timestamp: timer_time,
        };

        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(timer, &mut ctx)
        };

        assert_eq!(outputs.len(), 1);
        match &outputs[0] {
            Output::Event(e) => {
                assert_eq!(e.timestamp, 1500); // 500 + gap (1000)
                let result = e
                    .data
                    .column(2)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap()
                    .value(0);
                assert_eq!(result, 1);
            }
            _ => panic!("Expected Event output"),
        }

        // Session should be cleaned up
        assert_eq!(operator.active_session_count(), 0);
    }

    // ---------------------------------------------------------------
    //  Keyed session tracking
    // ---------------------------------------------------------------

    #[test]
    fn test_session_keyed_tracking() {
        let aggregator = SumAggregator::new(1); // Sum column 1 (value)
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );
        operator.set_key_column(0); // Key by column 0

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Events for key 1
        let event1 = create_keyed_event(100, 1, 10);
        let event2 = create_keyed_event(500, 1, 20);

        // Events for key 2
        let event3 = create_keyed_event(200, 2, 100);
        let event4 = create_keyed_event(600, 2, 200);

        for event in [event1, event2, event3, event4] {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }

        // Should have 2 active sessions (one per key)
        assert_eq!(operator.active_session_count(), 2);
    }

    // ---------------------------------------------------------------
    //  Late events
    // ---------------------------------------------------------------

    #[test]
    fn test_session_late_event_dropped() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(0);

        // Advance watermark far ahead
        let event1 = create_test_event(10000, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }

        // Process late event
        let late_event = create_test_event(100, 1);
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&late_event, &mut ctx)
        };

        let is_late = outputs.iter().any(|o| matches!(o, Output::LateEvent(_)));
        assert!(is_late);
        assert_eq!(operator.late_data_metrics().late_events_dropped(), 1);
    }

    #[test]
    fn test_session_late_event_side_output() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );
        operator.set_late_data_config(LateDataConfig::with_side_output("late".to_string()));

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(0);

        // Advance watermark
        let event1 = create_test_event(10000, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }

        // Process late event
        let late_event = create_test_event(100, 1);
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&late_event, &mut ctx)
        };

        let side_output = outputs.iter().find_map(|o| {
            if let Output::SideOutput(data) = o {
                Some(data.name.clone())
            } else {
                None
            }
        });
        assert_eq!(side_output.as_deref(), Some("late"));
        assert_eq!(operator.late_data_metrics().late_events_side_output(), 1);
    }

    // ---------------------------------------------------------------
    //  Emit strategies
    // ---------------------------------------------------------------

    #[test]
    fn test_session_emit_on_update() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::OnUpdate);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        let event = create_test_event(500, 1);
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx)
        };

        let event_count = outputs
            .iter()
            .filter(|o| matches!(o, Output::Event(_)))
            .count();
        assert_eq!(event_count, 1);
    }

    #[test]
    fn test_session_emit_changelog() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::Changelog);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        let event = create_test_event(500, 1);
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx)
        };

        let changelog_count = outputs
            .iter()
            .filter(|o| matches!(o, Output::Changelog(_)))
            .count();
        assert_eq!(changelog_count, 1);
    }

    #[test]
    fn test_session_emit_final_drops_late() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::Final);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(0);

        // Advance watermark
        let event1 = create_test_event(10000, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }

        // Process late event - should be silently dropped
        let late_event = create_test_event(100, 1);
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&late_event, &mut ctx)
        };

        assert!(outputs.is_empty());
        assert_eq!(operator.late_data_metrics().late_events_dropped(), 1);
    }

    // ---------------------------------------------------------------
    //  Checkpoint / restore
    // ---------------------------------------------------------------

    #[test]
    fn test_session_checkpoint_restore() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator.clone(),
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create some sessions
        for ts in [100, 500] {
            let event = create_test_event(ts, 1);
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }

        let counter_before = operator.session_id_counter.load(Ordering::Relaxed);

        // Checkpoint
        let checkpoint = operator.checkpoint();

        // Create new operator and restore
        let mut restored = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );
        restored.restore(checkpoint).unwrap();

        // Pending timers and counter should be restored
        assert_eq!(restored.pending_timers.len(), 1);
        assert_eq!(
            restored.session_id_counter.load(Ordering::Relaxed),
            counter_before,
        );
    }

    // ---------------------------------------------------------------
    //  Stale timer
    // ---------------------------------------------------------------

    #[test]
    fn test_session_stale_timer_ignored() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create session
        let event1 = create_test_event(500, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }

        let (sid, old_timer_time, _kh) = first_pending_timer(&operator);

        // Extend session — timer time changes
        let event2 = create_test_event(1200, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx);
        }

        // Fire the *old* timer (stale)
        let stale_timer = Timer {
            key: SessionWindowOperator::<CountAggregator>::timer_key(sid),
            timestamp: old_timer_time,
        };

        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(stale_timer, &mut ctx)
        };

        assert!(outputs.is_empty());
        assert_eq!(operator.active_session_count(), 1);
    }

    // ---------------------------------------------------------------
    //  Sum aggregation end-to-end
    // ---------------------------------------------------------------

    #[test]
    fn test_session_sum_aggregation() {
        let aggregator = SumAggregator::new(0);
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        for (ts, value) in [(100, 10), (500, 20), (800, 30)] {
            let event = create_test_event(ts, value);
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }

        // Fire timer
        let (sid, timer_time, _kh) = first_pending_timer(&operator);
        let timer = Timer {
            key: SessionWindowOperator::<SumAggregator>::timer_key(sid),
            timestamp: timer_time,
        };

        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(timer, &mut ctx)
        };

        match &outputs[0] {
            Output::Event(e) => {
                let result = e
                    .data
                    .column(2)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap()
                    .value(0);
                assert_eq!(result, 60); // 10 + 20 + 30
            }
            _ => panic!("Expected Event output"),
        }
    }

    // ---------------------------------------------------------------
    //  Output schema
    // ---------------------------------------------------------------

    #[test]
    fn test_session_output_schema() {
        use arrow_schema::Schema;
        let schema = Schema::new(vec![
            arrow_schema::Field::new("window_start", arrow_schema::DataType::Int64, false),
            arrow_schema::Field::new("window_end", arrow_schema::DataType::Int64, false),
            arrow_schema::Field::new("result", arrow_schema::DataType::Int64, false),
        ]);

        assert_eq!(schema.fields().len(), 3);
        assert_eq!(schema.field(0).name(), "window_start");
        assert_eq!(schema.field(1).name(), "window_end");
        assert_eq!(schema.field(2).name(), "result");
    }

    // ---------------------------------------------------------------
    //  NEW: Multi-session per key
    // ---------------------------------------------------------------

    #[test]
    fn test_multi_session_per_key() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // First event creates session at [100, 1100)
        let event1 = create_test_event(100, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 1);

        // Second event at 5000 — gap > 1000 ⇒ new session at [5000, 6000)
        let event2 = create_test_event(5000, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 2);

        // Both sessions exist in the same index
        let key_hash = SessionWindowOperator::<CountAggregator>::key_hash(&[]);
        let index = operator.load_session_index(key_hash, &state);
        assert_eq!(index.len(), 2);
        assert_eq!(index.sessions[0].start, 100);
        assert_eq!(index.sessions[1].start, 5000);
    }

    // ---------------------------------------------------------------
    //  NEW: Session persists across batch boundaries
    // ---------------------------------------------------------------

    #[test]
    fn test_session_persists_across_batch_boundary() {
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_op".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Batch 1: create session at [500, 1500)
        let event1 = create_test_event(500, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 1);

        // Simulate batch boundary: clear in-memory cache
        operator.session_indices.clear();
        assert_eq!(operator.active_session_count(), 0); // cache cleared

        // Batch 2: event at 800 should reload from state and extend
        let event2 = create_test_event(800, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx);
        }

        // Session should have been loaded from state store and extended
        assert_eq!(operator.active_session_count(), 1);

        let key_hash = SessionWindowOperator::<CountAggregator>::key_hash(&[]);
        let index = operator.load_session_index(key_hash, &state);
        assert_eq!(index.len(), 1);
        // start stays 500, end extends to max(1500, 800+1000=1800)
        assert_eq!(index.sessions[0].start, 500);
        assert_eq!(index.sessions[0].end, 1800);

        // Accumulator should reflect both events
        let sid = index.sessions[0].id;
        let acc = operator.load_accumulator(sid, &state);
        assert_eq!(acc.result(), 2);
    }

    // ---------------------------------------------------------------
    //  Session merging tests
    // ---------------------------------------------------------------

    #[test]
    fn test_two_way_session_merge() {
        // gap=1000. Create 2 sessions: [100, 1100) and [1200, 2200).
        // Bridge event at 1050 → potential [1050, 2050) overlaps both → merge.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_merge".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create first session [100, 1100)
        let event1 = create_test_event(100, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }

        // Create second session [1200, 2200) (gap of 100ms between sessions)
        let event2 = create_test_event(1200, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 2);

        // Bridge at 1050 → [1050, 2050) overlaps [100,1100) and [1200,2200)
        let bridge = create_test_event(1050, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&bridge, &mut ctx);
        }

        // Should have merged into a single session
        assert_eq!(operator.active_session_count(), 1);

        let key_hash = SessionWindowOperator::<CountAggregator>::key_hash(&[]);
        let index = operator.load_session_index(key_hash, &state);
        assert_eq!(index.len(), 1);

        // Merged bounds: start=100, end=max(1100, 2200, 2050)=2200
        assert_eq!(index.sessions[0].start, 100);
        assert_eq!(index.sessions[0].end, 2200);

        // Accumulator: 3 events total
        let sid = index.sessions[0].id;
        let acc = operator.load_accumulator(sid, &state);
        assert_eq!(acc.result(), 3);
    }

    #[test]
    fn test_three_way_session_merge() {
        // gap=500. Create 3 sessions with small gaps between them.
        // Use large allowed_lateness to prevent late data dropping on bridge events.
        // Session 1: [100, 600), Session 2: [700, 1200), Session 3: [1300, 1800)
        // Bridge at 550 → [550, 1050) overlaps sessions 1 and 2 → merge to [100, 1200)
        // Bridge at 1150 → [1150, 1650) overlaps [100, 1200) and [1300, 1800) → merge all
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(500),
            aggregator,
            Duration::from_millis(5000),
            "test_merge3".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create 3 sessions with 100ms gaps between them
        for ts in [100, 700, 1300] {
            let event = create_test_event(ts, 1);
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 3);

        // Bridge at 550 → [550, 1050) overlaps [100,600) and [700,1200)
        let bridge1 = create_test_event(550, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&bridge1, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 2);

        // Check intermediate state
        let key_hash = SessionWindowOperator::<CountAggregator>::key_hash(&[]);
        let index = operator.load_session_index(key_hash, &state);
        assert_eq!(index.len(), 2);

        // Bridge at 1150 → [1150, 1650) overlaps merged [100,1200) and [1300,1800)
        let bridge2 = create_test_event(1150, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&bridge2, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 1);

        let index = operator.load_session_index(key_hash, &state);
        assert_eq!(index.len(), 1);

        // Merged bounds: [100, 1800)
        assert_eq!(index.sessions[0].start, 100);
        assert_eq!(index.sessions[0].end, 1800);

        // 5 events total
        let sid = index.sessions[0].id;
        let acc = operator.load_accumulator(sid, &state);
        assert_eq!(acc.result(), 5);
    }

    #[test]
    fn test_merge_emits_changelog_retractions() {
        // Use Changelog strategy. Create 2 sessions (2 inserts emitted).
        // Merge them → should emit deletes for both old states + insert for merged.
        // gap=1000, sessions [100, 1100) and [1200, 2200) — 100ms gap between them.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_cl_merge".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::Changelog);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create first session [100, 1100) — emits Changelog insert
        let event1 = create_test_event(100, 1);
        let out1 = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx)
        };
        let insert_count_1 = out1
            .iter()
            .filter(|o| matches!(o, Output::Changelog(r) if r.weight == 1))
            .count();
        assert_eq!(insert_count_1, 1);

        // Create second session [1200, 2200) — emits Changelog insert
        let event2 = create_test_event(1200, 1);
        let out2 = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx)
        };
        let insert_count_2 = out2
            .iter()
            .filter(|o| matches!(o, Output::Changelog(r) if r.weight == 1))
            .count();
        assert_eq!(insert_count_2, 1);
        assert_eq!(operator.active_session_count(), 2);

        // Bridge at 1050 → [1050, 2050) overlaps both sessions → merge
        let bridge = create_test_event(1050, 1);
        let out3 = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&bridge, &mut ctx)
        };

        // Expect: 2 deletes (one for winner's old state, one for loser) + 1 insert (merged)
        let deletes: Vec<_> = out3
            .iter()
            .filter(|o| matches!(o, Output::Changelog(r) if r.weight == -1))
            .collect();
        let inserts: Vec<_> = out3
            .iter()
            .filter(|o| matches!(o, Output::Changelog(r) if r.weight == 1))
            .collect();

        assert_eq!(deletes.len(), 2, "Expected 2 delete retractions");
        assert_eq!(inserts.len(), 1, "Expected 1 insert for merged result");

        // Session count should be 1 after merge
        assert_eq!(operator.active_session_count(), 1);
    }

    #[test]
    fn test_merge_accumulator_correctness_sum() {
        // Use SumAggregator. Session A: events 10, 20 (sum=30).
        // Session B: event 50 (sum=50). Bridge with value 5.
        // gap=500 to keep sessions tight.
        // Session A: events at 100 and 300 → extends to [100, 800).
        // Session B: event at 900 → [900, 1400). Gap of 100ms.
        // Bridge at 750 → [750, 1250) overlaps both → merge.
        // Merged sum = 10 + 20 + 50 + 5 = 85.
        let aggregator = SumAggregator::new(0);
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(500),
            aggregator,
            Duration::from_millis(0),
            "test_sum_merge".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Session A: [100, 800) with sum=30
        for (ts, val) in [(100, 10), (300, 20)] {
            let event = create_test_event(ts, val);
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }

        // Session B: [900, 1400) with sum=50
        let event_b = create_test_event(900, 50);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event_b, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 2);

        // Bridge at 750 with value 5 → [750, 1250) overlaps both → merge
        let bridge = create_test_event(750, 5);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&bridge, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 1);

        // Verify merged sum
        let key_hash = SessionWindowOperator::<SumAggregator>::key_hash(&[]);
        let index = operator.load_session_index(key_hash, &state);
        let sid = index.sessions[0].id;
        let acc = operator.load_accumulator(sid, &state);
        assert_eq!(acc.result(), 85); // 10 + 20 + 50 + 5
    }

    #[test]
    fn test_merge_cleans_up_loser_timers() {
        // Create 2 sessions with pending timers. Merge them.
        // gap=1000, sessions [100, 1100) and [1200, 2200). Bridge at 1050.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_timer_merge".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create first session [100, 1100)
        let event1 = create_test_event(100, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }

        // Create second session [1200, 2200)
        let event2 = create_test_event(1200, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx);
        }
        assert_eq!(operator.pending_timers.len(), 2);

        // Bridge at 1050 → [1050, 2050) merges both sessions
        let bridge = create_test_event(1050, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&bridge, &mut ctx);
        }

        // Only 1 timer should remain (for the winner)
        assert_eq!(operator.pending_timers.len(), 1);
        assert_eq!(operator.active_session_count(), 1);
    }

    #[test]
    fn test_merge_no_retractions_for_on_watermark() {
        // OnWatermark strategy should NOT emit retractions on merge
        // (no intermediate output was emitted).
        // gap=1000, sessions [100, 1100) and [1200, 2200). Bridge at 1050.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_wm_merge".to_string(),
        );
        // Default strategy is OnWatermark

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create 2 sessions
        let event1 = create_test_event(100, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx);
        }
        let event2 = create_test_event(1200, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx);
        }

        // Bridge at 1050 → merge. No changelog outputs expected.
        let bridge = create_test_event(1050, 1);
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&bridge, &mut ctx)
        };

        let changelog_count = outputs
            .iter()
            .filter(|o| matches!(o, Output::Changelog(_)))
            .count();
        assert_eq!(changelog_count, 0);
        assert_eq!(operator.active_session_count(), 1);
    }

    // ---------------------------------------------------------------
    //  Emit strategy verification tests
    // ---------------------------------------------------------------

    #[test]
    fn test_on_watermark_no_intermediate_emits_on_timer() {
        // OnWatermark (default): process() should NOT emit Event/Changelog,
        // only on_timer() should emit.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_wm_emit".to_string(),
        );
        // Default strategy is OnWatermark

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Process 2 events in the same session
        let event1 = create_test_event(500, 1);
        let out1 = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx)
        };

        let event2 = create_test_event(800, 1);
        let out2 = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx)
        };

        // No Event or Changelog output from process() (only Watermark allowed)
        let event_or_changelog = |o: &&Output| matches!(o, Output::Event(_) | Output::Changelog(_));
        assert_eq!(out1.iter().filter(event_or_changelog).count(), 0);
        assert_eq!(out2.iter().filter(event_or_changelog).count(), 0);

        // Fire timer → should emit Event with count=2
        let (sid, timer_time, _kh) = first_pending_timer(&operator);
        let timer = Timer {
            key: SessionWindowOperator::<CountAggregator>::timer_key(sid),
            timestamp: timer_time,
        };
        let timer_out = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(timer, &mut ctx)
        };

        assert_eq!(timer_out.len(), 1);
        match &timer_out[0] {
            Output::Event(e) => {
                let result = e
                    .data
                    .column(2)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap()
                    .value(0);
                assert_eq!(result, 2);
            }
            _ => panic!("Expected Event output from on_timer"),
        }
    }

    #[test]
    fn test_on_window_close_no_intermediate_emits_on_timer() {
        // OnWindowClose: process() should NOT emit Event/Changelog,
        // only on_timer() should emit.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_owc_emit".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::OnWindowClose);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Process event
        let event = create_test_event(500, 1);
        let out = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx)
        };

        // No Event or Changelog from process()
        let event_or_changelog = |o: &&Output| matches!(o, Output::Event(_) | Output::Changelog(_));
        assert_eq!(out.iter().filter(event_or_changelog).count(), 0);

        // Fire timer → should emit Event
        let (sid, timer_time, _kh) = first_pending_timer(&operator);
        let timer = Timer {
            key: SessionWindowOperator::<CountAggregator>::timer_key(sid),
            timestamp: timer_time,
        };
        let timer_out = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(timer, &mut ctx)
        };

        assert_eq!(timer_out.len(), 1);
        assert!(matches!(&timer_out[0], Output::Event(_)));
    }

    #[test]
    fn test_changelog_timer_emits_changelog_record() {
        // Changelog strategy: on_timer() should emit Output::Changelog
        // with weight == 1 (insert).
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_cl_timer".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::Changelog);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Process event (emits intermediate Changelog insert)
        let event = create_test_event(500, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }

        // Fire timer → should emit Changelog insert
        let (sid, timer_time, _kh) = first_pending_timer(&operator);
        let timer = Timer {
            key: SessionWindowOperator::<CountAggregator>::timer_key(sid),
            timestamp: timer_time,
        };
        let timer_out = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(timer, &mut ctx)
        };

        assert_eq!(timer_out.len(), 1);
        match &timer_out[0] {
            Output::Changelog(record) => {
                assert_eq!(record.weight, 1, "Expected insert weight");
            }
            _ => panic!("Expected Changelog output from on_timer"),
        }
    }

    #[test]
    fn test_on_update_emits_updated_count() {
        // OnUpdate: each process() call should emit Event with the latest count.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_ou_emit".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::OnUpdate);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // First event → count=1
        let event1 = create_test_event(500, 1);
        let out1 = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event1, &mut ctx)
        };

        let emitted1: Vec<_> = out1
            .iter()
            .filter_map(|o| match o {
                Output::Event(e) => Some(e),
                _ => None,
            })
            .collect();
        assert_eq!(emitted1.len(), 1);
        let result1 = emitted1[0]
            .data
            .column(2)
            .as_any()
            .downcast_ref::<Int64Array>()
            .unwrap()
            .value(0);
        assert_eq!(result1, 1);

        // Second event in same session → count=2
        let event2 = create_test_event(800, 1);
        let out2 = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event2, &mut ctx)
        };

        let emitted2: Vec<_> = out2
            .iter()
            .filter_map(|o| match o {
                Output::Event(e) => Some(e),
                _ => None,
            })
            .collect();
        assert_eq!(emitted2.len(), 1);
        let result2 = emitted2[0]
            .data
            .column(2)
            .as_any()
            .downcast_ref::<Int64Array>()
            .unwrap()
            .value(0);
        assert_eq!(result2, 2);
    }

    // ---------------------------------------------------------------
    //  Timer persistence after checkpoint/restore
    // ---------------------------------------------------------------

    #[test]
    fn test_timer_reregistration_after_restore() {
        // Create session, checkpoint, restore to fresh operator + fresh TimerService.
        // Verify needs_timer_reregistration is true.
        // Process a new event (triggers lazy re-registration).
        // Verify TimerService has both the restored timer and the new timer.
        // Poll timers at restored timer time → verify timer fires.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator.clone(),
            Duration::from_millis(0),
            "test_restore_timers".to_string(),
        );

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Create a session at t=500 → timer at 500+1000+0 = 1500
        let event = create_test_event(500, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&event, &mut ctx);
        }
        assert_eq!(operator.pending_timers.len(), 1);
        let (_, original_timer_time, _) = first_pending_timer(&operator);
        assert_eq!(original_timer_time, 1500);

        // Checkpoint
        let checkpoint = operator.checkpoint();

        // Create fresh operator and fresh TimerService (simulating restart)
        let mut restored = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "test_restore_timers".to_string(),
        );
        let mut fresh_timers = TimerService::new();
        assert_eq!(fresh_timers.pending_count(), 0);

        // Restore
        restored.restore(checkpoint).unwrap();
        assert!(restored.needs_timer_reregistration);
        assert_eq!(restored.pending_timers.len(), 1);

        // Fresh TimerService still has no timers (not re-registered yet)
        assert_eq!(fresh_timers.pending_count(), 0);

        // Process a new event → triggers lazy re-registration + creates new session
        let event2 = create_test_event(5000, 1);
        {
            let mut ctx = create_test_context(&mut fresh_timers, &mut state, &mut watermark_gen);
            restored.process(&event2, &mut ctx);
        }

        // TimerService should now have at least 2 timers:
        // restored timer (1500) + new timer (5000+1000=6000)
        assert!(
            fresh_timers.pending_count() >= 2,
            "Expected at least 2 timers, got {}",
            fresh_timers.pending_count()
        );
        assert!(!restored.needs_timer_reregistration);

        // Poll timers at the restored timer time → should fire
        let fired = fresh_timers.poll_timers(1500);
        assert_eq!(fired.len(), 1);
        assert_eq!(fired[0].timestamp, 1500);

        // Use the fired timer to invoke on_timer
        let timer = Timer {
            key: fired[0].key.clone().unwrap(),
            timestamp: fired[0].timestamp,
        };
        let timer_out = {
            let mut ctx = create_test_context(&mut fresh_timers, &mut state, &mut watermark_gen);
            restored.on_timer(timer, &mut ctx)
        };

        // Should emit the restored session's result (count=1 from the original event)
        assert_eq!(timer_out.len(), 1);
        match &timer_out[0] {
            Output::Event(e) => {
                let result = e
                    .data
                    .column(2)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap()
                    .value(0);
                assert_eq!(result, 1);
            }
            _ => panic!("Expected Event output from restored timer"),
        }
    }

    // ========================================================================
    // EMIT ON WINDOW CLOSE (EOWC) — Session Window Tests (Issue #52)
    // ========================================================================

    #[test]
    fn test_eowc_session_basic_close() {
        // Create session with 30s gap, send events at t=0, 10000, 20000.
        // Session = [0, 50000). Fire timer. Verify single emission.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(30_000), // 30s gap
            aggregator,
            Duration::from_millis(0),
            "eowc_session".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::OnWindowClose);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(0);

        // Send 3 events within the same session
        for ts in [0, 10_000, 20_000] {
            let event = create_test_event(ts, 1);
            let outputs = {
                let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
                operator.process(&event, &mut ctx)
            };
            // No intermediate Event emissions
            let event_count = outputs
                .iter()
                .filter(|o| matches!(o, Output::Event(_)))
                .count();
            assert_eq!(
                event_count, 0,
                "EOWC session should not emit intermediate results"
            );
        }

        // Session should be [0, 50000) (last event at 20000 + gap 30000)
        assert_eq!(operator.active_session_count(), 1);

        // Get the pending timer and fire it
        let (sid, timer_time, _kh) = first_pending_timer(&operator);
        let timer = Timer {
            key: SessionWindowOperator::<CountAggregator>::timer_key(sid),
            timestamp: timer_time,
        };
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(timer, &mut ctx)
        };

        assert_eq!(outputs.len(), 1, "Should emit exactly once");
        match &outputs[0] {
            Output::Event(e) => {
                let result = e
                    .data
                    .column(2)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap();
                assert_eq!(result.value(0), 3, "Session count should be 3");
            }
            other => panic!("Expected Output::Event, got: {other:?}"),
        }

        // Session should be cleaned up
        assert_eq!(operator.active_session_count(), 0);
    }

    #[test]
    fn test_eowc_session_no_intermediate_on_extend() {
        // Extending a session with new events should not produce
        // intermediate emissions with OnWindowClose.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "eowc_extend".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::OnWindowClose);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Repeatedly extend the session
        for ts in (0..5).map(|i| i * 500) {
            let event = create_test_event(ts, 1);
            let outputs = {
                let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
                operator.process(&event, &mut ctx)
            };
            for output in &outputs {
                assert!(
                    !matches!(output, Output::Event(_)),
                    "No intermediate Event on session extend (ts={ts})"
                );
                assert!(
                    !matches!(output, Output::Changelog(_)),
                    "No intermediate Changelog on session extend (ts={ts})"
                );
            }
        }

        // Still 1 session, not yet emitted
        assert_eq!(operator.active_session_count(), 1);
    }

    #[test]
    fn test_eowc_session_merge_before_close() {
        // Create two separate sessions, then bridge them.
        // Merged session should emit once with combined count.
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000), // 1s gap
            aggregator,
            Duration::from_millis(0),
            "eowc_merge".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::OnWindowClose);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(100);

        // Session 1: event at t=100 → session [100, 1100)
        let e1 = create_test_event(100, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&e1, &mut ctx);
        }

        // Session 2: event at t=1200 → session [1200, 2200)
        let e2 = create_test_event(1200, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&e2, &mut ctx);
        }
        assert_eq!(
            operator.active_session_count(),
            2,
            "Should have 2 sessions before merge"
        );

        // Bridge: event at t=1050 → session [1050, 2050) overlaps both
        let bridge = create_test_event(1050, 1);
        let bridge_outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&bridge, &mut ctx)
        };

        // No intermediate Event from merge with OnWindowClose
        let event_count = bridge_outputs
            .iter()
            .filter(|o| matches!(o, Output::Event(_)))
            .count();
        assert_eq!(
            event_count, 0,
            "Merge should not emit intermediate results with EOWC"
        );
        assert_eq!(
            operator.active_session_count(),
            1,
            "Should have 1 merged session"
        );

        // Fire the merged session's timer
        let (sid, timer_time, _kh) = first_pending_timer(&operator);
        let timer = Timer {
            key: SessionWindowOperator::<CountAggregator>::timer_key(sid),
            timestamp: timer_time,
        };
        let outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(timer, &mut ctx)
        };

        assert_eq!(outputs.len(), 1, "Merged session should emit once");
        if let Output::Event(e) = &outputs[0] {
            let result = e
                .data
                .column(2)
                .as_any()
                .downcast_ref::<Int64Array>()
                .unwrap();
            assert_eq!(
                result.value(0),
                3,
                "Merged session count should be 3 (1+1+1)"
            );
        } else {
            panic!("Expected Event output for merged session");
        }
    }

    #[test]
    fn test_eowc_session_late_data_after_close() {
        // Close a session, then send late event for the same time range.
        // Late event should be dropped (or routed to side output).
        let aggregator = CountAggregator::new();
        let mut operator = SessionWindowOperator::with_id(
            Duration::from_millis(1000),
            aggregator,
            Duration::from_millis(0),
            "eowc_late".to_string(),
        );
        operator.set_emit_strategy(EmitStrategy::OnWindowClose);

        let mut timers = TimerService::new();
        let mut state = InMemoryStore::new();
        let mut watermark_gen = BoundedOutOfOrdernessGenerator::new(0);

        // Create session at t=500 → session [500, 1500)
        let e1 = create_test_event(500, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&e1, &mut ctx);
        }

        // Fire timer to close the session
        let (sid, timer_time, _kh) = first_pending_timer(&operator);
        let timer = Timer {
            key: SessionWindowOperator::<CountAggregator>::timer_key(sid),
            timestamp: timer_time,
        };
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.on_timer(timer, &mut ctx);
        }
        assert_eq!(operator.active_session_count(), 0);

        // Advance watermark past the session close
        let advance = create_test_event(5000, 1);
        {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&advance, &mut ctx);
        }

        // Send late event for the closed session's time range
        let late = create_test_event(600, 99);
        let late_outputs = {
            let mut ctx = create_test_context(&mut timers, &mut state, &mut watermark_gen);
            operator.process(&late, &mut ctx)
        };

        // Should be classified as late
        let is_late = late_outputs
            .iter()
            .any(|o| matches!(o, Output::LateEvent(_)));
        assert!(is_late, "Late event after session close should be detected");

        // Should NOT produce a new Event
        let is_event = late_outputs.iter().any(|o| matches!(o, Output::Event(_)));
        assert!(!is_event, "Late event should not re-open closed session");

        assert_eq!(operator.late_data_metrics().late_events_dropped(), 1);
    }
}