fv-streams-engine 0.6.0

//! THE STATEFUL OPERATORS AS TASKS (Phase 2c): the window, session, rank and ring operators
//! and the interval join behind the runtime's `Operator` trait, with the exact semantics one
//! consumer thread gives them today.
//!
//! - Every row carries its **origin** (`__fv_partition`: the Kafka partition it came from, or
//!   the task that produced it). An operator pushes each origin's rows separately, so a
//!   co-partitioned shard's own watermark — the minimum across origins of each origin's max event
//!   time, minus the allowed lateness, with bounded idleness — is what a consumer thread computes
//!   across its assigned partitions.
//! - **Sharding:** a keyBy stage holds one instance per task over its vnode range (the shuffle
//!   already sends a key's rows here); a stage that assumes co-partitioned input holds one
//!   instance per origin, as one consumer thread holds one per assigned partition.
//! - **A keyBy shard's time comes in band.** Behind a shuffle, a shard cannot derive time from
//!   the rows it receives: when one upstream task runs ahead, the maximum over the rows seen so
//!   far fires windows the slower upstream still has rows for (the crash e2e lost two-thirds of
//!   its input this way), and waiting on every upstream partition instead stalls a shard whose
//!   keys never touch one of them. So an event-time keyBy shard fires on the runtime's
//!   watermark — each source's minimum over its partitions, the minimum across every input edge,
//!   delivered behind the rows it covers — less the stage's allowed lateness, and the rows' own
//!   times never move it. A co-partitioned (per-origin) shard keeps deriving its partition's
//!   time from its rows, as the consumer thread does; so does a processing-time shard.
//! - Fired rows leave with their deterministic emission key in `__fv_key` (computed before the
//!   post steps, as the chain does) and this task's id as their origin.
//! - `on_barrier` returns the task's checkpoint object: every shard's own snapshot, framed;
//!   `with_state` restores it (Phase 2d).

use super::*;
use crate::dataflow::{EdgeId, OpSnapshot, Operator, Out, StateFile, TaskId};
use std::collections::{HashMap, HashSet};

/// How a stateful stage shards its state across a task.
#[derive(Clone, Copy)]
pub(super) enum Sharding {
    /// One instance per origin partition (a stage that assumes its input co-partitioned by the key).
    PerOrigin,
    /// One instance for the task's whole vnode range (a keyBy stage).
    Range,
}

/// Built once per shard; `Sync` so one factory can be shared across a stage's tasks.
pub(super) type OperatorFactory = Box<dyn Fn() -> Box<dyn fv_streams_ops::WindowOperator + Send> + Send + Sync>;

/// A window / session / topN / lastN stage as a task.
pub(super) struct StatefulOp {
    task: TaskId,
    mk: OperatorFactory,
    sharding: Sharding,
    tsrc: TimeSource,
    shape: EmitShape,
    group_by: Vec<String>,
    shards: HashMap<i32, Box<dyn fv_streams_ops::WindowOperator + Send>>,
    /// The origins each shard was told to expect (a range instance learns them as they appear).
    expected: HashMap<i32, Vec<i32>>,
    timing: Timing,
    /// The last watermark the runtime delivered (in-band time only).
    watermark: Option<i64>,
    post: crate::steps::BatchSteps,
    raw_post: Arc<Vec<fv_plan::inline::Step>>,
    dropped: Arc<AtomicU64>,
    /// File-backed checkpoints (a small head naming one immutable file per shard, so the barrier
    /// never materialises the whole state) — on the object store, and only for shard operators that
    /// support it (a keyBy window under a memory limit). `None` dir = keep the whole snapshot.
    checkpoint_files: bool,
    ckpt_dir: Option<std::path::PathBuf>,
}

/// How a stateful task takes and passes on time.
#[derive(Clone, Copy, Debug)]
pub(super) struct Timing {
    /// `true`: fire on the runtime's watermark (an event-time keyBy stage whose every input is
    /// stamped); `false`: on the time the rows carry.
    pub(super) in_band: bool,
    /// How far behind its input the task's own watermark runs: a window's size, so a stage
    /// downstream windowing on `windowStart` never sees a fired window as late (a window
    /// starting at `s` fires when time reaches `s + size`); zero for the other operators.
    pub(super) emit_hold_ms: i64,
}

/// What drives a firing: the shards' own time at a processing-time instant, the runtime's
/// watermark, or the end of the stream.
#[derive(Clone, Copy, Debug)]
enum Firing {
    Now(i64),
    Watermark(i64),
    Flush,
}

/// What one shard fired: a batch (the columnar operators) or keyed rows (the row seam).
enum Fired {
    Batch(arrow::array::RecordBatch),
    Rows(Vec<(String, fv_plan::row::Row)>),
}

impl StatefulOp {
    #[allow(clippy::too_many_arguments)] // one constructor, every field named at the call site
    pub(super) fn new(
        task: TaskId,
        mk: OperatorFactory,
        sharding: Sharding,
        tsrc: TimeSource,
        shape: EmitShape,
        group_by: Vec<String>,
        timing: Timing,
        post: Arc<Vec<fv_plan::inline::Step>>,
        dropped: Arc<AtomicU64>,
    ) -> Self {
        StatefulOp {
            task,
            mk,
            sharding,
            tsrc,
            shape,
            group_by,
            shards: HashMap::new(),
            expected: HashMap::new(),
            timing,
            watermark: None,
            post: crate::steps::BatchSteps::new(post.as_ref().clone()),
            raw_post: post,
            dropped,
            checkpoint_files: false,
            ckpt_dir: None,
        }
    }

    /// Enable file-backed checkpoints (a head plus one immutable file per shard) rooted at `dir`,
    /// on the object store. A no-op unless the operator's shards support it; the whole-snapshot path
    /// stays for the topic/memory backends and for operators without shard files.
    pub(super) fn with_checkpoint_files(mut self, on: bool, dir: Option<std::path::PathBuf>) -> Self {
        self.checkpoint_files = on;
        self.ckpt_dir = dir;
        self
    }

    /// Fire what the task's time allows: the runtime's watermark for an in-band task (nothing
    /// before the first one arrives), else the shards' own time at `now`.
    fn fire_due(&mut self, now: i64, out: &mut Out) {
        if self.timing.in_band {
            if let Some(wm) = self.watermark {
                self.fire(Firing::Watermark(wm), out);
            }
        } else {
            self.fire(Firing::Now(now), out);
        }
    }

