flow/batching_mode/

state.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Batching mode task state, which changes frequently
//!

use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::time::Duration;

use common_telemetry::debug;
use common_time::Timestamp;
use datatypes::value::Value;
use session::context::QueryContextRef;
use snafu::{OptionExt, ResultExt, ensure};
use tokio::sync::oneshot;
use tokio::time::Instant;

use crate::batching_mode::task::BatchingTask;
use crate::batching_mode::time_window::TimeWindowExpr;
use crate::error::{DatatypesSnafu, InternalSnafu, TimeSnafu, UnexpectedSnafu};
use crate::metrics::{
    METRIC_FLOW_BATCHING_ENGINE_QUERY_WINDOW_CNT, METRIC_FLOW_BATCHING_ENGINE_QUERY_WINDOW_SIZE,
    METRIC_FLOW_BATCHING_ENGINE_STALLED_WINDOW_SIZE,
};
use crate::{Error, FlowId};

/// The state of the [`BatchingTask`].
#[derive(Debug)]
pub struct TaskState {
    /// Query context
    pub(crate) query_ctx: QueryContextRef,
    /// last query complete time
    last_update_time: Instant,
    /// last time query duration
    last_query_duration: Duration,
    /// Last successful execution time in unix timestamp milliseconds.
    last_exec_time_millis: Option<i64>,
    /// First execution time in unix timestamp milliseconds, set once.
    start_time_millis: Option<i64>,
    /// Dirty Time windows need to be updated
    /// mapping of `start -> end` and non-overlapping
    pub(crate) dirty_time_windows: DirtyTimeWindows,
    checkpoint_mode: CheckpointMode,
    pending_fenced_repair: Option<FencedRepair>,
    /// Region id -> last consumed watermark sequence. Incremental scans use
    /// this as the next lower sequence bound for each source region.
    checkpoints: BTreeMap<u64, u64>,
    /// Once set, the task will never attempt incremental mode again.
    /// Set when the flow's query shape is deterministically incompatible
    /// with incremental execution (e.g. unsupported aggregate expressions).
    incremental_disabled: bool,
    exec_state: ExecState,
    /// Shutdown receiver
    pub(crate) shutdown_rx: oneshot::Receiver<()>,
    /// Task handle
    pub(crate) task_handle: Option<tokio::task::JoinHandle<()>>,
}
impl TaskState {
    pub fn new(query_ctx: QueryContextRef, shutdown_rx: oneshot::Receiver<()>) -> Self {
        Self::with_dirty_time_windows(query_ctx, shutdown_rx, DirtyTimeWindows::default())
    }

    pub fn with_dirty_time_windows(
        query_ctx: QueryContextRef,
        shutdown_rx: oneshot::Receiver<()>,
        dirty_time_windows: DirtyTimeWindows,
    ) -> Self {
        Self {
            query_ctx,
            last_update_time: Instant::now(),
            last_query_duration: Duration::from_secs(0),
            last_exec_time_millis: None,
            start_time_millis: None,
            dirty_time_windows,
            checkpoint_mode: CheckpointMode::FullSnapshot,
            pending_fenced_repair: None,
            checkpoints: Default::default(),
            incremental_disabled: false,
            exec_state: ExecState::Idle,
            shutdown_rx,
            task_handle: None,
        }
    }

    /// Record the first-execution start time. Call this once, just before
    /// the first frontend query is dispatched, not after it completes.
    pub fn record_start_time_if_first(&mut self) {
        if self.start_time_millis.is_none() {
            // start_time is recorded just before the first frontend query is dispatched
            // (pre-execution), so it may be marginally earlier than the streaming engine's
            // start_time which is set post-execution. Both are valid approximations of
            // "when this flow first ran".
            self.start_time_millis = Some(common_time::util::current_time_millis());
        }
    }

    pub fn after_query_exec(&mut self, elapsed: Duration, is_succ: bool) {
        self.exec_state = ExecState::Idle;
        self.last_query_duration = elapsed;
        self.last_update_time = Instant::now();
        if is_succ {
            self.last_exec_time_millis = Some(common_time::util::current_time_millis());
        }
    }

    pub fn last_execution_time_millis(&self) -> Option<i64> {
        self.last_exec_time_millis
    }

    /// First execution time in unix timestamp milliseconds, set once.
    pub fn start_time_millis(&self) -> Option<i64> {
        self.start_time_millis
    }

    pub fn checkpoint_mode(&self) -> CheckpointMode {
        self.checkpoint_mode
    }

    pub fn checkpoints(&self) -> &BTreeMap<u64, u64> {
        &self.checkpoints
    }

    /// Returns the in-progress fenced repair, if the task is repairing dirty
    /// windows under a frozen full-snapshot high watermark.
    pub fn pending_fenced_repair(&self) -> Option<&FencedRepair> {
        self.pending_fenced_repair.as_ref()
    }

    pub fn is_incremental_disabled(&self) -> bool {
        self.incremental_disabled
    }

    /// Permanently disable incremental mode for this task and
    /// immediately fall back to full snapshot for the current cycle.
    pub fn disable_incremental(&mut self) {
        self.incremental_disabled = true;
        self.mark_full_snapshot();
    }

    /// Move back to top-level FullSnapshot mode. If a fenced repair is active,
    /// restore its not-yet-in-flight pending windows to the live dirty queue so
    /// the moved backlog is not lost.
    pub fn mark_full_snapshot(&mut self) {
        self.abandon_fenced_repair();
    }

    /// Replace full-snapshot checkpoints with a complete watermark proof.
    /// Clears fenced repair state and enters Incremental unless disabled.
    pub fn advance_checkpoints(&mut self, watermark_map: HashMap<u64, u64>) {
        self.checkpoints = watermark_map.into_iter().collect();
        self.pending_fenced_repair = None;
        if !self.incremental_disabled {
            self.checkpoint_mode = CheckpointMode::Incremental;
        }
    }

    /// Advance only the participating regions for an incremental delta query.
    /// This also clears any stale fenced repair sub-state.
    pub fn advance_incremental_checkpoints_with_participation(
        &mut self,
        participating_regions: &BTreeSet<u64>,
        watermark_map: HashMap<u64, u64>,
    ) {
        for region_id in participating_regions {
            if let Some(seq) = watermark_map.get(region_id) {
                self.checkpoints.insert(*region_id, *seq);
            }
        }
        if !self.incremental_disabled {
            self.checkpoint_mode = CheckpointMode::Incremental;
        }
        self.pending_fenced_repair = None;
    }

    /// Start repairing the current live dirty windows under a frozen high `H`.
    /// The current live backlog is moved into the fenced repair so successful
    /// chunks are consumed from that backlog. New post-`H` dirty signals can
    /// still arrive in the live queue while the fenced repair is active.
    pub fn start_fenced_repair(&mut self, high: BTreeMap<u64, u64>) -> Option<&FencedRepair> {
        if self.dirty_time_windows.is_empty() {
            self.pending_fenced_repair = None;
            return None;
        }

        let pending_windows = self.dirty_time_windows.clone();
        self.dirty_time_windows.clean();
        self.pending_fenced_repair = Some(FencedRepair {
            high,
            pending_windows,
        });
        self.checkpoint_mode = CheckpointMode::FullSnapshot;
        self.pending_fenced_repair.as_ref()
    }