    /// Fire every shard (in id order, so emission order is stable), then route what fired.
    fn fire(&mut self, firing: Firing, out: &mut Out) {
        let mut ids: Vec<i32> = self.shards.keys().copied().collect();
        ids.sort_unstable();
        let mut fired: Vec<Fired> = Vec::new();
        for id in ids {
            let op = self.shards.get_mut(&id).expect("shard");
            let columnar = match firing {
                Firing::Flush => op.flush_batch(),
                Firing::Now(now) => op.advance_batch(now),
                Firing::Watermark(wm) => op.advance_batch_to(wm),
            };
            match columnar {
                Some(Some(b)) => fired.push(Fired::Batch(b)),
                Some(None) => {}
                None => {
                    let rows = match firing {
                        Firing::Flush => op.flush(),
                        Firing::Now(now) => op.advance(now),
                        Firing::Watermark(wm) => op.advance_to(wm),
                    };
                    if !rows.is_empty() {
                        fired.push(Fired::Rows(fired_to_keyed(rows, self.shape, &self.group_by)));
                    }
                }
            }
        }
        if env("STREAM_TRACE_TASKS", "0") == "1" {
            let n: usize = fired
                .iter()
                .map(|f| match f {
                    Fired::Batch(b) => b.num_rows(),
                    Fired::Rows(r) => r.len(),
                })
                .sum();
            if n > 0 || matches!(firing, Firing::Flush) {
                eprintln!("task {}: fired {n} row(s) on {firing:?}", self.task);
            }
        }
        for f in fired {
            let (batch, keys) = match f {
                Fired::Batch(b) => {
                    // keys from the fired columns BEFORE the post steps: the dedupe identity stays
                    // stable even if post renames or drops the window/group columns.
                    let keys = window_keys(&b, &self.group_by);
                    self.post_batch(b, keys)
                }
                Fired::Rows(keyed) => self.post_rows(keyed),
            };
            if batch.num_rows() > 0 {
                out.push(crate::decode::with_partition(
                    &crate::decode::with_keys(&batch, &keys),
                    self.task as i32,
                ));
            }
        }
    }

    /// The post steps on a fired batch (keys riding as a meta column), row by row when the
    /// translator cannot compile them for this schema.
    fn post_batch(
        &mut self,
        fired: arrow::array::RecordBatch,
        keys: Vec<String>,
    ) -> (arrow::array::RecordBatch, Vec<Option<String>>) {
        if self.post.is_empty() {
            return (fired, keys.into_iter().map(Some).collect());
        }
        let with_keys = with_key_column(&fired, &keys);
        let before = self.post.dropped;
        match self.post.apply(&with_keys) {
            Ok(o) => {
                let poison = self.post.dropped - before;
                if poison > 0 {
                    self.dropped.fetch_add(poison, Ordering::Relaxed);
                }
                let (b, _, k) = crate::decode::split_meta(&o);
                (b, k)
            }
            Err(_) => {
                let keyed: Vec<(String, fv_plan::row::Row)> = crate::rows::batch_to_rows(&fired)
                    .into_iter()
                    .zip(keys)
                    .map(|(row, key)| (key, row))
                    .collect();
                self.post_rows(keyed)
            }
        }
    }

    /// The post steps per row with poison isolation (the row seam's shape), then one batch.
    fn post_rows(
        &mut self,
        keyed: Vec<(String, fv_plan::row::Row)>,
    ) -> (arrow::array::RecordBatch, Vec<Option<String>>) {
        let mut keys = Vec::with_capacity(keyed.len());
        let mut rows = Vec::with_capacity(keyed.len());
        for (key, row) in keyed {
            if self.raw_post.is_empty() {
                keys.push(Some(key));
                rows.push(row);
                continue;
            }
            let (o, dropped) =
                fv_plan::inline::apply_steps_isolating(self.raw_post.as_slice(), std::slice::from_ref(&row));
            if dropped > 0 {
                self.dropped.fetch_add(dropped as u64, Ordering::Relaxed);
            }
            if let Some(r) = o.into_iter().next() {
                keys.push(Some(key));
                rows.push(r);
            }
        }
        (crate::rows::rows_to_batch(&rows), keys)
    }
}

impl Operator for StatefulOp {
    fn on_data(&mut self, _edge: EdgeId, batch: arrow::array::RecordBatch, out: &mut Out) {
        let now = now_ms();
        for (origin, part) in crate::decode::split_by_partition(&batch) {
            // rows without a usable event time are masked out (and counted), as the chain does.
            let (b, times) = match &self.tsrc {
                TimeSource::Column(c) => batch_with_times(&part, &[], None, Some(c), now),
                TimeSource::Ingest => batch_with_times(&part, &[], None, None, now),
                TimeSource::None => batch_with_times(&part, &[], None, None, 0),
            };
            let masked = part.num_rows() - b.num_rows();
            if masked > 0 {
                self.dropped.fetch_add(masked as u64, Ordering::Relaxed);
            }
            if b.num_rows() == 0 {
                continue;
            }
            let shard = match self.sharding {
                Sharding::PerOrigin => origin,
                Sharding::Range => 0,
            };
            let mk = &self.mk;
            let op = self.shards.entry(shard).or_insert_with(|| mk());
            let known = self.expected.entry(shard).or_default();
            if !known.contains(&origin) {
                known.push(origin);
                op.expect_partitions(&[origin], now);
            }
            if op.push_batch(origin, &b, &times).is_none() {
                // a row operator behind the batch flag: push its rows.
                for (row, t) in crate::rows::batch_to_rows(&b).into_iter().zip(&times) {
                    op.push(origin, *t, &row.0);
                }
            }
        }
        self.fire_due(now, out);
    }

    fn on_watermark(&mut self, wm: i64, _now_ms: i64, out: &mut Out) {
        if self.timing.in_band {
            self.watermark = Some(wm);
            self.fire(Firing::Watermark(wm), out);
        }
        // otherwise the shards' own per-origin time drives firing (see the module doc).
        // What leaves is the input time held back by the window size (`Timing::emit_hold_ms`).
        out.watermark = Some(Some(wm.saturating_sub(self.timing.emit_hold_ms)));
    }

    fn on_tick(&mut self, now_ms: i64, out: &mut Out) {
        self.fire_due(now_ms, out);
    }

    fn on_barrier(&mut self, epoch: u64, _out: &mut Out) -> Result<OpSnapshot, String> {
        let mut ids: Vec<i32> = self.shards.keys().copied().collect();
        ids.sort_unstable();

        // File-backed: on the object store, when every shard operator persists as files (a keyBy
        // window under a memory limit), the barrier writes ONE FILE PER SHARD instead of one giant
        // object — the peak in memory is a single shard's bytes, not the whole operator's state
        // (the fix for the gigabyte q5/window checkpoint). The head names the files; the files are
        // uploaded from a `.ckpt` staging dir off this thread and removed once durable.
        let ckpt_root = self.ckpt_dir.clone();
        if let (true, Some(root)) = (
            self.checkpoint_files && !ids.is_empty() && ids.iter().all(|id| self.shards[id].shard_files_supported()),
            ckpt_root,
        ) {
            let staging = root.join(".ckpt");
            let _ = std::fs::create_dir_all(&staging);
            let mut head = Vec::new();
            head.extend_from_slice(&(ids.len() as u32).to_le_bytes());
            let mut files: Vec<StateFile> = Vec::new();
            for id in &ids {
                let tag = format!("s{}-sh{}-e{}", self.task, id, epoch);
                // a state-tier fault (a full disk under the staging dir) fails this checkpoint,
                // naming the task — the coordinator ends the run with it.
                let shard_files = self
                    .shards
                    .get_mut(id)
                    .expect("shard")
                    .snapshot_shard_files(&staging, &tag)
                    .map_err(|e| format!("shard {id}: {e}"))?;
                head.extend_from_slice(&id.to_le_bytes());
                head.extend_from_slice(&(shard_files.len() as u32).to_le_bytes());
                for (sub, path) in shard_files {
                    let name = path
                        .file_name()
                        .and_then(|n| n.to_str())
                        .expect("shard file name")
                        .to_string();
                    head.extend_from_slice(&(sub as u32).to_le_bytes());
                    head.extend_from_slice(&(name.len() as u32).to_le_bytes());
                    head.extend_from_slice(name.as_bytes());
                    // A window shard is not time-range queried on restore (every file is read), so
                    // the stats are left at zero — the manifest carries them only for a join's reads.
                    files.push(StateFile {
                        name,
                        path,
                        min_time: 0,
                        max_time: 0,
                    });
                }
            }
            return Ok(OpSnapshot::head_with_files(head, files));
        }

        // Whole snapshot (topic/memory, or operators without shard files): every shard's own
        // snapshot framed per shard `[i32 shard][u32 len][bytes]…` — the object this task contributes.
        let mut bytes = Vec::new();
        let trace = env("STREAM_TRACE_TASKS", "0") == "1";
        for id in ids {
            let snap = self
                .shards
                .get_mut(&id)
                .expect("shard")
                .snapshot()
                .map_err(|e| format!("shard {id}: {e}"))?;
            if trace {
                // what the checkpoint holds for this shard, from the snapshot's own header.
                eprintln!(
                    "task {}: shard {id} snapshot {} B {}",
                    self.task,
                    snap.len(),
                    snapshot_header(&snap)
                );
            }
            bytes.extend_from_slice(&id.to_le_bytes());
            bytes.extend_from_slice(&(snap.len() as u32).to_le_bytes());
            bytes.extend_from_slice(&snap);
        }
        Ok(OpSnapshot::whole(bytes))
    }