    /// Finish the fenced repair and promote the frozen high watermark to the
    /// checkpoint map. Incremental-disabled flows stay in FullSnapshot mode.
    pub fn finish_fenced_repair(&mut self) -> Option<BTreeMap<u64, u64>> {
        let repair = self.pending_fenced_repair.take()?;
        self.checkpoints = repair.high;
        if !self.incremental_disabled {
            self.checkpoint_mode = CheckpointMode::Incremental;
        }
        Some(self.checkpoints.clone())
    }

    /// Abandon the current fenced repair and restore all not-yet-in-flight
    /// pending windows to the live dirty queue for a fresh scoped repair.
    pub fn abandon_fenced_repair(&mut self) -> bool {
        self.checkpoint_mode = CheckpointMode::FullSnapshot;
        let Some(repair) = self.pending_fenced_repair.take() else {
            return false;
        };

        self.dirty_time_windows
            .add_dirty_windows(&repair.pending_windows);
        true
    }

    /// Restore a scoped query's windows after a failed or unproven run. During
    /// an active fenced repair this requeues into `pending_windows`; otherwise
    /// it restores to the live dirty queue.
    pub fn restore_scoped_windows(&mut self, filter: &FilterExprInfo) {
        if let Some(repair) = self.pending_fenced_repair.as_mut() {
            repair
                .pending_windows
                .add_windows(filter.time_ranges.clone());
            return;
        }

        self.dirty_time_windows
            .add_windows(filter.time_ranges.clone());
    }

    /// Generate the next scoped filter from the fenced-repair queue when active;
    /// otherwise consume windows from the live dirty queue.
    pub fn gen_scoped_filter_exprs(
        &mut self,
        col_name: &str,
        expire_lower_bound: Option<Timestamp>,
        window_size: chrono::Duration,
        window_cnt: usize,
        flow_id: FlowId,
        task_ctx: Option<&BatchingTask>,
    ) -> Result<Option<FilterExprInfo>, Error> {
        if let Some(repair) = self.pending_fenced_repair.as_mut() {
            let expr = repair.pending_windows.gen_filter_exprs(
                col_name,
                expire_lower_bound,
                window_size,
                window_cnt,
                flow_id,
                task_ctx,
            )?;
            if expr.is_some() || !repair.pending_windows.is_empty() {
                return Ok(expr);
            }

            // All pending repair windows may have expired during merge. Clear
            // the empty repair so this call can fall back to live dirty windows
            // instead of routing future executions to an empty queue forever.
            self.pending_fenced_repair = None;
        }

        self.dirty_time_windows.gen_filter_exprs(
            col_name,
            expire_lower_bound,
            window_size,
            window_cnt,
            flow_id,
            task_ctx,
        )
    }

    /// Returns true only when the query result's participating regions and
    /// terminal watermarks exactly match the fenced repair's frozen high `H`.
    pub fn fenced_repair_watermarks_match_high(
        &self,
        participating_regions: &BTreeSet<u64>,
        watermark_map: &HashMap<u64, u64>,
    ) -> bool {
        let Some(repair) = self.pending_fenced_repair.as_ref() else {
            return false;
        };

        !participating_regions.is_empty()
            && participating_regions.len() == repair.high.len()
            && watermark_map.len() == repair.high.len()
            && participating_regions.iter().all(|region_id| {
                repair
                    .high
                    .get(region_id)
                    .zip(watermark_map.get(region_id))
                    .is_some_and(|(high, watermark)| high == watermark)
            })
    }

    /// Whether the active fenced repair has drained all pending windows.
    pub fn fenced_repair_pending_is_empty(&self) -> bool {
        self.pending_fenced_repair
            .as_ref()
            .is_some_and(|repair| repair.pending_windows.is_empty())
    }

    /// Full-snapshot checkpoint advances require a watermark for every region
    /// that participated in the query.
    pub fn can_advance_full_snapshot_checkpoints(
        &self,
        participating_regions: &BTreeSet<u64>,
        watermark_map: &HashMap<u64, u64>,
    ) -> bool {
        !participating_regions.is_empty()
            && participating_regions.len() == watermark_map.len()
            && participating_regions
                .iter()
                .all(|region_id| watermark_map.contains_key(region_id))
    }

    /// Incremental advances are limited to participating regions whose returned
    /// watermark is not older than the stored checkpoint.
    pub fn can_advance_incremental_checkpoints_with_participation(
        &self,
        participating_regions: &BTreeSet<u64>,
        watermark_map: &HashMap<u64, u64>,
    ) -> bool {
        !self.incremental_disabled
            && !self.checkpoints.is_empty()
            && !participating_regions.is_empty()
            && participating_regions.len() == watermark_map.len()
            && participating_regions
                .iter()
                .all(|region_id| self.checkpoints.contains_key(region_id))
            && participating_regions.iter().all(|region_id| {
                let checkpoint = self.checkpoints.get(region_id);
                watermark_map
                    .get(region_id)
                    .zip(checkpoint)
                    .is_some_and(|(seq, checkpoint)| seq >= checkpoint)
            })
    }

    /// Compute the next query delay based on the time window size or the last query duration.
    /// Aiming to avoid too frequent queries. But also not too long delay.
    ///
    /// next wait time is calculated as:
    /// last query duration, capped by [max(min_run_interval, time_window_size), max_timeout],
    /// note at most wait for `max_timeout`.
    ///
    /// if current the dirty time range is longer than one query can handle,
    /// execute immediately to faster clean up dirty time windows.
    /// Active fenced repairs also execute immediately while pending windows
    /// remain: the current backlog has moved out of live dirty windows and into
    /// `pending_fenced_repair.pending_windows`.
    ///
    /// If `prefer_short_incremental_cadence` is true, run incremental queries
    /// more often when there is no large dirty backlog. This only reduces the
    /// chance of hitting a stale cursor after flush; it is not required for
    /// correctness.
    pub fn get_next_start_query_time(
        &self,
        flow_id: FlowId,
        time_window_size: &Option<Duration>,
        min_refresh_duration: Duration,
        max_timeout: Option<Duration>,
        max_filter_num_per_query: usize,
        prefer_short_incremental_cadence: bool,
    ) -> Instant {
        // = last query duration, capped by [max(min_run_interval, time_window_size), max_timeout], note at most `max_timeout`
        let lower = time_window_size.unwrap_or(min_refresh_duration);
        let next_duration = self.last_query_duration.max(lower);
        let next_duration = if let Some(max_timeout) = max_timeout {
            next_duration.min(max_timeout)
        } else {
            next_duration
        };

        if self
            .pending_fenced_repair
            .as_ref()
            .is_some_and(|repair| !repair.pending_windows().is_empty())
        {
            debug!(
                "Flow id = {}, active fenced repair still has pending windows, execute immediately",
                flow_id,
            );
            return Instant::now();
        }

        let cur_dirty_window_size = self.dirty_time_windows.window_size();
        // compute how much time range can be handled in one query
        let max_query_update_range = (*time_window_size)
            .unwrap_or_default()
            .mul_f64(max_filter_num_per_query as f64);
        // if dirty time range is more than one query can handle, execute immediately
        // to faster clean up dirty time windows
        if cur_dirty_window_size < max_query_update_range {
            if prefer_short_incremental_cadence {
                // Run incremental queries sooner than the normal time-window
                // cadence, while still backing off by at least the previous
                // query duration and respecting the max-timeout cap.
                let next_duration = self.last_query_duration.max(min_refresh_duration);
                let next_duration = if let Some(max_timeout) = max_timeout {
                    next_duration.min(max_timeout)
                } else {
                    next_duration
                };
                self.last_update_time + next_duration
            } else {
                self.last_update_time + next_duration
            }
        } else {
            // if dirty time windows can't be clean up in one query, execute immediately to faster
            // clean up dirty time windows
            debug!(
                "Flow id = {}, still have too many {} dirty time window({:?}), execute immediately",
                flow_id,
                self.dirty_time_windows.windows.len(),
                self.dirty_time_windows.windows
            );
            Instant::now()
        }
    }
}

/// For keep recording of dirty time windows, which is time window that have new data inserted
/// since last query.
#[derive(Debug, Clone)]
pub struct DirtyTimeWindows {
    /// windows's `start -> end` and non-overlapping
    /// `end` is exclusive(and optional)
    windows: BTreeMap<Timestamp, Option<Timestamp>>,
    /// Maximum number of filters allowed in a single query
    max_filter_num_per_query: usize,
    /// Time window merge distance
    ///
    time_window_merge_threshold: usize,
}

impl DirtyTimeWindows {
    pub fn new(max_filter_num_per_query: usize, time_window_merge_threshold: usize) -> Self {
        Self {
            windows: BTreeMap::new(),
            max_filter_num_per_query,
            time_window_merge_threshold,
        }
    }

    #[cfg(test)]
    pub(crate) fn max_filter_num_per_query(&self) -> usize {
        self.max_filter_num_per_query
    }

    #[cfg(test)]
    pub(crate) fn time_window_merge_threshold(&self) -> usize {
        self.time_window_merge_threshold
    }
}

impl Default for DirtyTimeWindows {
    fn default() -> Self {
        Self {
            windows: BTreeMap::new(),
            max_filter_num_per_query: 20,
            time_window_merge_threshold: 3,
        }
    }
}

impl DirtyTimeWindows {
    /// Time window merge distance
    ///
    /// TODO(discord9): make those configurable
    pub const MERGE_DIST: i32 = 3;

    /// Add lower bounds to the dirty time windows. Upper bounds are ignored.
    ///
    /// # Arguments
    ///
    /// * `lower_bounds` - An iterator of lower bounds to be added.
    pub fn add_lower_bounds(&mut self, lower_bounds: impl Iterator<Item = Timestamp>) {
        for lower_bound in lower_bounds {
            let entry = self.windows.entry(lower_bound);
            entry.or_insert(None);
        }
    }

    pub fn window_size(&self) -> Duration {
        let mut ret = Duration::from_secs(0);
        for (start, end) in &self.windows {
            if let Some(end) = end
                && let Some(duration) = end.sub(start)
            {
                ret += duration.to_std().unwrap_or_default();
            }
        }
        ret
    }

    pub fn add_window(&mut self, start: Timestamp, end: Option<Timestamp>) {
        self.add_or_merge_window(start, end);
    }

    pub fn add_windows(&mut self, time_ranges: Vec<(Timestamp, Timestamp)>) {
        for (start, end) in time_ranges {
            self.add_or_merge_window(start, Some(end));
        }
    }

    /// Add all dirty markers from another dirty-window set.
    pub fn add_dirty_windows(&mut self, dirty_windows: &DirtyTimeWindows) {
        for (start, end) in &dirty_windows.windows {
            self.add_or_merge_window(*start, *end);
        }
    }

    fn add_or_merge_window(&mut self, start: Timestamp, end: Option<Timestamp>) {
        self.windows
            .entry(start)
            .and_modify(|current_end| {
                *current_end = Self::union_window_end(*current_end, end);
            })
            .or_insert(end);
    }

    fn union_window_end(
        current_end: Option<Timestamp>,
        incoming_end: Option<Timestamp>,
    ) -> Option<Timestamp> {
        match (current_end, incoming_end) {
            (Some(current), Some(incoming)) => Some(current.max(incoming)),
            // `None` is a dirty marker without a known upper bound.  When one
            // side has a concrete end, keep it so merging a restored snapshot
            // never shrinks an already-known dirty range with the same start.
            (Some(end), None) | (None, Some(end)) => Some(end),
            (None, None) => None,
        }
    }

    /// Clean all dirty time windows, useful when can't found time window expr
    pub fn clean(&mut self) {
        self.windows.clear();
    }

    /// Set windows to be dirty, only useful for full aggr without time window
    /// to mark some new data is inserted
    pub fn set_dirty(&mut self) {
        self.add_or_merge_window(Timestamp::new_second(0), None);
    }

    /// Number of dirty windows.
    pub fn len(&self) -> usize {
        self.windows.len()
    }

    pub fn is_empty(&self) -> bool {
        self.windows.is_empty()
    }

    /// Get the effective count of time windows, which is the number of time windows that can be
    /// used for query, compute from total time window range divided by `window_size`.
    pub fn effective_count(&self, window_size: &Duration) -> usize {
        if self.windows.is_empty() {
            return 0;
        }
        let window_size =
            chrono::Duration::from_std(*window_size).unwrap_or(chrono::Duration::zero());
        let total_window_time_range =
            self.windows
                .iter()
                .fold(chrono::Duration::zero(), |acc, (start, end)| {
                    if let Some(end) = end {
                        acc + end.sub(start).unwrap_or(chrono::Duration::zero())
                    } else {
                        acc + window_size
                    }
                });

        // not sure window_size is zero have any meaning, but just in case
        if window_size.num_seconds() == 0 {
            0
        } else {
            (total_window_time_range.num_seconds() / window_size.num_seconds()) as usize
        }
    }