    fn on_eos(&mut self, out: &mut Out) {
        self.fire(Firing::Flush, out);
    }
}

/// The metadata of a snapshot's first frame (`[u32 len][IPC stream]`), for the trace: the
/// columnar window's header carries `fired_through`, the watermark state and the pane count; a
/// row operator's JSON snapshot has no frame and is described by its length only.
fn snapshot_header(snap: &[u8]) -> String {
    let Some(len) = snap
        .get(..4)
        .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as usize)
    else {
        return String::new();
    };
    let Some(frame) = snap.get(4..4 + len) else {
        return String::new();
    };
    match arrow::ipc::reader::StreamReader::try_new(std::io::Cursor::new(frame), None) {
        Ok(r) => {
            let md = r.schema().metadata().clone();
            let mut parts: Vec<String> = md.iter().map(|(k, v)| format!("{k}={v}")).collect();
            parts.sort();
            parts.join(" ")
        }
        Err(_) => String::new(),
    }
}

impl StatefulOp {
    /// Restore the shards from a checkpoint object written by `on_barrier`: each shard is rebuilt
    /// by the factory and loaded (the operators refuse a snapshot of a different configuration).
    pub(super) fn with_state(mut self, bytes: &[u8]) -> Result<Self, String> {
        let mut at = 0usize;
        while at < bytes.len() {
            let id = i32::from_le_bytes(
                bytes
                    .get(at..at + 4)
                    .ok_or("shard frame truncated")?
                    .try_into()
                    .unwrap(),
            );
            let len = u32::from_le_bytes(
                bytes
                    .get(at + 4..at + 8)
                    .ok_or("shard frame truncated")?
                    .try_into()
                    .unwrap(),
            ) as usize;
            let snap = bytes.get(at + 8..at + 8 + len).ok_or("shard bytes truncated")?;
            let mut op = (self.mk)();
            op.load_snapshot(snap).map_err(|e| format!("shard {id}: {e}"))?;
            self.shards.insert(id, op);
            at += 8 + len;
        }
        Ok(self)
    }

    /// Restore from a file-backed head (written by `on_barrier`) plus the shard files it names,
    /// each already downloaded to a local path keyed by file name. The head is
    /// `[u32 n_shards]` then per engine shard `[i32 id][u32 n_files]` then per file
    /// `[u32 sub_index][u32 name_len][name]`. Bounded to one shard file at a time.
    pub(super) fn with_file_state(
        mut self,
        head: &[u8],
        files: &std::collections::HashMap<String, std::path::PathBuf>,
    ) -> Result<Self, String> {
        let mut at = 0usize;
        let read_u32 = |b: &[u8], at: &mut usize| -> Result<u32, String> {
            let v = b.get(*at..*at + 4).ok_or("stateful head truncated")?;
            *at += 4;
            Ok(u32::from_le_bytes(v.try_into().unwrap()))
        };
        let n_shards = read_u32(head, &mut at)?;
        for _ in 0..n_shards {
            let id = read_u32(head, &mut at)? as i32;
            let n_files = read_u32(head, &mut at)?;
            let mut shard_files: Vec<(usize, std::path::PathBuf)> = Vec::with_capacity(n_files as usize);
            for _ in 0..n_files {
                let sub = read_u32(head, &mut at)? as usize;
                let name_len = read_u32(head, &mut at)? as usize;
                let name = head.get(at..at + name_len).ok_or("stateful head: name truncated")?;
                at += name_len;
                let name = std::str::from_utf8(name).map_err(|_| "stateful head: name utf8")?;
                let path = files
                    .get(name)
                    .ok_or_else(|| format!("shard file {name} missing from the restored set"))?;
                shard_files.push((sub, path.clone()));
            }
            let mut op = (self.mk)();
            op.load_shard_files(&shard_files)
                .map_err(|e| format!("shard {id}: {e}"))?;
            self.shards.insert(id, op);
        }
        Ok(self)
    }
}

/// A `streamJoin` stage as a task: one `StreamJoinBatch` over the task's rows of both sides.
pub(super) struct JoinOp {
    task: TaskId,
    spec: JoinSpec,
    join: fv_streams_ops::StreamJoinBatch,
    /// Which input edges carry the left side (the rest are the right).
    left_edges: HashSet<EdgeId>,
    /// `true`: the join's time is the runtime's watermark (a keyBy join whose every input is
    /// stamped), so a fast side never evicts what a slow side still has partners for.
    in_band: bool,
    post: crate::steps::BatchSteps,
    raw_post: Arc<Vec<fv_plan::inline::Step>>,
    dropped: Arc<AtomicU64>,
    /// `true`: checkpoint as a small head plus upload-once frozen files (object store); `false`: a
    /// whole snapshot (topic/memory).
    checkpoint_files: bool,
}

impl JoinOp {
    #[allow(clippy::too_many_arguments)]
    pub(super) fn new(
        task: TaskId,
        spec: JoinSpec,
        in_edges: &[(EdgeId, TaskId)],
        left_tasks: &HashSet<TaskId>,
        in_band: bool,
        post: Arc<Vec<fv_plan::inline::Step>>,
        dropped: Arc<AtomicU64>,
        spill: Option<(
            std::path::PathBuf,
            Arc<dyn datafusion::execution::memory_pool::MemoryPool>,
        )>,
    ) -> Self {
        let mut join =
            fv_streams_ops::StreamJoinBatch::new(spec.join_key.clone(), spec.window_ms, spec.allowed_lateness_ms)
                .with_idle_timeout(spec.idle_timeout_ms);
        if in_band {
            join = join.with_in_band_time();
        }
        if let Some((dir, pool)) = spill {
            join = join.with_spill(&dir.join(format!("join-task{task}")), &pool, 8);
        }
        // both sides must report before eviction advances (the two-input min).
        join.expect_partitions(&[JOIN_SIDE_LEFT, JOIN_SIDE_RIGHT], now_ms());
        let left_edges = in_edges
            .iter()
            .filter(|(_, upstream)| left_tasks.contains(upstream))
            .map(|(e, _)| *e)
            .collect();
        JoinOp {
            task,
            spec,
            join,
            left_edges,
            in_band,
            post: crate::steps::BatchSteps::new(post.as_ref().clone()),
            raw_post: post,
            dropped,
            checkpoint_files: false,
        }
    }

    /// Checkpoint file-backed (a small head + upload-once frozen files) — set for the object store.
    pub(super) fn with_checkpoint_files(mut self, on: bool) -> Self {
        self.checkpoint_files = on;
        self
    }
}

impl Operator for JoinOp {
    fn on_data(&mut self, edge: EdgeId, batch: arrow::array::RecordBatch, out: &mut Out) {
        let is_left = self.left_edges.contains(&edge);
        let (b, times) = batch_with_times(&batch, &[], None, Some(&self.spec.time_column), 0);
        let masked = batch.num_rows() - b.num_rows();
        if masked > 0 {
            self.dropped.fetch_add(masked as u64, Ordering::Relaxed);
        }
        let now = now_ms();
        if b.num_rows() > 0 {
            let (side_id, side) = if is_left {
                (JOIN_SIDE_LEFT, fv_streams_ops::Side::Left)
            } else {
                (JOIN_SIDE_RIGHT, fv_streams_ops::Side::Right)
            };
            if let Some(joined) = self.join.push_batch(side_id, side, &b, &times, now) {
                // the post steps on the joined batch, then keys from the post-shaped batch
                // (replay produces identical rows ⇒ identical keys), as the chain does.
                let stepped = if self.post.is_empty() {
                    Ok(joined.clone())
                } else {
                    let before = self.post.dropped;
                    let r = self.post.apply(&joined);
                    let poison = self.post.dropped - before;
                    if poison > 0 {
                        self.dropped.fetch_add(poison, Ordering::Relaxed);
                    }
                    r
                };
                let (outb, keys): (arrow::array::RecordBatch, Vec<Option<String>>) = match stepped {
                    Ok(b) => {
                        let keys = join_keys(&b, &self.spec.join_key).into_iter().map(Some).collect();
                        (b, keys)
                    }
                    Err(_) => {
                        let mut keys = Vec::new();
                        let mut rows = Vec::new();
                        for r in crate::rows::batch_to_rows(&joined) {
                            let (o, d) = fv_plan::inline::apply_steps_isolating(
                                self.raw_post.as_slice(),
                                std::slice::from_ref(&r),
                            );
                            if d > 0 {
                                self.dropped.fetch_add(d as u64, Ordering::Relaxed);
                            }
                            if let Some(r) = o.into_iter().next() {
                                keys.push(Some(join_emit_key(&r, &self.spec.join_key)));
                                rows.push(r);
                            }
                        }
                        (crate::rows::rows_to_batch(&rows), keys)
                    }
                };
                if outb.num_rows() > 0 {
                    out.push(crate::decode::with_partition(
                        &crate::decode::with_keys(&outb, &keys),
                        self.task as i32,
                    ));
                }
            }
        }
        self.join.evict(now); // bound the state to ~one window behind the two-side watermark
    }

    fn on_watermark(&mut self, wm: i64, now_ms: i64, _out: &mut Out) {
        if self.in_band {
            self.join.set_watermark(wm);
            self.join.evict(now_ms);
        }
        // otherwise the two sides' own maxima drive it (a co-partitioned join). What leaves is
        // the input time: a joined row carries its inputs' own times.
    }

    fn on_tick(&mut self, now_ms: i64, _out: &mut Out) {
        self.join.evict(now_ms);
    }

    fn on_barrier(&mut self, _epoch: u64, _out: &mut Out) -> Result<OpSnapshot, String> {
        if self.checkpoint_files {
            // a small head (metadata + hot batches) plus the immutable frozen files, uploaded once.
            // Hardlink each frozen file into a per-task staging dir SYNCHRONOUSLY here, before the
            // task resumes and its eviction can delete the original — the checkpoint (which runs
            // off this thread) then uploads from the staging link, immune to the eviction race. The
            // link is cheap (same inode) and the checkpointer removes it once the file is durable.
            let files = self
                .join
                .frozen_files()
                .into_iter()
                .map(|f| {
                    let path = match f.path.parent() {
                        Some(dir) => {
                            let staging = dir.join(".ckpt");
                            let _ = std::fs::create_dir_all(&staging);
                            let link = staging.join(&f.name);
                            if !link.exists() && std::fs::hard_link(&f.path, &link).is_err() {
                                // fall back to the original path (best effort) if the link fails
                                f.path.clone()
                            } else {
                                link
                            }
                        }
                        None => f.path.clone(),
                    };
                    StateFile {
                        name: f.name,
                        path,
                        min_time: f.min_time,
                        max_time: f.max_time,
                    }
                })
                .collect();
            Ok(OpSnapshot::head_with_files(self.join.snapshot_head(), files))
        } else {
            Ok(OpSnapshot::whole(self.join.snapshot()))
        }
    }

    fn on_eos(&mut self, _out: &mut Out) {}
}

impl JoinOp {
    /// Restore the join from a whole checkpoint object written by `on_barrier` (whole mode).
    pub(super) fn with_state(mut self, bytes: &[u8]) -> Result<Self, String> {
        self.join.load_snapshot(bytes)?;
        Ok(self)
    }

    /// Restore from a head and the frozen files it references (file-backed mode): the engine has
    /// placed each file at a local path, keyed by name.
    pub(super) fn with_file_state(
        mut self,
        head: &[u8],
        files: &std::collections::HashMap<String, std::path::PathBuf>,
    ) -> Result<Self, String> {
        self.join.restore_incremental(head, files)?;
        Ok(self)
    }
}

/// The LOOKUP JOIN driver (5b, Nexmark q13): enrich a STREAM against a BOUNDED side TABLE.
///
/// Two inputs, distinguished by edge (like [`JoinOp`]): the **table** edges upsert the keyed lookup
/// table; every other edge is the **stream**, enriched row-by-row and emitted. Processing-time — the
/// operator ignores watermarks (no eviction, no gating), so a broadcast table's clock never pins the
/// stream. Like a join, the output is a NEW keyed stream, so it stamps its own task id as the
/// partition and keys each enriched row by the join value + content hash (`join_keys`).
pub(super) struct LookupJoinOp {
    task: TaskId,
    join_key: String,
    lookup: fv_streams_ops::LookupJoinBatch,
    /// Which input edges carry the TABLE side (the rest are the stream).
    table_edges: HashSet<EdgeId>,
    post: crate::steps::BatchSteps,
    raw_post: Arc<Vec<fv_plan::inline::Step>>,
    dropped: Arc<AtomicU64>,
}

impl LookupJoinOp {
    pub(super) fn new(
        task: TaskId,
        join_key: String,
        in_edges: &[(EdgeId, TaskId)],
        table_tasks: &HashSet<TaskId>,
        post: Arc<Vec<fv_plan::inline::Step>>,
        dropped: Arc<AtomicU64>,
    ) -> Self {
        let table_edges = in_edges
            .iter()
            .filter(|(_, upstream)| table_tasks.contains(upstream))
            .map(|(e, _)| *e)
            .collect();
        LookupJoinOp {
            task,
            join_key: join_key.clone(),
            lookup: fv_streams_ops::LookupJoinBatch::new(join_key),
            table_edges,
            post: crate::steps::BatchSteps::new(post.as_ref().clone()),
            raw_post: post,
            dropped,
        }
    }