    /// Generate all filter expressions consuming all time windows
    ///
    /// there is two limits:
    /// - shouldn't return a too long time range(<=`window_size * window_cnt`), so that the query can be executed in a reasonable time
    /// - shouldn't return too many time range exprs, so that the query can be parsed properly instead of causing parser to overflow
    pub fn gen_filter_exprs(
        &mut self,
        col_name: &str,
        expire_lower_bound: Option<Timestamp>,
        window_size: chrono::Duration,
        window_cnt: usize,
        flow_id: FlowId,
        task_ctx: Option<&BatchingTask>,
    ) -> Result<Option<FilterExprInfo>, Error> {
        ensure!(
            window_size.num_seconds() > 0,
            UnexpectedSnafu {
                reason: "window_size is zero, can't generate filter exprs",
            }
        );

        debug!(
            "expire_lower_bound: {:?}, window_size: {:?}",
            expire_lower_bound.map(|t| t.to_iso8601_string()),
            window_size
        );
        self.merge_dirty_time_windows(window_size, expire_lower_bound)?;

        if self.windows.len() > window_cnt {
            let first_time_window = self.windows.first_key_value();
            let last_time_window = self.windows.last_key_value();

            if let Some(task_ctx) = task_ctx {
                debug!(
                    "Flow id = {:?}, too many time windows: {}, only the first {} are taken for this query, the group by expression might be wrong. Time window expr={:?}, expire_after={:?}, first_time_window={:?}, last_time_window={:?}, the original query: {:?}",
                    task_ctx.config.flow_id,
                    self.windows.len(),
                    window_cnt,
                    task_ctx.config.time_window_expr,
                    task_ctx.config.expire_after,
                    first_time_window,
                    last_time_window,
                    task_ctx.config.query
                );
            } else {
                debug!(
                    "Flow id = {:?}, too many time windows: {}, only the first {} are taken for this query, the group by expression might be wrong. first_time_window={:?}, last_time_window={:?}",
                    flow_id,
                    self.windows.len(),
                    window_cnt,
                    first_time_window,
                    last_time_window
                )
            }
        }

        // get the first `window_cnt` time windows
        let max_time_range = window_size * window_cnt as i32;

        let mut to_be_query = BTreeMap::new();
        let mut new_windows = self.windows.clone();
        let mut cur_time_range = chrono::Duration::zero();
        for (idx, (start, end)) in self.windows.iter().enumerate() {
            let first_end = start
                .add_duration(window_size.to_std().unwrap())
                .context(TimeSnafu)?;
            let end = end.unwrap_or(first_end);

            // if time range is too long, stop
            if cur_time_range >= max_time_range {
                break;
            }

            // if we have enough time windows, stop
            if idx >= window_cnt {
                break;
            }

            let Some(x) = end.sub(start) else {
                continue;
            };
            if cur_time_range + x <= max_time_range {
                to_be_query.insert(*start, Some(end));
                new_windows.remove(start);
                cur_time_range += x;
            } else {
                // too large a window, split it
                // split at window_size * times
                let surplus = max_time_range - cur_time_range;
                if surplus.num_seconds() <= window_size.num_seconds() {
                    // Skip splitting if surplus is smaller than window_size
                    break;
                }
                let times = surplus.num_seconds() / window_size.num_seconds();

                let split_offset = window_size * times as i32;
                let split_at = start
                    .add_duration(split_offset.to_std().unwrap())
                    .context(TimeSnafu)?;
                to_be_query.insert(*start, Some(split_at));

                // remove the original window
                new_windows.remove(start);
                new_windows.insert(split_at, Some(end));
                cur_time_range += split_offset;
                break;
            }
        }

        self.windows = new_windows;

        METRIC_FLOW_BATCHING_ENGINE_QUERY_WINDOW_CNT
            .with_label_values(&[flow_id.to_string().as_str()])
            .observe(to_be_query.len() as f64);

        let full_time_range = to_be_query
            .iter()
            .fold(chrono::Duration::zero(), |acc, (start, end)| {
                if let Some(end) = end {
                    acc + end.sub(start).unwrap_or(chrono::Duration::zero())
                } else {
                    acc + window_size
                }
            })
            .num_seconds() as f64;
        METRIC_FLOW_BATCHING_ENGINE_QUERY_WINDOW_SIZE
            .with_label_values(&[flow_id.to_string().as_str()])
            .observe(full_time_range);

        let stalled_time_range =
            self.windows
                .iter()
                .fold(chrono::Duration::zero(), |acc, (start, end)| {
                    if let Some(end) = end {
                        acc + end.sub(start).unwrap_or(chrono::Duration::zero())
                    } else {
                        acc + window_size
                    }
                });

        METRIC_FLOW_BATCHING_ENGINE_STALLED_WINDOW_SIZE
            .with_label_values(&[flow_id.to_string().as_str()])
            .observe(stalled_time_range.num_seconds() as f64);

        let std_window_size = window_size.to_std().map_err(|e| {
            InternalSnafu {
                reason: e.to_string(),
            }
            .build()
        })?;

        let mut expr_lst = vec![];
        let mut time_ranges = vec![];
        for (start, end) in to_be_query.into_iter() {
            // align using time window exprs
            let (start, end) = if let Some(ctx) = task_ctx {
                let Some(time_window_expr) = &ctx.config.time_window_expr else {
                    UnexpectedSnafu {
                        reason: "time_window_expr is not set",
                    }
                    .fail()?
                };
                Self::align_time_window(start, end, time_window_expr)?
            } else {
                (start, end)
            };
            let end = end.unwrap_or(start.add_duration(std_window_size).context(TimeSnafu)?);
            time_ranges.push((start, end));

            debug!(
                "Time window start: {:?}, end: {:?}",
                start.to_iso8601_string(),
                end.to_iso8601_string()
            );

            use datafusion_expr::{col, lit};
            let lower = to_df_literal(start)?;
            let upper = to_df_literal(end)?;
            let expr = col(col_name)
                .gt_eq(lit(lower))
                .and(col(col_name).lt(lit(upper)));
            expr_lst.push(expr);
        }
        let expr = expr_lst.into_iter().reduce(|a, b| a.or(b));
        let ret = expr.map(|expr| FilterExprInfo {
            expr,
            col_name: col_name.to_string(),
            time_ranges,
            window_size,
        });
        Ok(ret)
    }

    /// Align a time range `[start, end)` (end is optional and exclusive) to
    /// time window boundaries defined by the time window expr.
    pub(crate) fn align_time_window(
        start: Timestamp,
        end: Option<Timestamp>,
        time_window_expr: &TimeWindowExpr,
    ) -> Result<(Timestamp, Option<Timestamp>), Error> {
        let align_start = time_window_expr.eval(start)?.0.context(UnexpectedSnafu {
            reason: format!(
                "Failed to align start time {:?} with time window expr {:?}",
                start, time_window_expr
            ),
        })?;
        let align_end = end
            .and_then(|end| {
                time_window_expr
                    .eval(end)
                    // if after aligned, end is the same, then use end(because it's already aligned) else use aligned end
                    .map(|r| if r.0 == Some(end) { r.0 } else { r.1 })
                    .transpose()
            })
            .transpose()?;
        Ok((align_start, align_end))
    }

    /// Merge time windows that overlaps or get too close
    ///
    /// TODO(discord9): not merge and prefer to send smaller time windows? how?
    pub fn merge_dirty_time_windows(
        &mut self,
        window_size: chrono::Duration,
        expire_lower_bound: Option<Timestamp>,
    ) -> Result<(), Error> {
        if self.windows.is_empty() {
            return Ok(());
        }

        let mut new_windows = BTreeMap::new();

        let std_window_size = window_size.to_std().map_err(|e| {
            InternalSnafu {
                reason: e.to_string(),
            }
            .build()
        })?;

        // previous time window
        let mut prev_tw = None;
        for (mut lower_bound, upper_bound) in std::mem::take(&mut self.windows) {
            // filter out expired time window
            if let Some(expire_lower_bound) = expire_lower_bound {
                match upper_bound {
                    // A bounded range ending at or before the expire bound is
                    // fully expired, drop it.
                    Some(upper_bound) if upper_bound <= expire_lower_bound => continue,
                    // A bounded range crossing the expire bound keeps its
                    // still-live suffix. The expire bound is aligned to the
                    // time window boundary by the caller, so the clipped start
                    // stays aligned.
                    Some(_) if lower_bound < expire_lower_bound => {
                        lower_bound = expire_lower_bound;
                    }
                    // Unbounded windows keep the start-based behavior.
                    None if lower_bound < expire_lower_bound => continue,
                    _ => {}
                }
            }

            let Some(prev_tw) = &mut prev_tw else {
                prev_tw = Some((lower_bound, upper_bound));
                continue;
            };

            // if cur.lower - prev.upper <= window_size * MERGE_DIST, merge
            // this also deal with overlap windows because cur.lower > prev.lower is always true
            let prev_upper = prev_tw
                .1
                .unwrap_or(prev_tw.0.add_duration(std_window_size).context(TimeSnafu)?);
            prev_tw.1 = Some(prev_upper);

            let cur_upper = upper_bound.unwrap_or(
                lower_bound
                    .add_duration(std_window_size)
                    .context(TimeSnafu)?,
            );

            if lower_bound
                .sub(&prev_upper)
                .map(|dist| dist <= window_size * self.time_window_merge_threshold as i32)
                .unwrap_or(false)
            {
                // Union the two windows: the current window may be contained
                // in the previous one, so keep the larger upper bound.
                prev_tw.1 = Some(prev_upper.max(cur_upper));
            } else {
                new_windows.insert(prev_tw.0, prev_tw.1);
                *prev_tw = (lower_bound, Some(cur_upper));
            }
        }

        if let Some(prev_tw) = prev_tw {
            new_windows.insert(prev_tw.0, prev_tw.1);
        }

        self.windows = new_windows;

        Ok(())
    }
}

pub(crate) fn to_df_literal(value: Timestamp) -> Result<datafusion_common::ScalarValue, Error> {
    let value = Value::from(value);
    let value = value
        .try_to_scalar_value(&value.data_type())
        .with_context(|_| DatatypesSnafu {
            extra: format!("Failed to convert to scalar value: {}", value),
        })?;
    Ok(value)
}

#[derive(Debug, Clone)]
enum ExecState {
    Idle,
    Executing,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CheckpointMode {
    FullSnapshot,
    Incremental,
}

/// Dirty windows that must be repaired under a frozen full-snapshot watermark.
/// This is a FullSnapshot sub-state, not a separate checkpoint mode.
#[derive(Debug, Clone)]
pub struct FencedRepair {
    high: BTreeMap<u64, u64>,
    pending_windows: DirtyTimeWindows,
}

impl FencedRepair {
    /// Frozen high watermark `H` used as the snapshot upper bound for chunks.
    pub fn high(&self) -> &BTreeMap<u64, u64> {
        &self.high
    }

    /// Dirty windows still waiting to be repaired under `high`.
    pub fn pending_windows(&self) -> &DirtyTimeWindows {
        &self.pending_windows
    }
}

/// Filter Expression's information
#[derive(Debug, Clone)]
pub struct FilterExprInfo {
    pub expr: datafusion_expr::Expr,
    pub col_name: String,
    pub time_ranges: Vec<(Timestamp, Timestamp)>,
    pub window_size: chrono::Duration,
}

impl FilterExprInfo {
    pub fn total_window_length(&self) -> chrono::Duration {
        self.time_ranges
            .iter()
            .fold(chrono::Duration::zero(), |acc, (start, end)| {
                acc + end.sub(start).unwrap_or(chrono::Duration::zero())
            })
    }

    pub fn predicate_for_col(
        &self,
        col_name: &str,
    ) -> Result<Option<datafusion_expr::Expr>, Error> {
        use datafusion_common::Column;
        use datafusion_expr::{Expr, lit};

        let mut expr_lst = Vec::with_capacity(self.time_ranges.len());
        for (start, end) in &self.time_ranges {
            let lower = to_df_literal(*start)?;
            let upper = to_df_literal(*end)?;
            let filter_col = || Expr::Column(Column::new_unqualified(col_name));
            expr_lst.push(
                filter_col()
                    .gt_eq(lit(lower))
                    .and(filter_col().lt(lit(upper))),
            );
        }

        Ok(expr_lst.into_iter().reduce(|a, b| a.or(b)))
    }
}

#[cfg(test)]
mod test {
    use pretty_assertions::assert_eq;
    use session::context::QueryContext;

    use super::*;
    use crate::batching_mode::time_window::find_time_window_expr;
    use crate::batching_mode::utils::sql_to_df_plan;
    use crate::test_utils::create_test_query_engine;

    #[test]
    fn test_task_state_records_last_execution_time() {
        let query_ctx = QueryContext::arc();
        let (_tx, rx) = tokio::sync::oneshot::channel();
        let mut state = TaskState::new(query_ctx, rx);

        assert_eq!(None, state.last_execution_time_millis());
        state.after_query_exec(std::time::Duration::from_millis(1), false);
        assert_eq!(None, state.last_execution_time_millis());

        state.after_query_exec(std::time::Duration::from_millis(1), true);
        assert!(state.last_execution_time_millis().is_some());
    }

    #[test]
    fn test_merge_dirty_time_windows() {
        let merge_dist = DirtyTimeWindows::default().time_window_merge_threshold;
        let testcases = vec![
            // just enough to merge
            (
                vec![
                    Timestamp::new_second(0),
                    Timestamp::new_second((1 + merge_dist as i64) * 5 * 60),
                ],
                (chrono::Duration::seconds(5 * 60), None),
                BTreeMap::from([(
                    Timestamp::new_second(0),
                    Some(Timestamp::new_second((2 + merge_dist as i64) * 5 * 60)),
                )]),
                Some(
                    "((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:25:00' AS TIMESTAMP)))",
                ),
            ),
            // separate time window
            (
                vec![
                    Timestamp::new_second(0),
                    Timestamp::new_second((2 + merge_dist as i64) * 5 * 60),
                ],
                (chrono::Duration::seconds(5 * 60), None),
                BTreeMap::from([
                    (
                        Timestamp::new_second(0),
                        Some(Timestamp::new_second(5 * 60)),
                    ),
                    (
                        Timestamp::new_second((2 + merge_dist as i64) * 5 * 60),
                        Some(Timestamp::new_second((3 + merge_dist as i64) * 5 * 60)),
                    ),
                ]),
                Some(
                    "(((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:05:00' AS TIMESTAMP))) OR ((ts >= CAST('1970-01-01 00:25:00' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:30:00' AS TIMESTAMP))))",
                ),
            ),
            // overlapping
            (
                vec![
                    Timestamp::new_second(0),
                    Timestamp::new_second((merge_dist as i64) * 5 * 60),
                ],
                (chrono::Duration::seconds(5 * 60), None),
                BTreeMap::from([(
                    Timestamp::new_second(0),
                    Some(Timestamp::new_second((1 + merge_dist as i64) * 5 * 60)),
                )]),
                Some(
                    "((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:20:00' AS TIMESTAMP)))",
                ),
            ),
            // complex overlapping
            (
                vec![
                    Timestamp::new_second(0),
                    Timestamp::new_second((merge_dist as i64) * 3),
                    Timestamp::new_second((merge_dist as i64) * 3 * 2),
                ],
                (chrono::Duration::seconds(3), None),
                BTreeMap::from([(
                    Timestamp::new_second(0),
                    Some(Timestamp::new_second((merge_dist as i64) * 7)),
                )]),
                Some(
                    "((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:00:21' AS TIMESTAMP)))",
                ),
            ),
            // split range
            (
                Vec::from_iter((0..20).map(|i| Timestamp::new_second(i * 3)).chain(
                    std::iter::once(Timestamp::new_second(
                        60 + 3 * (DirtyTimeWindows::MERGE_DIST as i64 + 1),
                    )),
                )),
                (chrono::Duration::seconds(3), None),
                BTreeMap::from([
                    (Timestamp::new_second(0), Some(Timestamp::new_second(60))),
                    (
                        Timestamp::new_second(60 + 3 * (DirtyTimeWindows::MERGE_DIST as i64 + 1)),
                        Some(Timestamp::new_second(
                            60 + 3 * (DirtyTimeWindows::MERGE_DIST as i64 + 1) + 3,
                        )),
                    ),
                ]),
                Some(
                    "((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:01:00' AS TIMESTAMP)))",
                ),
            ),
            // split 2 min into 1 min
            (
                Vec::from_iter((0..40).map(|i| Timestamp::new_second(i * 3))),
                (chrono::Duration::seconds(3), None),
                BTreeMap::from([(
                    Timestamp::new_second(0),
                    Some(Timestamp::new_second(40 * 3)),
                )]),
                Some(
                    "((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:01:00' AS TIMESTAMP)))",
                ),
            ),
            // split 3s + 1min into 3s + 57s
            (
                Vec::from_iter(
                    std::iter::once(Timestamp::new_second(0))
                        .chain((0..40).map(|i| Timestamp::new_second(20 + i * 3))),
                ),
                (chrono::Duration::seconds(3), None),
                BTreeMap::from([
                    (Timestamp::new_second(0), Some(Timestamp::new_second(3))),
                    (Timestamp::new_second(20), Some(Timestamp::new_second(140))),
                ]),
                Some(
                    "(((ts >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:00:03' AS TIMESTAMP))) OR ((ts >= CAST('1970-01-01 00:00:20' AS TIMESTAMP)) AND (ts < CAST('1970-01-01 00:01:17' AS TIMESTAMP))))",
                ),
            ),
            // expired
            (
                vec![
                    Timestamp::new_second(0),
                    Timestamp::new_second((merge_dist as i64) * 5 * 60),
                ],
                (
                    chrono::Duration::seconds(5 * 60),
                    Some(Timestamp::new_second((merge_dist as i64) * 6 * 60)),
                ),
                BTreeMap::from([]),
                None,
            ),
        ];
        // let len = testcases.len();
        // let testcases = testcases[(len - 2)..(len - 1)].to_vec();
        for (lower_bounds, (window_size, expire_lower_bound), expected, expected_filter_expr) in
            testcases
        {
            let mut dirty = DirtyTimeWindows::default();
            dirty.add_lower_bounds(lower_bounds.into_iter());
            dirty
                .merge_dirty_time_windows(window_size, expire_lower_bound)
                .unwrap();
            assert_eq!(expected, dirty.windows);
            let filter_expr = dirty
                .gen_filter_exprs(
                    "ts",
                    expire_lower_bound,
                    window_size,
                    dirty.max_filter_num_per_query,
                    0,
                    None,
                )
                .unwrap()
                .map(|e| e.expr);

            let unparser = datafusion::sql::unparser::Unparser::default();
            let to_sql = filter_expr
                .as_ref()
                .map(|e| unparser.expr_to_sql(e).unwrap().to_string());
            assert_eq!(expected_filter_expr, to_sql.as_deref());
        }
    }