    pub(super) fn with_state(mut self, bytes: &[u8]) -> Result<Self, String> {
        self.lookup.load_snapshot(bytes)?;
        Ok(self)
    }

    /// Post-step an enriched batch, key it by the join value + content hash, stamp this task's id as
    /// the partition (the join output is a new keyed stream), and push it.
    fn emit(&mut self, enriched: arrow::array::RecordBatch, out: &mut Out) {
        if enriched.num_rows() == 0 {
            return;
        }
        let (outb, keys): (arrow::array::RecordBatch, Vec<Option<String>>) = if self.post.is_empty() {
            let keys = join_keys(&enriched, &self.join_key).into_iter().map(Some).collect();
            (enriched, keys)
        } else {
            let before = self.post.dropped;
            match self.post.apply(&enriched) {
                Ok(b) => {
                    let poison = self.post.dropped - before;
                    if poison > 0 {
                        self.dropped.fetch_add(poison, Ordering::Relaxed);
                    }
                    let keys = join_keys(&b, &self.join_key).into_iter().map(Some).collect();
                    (b, keys)
                }
                Err(_) => {
                    // per-row isolation, exactly as JoinOp's post-step fallback.
                    let mut keys = Vec::new();
                    let mut rows = Vec::new();
                    for r in crate::rows::batch_to_rows(&enriched) {
                        let (o, d) =
                            fv_plan::inline::apply_steps_isolating(self.raw_post.as_slice(), std::slice::from_ref(&r));
                        if d > 0 {
                            self.dropped.fetch_add(d as u64, Ordering::Relaxed);
                        }
                        if let Some(r) = o.into_iter().next() {
                            keys.push(Some(join_emit_key(&r, &self.join_key)));
                            rows.push(r);
                        }
                    }
                    (crate::rows::rows_to_batch(&rows), keys)
                }
            }
        };
        if outb.num_rows() > 0 {
            out.push(crate::decode::with_partition(
                &crate::decode::with_keys(&outb, &keys),
                self.task as i32,
            ));
        }
    }
}

impl Operator for LookupJoinOp {
    fn on_data(&mut self, edge: EdgeId, batch: arrow::array::RecordBatch, out: &mut Out) {
        let (data, _, _) = crate::decode::split_meta(&batch);
        if self.table_edges.contains(&edge) {
            // TABLE side: upsert; the first table batch may release stream rows buffered before it.
            for released in self.lookup.upsert_table(&data) {
                self.emit(released, out);
            }
        } else if let Some(enriched) = self.lookup.enrich(&data) {
            // STREAM side: enrich against the current table (null table columns on a miss = LEFT).
            self.emit(enriched, out);
        }
    }

    // Processing-time: the lookup join does not gate on the watermark or tick (no eviction).
    fn on_watermark(&mut self, _wm: i64, _now_ms: i64, _out: &mut Out) {}

    fn on_barrier(&mut self, _epoch: u64, _out: &mut Out) -> Result<OpSnapshot, String> {
        Ok(OpSnapshot::whole(self.lookup.snapshot()))
    }

    fn on_eos(&mut self, _out: &mut Out) {}
}

/// A stateless stage fed by another stage: its inline steps on the batch (per row with poison
/// isolation when they do not compile for this schema), keyed by the row's own `rid`, else the
/// upstream key it arrived under.
pub(super) struct StepsOp {
    steps: crate::steps::BatchSteps,
    raw_steps: Arc<Vec<fv_plan::inline::Step>>,
    /// `true` when the batch goes straight to a sink: `__fv_key` becomes the emission key.
    finalize_keys: bool,
    dropped: Arc<AtomicU64>,
}

impl StepsOp {
    pub(super) fn new(steps: Arc<Vec<fv_plan::inline::Step>>, finalize_keys: bool, dropped: Arc<AtomicU64>) -> Self {
        StepsOp {
            steps: crate::steps::BatchSteps::new(steps.as_ref().clone()),
            raw_steps: steps,
            finalize_keys,
            dropped,
        }
    }
}

impl Operator for StepsOp {
    fn on_data(&mut self, _edge: EdgeId, batch: arrow::array::RecordBatch, out: &mut Out) {
        // The columnar path: every source writes a row's key and origin as meta columns, and the
        // compiled steps keep meta columns aligned through a filter — so the steps apply to the
        // batch as it arrived and the key is finalized IN PLACE (the row's own `rid`, else the key
        // it came with, else the batch's fallback) with one `zip` per column: no per-row strings,
        // no split / rebuild of the meta columns on either side of the steps.
        let before = self.steps.dropped;
        // the batch's fallback key is its first CONSUMED key, read before the steps (a row the
        // filter drops still supplies it; a rid never does) — the fused source's rule.
        let fallback = first_key(&batch);
        match self.steps.apply(&batch) {
            Ok(o) => {
                let poison = self.steps.dropped - before;
                if poison > 0 {
                    self.dropped.fetch_add(poison, Ordering::Relaxed);
                }
                if o.num_rows() > 0 {
                    out.push(finalize_key_column(&o, self.finalize_keys, fallback.as_deref()));
                }
            }
            Err(_) => {
                // this schema cannot take the batch path: per row with poison isolation, as the
                // chain does, per origin (the consumed keys do not survive a row path; the fallback
                // does).
                for (origin, part) in crate::decode::split_by_partition(&batch) {
                    let (data, _, keys) = crate::decode::split_meta(&part);
                    let fallback = keys.iter().find_map(|k| k.clone());
                    let rows = crate::rows::batch_to_rows(&data);
                    let (o, dropped) = fv_plan::inline::apply_steps_isolating(self.raw_steps.as_slice(), &rows);
                    if dropped > 0 {
                        self.dropped.fetch_add(dropped as u64, Ordering::Relaxed);
                    }
                    let b = crate::rows::rows_to_batch(&o);
                    let keys = if self.finalize_keys {
                        rid_keys(&b, &[], fallback.as_deref())
                    } else {
                        vec![fallback.clone(); b.num_rows()]
                    };
                    if b.num_rows() > 0 {
                        out.push(crate::decode::with_partition(
                            &crate::decode::with_keys(&b, &keys),
                            origin,
                        ));
                    }
                }
            }
        }
    }

    fn on_watermark(&mut self, _wm: i64, _now_ms: i64, _out: &mut Out) {}

    fn on_barrier(&mut self, _epoch: u64, _out: &mut Out) -> Result<OpSnapshot, String> {
        Ok(OpSnapshot::whole(Vec::new()))
    }

    fn on_eos(&mut self, _out: &mut Out) {}
}

/// The batch's first consumed key (its `__fv_key` column's first non-null value): the fallback for
/// rows that have none.
fn first_key(batch: &arrow::array::RecordBatch) -> Option<String> {
    use arrow::array::{Array, StringArray};
    let col = batch.column_by_name(crate::decode::META_KEY)?;
    let strings = col.as_any().downcast_ref::<StringArray>()?;
    (0..strings.len())
        .find(|i| strings.is_valid(*i))
        .map(|i| strings.value(i).to_string())
}