    #[test]
    fn test_merge_dirty_time_windows_with_bounded_ranges() {
        let window_size = chrono::Duration::seconds(5);
        let testcases = vec![
            // A contained bounded range must not shrink the containing window:
            // [0s, 15s) merged with nested [5s, 10s) stays [0s, 15s).
            (
                vec![
                    (Timestamp::new_second(0), Some(Timestamp::new_second(15))),
                    (Timestamp::new_second(5), Some(Timestamp::new_second(10))),
                ],
                BTreeMap::from([(Timestamp::new_second(0), Some(Timestamp::new_second(15)))]),
            ),
            // An unbounded dirty window nested in a bounded range must not
            // shrink the range either: [0s, 15s) merged with 3s (window end
            // 8s) stays [0s, 15s).
            (
                vec![
                    (Timestamp::new_second(0), Some(Timestamp::new_second(15))),
                    (Timestamp::new_second(3), None),
                ],
                BTreeMap::from([(Timestamp::new_second(0), Some(Timestamp::new_second(15)))]),
            ),
            // Disjoint bounded ranges far apart are kept separate.
            (
                vec![
                    (Timestamp::new_second(0), Some(Timestamp::new_second(5))),
                    (Timestamp::new_second(100), Some(Timestamp::new_second(110))),
                ],
                BTreeMap::from([
                    (Timestamp::new_second(0), Some(Timestamp::new_second(5))),
                    (Timestamp::new_second(100), Some(Timestamp::new_second(110))),
                ]),
            ),
            // Overlapping bounded ranges are unioned: [0s, 10s) and [5s, 20s)
            // become [0s, 20s).
            (
                vec![
                    (Timestamp::new_second(0), Some(Timestamp::new_second(10))),
                    (Timestamp::new_second(5), Some(Timestamp::new_second(20))),
                ],
                BTreeMap::from([(Timestamp::new_second(0), Some(Timestamp::new_second(20)))]),
            ),
        ];

        for (windows, expected) in testcases {
            let mut dirty = DirtyTimeWindows::default();
            for (start, end) in windows {
                dirty.add_window(start, end);
            }
            dirty.merge_dirty_time_windows(window_size, None).unwrap();
            assert_eq!(expected, dirty.windows);
        }

        // Expire bound handling for bounded ranges vs unbounded windows.
        let expire_testcases = vec![
            // A bounded range ending at the expire bound is fully expired.
            (
                vec![(Timestamp::new_second(0), Some(Timestamp::new_second(10)))],
                BTreeMap::from([]),
            ),
            // A bounded range ending before the expire bound is fully expired.
            (
                vec![(Timestamp::new_second(0), Some(Timestamp::new_second(5)))],
                BTreeMap::from([]),
            ),
            // A bounded range crossing the expire bound keeps its live
            // suffix: [0s, 15s) with expire 10s becomes [10s, 15s).
            (
                vec![(Timestamp::new_second(0), Some(Timestamp::new_second(15)))],
                BTreeMap::from([(Timestamp::new_second(10), Some(Timestamp::new_second(15)))]),
            ),
            // A bounded range starting at the expire bound is kept intact.
            (
                vec![(Timestamp::new_second(10), Some(Timestamp::new_second(15)))],
                BTreeMap::from([(Timestamp::new_second(10), Some(Timestamp::new_second(15)))]),
            ),
            // An unbounded window starting before the expire bound is
            // dropped, preserving the existing start-based behavior.
            (vec![(Timestamp::new_second(5), None)], BTreeMap::from([])),
            // An unbounded window starting at the expire bound is kept.
            (
                vec![(Timestamp::new_second(10), None)],
                BTreeMap::from([(Timestamp::new_second(10), None)]),
            ),
        ];

        for (windows, expected) in expire_testcases {
            let mut dirty = DirtyTimeWindows::default();
            for (start, end) in windows {
                dirty.add_window(start, end);
            }
            dirty
                .merge_dirty_time_windows(window_size, Some(Timestamp::new_second(10)))
                .unwrap();
            assert_eq!(expected, dirty.windows);
        }
    }

    #[tokio::test]
    async fn test_align_time_window() {
        type TimeWindow = (Timestamp, Option<Timestamp>);
        struct TestCase {
            sql: String,
            aligns: Vec<(TimeWindow, TimeWindow)>,
        }
        let testcases: Vec<TestCase> = vec![TestCase{
            sql: "SELECT date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window;".to_string(),
            aligns: vec![
                ((Timestamp::new_second(3), None), (Timestamp::new_second(0), None)),
                ((Timestamp::new_second(8), None), (Timestamp::new_second(5), None)),
                ((Timestamp::new_second(8), Some(Timestamp::new_second(10))), (Timestamp::new_second(5), Some(Timestamp::new_second(10)))),
                ((Timestamp::new_second(8), Some(Timestamp::new_second(9))), (Timestamp::new_second(5), Some(Timestamp::new_second(10)))),
            ],
        }];

        let query_engine = create_test_query_engine();
        let ctx = QueryContext::arc();
        for TestCase { sql, aligns } in testcases {
            let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), &sql, true)
                .await
                .unwrap();