/// The batch with its `__fv_key` meta column finalized in place, columnar: when `finalize` (the
/// batch goes straight to a sink) the row's own `rid` when it is a non-null string, else the key
/// the row came with, else `fallback`; when not, the key it came with, else `fallback`. A batch
/// without a key column gets one. One `zip` per layer; no per-row strings.
fn finalize_key_column(
    batch: &arrow::array::RecordBatch,
    finalize: bool,
    fallback: Option<&str>,
) -> arrow::array::RecordBatch {
    use arrow::array::{Array, ArrayRef, Scalar, StringArray};
    use arrow::compute::kernels::zip::zip;
    use arrow::compute::{is_not_null, is_null};
    use arrow::datatypes::{DataType, Field, Schema};
    let n = batch.num_rows();
    let key_idx = batch.schema().index_of(crate::decode::META_KEY).ok();
    let existing: ArrayRef = match key_idx {
        Some(i) => Arc::clone(batch.column(i)),
        None => Arc::new(StringArray::new_null(n)),
    };
    let existing = if existing.data_type() == &DataType::Utf8 {
        existing
    } else {
        arrow::compute::cast(&existing, &DataType::Utf8).unwrap_or_else(|_| Arc::new(StringArray::new_null(n)))
    };
    let mut key = existing;
    if finalize {
        if let Some(rid) = batch.column_by_name("rid").filter(|c| c.data_type() == &DataType::Utf8) {
            if rid.null_count() < n {
                key = zip(&is_not_null(rid.as_ref()).expect("mask"), rid, &key).expect("rid over key");
            }
        }
    }
    if key.null_count() > 0 {
        if let Some(fallback) = fallback {
            let fallback = Scalar::new(StringArray::from(vec![fallback]));
            key = zip(&is_null(key.as_ref()).expect("mask"), &fallback, &key).expect("fallback under key");
        }
    }
    let mut fields: Vec<Field> = batch.schema().fields().iter().map(|f| f.as_ref().clone()).collect();
    let mut cols: Vec<ArrayRef> = batch.columns().to_vec();
    match key_idx {
        Some(i) => cols[i] = key,
        None => {
            fields.push(Field::new(crate::decode::META_KEY, DataType::Utf8, true));
            cols.push(key);
        }
    }
    arrow::array::RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).expect("key column matches the batch")
}

#[cfg(test)]
mod key_column_tests {
    use super::*;
    use arrow::array::{Array, Int64Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema};