            let (column_name, time_window_expr, _, df_schema) = find_time_window_expr(
                &plan,
                query_engine.engine_state().catalog_manager().clone(),
                ctx.clone(),
            )
            .await
            .unwrap();

            let time_window_expr = time_window_expr
                .map(|expr| {
                    TimeWindowExpr::from_expr(
                        &expr,
                        &column_name,
                        &df_schema,
                        &query_engine.engine_state().session_state(),
                    )
                })
                .transpose()
                .unwrap()
                .unwrap();

            for (before_align, expected_after_align) in aligns {
                let after_align = DirtyTimeWindows::align_time_window(
                    before_align.0,
                    before_align.1,
                    &time_window_expr,
                )
                .unwrap();
                assert_eq!(expected_after_align, after_align);
            }
        }
    }

    #[test]
    fn test_task_state_checkpoint_mode_and_advancement() {
        let query_ctx = QueryContext::arc();
        let (_tx, rx) = tokio::sync::oneshot::channel();
        let mut state = TaskState::new(query_ctx, rx);

        assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot);
        assert!(state.checkpoints().is_empty());

        state.advance_checkpoints(HashMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]));
        assert_eq!(state.checkpoint_mode(), CheckpointMode::Incremental);
        assert_eq!(
            state.checkpoints(),
            &BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)])
        );

        state.mark_full_snapshot();
        assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot);
        assert_eq!(
            state.checkpoints(),
            &BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)])
        );
    }

    #[test]
    fn test_mark_full_snapshot_restores_pending_fenced_repair_windows() {
        let query_ctx = QueryContext::arc();
        let (_tx, rx) = tokio::sync::oneshot::channel();
        let mut state = TaskState::new(query_ctx, rx);
        state
            .dirty_time_windows
            .add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15)));
        state
            .dirty_time_windows
            .add_window(Timestamp::new_second(100), Some(Timestamp::new_second(105)));

        state
            .start_fenced_repair(BTreeMap::from([(1_u64, 10_u64)]))
            .unwrap();
        assert!(state.dirty_time_windows.is_empty());
        assert_eq!(
            state
                .pending_fenced_repair()
                .unwrap()
                .pending_windows()
                .len(),
            2
        );

        state.mark_full_snapshot();

        assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot);
        assert!(state.pending_fenced_repair().is_none());
        assert_eq!(state.dirty_time_windows.len(), 2);
    }

    #[test]
    fn test_disable_incremental_persists_full_snapshot_mode() {
        let query_ctx = QueryContext::arc();
        let (_tx, rx) = tokio::sync::oneshot::channel();
        let mut state = TaskState::new(query_ctx, rx);

        assert!(!state.is_incremental_disabled());

        // After disable, mode becomes FullSnapshot and flag is set.
        state.disable_incremental();
        assert!(state.is_incremental_disabled());
        assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot);

        // `advance_checkpoints` will NOT transition to Incremental when disabled.
        state.advance_checkpoints(HashMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]));
        assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot);
        assert_eq!(
            state.checkpoints(),
            &BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)])
        );

        // `mark_full_snapshot` does not re-enable incremental.
        state.mark_full_snapshot();
        assert!(state.is_incremental_disabled());
        assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot);
    }

    #[test]
    fn test_full_snapshot_checkpoint_advancement_requires_participating_regions() {
        let query_ctx = QueryContext::arc();
        let (_tx, rx) = tokio::sync::oneshot::channel();
        let state = TaskState::new(query_ctx, rx);

        assert!(!state.can_advance_full_snapshot_checkpoints(&BTreeSet::new(), &HashMap::new()));
        assert!(!state.can_advance_full_snapshot_checkpoints(
            &BTreeSet::from([1_u64, 2_u64]),
            &HashMap::from([(1_u64, 10_u64)]),
        ));
        assert!(state.can_advance_full_snapshot_checkpoints(
            &BTreeSet::from([1_u64, 2_u64]),
            &HashMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]),
        ));
    }

    #[test]
    fn test_incremental_checkpoint_advancement_requires_participation_alignment() {
        let query_ctx = QueryContext::arc();
        let (_tx, rx) = tokio::sync::oneshot::channel();
        let mut state = TaskState::new(query_ctx, rx);
        state.advance_checkpoints(HashMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]));

        assert!(
            state.can_advance_incremental_checkpoints_with_participation(
                &BTreeSet::from([1_u64]),
                &HashMap::from([(1_u64, 11_u64)]),
            )
        );
        assert!(
            !state.can_advance_incremental_checkpoints_with_participation(
                &BTreeSet::from([1_u64, 2_u64]),
                &HashMap::from([(1_u64, 11_u64)]),
            )
        );
        assert!(
            !state.can_advance_incremental_checkpoints_with_participation(
                &BTreeSet::from([3_u64]),
                &HashMap::from([(3_u64, 11_u64)]),
            )
        );
        assert!(
            !state.can_advance_incremental_checkpoints_with_participation(
                &BTreeSet::from([1_u64]),
                &HashMap::from([(1_u64, 9_u64)]),
            )
        );
        assert!(
            state.can_advance_incremental_checkpoints_with_participation(
                &BTreeSet::from([1_u64, 2_u64]),
                &HashMap::from([(1_u64, 11_u64), (2_u64, 21_u64)]),
            )
        );

        state.disable_incremental();
        assert!(
            !state.can_advance_incremental_checkpoints_with_participation(
                &BTreeSet::from([1_u64, 2_u64]),
                &HashMap::from([(1_u64, 12_u64), (2_u64, 22_u64)]),
            )
        );
    }

    #[test]
    fn test_incremental_checkpoint_advancement_merges_participating_subset() {
        let query_ctx = QueryContext::arc();
        let (_tx, rx) = tokio::sync::oneshot::channel();
        let mut state = TaskState::new(query_ctx, rx);
        state.advance_checkpoints(HashMap::from([
            (1_u64, 10_u64),
            (2_u64, 20_u64),
            (3_u64, 30_u64),
        ]));

        state.advance_incremental_checkpoints_with_participation(
            &BTreeSet::from([1_u64, 3_u64]),
            HashMap::from([(1_u64, 12_u64), (3_u64, 35_u64)]),
        );

        assert_eq!(state.checkpoint_mode(), CheckpointMode::Incremental);
        assert_eq!(
            state.checkpoints(),
            &BTreeMap::from([(1_u64, 12_u64), (2_u64, 20_u64), (3_u64, 35_u64)])
        );
    }

    #[test]
    fn test_filter_expr_info_predicate_for_col_empty_ranges() {
        let filter = FilterExprInfo {
            expr: datafusion_expr::col("ts"),
            col_name: "ts".to_string(),
            time_ranges: vec![],
            window_size: chrono::Duration::seconds(1),
        };

        assert!(filter.predicate_for_col("time_window").unwrap().is_none());
    }

    #[test]
    fn test_filter_expr_info_predicate_for_col_single_range() {
        let filter = FilterExprInfo {
            expr: datafusion_expr::col("ts"),
            col_name: "ts".to_string(),
            time_ranges: vec![(Timestamp::new_second(0), Timestamp::new_second(1))],
            window_size: chrono::Duration::seconds(1),
        };

        let predicate = filter.predicate_for_col("time_window").unwrap().unwrap();
        let unparser = datafusion::sql::unparser::Unparser::default();
        assert_eq!(
            "((time_window >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (time_window < CAST('1970-01-01 00:00:01' AS TIMESTAMP)))",
            unparser.expr_to_sql(&predicate).unwrap().to_string()
        );
    }

    #[test]
    fn test_filter_expr_info_predicate_for_col_multiple_ranges() {
        let filter = FilterExprInfo {
            expr: datafusion_expr::col("ts"),
            col_name: "ts".to_string(),
            time_ranges: vec![
                (Timestamp::new_second(0), Timestamp::new_second(1)),
                (Timestamp::new_second(10), Timestamp::new_second(11)),
            ],
            window_size: chrono::Duration::seconds(1),
        };

        let predicate = filter.predicate_for_col("time_window").unwrap().unwrap();
        let unparser = datafusion::sql::unparser::Unparser::default();
        assert_eq!(
            "(((time_window >= CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AND (time_window < CAST('1970-01-01 00:00:01' AS TIMESTAMP))) OR ((time_window >= CAST('1970-01-01 00:00:10' AS TIMESTAMP)) AND (time_window < CAST('1970-01-01 00:00:11' AS TIMESTAMP))))",
            unparser.expr_to_sql(&predicate).unwrap().to_string()
        );
    }

    /// Helper: create a `TaskState` whose `last_update_time` is a known duration in the past.
    fn state_with_past_update(age: Duration) -> TaskState {
        let query_ctx = QueryContext::arc();
        let (_tx, rx) = tokio::sync::oneshot::channel();
        let mut state = TaskState::new(query_ctx, rx);
        state.last_update_time = Instant::now() - age;
        state
    }

    #[test]
    fn test_short_incremental_cadence_uses_min_refresh() {
        // When prefer_short_incremental_cadence is true and dirty backlog is manageable,
        // the next start time should be last_update_time + min_refresh (short cadence),
        // ignoring the longer time_window_size.
        let state = state_with_past_update(Duration::from_secs(10));

        let time_window_size = Some(Duration::from_secs(60)); // large window
        let min_refresh = Duration::from_secs(5);
        let flow_id = 1;

        let result = state.get_next_start_query_time(
            flow_id,
            &time_window_size,
            min_refresh,
            None,
            20,
            true, // prefer_short_incremental_cadence
        );

        // With short cadence, result should be last_update_time + min_refresh.
        let expected = state.last_update_time + min_refresh;
        assert_eq!(result, expected);
    }

    #[test]
    fn test_short_incremental_cadence_respects_last_query_duration() {
        let mut state = state_with_past_update(Duration::from_secs(10));
        state.last_query_duration = Duration::from_secs(20);

        let time_window_size = Some(Duration::from_secs(60));
        let min_refresh = Duration::from_secs(5);
        let flow_id = 1;

        let result = state.get_next_start_query_time(
            flow_id,
            &time_window_size,
            min_refresh,
            None,
            20,
            true,
        );

        assert_eq!(result, state.last_update_time + state.last_query_duration);
    }

    #[test]
    fn test_short_incremental_cadence_respects_max_timeout() {
        let mut state = state_with_past_update(Duration::from_secs(10));
        state.last_query_duration = Duration::from_secs(20);

        let time_window_size = Some(Duration::from_secs(60));
        let min_refresh = Duration::from_secs(30);
        let max_timeout = Duration::from_secs(5);
        let flow_id = 1;

        let result = state.get_next_start_query_time(
            flow_id,
            &time_window_size,
            min_refresh,
            Some(max_timeout),
            20,
            true,
        );

        assert_eq!(result, state.last_update_time + max_timeout);
    }

    #[test]
    fn test_full_snapshot_ignores_short_cadence() {
        // When prefer_short_incremental_cadence is false (full snapshot mode),
        // the normal long-cadence based on time_window_size applies.
        let mut state = state_with_past_update(Duration::from_secs(10));
        // Make last_query_duration small so the lower bound (time_window_size) dominates.
        state.last_query_duration = Duration::from_secs(1);

        let time_window_size = Some(Duration::from_secs(60)); // large window
        let min_refresh = Duration::from_secs(5);
        let flow_id = 1;

        let result = state.get_next_start_query_time(
            flow_id,
            &time_window_size,
            min_refresh,
            None,
            20,
            false, // prefer_short_incremental_cadence = false
        );

        // With normal cadence, result should be last_update_time + time_window_size
        // (since last_query_duration < time_window_size).
        let expected = state.last_update_time + Duration::from_secs(60);
        assert_eq!(result, expected);
    }

    #[test]
    fn test_dirty_window_overflow_schedules_immediately_even_with_short_cadence() {
        // Dirty-window overflow must always schedule immediately,
        // regardless of prefer_short_incremental_cadence.
        let mut state = state_with_past_update(Duration::from_secs(10));
        // Create a very large dirty backlog.
        state
            .dirty_time_windows
            .add_window(Timestamp::new_second(0), Some(Timestamp::new_second(3600)));

        let time_window_size = Some(Duration::from_secs(1)); // tiny window => overflow
        let min_refresh = Duration::from_secs(5);
        let flow_id = 1;

        // With short cadence flag.
        let result = state.get_next_start_query_time(
            flow_id,
            &time_window_size,
            min_refresh,
            None,
            1, // max 1 filter => tiny capacity
            true,
        );
        assert!(
            result <= Instant::now(),
            "dirty overflow should schedule immediately"
        );

        // Without short cadence flag — same behavior.
        let result2 = state.get_next_start_query_time(
            flow_id,
            &time_window_size,
            min_refresh,
            None,
            1,
            false,
        );
        assert!(
            result2 <= Instant::now(),
            "dirty overflow should schedule immediately"
        );
    }

    #[test]
    fn test_pending_fenced_repair_schedules_immediately() {
        let mut state = state_with_past_update(Duration::from_secs(10));
        state
            .dirty_time_windows
            .add_window(Timestamp::new_second(0), Some(Timestamp::new_second(5)));
        state
            .start_fenced_repair(BTreeMap::from([(1_u64, 10_u64)]))
            .unwrap();
        assert!(state.dirty_time_windows.is_empty());
        assert!(!state.fenced_repair_pending_is_empty());

        let result = state.get_next_start_query_time(
            1,
            &Some(Duration::from_secs(60)),
            Duration::from_secs(5),
            None,
            20,
            false,
        );

        assert!(
            result <= Instant::now(),
            "pending fenced repair backlog should schedule immediately"
        );
    }

    #[test]
    fn test_incremental_disabled_ignores_short_cadence() {
        // When prefer_short_incremental_cadence is true but the dirty backlog is
        // manageable, the short cadence is applied. This test verifies that the
        // caller-side guard (checkpoint_mode + !is_incremental_disabled) controls
        // whether short cadence is requested at all — when incremental is disabled,
        // the flag is false, and the long cadence applies.
        //
        // This simulates the case where the caller computed
        // prefer_short_incremental_cadence = false (e.g. incremental disabled
        // or FullSnapshot mode), so the long cadence is used.
        let mut state = state_with_past_update(Duration::from_secs(10));
        state.last_query_duration = Duration::from_secs(1);

        let time_window_size = Some(Duration::from_secs(60));
        let min_refresh = Duration::from_secs(5);
        let flow_id = 1;

        let result = state.get_next_start_query_time(
            flow_id,
            &time_window_size,
            min_refresh,
            None,
            20,
            false, // prefer_short_incremental_cadence = false
        );

        // With normal cadence, result should be last_update_time + time_window_size.
        let expected = state.last_update_time + Duration::from_secs(60);
        assert_eq!(result, expected);
    }
}