    fn batch(rid: &[Option<&str>], key: Option<&[Option<&str>]>) -> arrow::array::RecordBatch {
        let n = rid.len();
        let mut fields = vec![
            Field::new("rid", DataType::Utf8, true),
            Field::new("n", DataType::Int64, false),
        ];
        let mut cols: Vec<arrow::array::ArrayRef> = vec![
            Arc::new(StringArray::from(rid.to_vec())),
            Arc::new(Int64Array::from((0..n as i64).collect::<Vec<_>>())),
        ];
        if let Some(k) = key {
            fields.push(Field::new(crate::decode::META_KEY, DataType::Utf8, true));
            cols.push(Arc::new(StringArray::from(k.to_vec())));
        }
        arrow::array::RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).unwrap()
    }

    fn keys(b: &arrow::array::RecordBatch) -> Vec<Option<String>> {
        let (_, _, k) = crate::decode::split_meta(b);
        k
    }

    #[test]
    fn finalize_takes_rid_then_the_consumed_key_then_the_batch_fallback() {
        let b = batch(&[Some("r0"), None, None], Some(&[Some("k0"), Some("k1"), None]));
        let fallback = first_key(&b);
        assert_eq!(fallback.as_deref(), Some("k0"));
        assert_eq!(
            keys(&finalize_key_column(&b, true, fallback.as_deref())),
            vec![Some("r0".into()), Some("k1".into()), Some("k0".into())],
            "rid first; else the consumed key; else the batch's first consumed key"
        );
        assert_eq!(
            keys(&finalize_key_column(&b, false, fallback.as_deref())),
            vec![Some("k0".into()), Some("k1".into()), Some("k0".into())],
            "not finalizing leaves rid alone: the consumed key, else the fallback"
        );
    }

    #[test]
    fn a_batch_without_keys_gets_a_key_column_from_rid_or_nulls() {
        let b = batch(&[Some("r0"), None], None);
        assert_eq!(first_key(&b), None, "no key column: no fallback (a rid is never one)");
        let f = finalize_key_column(&b, true, None);
        assert_eq!(
            keys(&f),
            vec![Some("r0".into()), None],
            "rid where present, nothing else"
        );
        let none = finalize_key_column(&batch(&[None, None], None), true, None);
        assert_eq!(keys(&none), vec![None, None]);
        assert_eq!(none.num_columns(), 3, "the key column is added once");
    }

    #[test]
    fn the_other_columns_and_the_row_count_are_untouched() {
        let b = batch(&[Some("a"), Some("b")], Some(&[None, None]));
        let f = finalize_key_column(&b, true, first_key(&b).as_deref());
        assert_eq!(f.num_rows(), 2);
        assert_eq!(f.column_by_name("n").unwrap().len(), 2);
        assert_eq!(keys(&f), vec![Some("a".into()), Some("b".into())]);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::{Float64Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema};

    fn events(keys: &[&str], ts: &[i64], partition: i32) -> arrow::array::RecordBatch {
        let schema = Arc::new(Schema::new(vec![
            Field::new("k", DataType::Utf8, true),
            Field::new("ts", DataType::Float64, true),
            Field::new("v", DataType::Float64, true),
        ]));
        let k: StringArray = keys.to_vec().into();
        let t: Float64Array = ts.iter().map(|x| *x as f64).collect::<Vec<_>>().into();
        let v: Float64Array = ts.iter().map(|x| *x as f64 / 10.0).collect::<Vec<_>>().into();
        let b = arrow::array::RecordBatch::try_new(schema, vec![Arc::new(k), Arc::new(t), Arc::new(v)]).unwrap();
        crate::decode::with_partition(&b, partition)
    }

    fn window_op(sharding: Sharding) -> StatefulOp {
        window_op_with(sharding, 0)
    }

    fn window_op_with(sharding: Sharding, lateness_ms: i64) -> StatefulOp {
        let timing = Timing {
            in_band: matches!(sharding, Sharding::Range),
            emit_hold_ms: 1000,
        };
        StatefulOp::new(
            7,
            Box::new(move || {
                Box::new(fv_streams_ops::WindowAggBatch::tumbling(
                    1000,
                    lateness_ms,
                    vec!["k".into()],
                    vec![fv_streams_ops::Agg::count("n")],
                ))
            }),
            sharding,
            TimeSource::Column("ts".into()),
            EmitShape::Window,
            vec!["k".into()],
            timing,
            Arc::new(Vec::new()),
            Arc::new(AtomicU64::new(0)),
        )
    }

    #[test]
    fn a_window_task_passes_its_time_on_held_back_by_its_window_size() {
        // downstream, a stage windowing on `windowStart` must never see a fired window as late:
        // a window starting at s fires at s + size, so the time this task passes on is the input
        // time less its size — the consumer thread's latency for such a stage, exactly.
        let mut op = window_op(Sharding::Range);
        let mut out = Out::default();
        op.on_data(0, events(&["a"], &[100], 0), &mut out);
        op.on_watermark(1_000, 0, &mut out);
        assert_eq!(out.batches.len(), 1, "window [0,1000) fired at watermark 1000");
        assert_eq!(
            out.watermark,
            Some(Some(0)),
            "the time passed on is 1000 − 1000: the fired windowStart is not late"
        );
        let mut per = window_op(Sharding::PerOrigin);
        let mut out = Out::default();
        per.on_watermark(5_000, 0, &mut out);
        assert_eq!(
            out.watermark,
            Some(Some(4_000)),
            "a per-origin task passes held-back time on too"
        );
    }

    #[test]
    fn a_keyby_shard_fires_on_the_runtime_watermark_never_on_its_rows() {
        // REGRESSION (the crash e2e): four source tasks feed a keyBy shard through the shuffle;
        // one runs far ahead. Deriving time from the rows fired window 0 on the first origin
        // alone and dropped the slower origins' rows as late. In-band time: the rows never move
        // it, the runtime's watermark (the minimum across every input) does.
        let mut op = window_op_with(Sharding::Range, 500);
        let mut out = Out::default();
        op.on_data(0, events(&["a", "a"], &[100, 9_900], 0), &mut out);
        op.on_data(0, events(&["b"], &[9_900], 1), &mut out);
        op.on_tick(1_000_000, &mut out);
        assert!(
            out.batches.is_empty(),
            "rows at 9900 and a tick: nothing fires without a watermark"
        );
        op.on_watermark(1_400, 0, &mut out);
        assert!(
            out.batches.is_empty(),
            "watermark 1400 less 500 ms lateness is 900: window [0,1000) is still open"
        );
        op.on_data(0, events(&["d", "d"], &[200, 9_900], 3), &mut out);
        op.on_watermark(1_500, 0, &mut out);
        let mut keys: Vec<String> = out.batches.iter().flat_map(keys_of).collect();
        keys.sort();
        assert_eq!(
            keys,
            vec!["w|n0|n1000|sa", "w|n0|n1000|sd"],
            "watermark 1500: window [0,1000) fires with the late-arriving origin's row in it"
        );
        // a tick after a watermark fires nothing new; the next watermark does.
        let mut out = Out::default();
        op.on_tick(2_000_000, &mut out);
        assert!(out.batches.is_empty());
        op.on_watermark(20_000, 0, &mut out);
        let mut keys: Vec<String> = out.batches.iter().flat_map(keys_of).collect();
        keys.sort();
        assert_eq!(
            keys,
            vec!["w|n9000|n10000|sa", "w|n9000|n10000|sb", "w|n9000|n10000|sd"]
        );
        // a per-origin shard keeps its rows' own time: origin 0 past 1000 fires its window.
        let mut per = window_op(Sharding::PerOrigin);
        let mut out = Out::default();
        per.on_data(0, events(&["a", "a"], &[100, 1_500], 0), &mut out);
        assert_eq!(out.batches.len(), 1, "co-partitioned: the partition's own time fires");
    }

    fn keys_of(b: &arrow::array::RecordBatch) -> Vec<String> {
        crate::decode::split_meta(b).2.into_iter().flatten().collect()
    }

    #[test]
    fn a_range_shard_fires_on_the_delivered_watermark_and_keys_its_windows() {
        // one keyBy instance fed by two origins through the shuffle: the runtime's watermark (the
        // minimum across both) is what closes window [0,1000), not the rows' own times.
        let mut op = window_op(Sharding::Range);
        let mut out = Out::default();
        op.on_data(0, events(&["a", "a"], &[100, 900], 0), &mut out);
        op.on_data(0, events(&["b"], &[300], 1), &mut out);
        assert!(out.batches.is_empty(), "no watermark yet");
        op.on_data(0, events(&["a"], &[2500], 0), &mut out);
        op.on_watermark(300, 0, &mut out);
        assert!(
            out.batches.is_empty(),
            "origin 0 is at 2500 but the runtime's minimum is 300: window 0 holds"
        );
        op.on_data(0, events(&["b"], &[1700], 1), &mut out);
        op.on_watermark(1700, 0, &mut out);
        assert_eq!(out.batches.len(), 1, "the watermark passed 1000: window [0,1000) fires");
        let fired = &out.batches[0];
        let mut keys = keys_of(fired);
        keys.sort();
        assert_eq!(
            keys,
            vec!["w|n0|n1000|sa", "w|n0|n1000|sb"],
            "deterministic window keys, from the columns"
        );
        let origin = crate::decode::split_by_partition(fired);
        assert_eq!(origin[0].0, 7, "fired rows carry this task as their origin");
        // end of stream flushes what is left.
        let mut out = Out::default();
        op.on_eos(&mut out);
        let mut left: Vec<String> = out.batches.iter().flat_map(keys_of).collect();
        left.sort();
        assert_eq!(left, vec!["w|n1000|n2000|sb", "w|n2000|n3000|sa"]);
    }

    #[test]
    fn a_stateful_task_snapshots_its_shards_and_restores_them_exactly() {
        // two origins, two shards, a window open in each; the snapshot restores into a fresh task
        // that then fires the same windows with the same counts.
        let mut op = window_op(Sharding::PerOrigin);
        let mut out = Out::default();
        op.on_data(0, events(&["a", "a"], &[100, 900], 0), &mut out);
        op.on_data(0, events(&["b"], &[300], 1), &mut out);
        assert!(out.batches.is_empty());
        let snap = op.on_barrier(1, &mut out).unwrap();
        assert!(!snap.head.is_empty());
        let bytes = snap.head;
        let mut restored = window_op(Sharding::PerOrigin).with_state(&bytes).unwrap();
        assert_eq!(restored.shards.len(), 2);
        let mut out = Out::default();
        restored.on_eos(&mut out);
        let mut keys: Vec<String> = out.batches.iter().flat_map(keys_of).collect();
        keys.sort();
        assert_eq!(keys, vec!["w|n0|n1000|sa", "w|n0|n1000|sb"]);
        let counts: Vec<i64> = out
            .batches
            .iter()
            .flat_map(|b| {
                let n = b.column_by_name("n").unwrap();
                let n = n.as_any().downcast_ref::<arrow::array::Float64Array>().unwrap();
                n.values().iter().map(|v| *v as i64).collect::<Vec<i64>>()
            })
            .collect();
        let mut counts = counts;
        counts.sort();
        assert_eq!(counts, vec![1, 2], "a's two rows and b's one survived the round trip");
        assert!(
            window_op(Sharding::Range).with_state(&[1, 2, 3]).is_err(),
            "a truncated object is refused"
        );
    }

    #[test]
    fn a_snapshot_between_pushing_and_firing_loses_nothing() {
        // REGRESSION (the crash e2e): windows the checkpoint marked as fired never reached the
        // sink. The barrier's snapshot must leave every open pane exactly as it was.
        let mut op = window_op(Sharding::Range);
        let mut out = Out::default();
        // 50 keys over 4 origins, monotonic time over 9 windows; every window gets 10 rows per
        // key... but with the keys interleaved so panes and key counts grow between pushes.
        let keys: Vec<String> = (0..50).map(|k| format!("k{k}")).collect();
        for step in 0..90i64 {
            let origin = (step % 4) as i32;
            let ts = step * 100; // 0..9000: window w = step / 10
            let ks: Vec<&str> = keys.iter().map(String::as_str).collect();
            let tss: Vec<i64> = vec![ts; 50];
            op.on_data(0, events(&ks, &tss, origin), &mut out);
            if step == 37 {
                // a barrier mid-stream: snapshot (and keep going with the live state).
                let snap = op.on_barrier(1, &mut out).unwrap();
                assert!(!snap.head.is_empty());
            }
        }
        // the watermark is 8900: windows 0..8 (end <= 9000) fire; window 9 (start 9000) has no rows.
        let mut fired: Vec<(String, i64)> = Vec::new();
        for b in &out.batches {
            let n = b.column_by_name("n").unwrap();
            let n = n.as_any().downcast_ref::<arrow::array::Float64Array>().unwrap();
            for (i, key) in keys_of(b).iter().enumerate() {
                fired.push((key.clone(), n.value(i) as i64));
            }
        }
        // in-band time: the runtime's watermark at the last data closes windows 0..8.
        let mut out2 = Out::default();
        op.on_watermark(8_900, 0, &mut out2);
        for b in &out2.batches {
            let n = b.column_by_name("n").unwrap();
            let n = n.as_any().downcast_ref::<arrow::array::Float64Array>().unwrap();
            for (i, key) in keys_of(b).iter().enumerate() {
                fired.push((key.clone(), n.value(i) as i64));
            }
        }
        // 8 closed windows (0..7) x 50 keys before the tick's watermark... count what fired:
        // every fired (window, key) must carry n = 10 (10 steps per window, each with every key).
        let bad: Vec<&(String, i64)> = fired.iter().filter(|(_, n)| *n != 10).collect();
        assert!(bad.is_empty(), "windows fired with lost rows: {bad:?}");
        let windows: std::collections::BTreeSet<String> = fired
            .iter()
            .map(|(k, _)| k.split('|').nth(1).unwrap().to_string())
            .collect();
        assert_eq!(
            fired.len(),
            windows.len() * 50,
            "every key of every fired window: {} rows over {:?}",
            fired.len(),
            windows
        );
        assert!(
            windows.len() >= 8,
            "windows 0..7 closed by the last data (ts 8900): {windows:?}"
        );
    }

    #[test]
    fn per_origin_sharding_keeps_one_instance_per_partition() {
        // two origins, each its own instance: each fires on its own time, like one consumer
        // thread's per-partition instances.
        let mut op = window_op(Sharding::PerOrigin);
        let mut out = Out::default();
        op.on_data(0, events(&["a"], &[100], 0), &mut out);
        op.on_data(0, events(&["a"], &[2500], 0), &mut out);
        assert_eq!(out.batches.len(), 1, "origin 0's window fired on its own watermark");
        assert_eq!(keys_of(&out.batches[0]), vec!["w|n0|n1000|sa"]);
        let mut out = Out::default();
        op.on_data(0, events(&["b"], &[100], 1), &mut out);
        assert!(out.batches.is_empty(), "origin 1's instance is separate and still open");
        assert_eq!(op.shards.len(), 2);
    }

    #[test]
    fn a_join_task_pairs_rows_across_its_sides_and_keys_them() {
        let spec = JoinSpec {
            join_key: "k".into(),
            time_column: "ts".into(),
            window_ms: 1000,
            allowed_lateness_ms: 0,
            idle_timeout_ms: 0,
            key_by: true,
        };
        let left_tasks: HashSet<TaskId> = [1].into_iter().collect();
        let mut op = JoinOp::new(
            9,
            spec,
            &[(10, 1), (11, 2)],
            &left_tasks,
            false,
            Arc::new(Vec::new()),
            Arc::new(AtomicU64::new(0)),
            None,
        );
        let mut out = Out::default();
        op.on_data(10, events(&["a", "b"], &[100, 200], 0), &mut out);
        assert!(out.batches.is_empty(), "no partner yet");
        op.on_data(11, events(&["a"], &[500], 0), &mut out);
        assert_eq!(
            out.batches.len(),
            1,
            "a's right row meets a's left row within the window"
        );
        let joined = &out.batches[0];
        assert_eq!(joined.num_rows(), 1);
        let keys = keys_of(joined);
        assert!(
            keys[0].starts_with("j|sa|"),
            "join keys lead with the join-key value: {keys:?}"
        );
        assert_eq!(crate::decode::split_by_partition(joined)[0].0, 9);
    }

    #[test]
    fn a_keyby_join_task_evicts_on_the_runtime_watermark_not_on_a_fast_side() {
        // the left side runs to 9000 while the right is still at 100: on the sides' own maxima
        // the join would drop the right's late-arriving partner; on the delivered watermark it
        // waits for the runtime's minimum.
        let spec = JoinSpec {
            join_key: "k".into(),
            time_column: "ts".into(),
            window_ms: 1000,
            allowed_lateness_ms: 0,
            idle_timeout_ms: 0,
            key_by: true,
        };
        let left_tasks: HashSet<TaskId> = [1].into_iter().collect();
        let mut op = JoinOp::new(
            9,
            spec,
            &[(10, 1), (11, 2)],
            &left_tasks,
            true,
            Arc::new(Vec::new()),
            Arc::new(AtomicU64::new(0)),
            None,
        );
        let mut out = Out::default();
        op.on_data(10, events(&["a", "a"], &[100, 9_000], 0), &mut out);
        op.on_watermark(150, 0, &mut out);
        op.on_tick(1_000_000, &mut out);
        op.on_data(11, events(&["a"], &[300], 0), &mut out);
        assert_eq!(out.batches.len(), 1, "the right's 300 still meets the left's 100");
        assert_eq!(out.batches[0].num_rows(), 1);
        op.on_watermark(9_000, 0, &mut out);
        let mut out = Out::default();
        op.on_data(11, events(&["a"], &[300], 0), &mut out);
        assert!(out.batches.is_empty(), "after the watermark passed 9000, a 300 is late");
    }

    #[test]
    fn a_steps_task_transforms_and_keys_by_rid_then_upstream_key() {
        let steps = Arc::new(vec![
            fv_plan::inline::Step::Filter {
                expression: "v > 20".into(),
            },
            fv_plan::inline::Step::ApplyExpression {
                column: "double".into(),
                expression: "v * 2".into(),
            },
        ]);
        let mut op = StepsOp::new(steps, true, Arc::new(AtomicU64::new(0)));
        let b = events(&["a", "b", "c"], &[100, 300, 500], 4); // v = 10, 30, 50
        let b = crate::decode::with_keys(&b, &[Some("ka".into()), Some("kb".into()), None]);
        let mut out = Out::default();
        op.on_data(0, b, &mut out);
        assert_eq!(out.batches.len(), 1);
        let o = &out.batches[0];
        assert_eq!(o.num_rows(), 2, "the filter kept v > 20");
        let (data, _, keys) = crate::decode::split_meta(o);
        assert!(data.column_by_name("double").is_some());
        assert_eq!(
            keys,
            vec![Some("kb".to_string()), Some("ka".to_string())],
            "no rid column: the upstream key, else the batch's fallback key"
        );
        assert_eq!(
            crate::decode::split_by_partition(o)[0].0,
            4,
            "the origin passes through"
        );
    }
}