mito2/read/

range_cache.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.

//! Utilities for the partition range scan result cache.

use std::future::Future;
use std::mem;
use std::sync::Arc;

use async_stream::try_stream;
use common_telemetry::warn;
use common_time::Timestamp;
use common_time::range::TimestampRange;
use common_time::timestamp::TimeUnit;
use datafusion_expr::expr::Expr;
use datafusion_expr::{Between, BinaryExpr, Operator};
use datatypes::arrow::compute::concat_batches;
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::prelude::ConcreteDataType;
use datatypes::value::scalar_value_to_timestamp;
use futures::TryStreamExt;
use snafu::ResultExt;
use store_api::region_engine::PartitionRange;
use store_api::storage::{FileId, RegionId, SequenceRange, TimeSeriesRowSelector};
use table::predicate::is_string_timestamp_literal;
use tokio::sync::{mpsc, oneshot};

use crate::cache::CacheStrategy;
use crate::cache::range_cache::{RangeCacheDelivery, RangeCacheLookup, RangeCacheWorker};
use crate::error::{ComputeArrowSnafu, Result};
use crate::read::BoxedRecordBatchStream;
use crate::read::read_columns::ReadColumns;
use crate::read::scan_region::StreamContext;
use crate::read::scan_util::PartitionMetrics;
use crate::read::series_reader::SeriesRange;
use crate::region::options::MergeMode;
use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;

const RANGE_CACHE_COMPACT_THRESHOLD_BYTES: usize = 8 * 1024 * 1024;

/// Fingerprint of the scan request fields that affect partition range cache reuse.
///
/// It records a normalized view of the projected columns and filters, plus
/// scan options that can change the returned rows. Schema-dependent metadata
/// and the partition expression version are included so cached results are not
/// reused across incompatible schema or partitioning changes.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct ScanRequestFingerprint {
    /// Projection and filters without the time index and partition exprs.
    inner: Arc<SharedScanRequestFingerprint>,
    /// Filters with the time index column.
    time_filters: Option<Arc<Vec<String>>>,
    series_row_selector: Option<TimeSeriesRowSelector>,
    append_mode: bool,
    filter_deleted: bool,
    merge_mode: MergeMode,
    /// Exact sequence range applied row-level on SST reads, or `None` for scans
    /// without row-level sequence filtering. Kept in the cache key so filtered
    /// range results are never reused for scans with a different (or no) range.
    sequence_range: Option<SequenceRange>,
    stage: RangeScanStage,
    /// We keep the partition expr version to ensure we won't reuse the fingerprint after we change the partition expr.
    /// We store the version instead of the whole partition expr or partition expr filters.
    partition_expr_version: u64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum RangeScanStage {
    Data,
    /// Source batches may be split for a later cross-range merge. Keep their
    /// cached batch boundaries separate from ordinary sequential scans.
    DataForCrossRangeMerge,
    CandidateSeries,
    SeriesData(SeriesRange),
}

#[derive(Debug)]
pub(crate) struct ScanRequestFingerprintBuilder {
    pub(crate) read_columns: ReadColumns,
    pub(crate) read_column_types: Vec<Option<ConcreteDataType>>,
    pub(crate) filters: Vec<String>,
    pub(crate) time_filters: Vec<String>,
    pub(crate) series_row_selector: Option<TimeSeriesRowSelector>,
    pub(crate) append_mode: bool,
    pub(crate) filter_deleted: bool,
    pub(crate) merge_mode: MergeMode,
    pub(crate) sequence_range: Option<SequenceRange>,
    pub(crate) partition_expr_version: u64,
}

impl ScanRequestFingerprintBuilder {
    pub(crate) fn build(self) -> ScanRequestFingerprint {
        let Self {
            read_columns,
            read_column_types,
            filters,
            time_filters,
            series_row_selector,
            append_mode,
            filter_deleted,
            merge_mode,
            sequence_range,
            partition_expr_version,
        } = self;

        ScanRequestFingerprint {
            inner: Arc::new(SharedScanRequestFingerprint {
                read_columns,
                read_column_types,
                filters,
            }),
            time_filters: (!time_filters.is_empty()).then(|| Arc::new(time_filters)),
            series_row_selector,
            append_mode,
            filter_deleted,
            merge_mode,
            sequence_range,
            stage: RangeScanStage::Data,
            partition_expr_version,
        }
    }
}

/// Non-copiable struct of the fingerprint.
#[derive(Debug, PartialEq, Eq, Hash)]
struct SharedScanRequestFingerprint {
    /// Logical columns of the projection.
    read_columns: ReadColumns,
    /// Column types of the projection.
    /// We keep this to ensure we won't reuse the fingerprint after a schema change.
    read_column_types: Vec<Option<ConcreteDataType>>,
    /// Filters without the time index column and region partition exprs.
    filters: Vec<String>,
}

impl ScanRequestFingerprint {
    #[cfg(test)]
    pub(crate) fn read_columns(&self) -> &ReadColumns {
        &self.inner.read_columns
    }

    #[cfg(test)]
    pub(crate) fn read_column_types(&self) -> &[Option<ConcreteDataType>] {
        &self.inner.read_column_types
    }

    #[cfg(test)]
    pub(crate) fn filters(&self) -> &[String] {
        &self.inner.filters
    }

    #[cfg(test)]
    pub(crate) fn time_filters(&self) -> &[String] {
        self.time_filters
            .as_deref()
            .map(Vec::as_slice)
            .unwrap_or(&[])
    }

    pub(crate) fn without_time_filters(&self) -> Self {
        Self {
            inner: Arc::clone(&self.inner),
            time_filters: None,
            series_row_selector: self.series_row_selector,
            append_mode: self.append_mode,
            filter_deleted: self.filter_deleted,
            merge_mode: self.merge_mode,
            sequence_range: self.sequence_range,
            stage: self.stage,
            partition_expr_version: self.partition_expr_version,
        }
    }

    fn for_candidate_series(&self) -> Self {
        Self {
            inner: Arc::clone(&self.inner),
            time_filters: self.time_filters.clone(),
            series_row_selector: self.series_row_selector,
            append_mode: self.append_mode,
            filter_deleted: self.filter_deleted,
            merge_mode: self.merge_mode,
            sequence_range: self.sequence_range,
            stage: RangeScanStage::CandidateSeries,
            partition_expr_version: self.partition_expr_version,
        }
    }

    fn for_series_data(&self, range: SeriesRange) -> Self {
        Self {
            inner: Arc::clone(&self.inner),
            time_filters: self.time_filters.clone(),
            series_row_selector: self.series_row_selector,
            append_mode: self.append_mode,
            filter_deleted: self.filter_deleted,
            merge_mode: self.merge_mode,
            sequence_range: self.sequence_range,
            stage: RangeScanStage::SeriesData(range),
            partition_expr_version: self.partition_expr_version,
        }
    }

    pub(crate) fn estimated_size(&self) -> usize {
        mem::size_of::<SharedScanRequestFingerprint>()
            + self.inner.read_columns.estimated_size()
            + self.inner.read_column_types.capacity() * mem::size_of::<Option<ConcreteDataType>>()
            + self.inner.filters.capacity() * mem::size_of::<String>()
            + self
                .inner
                .filters
                .iter()
                .map(|filter| filter.capacity())
                .sum::<usize>()
            + self.time_filters.as_ref().map_or(0, |filters| {
                mem::size_of::<Vec<String>>()
                    + filters.capacity() * mem::size_of::<String>()
                    + filters
                        .iter()
                        .map(|filter| filter.capacity())
                        .sum::<usize>()
            })
    }
}

/// Cache key for range scan outputs.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct RangeScanCacheKey {
    pub(crate) region_id: RegionId,
    /// Sorted (file_id, row_group_index) pairs that uniquely identify the data this range covers.
    pub(crate) row_groups: Vec<(FileId, i64)>,
    pub(crate) scan: ScanRequestFingerprint,
}

impl RangeScanCacheKey {
    pub(crate) fn for_cross_range_merge(mut self) -> Self {
        debug_assert_eq!(self.scan.stage, RangeScanStage::Data);
        self.scan.stage = RangeScanStage::DataForCrossRangeMerge;
        self
    }

    pub(crate) fn estimated_size(&self) -> usize {
        mem::size_of::<Self>()
            + self.row_groups.capacity() * mem::size_of::<(FileId, i64)>()
            + self.scan.estimated_size()
    }
}

/// Cached result for one range scan.
#[derive(Debug)]
pub(crate) struct CachedBatchSlice {
    batch: RecordBatch,
    slice_lengths: Vec<usize>,
}

impl CachedBatchSlice {
    fn metadata_size(&self) -> usize {
        self.slice_lengths.capacity() * mem::size_of::<usize>()
    }
}

pub(crate) struct RangeScanCacheValue {
    cached_batches: Vec<CachedBatchSlice>,
    /// Precomputed size of all compacted batches.
    estimated_batches_size: usize,
}

impl RangeScanCacheValue {
    pub(crate) fn new(
        cached_batches: Vec<CachedBatchSlice>,
        estimated_batches_size: usize,
    ) -> Self {
        Self {
            cached_batches,
            estimated_batches_size,
        }
    }

    pub(crate) fn estimated_size(&self) -> usize {
        mem::size_of::<Self>()
            + self.cached_batches.capacity() * mem::size_of::<CachedBatchSlice>()
            + self
                .cached_batches
                .iter()
                .map(CachedBatchSlice::metadata_size)
                .sum::<usize>()
            + self.estimated_batches_size
    }
}

/// Row groups and whether all sources are file-only for a partition range.
pub(crate) struct PartitionRangeRowGroups {
    /// Sorted (file_id, row_group_index) pairs.
    pub(crate) row_groups: Vec<(FileId, i64)>,
    pub(crate) only_file_sources: bool,
}

/// Collects (file_id, row_group_index) pairs from a partition range's row group indices.
pub(crate) fn collect_partition_range_row_groups(
    stream_ctx: &StreamContext,
    part_range: &PartitionRange,
) -> PartitionRangeRowGroups {
    let range_meta = &stream_ctx.ranges[part_range.identifier];
    let mut row_groups = Vec::new();
    let mut only_file_sources = true;

    for index in &range_meta.row_group_indices {
        if stream_ctx.is_file_range_index(*index) {
            let file_id = stream_ctx.input.file_from_index(*index).file_id().file_id();
            row_groups.push((file_id, index.row_group_index));
        } else {
            only_file_sources = false;
        }
    }

    row_groups.sort_unstable_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()).then(a.1.cmp(&b.1)));

    PartitionRangeRowGroups {
        row_groups,
        only_file_sources,
    }
}

/// Returns the timestamp range where all time-only predicates are guaranteed true.
///
/// Returns `None` for empty input because there is no time-filter implication
/// that cache-key normalization or prefilter postponement can use. It also
/// returns `None` if any expression contains an unsupported shape: `OR`, `NOT`,
/// `IN`, non-literal RHS, unsupported operator, column-name mismatch, an `=`
/// literal that cannot be represented exactly in the column unit, or overflow
/// during bound adjustment.
///
/// This is intentionally stricter than `extract_time_range_from_expr` in
/// `table::predicate`: lower bounds round up and upper bounds round down. If a
/// partition's file-time range is contained by the returned range, every row in
/// that partition satisfies the original time predicates.
///
/// `IsNull`/`IsNotNull` on the time index are not routed into `time_filters`
/// today. If that changes, handle them here before stripping time filters from
/// the cache key.
pub(crate) fn implied_time_range_from_exprs(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    exprs: &[&Expr],
) -> Option<TimestampRange> {
    if exprs.is_empty() {
        return None;
    }

    let mut acc = TimestampRange::min_to_max();
    for expr in exprs {
        let r = implied_time_range_from_expr(ts_col_name, ts_col_unit, expr)?;
        acc = acc.and(&r);
    }
    Some(acc)
}

fn implied_time_range_from_expr(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    expr: &Expr,
) -> Option<TimestampRange> {
    match expr {
        Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
            Operator::And => {
                let l = implied_time_range_from_expr(ts_col_name, ts_col_unit, left)?;
                let r = implied_time_range_from_expr(ts_col_name, ts_col_unit, right)?;
                Some(l.and(&r))
            }
            Operator::Eq | Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => {
                implied_from_cmp(ts_col_name, ts_col_unit, left, *op, right)
            }
            // `OR` would require a strict intersection over a union of half-planes
            // (not the loose-span union provided by `TimestampRange::or`), so we
            // refuse it. Any other operator is unsupported.
            _ => None,
        },
        Expr::Between(Between {
            expr,
            negated,
            low,
            high,
        }) => {
            if *negated {
                return None;
            }
            implied_from_between(ts_col_name, ts_col_unit, expr, low, high)
        }
        // Includes `IsNull`, `IsNotNull`, `Not`, `InList`, function calls, etc.
        _ => None,
    }
}

fn match_ts_column_literal<'a>(
    ts_col_name: &str,
    left: &'a Expr,
    right: &'a Expr,
) -> Option<(Timestamp, bool)> {
    let (col, scalar, reverse) = match (left, right) {
        (Expr::Column(c), Expr::Literal(s, _)) => (c, s, false),
        (Expr::Literal(s, _), Expr::Column(c)) => (c, s, true),
        _ => return None,
    };
    if col.name != ts_col_name {
        return None;
    }
    // Reject string literals: their conversion needs a timezone we do not have,
    // and the existing extractor in `table::predicate` rejects them too.
    if is_string_timestamp_literal(scalar) {
        return None;
    }
    scalar_value_to_timestamp(scalar, None).map(|t| (t, reverse))
}

fn implied_from_cmp(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    left: &Expr,
    op: Operator,
    right: &Expr,
) -> Option<TimestampRange> {
    let (ts, reverse) = match_ts_column_literal(ts_col_name, left, right)?;
    // Normalize to "column OP literal".
    let op = if reverse {
        match op {
            Operator::Lt => Operator::Gt,
            Operator::LtEq => Operator::GtEq,
            Operator::Gt => Operator::Lt,
            Operator::GtEq => Operator::LtEq,
            Operator::Eq => Operator::Eq,
            _ => return None,
        }
    } else {
        op
    };

    match op {
        Operator::GtEq => {
            // ts >= L. Round the lower bound up in the column unit.
            let b = ts.convert_to_ceil(ts_col_unit)?;
            Some(TimestampRange::from_start(b))
        }
        Operator::Gt => {
            // ts > L. floor(L) + 1 is the tight lower bound in the column unit.
            let v = ts.convert_to(ts_col_unit)?.value().checked_add(1)?;
            Some(TimestampRange::from_start(Timestamp::new(v, ts_col_unit)))
        }
        Operator::LtEq => {
            // ts <= L. Round the upper bound down in the column unit.
            let b = ts.convert_to(ts_col_unit)?;
            Some(TimestampRange::until_end(b, true))
        }
        Operator::Lt => {
            // ts < L. `ts < ceil(L)` is the tight bound: equal to `ts < L` when
            // L is exactly representable, and `ts <= floor(L)` otherwise.
            let b = ts.convert_to_ceil(ts_col_unit)?;
            Some(TimestampRange::until_end(b, false))
        }
        Operator::Eq => {
            // ts = L. Only provable when L is exactly representable.
            let f = ts.convert_to(ts_col_unit)?;
            let c = ts.convert_to_ceil(ts_col_unit)?;
            if f.value() != c.value() {
                return None;
            }
            Some(TimestampRange::single(f))
        }
        _ => None,
    }
}

fn implied_from_between(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    expr: &Expr,
    low: &Expr,
    high: &Expr,
) -> Option<TimestampRange> {
    let Expr::Column(c) = expr else {
        return None;
    };
    if c.name != ts_col_name {
        return None;
    }
    let (low_s, high_s) = match (low, high) {
        (Expr::Literal(l, _), Expr::Literal(h, _)) => (l, h),
        _ => return None,
    };
    if is_string_timestamp_literal(low_s) || is_string_timestamp_literal(high_s) {
        return None;
    }
    let low_ts = scalar_value_to_timestamp(low_s, None)?;
    let high_ts = scalar_value_to_timestamp(high_s, None)?;
    // BETWEEN low AND high is equivalent to ts >= low AND ts <= high.
    let lo = low_ts.convert_to_ceil(ts_col_unit)?;
    let hi = high_ts.convert_to(ts_col_unit)?;
    Some(TimestampRange::new_inclusive(Some(lo), Some(hi)))
}

/// Builds a cache key for the given partition range if it is eligible for caching.
pub(crate) fn build_range_cache_key(
    stream_ctx: &StreamContext,
    part_range: &PartitionRange,
) -> Option<RangeScanCacheKey> {
    build_range_cache_key_inner(stream_ctx, part_range, None)
}

/// Builds a cache key for a candidate-series partition-range result.
pub(crate) fn build_candidate_range_cache_key(
    stream_ctx: &StreamContext,
    part_range: &PartitionRange,
) -> Option<RangeScanCacheKey> {
    build_range_cache_key_inner(
        stream_ctx,
        part_range,
        Some(RangeScanStage::CandidateSeries),
    )
}

/// Builds a cache key for a two-phase series-data partition-range result.
pub(crate) fn build_series_range_cache_key(
    stream_ctx: &StreamContext,
    part_range: &PartitionRange,
    range: SeriesRange,
) -> Option<RangeScanCacheKey> {
    build_range_cache_key_inner(
        stream_ctx,
        part_range,
        Some(RangeScanStage::SeriesData(range)),
    )
}

fn build_range_cache_key_inner(
    stream_ctx: &StreamContext,
    part_range: &PartitionRange,
    stage: Option<RangeScanStage>,
) -> Option<RangeScanCacheKey> {
    if !stream_ctx.input.cache_strategy.has_range_result_cache() {
        return None;
    }

    let fingerprint = stream_ctx.input.scan_fingerprint()?;

    // Dyn filters can change at runtime, so we can't cache when they're present.
    let has_dyn_filters = stream_ctx
        .input
        .predicate_group()
        .predicate_without_region()
        .is_some_and(|p| !p.dyn_filters().is_empty());
    if has_dyn_filters {
        return None;
    }

    let rg = collect_partition_range_row_groups(stream_ctx, part_range);
    if !rg.only_file_sources || rg.row_groups.is_empty() {
        return None;
    }

    // If the implied range covers this partition's `FileTimeRange`, drop
    // time-only predicates from the cache key so that queries with different
    // but equally-covering time bounds share an entry. `None` means there is no
    // analyzable time-only predicate or some predicate had an unsupported shape
    // (e.g. `OR`), so we keep the fingerprint unchanged. When `time_filters` is
    // already empty, cloning the fingerprint is equivalent to stripping them.
    let range_meta = &stream_ctx.ranges[part_range.identifier];
    let (file_min, file_max) = range_meta.time_range;
    let covers = match stream_ctx.input.implied_time_range() {
        // An empty implied range can never cover a non-empty file range, so
        // short-circuit.
        Some(implied) if !implied.is_empty() => {
            // `file_min`/`file_max` may use a different unit than the implied
            // range: after the time index unit was widened, pre-alter files
            // still carry the old unit while the implied range uses the
            // current column unit. `Timestamp` comparisons normalize units,
            // so `contains` compares instants and stays sound across units.
            implied.contains(&file_min) && implied.contains(&file_max)
        }
        _ => false,
    };
    let scan = if covers {
        fingerprint.without_time_filters()
    } else {
        fingerprint.clone()
    };
    let scan = match stage {
        Some(RangeScanStage::CandidateSeries) => scan.for_candidate_series(),
        Some(RangeScanStage::SeriesData(assignment)) => scan.for_series_data(assignment),
        Some(RangeScanStage::DataForCrossRangeMerge) => {
            let mut scan = scan;
            scan.stage = RangeScanStage::DataForCrossRangeMerge;
            scan
        }
        Some(RangeScanStage::Data) | None => scan,
    };

    Some(RangeScanCacheKey {
        region_id: stream_ctx.input.region_metadata().region_id,
        row_groups: rg.row_groups,
        scan,
    })
}

/// Returns a stream that replays cached record batches.
pub(crate) fn cached_flat_range_stream(
    delivery: Arc<RangeCacheDelivery>,
) -> BoxedRecordBatchStream {
    Box::pin(try_stream! {
        // Borrow through the delivery guard so it survives every yield. Batches
        // handed to downstream operators are then owned by the query, not by
        // this completed-result backlog.
        for cached_batch in &delivery.value().cached_batches {
            let mut offset = 0;
            for &len in &cached_batch.slice_lengths {
                yield cached_batch.batch.slice(offset, len);
                offset += len;
            }
        }
    })
}

/// Builds only one decoded result per full cache key. A load must be driven
/// independently of any caller: merge readers can wait for another stream with
/// the same key before polling their first stream to completion.
pub(crate) struct CachedRangeReader<'a> {
    pub(crate) stream_ctx: &'a StreamContext,
    pub(crate) part_range: &'a PartitionRange,
    pub(crate) key: Option<RangeScanCacheKey>,
    pub(crate) metrics: &'a PartitionMetrics,
}

impl CachedRangeReader<'_> {
    pub(crate) async fn build<F, Fut>(self, build: F) -> Result<(BoxedRecordBatchStream, usize)>
    where
        F: FnOnce() -> Fut + Clone + Send + 'static,
        Fut: Future<Output = Result<(BoxedRecordBatchStream, usize)>> + Send + 'static,
    {
        let Self {
            stream_ctx,
            part_range,
            key,
            metrics,
        } = self;
        let strategy = &stream_ctx.input.cache_strategy;
        let (Some(key), Some(worker)) = (key, strategy.range_cache_worker()) else {
            return build().await;
        };
        let files = stream_ctx.ranges[part_range.identifier]
            .row_group_indices
            .iter()
            .map(|index| stream_ctx.input.file_from_index(*index).clone())
            .collect();
        let loader = build.clone();
        let load_worker = worker.clone();
        let load_strategy = strategy.clone();
        let load_metrics = metrics.clone();
        let key_size = key.estimated_size();
        match worker.get_or_load(key, files, async move {
            let result = async {
                let (stream, _) = loader().await?;
                collect_range_result(stream, &load_worker, &load_strategy).await
            }
            .await;
            match result {
                Ok(Some(value)) => {
                    load_metrics.inc_range_cache_size(key_size + value.value().estimated_size());
                    Some(value)
                }
                Ok(None) => None,
                Err(error) => {
                    warn!(error; "Failed to load range cache; falling back to an uncached scan");
                    None
                }
            }
        }) {
            RangeCacheLookup::Hit(value) => {
                metrics.inc_range_cache_hit();
                Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE))
            }
            RangeCacheLookup::Loading(receiver) => {
                metrics.inc_range_cache_miss();
                Ok((
                    shared_range_stream(receiver, build),
                    DEFAULT_READ_BATCH_SIZE,
                ))
            }
            RangeCacheLookup::Bypass => {
                metrics.inc_range_cache_miss();
                build().await
            }
        }
    }
}

/// Waiting is lazy, but the producer is independent of every consumer.
fn shared_range_stream<F, Fut>(
    mut receiver: tokio::sync::watch::Receiver<Option<Arc<RangeCacheDelivery>>>,
    build: F,
) -> BoxedRecordBatchStream
where
    F: FnOnce() -> Fut + Send + 'static,
    Fut: Future<Output = Result<(BoxedRecordBatchStream, usize)>> + Send + 'static,
{
    Box::pin(try_stream! {
        let cached = receiver.wait_for(|cached| cached.is_some()).await
            .ok().and_then(|cached| cached.clone());
        drop(receiver);
        let mut stream = if let Some(cached) = cached {
            cached_flat_range_stream(cached)
        } else {
            // A failed/oversized/cancelled fill is never a query result.
            // Retry once through the ordinary streaming path, not the worker.
            build().await?.0
        };
        while let Some(batch) = stream.try_next().await? {
            yield batch;
        }
    })
}

/// The load budget covers retained input buffers and completed deliveries. Acquisition
/// must be nonblocking: waiting with a partially accumulated result deadlocks
/// when several loads consume all permits. Oversized ranges retain streaming
/// query semantics by falling back to an uncached scan.
async fn collect_range_result(
    mut stream: BoxedRecordBatchStream,
    worker: &RangeCacheWorker,
    strategy: &CacheStrategy,
) -> Result<Option<Arc<RangeCacheDelivery>>> {
    let mut permit: Option<tokio::sync::OwnedSemaphorePermit> = None;
    let mut buffer = CacheBatchBuffer::new(strategy);
    while let Some(batch) = stream.try_next().await? {
        let Some(next_permit) = worker
            .load_memory
            .try_acquire(batch.get_array_memory_size())
        else {
            return Ok(None);
        };
        if let Some(permit) = &mut permit {
            permit.merge(next_permit);
        } else {
            permit = Some(next_permit);
        }
        buffer.push(batch)?;
    }
    let Some(value) = buffer.finish().await.transpose()? else {
        return Ok(None);
    };
    let Some(permit) = worker
        .load_memory
        .try_resize(permit, value.estimated_size())
    else {
        return Ok(None);
    };
    Ok(Some(RangeCacheDelivery::new(value, permit)))
}

enum CacheConcatCommand {
    Compact(Vec<RecordBatch>),
    Finish {
        pending: Vec<RecordBatch>,
        result_tx: oneshot::Sender<Result<Arc<RangeScanCacheValue>>>,
    },
}

#[derive(Default)]
struct CacheConcatState {
    cached_batches: Vec<CachedBatchSlice>,
    estimated_size: usize,
}

impl CacheConcatState {
    async fn compact(
        &mut self,
        batches: Vec<RecordBatch>,
        limiter: &crate::cache::RangeResultMemoryLimiter,
    ) -> Result<()> {
        if batches.is_empty() {
            return Ok(());
        }

        let input_size = batches
            .iter()
            .map(RecordBatch::get_array_memory_size)
            .sum::<usize>();
        let _permit = limiter.acquire(input_size).await?;

        let compacted = compact_record_batches(batches)?;
        self.estimated_size += compacted.batch.get_array_memory_size();
        self.cached_batches.push(compacted);
        Ok(())
    }

    fn finish(self) -> RangeScanCacheValue {
        RangeScanCacheValue::new(self.cached_batches, self.estimated_size)
    }
}

fn compact_record_batches(batches: Vec<RecordBatch>) -> Result<CachedBatchSlice> {
    debug_assert!(!batches.is_empty());

    let slice_lengths = batches.iter().map(RecordBatch::num_rows).collect();
    build_cached_batch_slice(batches, slice_lengths)
}

fn build_cached_batch_slice(
    batches: Vec<RecordBatch>,
    slice_lengths: Vec<usize>,
) -> Result<CachedBatchSlice> {
    let batch = if batches.len() == 1 {
        batches.into_iter().next().unwrap()
    } else {
        let schema = batches[0].schema();
        concat_batches(&schema, &batches).context(ComputeArrowSnafu)?
    };

    Ok(CachedBatchSlice {
        batch,
        slice_lengths,
    })
}

async fn run_cache_concat_task(
    mut rx: mpsc::UnboundedReceiver<CacheConcatCommand>,
    limiter: Arc<crate::cache::RangeResultMemoryLimiter>,
    skip_threshold_bytes: usize,
) {
    let mut state = CacheConcatState::default();

    while let Some(cmd) = rx.recv().await {
        match cmd {
            CacheConcatCommand::Compact(batches) => {
                if let Err(err) = state.compact(batches, &limiter).await {
                    warn!(err; "Failed to compact range cache batches");
                    return;
                }
                // Close the channel to stop further work as soon as the cached
                // size exceeds the configured cache budget.
                if state.estimated_size > skip_threshold_bytes {
                    return;
                }
            }
            CacheConcatCommand::Finish { pending, result_tx } => {
                let compact_result = state
                    .compact(pending, &limiter)
                    .await
                    .map(|()| state.finish());
                let result = match compact_result {
                    Ok(v) => {
                        let value = Arc::new(v);
                        Ok(value)
                    }
                    Err(e) => {
                        warn!(e; "Failed to finalize range cache batches");

                        Err(e)
                    }
                };

                let _ = result_tx.send(result);

                break;
            }
        }
    }
}

struct CacheBatchBuffer {
    buffered_batches: Vec<RecordBatch>,
    buffered_rows: usize,
    buffered_size: usize,
    sender: Option<mpsc::UnboundedSender<CacheConcatCommand>>,
    /// Kept until concat finishes so abandoning the shared load cancels concat too.
    concat_task: Option<common_runtime::JoinHandle<()>>,
}

impl CacheBatchBuffer {
    fn new(cache_strategy: &CacheStrategy) -> Self {
        let (sender, concat_task) = cache_strategy
            .range_result_memory_limiter()
            .map(|limiter| {
                let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0);
                let (tx, rx) = mpsc::unbounded_channel();
                let task = common_runtime::spawn_query(run_cache_concat_task(
                    rx,
                    limiter.clone(),
                    skip_threshold_bytes,
                ));
                (tx, task)
            })
            .unzip();

        Self {
            buffered_batches: Vec::new(),
            buffered_rows: 0,
            buffered_size: 0,
            sender,
            concat_task,
        }
    }

    fn push(&mut self, batch: RecordBatch) -> Result<()> {
        if self.sender.is_none() {
            return Ok(());
        }

        self.buffered_rows += batch.num_rows();
        self.buffered_size += batch.get_array_memory_size();
        self.buffered_batches.push(batch);

        if self.buffered_batches.len() > 1
            && (self.buffered_rows > DEFAULT_READ_BATCH_SIZE
                || self.buffered_size > RANGE_CACHE_COMPACT_THRESHOLD_BYTES)
        {
            self.notify_compact();
        }

        Ok(())
    }

    fn notify_compact(&mut self) {
        if self.buffered_batches.is_empty() || self.sender.is_none() {
            return;
        }

        let batches = mem::take(&mut self.buffered_batches);
        self.buffered_rows = 0;
        self.buffered_size = 0;

        let Some(sender) = &self.sender else {
            return;
        };
        if sender.send(CacheConcatCommand::Compact(batches)).is_err() {
            self.sender = None;
        }
    }

    async fn finish(mut self) -> Option<Result<Arc<RangeScanCacheValue>>> {
        let sender = self.sender.take()?;
        let (result_tx, result_rx) = oneshot::channel();
        sender
            .send(CacheConcatCommand::Finish {
                pending: mem::take(&mut self.buffered_batches),
                result_tx,
            })
            .ok()?;
        // Keep ownership until the result arrives so cancellation also aborts
        // concat while it waits for its temporary memory permits.
        let result = result_rx.await.ok();
        self.concat_task = None;
        result
    }
}

impl Drop for CacheBatchBuffer {
    fn drop(&mut self) {
        // Aborting releases queued batches even while concat waits for a permit.
        if let Some(task) = &self.concat_task {
            task.abort();
        }
    }
}

/// Wraps a stream to cache its output for future range cache hits.
#[cfg(feature = "test")]
pub(crate) fn cache_flat_range_stream(
    mut stream: BoxedRecordBatchStream,
    cache_strategy: CacheStrategy,
    key: RangeScanCacheKey,
    part_metrics: PartitionMetrics,
) -> BoxedRecordBatchStream {
    Box::pin(try_stream! {
        let mut buffer = CacheBatchBuffer::new(&cache_strategy);
        while let Some(batch) = stream.try_next().await? {
            buffer.push(batch.clone())?;
            yield batch;
        }

        if let Some(Ok(value)) = buffer.finish().await {
            part_metrics.inc_range_cache_size(key.estimated_size() + value.estimated_size());
            cache_strategy.put_range_result(key, value);
        }
    })
}

/// Creates a `cache_flat_range_stream` with dummy internals for benchmarking.
///
/// This avoids exposing `RangeScanCacheKey`, `ScanRequestFingerprint`, and
/// `PartitionMetrics` publicly.
#[cfg(feature = "test")]
pub fn bench_cache_flat_range_stream(
    stream: BoxedRecordBatchStream,
    cache_size_bytes: u64,
    region_id: RegionId,
) -> BoxedRecordBatchStream {
    use std::time::Instant;

    use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;

    use crate::region::options::MergeMode;

    let cache_manager = Arc::new(
        crate::cache::CacheManager::builder()
            .range_result_cache_size(cache_size_bytes)
            .build(),
    );
    let cache_strategy = CacheStrategy::EnableAll(cache_manager);

    let fingerprint = ScanRequestFingerprintBuilder {
        read_columns: ReadColumns::new(std::iter::empty()),
        read_column_types: vec![],
        filters: vec![],
        time_filters: vec![],
        series_row_selector: None,
        append_mode: false,
        filter_deleted: false,
        merge_mode: MergeMode::LastRow,
        sequence_range: None,
        partition_expr_version: 0,
    }
    .build();

    let key = RangeScanCacheKey {
        region_id,
        row_groups: vec![],
        scan: fingerprint,
    };

    let metrics_set = ExecutionPlanMetricsSet::new();
    let part_metrics =
        PartitionMetrics::new(region_id, 0, "bench", Instant::now(), false, &metrics_set);

    cache_flat_range_stream(stream, cache_strategy, key, part_metrics)
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::time::Instant;

    use common_time::Timestamp;
    use common_time::range::TimestampRange;
    use common_time::timestamp::TimeUnit;
    use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
    use datafusion_common::ScalarValue;
    use datafusion_expr::{Expr, col, lit};
    use smallvec::smallvec;
    use store_api::storage::{FileId, RegionId};
    use tokio::sync::watch;

    use super::*;
    use crate::cache::CacheManager;
    use crate::cache::test_util::range_cache_key;
    use crate::read::flat_projection::FlatProjectionMapper;
    use crate::read::range::{RangeMeta, RowGroupIndex, SourceIndex};
    use crate::read::scan_region::{PredicateGroup, ScanInput};
    use crate::sst::file::FileTimeRange;
    use crate::test_util::memtable_util::metadata_with_primary_key;
    use crate::test_util::scheduler_util::SchedulerEnv;
    use crate::test_util::sst_util::sst_file_handle_with_file_id;

    fn test_cache_strategy() -> CacheStrategy {
        CacheStrategy::EnableAll(Arc::new(
            CacheManager::builder()
                .range_result_cache_size(1024 * 1024)
                .build(),
        ))
    }

    fn test_scan_fingerprint(
        filters: Vec<String>,
        time_filters: Vec<String>,
        series_row_selector: Option<TimeSeriesRowSelector>,
        filter_deleted: bool,
        partition_expr_version: u64,
    ) -> ScanRequestFingerprint {
        let read_columns = ReadColumns::new([1, 2]);
        ScanRequestFingerprintBuilder {
            read_columns,
            read_column_types: vec![None, None],
            filters,
            time_filters,
            series_row_selector,
            append_mode: false,
            filter_deleted,
            merge_mode: MergeMode::LastRow,
            sequence_range: None,
            partition_expr_version,
        }
        .build()
    }

    fn test_cache_context(strategy: &CacheStrategy) -> (RangeScanCacheKey, PartitionMetrics) {
        let region_id = RegionId::new(1, 1);
        let key = RangeScanCacheKey {
            region_id,
            row_groups: vec![],
            scan: test_scan_fingerprint(vec![], vec![], None, false, 0),
        };

        let metrics_set = ExecutionPlanMetricsSet::new();
        let part_metrics =
            PartitionMetrics::new(region_id, 0, "test", Instant::now(), false, &metrics_set);

        assert!(strategy.get_range_result(&key).is_none());
        (key, part_metrics)
    }

    async fn new_stream_context(
        filters: Vec<Expr>,
        query_time_range: Option<TimestampRange>,
        partition_time_range: FileTimeRange,
    ) -> (StreamContext, PartitionRange) {
        let env = SchedulerEnv::new().await;
        let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false));
        let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap();
        let predicate = PredicateGroup::new(metadata.as_ref(), &filters).unwrap();
        let file_id = FileId::random();
        let file = sst_file_handle_with_file_id(
            file_id,
            partition_time_range.0.value(),
            partition_time_range.1.value(),
        );
        let input = ScanInput::builder(env.access_layer.clone(), mapper)
            .with_predicate(predicate)
            .with_time_range(query_time_range)
            .with_files(vec![file])
            .with_cache(test_cache_strategy())
            .build();
        let range_meta = RangeMeta {
            time_range: partition_time_range,
            indices: smallvec![SourceIndex {
                index: 0,
                num_row_groups: 1,
            }],
            row_group_indices: smallvec![RowGroupIndex {
                index: 0,
                row_group_index: 0,
            }],
            num_rows: 10,
        };
        let partition_range = range_meta.new_partition_range(0);
        let stream_ctx = StreamContext {
            input,
            ranges: vec![range_meta],
            query_start: Instant::now(),
        };

        (stream_ctx, partition_range)
    }

    #[rstest::rstest]
    #[case(1024 * 1024, true)]
    #[case(1, false)]
    #[tokio::test]
    async fn shared_range_read_preserves_batches_or_falls_back_on_budget_exhaustion(
        #[case] capacity: u64,
        #[case] cached: bool,
    ) {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let (mut ctx, range) =
            new_stream_context(vec![col("k0").eq(lit("a"))], None, FileTimeRange::default()).await;
        ctx.input.cache_strategy = CacheStrategy::EnableAll(Arc::new(
            CacheManager::builder()
                .range_result_cache_size(capacity)
                .build(),
        ));
        let key = build_range_cache_key(&ctx, &range).unwrap();
        let (_, metrics) = test_cache_context(&ctx.input.cache_strategy);
        let calls = Arc::new(AtomicUsize::new(0));
        let expected = vec![make_batch(&[1, 2]), make_batch(&[3, 4, 5])];
        let batches = expected.clone();
        let count = calls.clone();
        let (stream, _) = CachedRangeReader {
            stream_ctx: &ctx,
            part_range: &range,
            key: Some(key.clone()),
            metrics: &metrics,
        }
        .build(move || async move {
            count.fetch_add(1, Ordering::Relaxed);
            Ok((
                Box::pin(futures::stream::iter(batches.into_iter().map(Ok)))
                    as BoxedRecordBatchStream,
                DEFAULT_READ_BATCH_SIZE,
            ))
        })
        .await
        .unwrap();
        let actual: Vec<_> = stream.try_collect().await.unwrap();
        assert_eq!(actual, expected);
        assert_eq!(calls.load(Ordering::Relaxed), if cached { 1 } else { 2 });
        assert_eq!(
            ctx.input.cache_strategy.get_range_result(&key).is_some(),
            cached
        );
        assert_eq!(
            ctx.input
                .cache_strategy
                .range_cache_worker()
                .unwrap()
                .load_memory
                .available_permits(),
            (capacity as usize).div_ceil(1024)
        );
    }

    fn start_shared_load(
        strategy: &CacheStrategy,
        key: RangeScanCacheKey,
        batches: Vec<RecordBatch>,
    ) -> watch::Receiver<Option<Arc<RangeCacheDelivery>>> {
        let worker = strategy.range_cache_worker().unwrap();
        let loader_worker = worker.clone();
        let loader_strategy = strategy.clone();
        let lookup = worker.get_or_load(key, vec![], async move {
            collect_range_result(
                Box::pin(futures::stream::iter(batches.into_iter().map(Ok))),
                &loader_worker,
                &loader_strategy,
            )
            .await
            .unwrap()
        });
        let RangeCacheLookup::Loading(receiver) = lookup else {
            panic!("expected a shared load");
        };
        receiver
    }

    async fn wait_for_shared_load(
        receiver: &mut watch::Receiver<Option<Arc<RangeCacheDelivery>>>,
    ) -> Option<Arc<RangeCacheDelivery>> {
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            receiver
                .wait_for(|value| value.is_some())
                .await
                .ok()
                .and_then(|value| value.clone())
        })
        .await
        .expect("a shared load must finish or fall back without waiting for consumers")
    }

    #[tokio::test]
    async fn completed_load_keeps_memory_until_all_subscribers_release() {
        let strategy = test_cache_strategy();
        let worker = strategy.range_cache_worker().unwrap();
        let capacity = worker.load_memory.available_permits();
        let (key, _) = test_cache_context(&strategy);
        let mut first = start_shared_load(&strategy, key.clone(), vec![make_batch(&[1, 2, 3])]);
        let second = first.clone();
        drop(wait_for_shared_load(&mut first).await.unwrap());

        // A completed flight no longer counts toward the flight limit, but
        // unpolled subscribers must still reserve its decoded result memory.
        assert!(worker.load_memory.available_permits() < capacity);
        drop(first);
        assert!(worker.load_memory.available_permits() < capacity);
        drop(second);
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            while worker.load_memory.available_permits() != capacity {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("the last subscriber must release the delivery reservation");
        assert!(
            worker.get(&key).is_some(),
            "the cache still owns the result"
        );
    }

    #[rstest::rstest]
    #[tokio::test]
    async fn successive_loads_with_slow_consumers_exhaust_delivery_budget(
        #[values(false, true)] invalidate: bool,
    ) {
        let strategy = CacheStrategy::EnableAll(Arc::new(
            CacheManager::builder()
                .range_result_cache_size(8192)
                .build(),
        ));
        let worker = strategy.range_cache_worker().unwrap();
        let capacity = worker.load_memory.available_permits();
        let mut held_streams = Vec::new();
        let mut fallbacks = 0;

        // Each flight finishes before the next starts, exceeding the 32-flight
        // limit cumulatively. Only the delivery budget can bound this backlog.
        for _ in 0..40 {
            let file = FileId::random();
            let key = range_cache_key(&[file]);
            let expected = vec![make_large_binary_batch(1, 2048)];
            let mut receiver = start_shared_load(&strategy, key, expected.clone());
            let shared = wait_for_shared_load(&mut receiver).await.is_some();
            if invalidate {
                worker.invalidate_file(file);
            }
            let fallback_batches = expected.clone();
            let stream = shared_range_stream(receiver, move || async move {
                Ok((
                    Box::pin(futures::stream::iter(fallback_batches.into_iter().map(Ok)))
                        as BoxedRecordBatchStream,
                    DEFAULT_READ_BATCH_SIZE,
                ))
            });
            if shared {
                held_streams.push(stream);
            } else {
                fallbacks += 1;
                let actual: Vec<_> =
                    tokio::time::timeout(std::time::Duration::from_secs(5), stream.try_collect())
                        .await
                        .expect("fallback must not wait for an unpolled range")
                        .unwrap();
                assert_eq!(actual, expected);
            }
        }

        assert!(!held_streams.is_empty());
        assert!(held_streams.len() < 32);
        assert!(
            fallbacks > 0,
            "completed results must not escape the budget"
        );
        assert!(worker.load_memory.available_permits() < capacity);

        // Transfer ownership out of watch receivers into actively polled replay
        // streams, then leave them suspended at their first yield.
        for stream in &mut held_streams {
            assert!(stream.try_next().await.unwrap().is_some());
        }
        assert!(worker.load_memory.available_permits() < capacity);
        drop(held_streams);
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            while worker.load_memory.available_permits() != capacity {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("dropping slow readers must unblock new shared loads");

        let mut retry = start_shared_load(
            &strategy,
            range_cache_key(&[FileId::random()]),
            vec![make_large_binary_batch(1, 2048)],
        );
        assert!(wait_for_shared_load(&mut retry).await.is_some());
    }

    #[tokio::test]
    async fn cancelling_a_partial_load_releases_its_reservation() {
        let strategy = test_cache_strategy();
        let worker = strategy.range_cache_worker().unwrap().clone();
        let capacity = worker.load_memory.available_permits();
        let (key, _) = test_cache_context(&strategy);
        let loader_worker = worker.clone();
        let lookup = worker.get_or_load(key.clone(), vec![], async move {
            let stream = Box::pin(try_stream! {
                yield make_batch(&[1, 2, 3]);
                std::future::pending::<()>().await;
            });
            collect_range_result(stream, &loader_worker, &strategy)
                .await
                .unwrap()
        });
        let RangeCacheLookup::Loading(receiver) = lookup else {
            panic!("expected a shared load");
        };
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            while worker.load_memory.available_permits() == capacity {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("the partial result must acquire its input reservation");
        drop(receiver);
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            while worker.load_memory.available_permits() != capacity {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("cancellation must release the partial result reservation");
        assert!(worker.get(&key).is_none());
    }

    #[tokio::test]
    async fn follower_can_finish_without_polling_the_first_range_stream() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let (ctx, range) =
            new_stream_context(vec![col("k0").eq(lit("a"))], None, FileTimeRange::default()).await;
        let key = build_range_cache_key(&ctx, &range).unwrap();
        let (_, metrics) = test_cache_context(&ctx.input.cache_strategy);
        let gate = Arc::new(tokio::sync::Semaphore::new(0));
        let calls = Arc::new(AtomicUsize::new(0));
        let count = calls.clone();
        let loader_gate = gate.clone();
        let build = move || async move {
            count.fetch_add(1, Ordering::Relaxed);
            let _permit = loader_gate.acquire().await.unwrap();
            Ok((
                Box::pin(futures::stream::iter([Ok(make_batch(&[1, 2, 3]))]))
                    as BoxedRecordBatchStream,
                DEFAULT_READ_BATCH_SIZE,
            ))
        };
        let (first, _) = CachedRangeReader {
            stream_ctx: &ctx,
            part_range: &range,
            key: Some(key.clone()),
            metrics: &metrics,
        }
        .build(build.clone())
        .await
        .unwrap();
        let (second, _) = CachedRangeReader {
            stream_ctx: &ctx,
            part_range: &range,
            key: Some(key),
            metrics: &metrics,
        }
        .build(build)
        .await
        .unwrap();
        gate.add_permits(1);
        let second_batches: Vec<_> =
            tokio::time::timeout(std::time::Duration::from_secs(5), second.try_collect())
                .await
                .expect("a follower must not depend on polling the first query")
                .unwrap();
        let first_batches: Vec<_> = first.try_collect().await.unwrap();
        assert_eq!(first_batches, second_batches);
        assert_eq!(second_batches, vec![make_batch(&[1, 2, 3])]);
        assert_eq!(calls.load(Ordering::Relaxed), 1);
    }

    /// Helper to create a timestamp millisecond literal.
    fn ts_lit(val: i64) -> Expr {
        lit(ScalarValue::TimestampMillisecond(Some(val), None))
    }

    fn normalized_exprs(exprs: impl IntoIterator<Item = Expr>) -> Vec<String> {
        let mut exprs = exprs
            .into_iter()
            .map(|expr| expr.to_string())
            .collect::<Vec<_>>();
        exprs.sort_unstable();
        exprs
    }

    async fn assert_range_cache_filters(
        filters: Vec<Expr>,
        query_time_range: Option<TimestampRange>,
        partition_time_range: FileTimeRange,
        expected_filters: Vec<Expr>,
        expected_time_filters: Vec<Expr>,
    ) {
        let (stream_ctx, part_range) =
            new_stream_context(filters, query_time_range, partition_time_range).await;

        let key = build_range_cache_key(&stream_ctx, &part_range).unwrap();

        assert_eq!(
            key.scan.filters(),
            normalized_exprs(expected_filters).as_slice()
        );
        assert_eq!(
            key.scan.time_filters(),
            normalized_exprs(expected_time_filters).as_slice()
        );
    }

    #[tokio::test]
    async fn range_cache_time_filter_key_cases() {
        let partition = (
            Timestamp::new_millisecond(1000),
            Timestamp::new_millisecond(2000),
        );

        struct Case {
            filters: Vec<Expr>,
            query_time_range: Option<TimestampRange>,
            expected_filters: Vec<Expr>,
            expected_time_filters: Vec<Expr>,
        }

        // Time filters are stripped only when their implied range fully covers
        // the partition's file-time range. `is_not_null(ts)` stays in regular
        // filters because it is not routed into `time_filters`.
        for case in [
            Case {
                filters: vec![
                    col("ts").gt_eq(ts_lit(1000)),
                    col("ts").lt(ts_lit(2001)),
                    col("ts").is_not_null(),
                    col("k0").eq(lit("foo")),
                ],
                query_time_range: TimestampRange::with_unit(1000, 2002, TimeUnit::Millisecond),
                expected_filters: vec![col("k0").eq(lit("foo")), col("ts").is_not_null()],
                expected_time_filters: vec![],
            },
            Case {
                filters: vec![
                    col("ts").gt_eq(ts_lit(500)),
                    col("ts").lt(ts_lit(3000)),
                    col("k0").eq(lit("foo")),
                ],
                query_time_range: TimestampRange::with_unit(500, 3000, TimeUnit::Millisecond),
                expected_filters: vec![col("k0").eq(lit("foo"))],
                expected_time_filters: vec![],
            },
            Case {
                filters: vec![
                    col("ts").gt_eq(ts_lit(1000)),
                    col("ts").lt_eq(ts_lit(2000)),
                    col("k0").eq(lit("foo")),
                ],
                query_time_range: TimestampRange::with_unit(1000, 2001, TimeUnit::Millisecond),
                expected_filters: vec![col("k0").eq(lit("foo"))],
                expected_time_filters: vec![],
            },
            Case {
                filters: vec![
                    col("ts").between(ts_lit(1000), ts_lit(2000)),
                    col("k0").eq(lit("foo")),
                ],
                query_time_range: TimestampRange::with_unit(1000, 2001, TimeUnit::Millisecond),
                expected_filters: vec![col("k0").eq(lit("foo"))],
                expected_time_filters: vec![],
            },
            Case {
                filters: vec![col("ts").gt_eq(ts_lit(1200)), col("k0").eq(lit("foo"))],
                query_time_range: TimestampRange::with_unit(1200, 2001, TimeUnit::Millisecond),
                expected_filters: vec![col("k0").eq(lit("foo"))],
                expected_time_filters: vec![col("ts").gt_eq(ts_lit(1200))],
            },
            Case {
                filters: vec![
                    col("ts").gt_eq(ts_lit(1500)),
                    col("ts").is_not_null(),
                    col("k0").eq(lit("foo")),
                ],
                query_time_range: None,
                expected_filters: vec![col("k0").eq(lit("foo")), col("ts").is_not_null()],
                expected_time_filters: vec![col("ts").gt_eq(ts_lit(1500))],
            },
        ] {
            assert_range_cache_filters(
                case.filters,
                case.query_time_range,
                partition,
                case.expected_filters,
                case.expected_time_filters,
            )
            .await;
        }
    }

    #[tokio::test]
    async fn two_distinct_queries_share_cache_key_when_both_cover() {
        let partition_range = (
            Timestamp::new_millisecond(1000),
            Timestamp::new_millisecond(2000),
        );

        let (ctx_a, part_a) = new_stream_context(
            vec![
                col("ts").gt_eq(ts_lit(500)),
                col("ts").lt(ts_lit(3000)),
                col("k0").eq(lit("foo")),
            ],
            TimestampRange::with_unit(500, 3000, TimeUnit::Millisecond),
            partition_range,
        )
        .await;
        let (ctx_b, part_b) = new_stream_context(
            vec![
                col("ts").gt_eq(ts_lit(100)),
                col("ts").lt(ts_lit(5000)),
                col("k0").eq(lit("foo")),
            ],
            TimestampRange::with_unit(100, 5000, TimeUnit::Millisecond),
            partition_range,
        )
        .await;

        let key_a = build_range_cache_key(&ctx_a, &part_a).unwrap();
        let key_b = build_range_cache_key(&ctx_b, &part_b).unwrap();
        assert_eq!(key_a.scan, key_b.scan);
        assert!(key_a.scan.time_filters().is_empty());
    }

    #[tokio::test]
    async fn scan_stages_have_separate_cache_keys() {
        let partition_range = (
            Timestamp::new_millisecond(1000),
            Timestamp::new_millisecond(2000),
        );
        let (ctx, part_range) =
            new_stream_context(vec![col("k0").eq(lit("foo"))], None, partition_range).await;

        let data_key = build_range_cache_key(&ctx, &part_range).unwrap();
        let candidate_key = build_candidate_range_cache_key(&ctx, &part_range).unwrap();
        let range_0 = SeriesRange::new(0, 2).unwrap();
        let range_1 = SeriesRange::new(1, 2).unwrap();
        let series_key_0 = build_series_range_cache_key(&ctx, &part_range, range_0).unwrap();
        let series_key_1 = build_series_range_cache_key(&ctx, &part_range, range_1).unwrap();

        let merge_key = data_key.clone().for_cross_range_merge();
        assert_ne!(merge_key, data_key);
        assert_ne!(merge_key.scan, candidate_key.scan);
        assert_ne!(merge_key.scan, series_key_0.scan);
        assert_eq!(merge_key.row_groups, data_key.row_groups);
        assert_eq!(merge_key, data_key.clone().for_cross_range_merge());
        assert_ne!(data_key.scan, candidate_key.scan);
        assert_ne!(data_key.scan, series_key_0.scan);
        assert_ne!(candidate_key.scan, series_key_0.scan);
        assert_ne!(series_key_0.scan, series_key_1.scan);
        assert_eq!(data_key.row_groups, candidate_key.row_groups);
        assert_eq!(data_key.row_groups, series_key_0.row_groups);
    }

    #[tokio::test]
    async fn disables_optimization_on_or_clause() {
        let partition_range = (
            Timestamp::new_millisecond(1000),
            Timestamp::new_millisecond(2000),
        );

        let or_a = col("ts").gt_eq(ts_lit(1000)).or(col("ts").lt(ts_lit(500)));
        let or_b = col("ts").gt_eq(ts_lit(900)).or(col("ts").lt(ts_lit(400)));

        let (ctx_a, part_a) = new_stream_context(
            vec![or_a.clone(), col("k0").eq(lit("foo"))],
            None,
            partition_range,
        )
        .await;
        let (ctx_b, part_b) = new_stream_context(
            vec![or_b.clone(), col("k0").eq(lit("foo"))],
            None,
            partition_range,
        )
        .await;

        assert!(ctx_a.input.implied_time_range().is_none());
        let key_a = build_range_cache_key(&ctx_a, &part_a).unwrap();
        let key_b = build_range_cache_key(&ctx_b, &part_b).unwrap();
        assert_ne!(key_a.scan, key_b.scan);
        assert_eq!(
            key_a.scan.time_filters(),
            normalized_exprs([or_a]).as_slice()
        );
    }

    #[tokio::test]
    async fn empty_implied_range_does_not_panic_on_non_ms_file_range() {
        // Contradictory time predicates make the implied range empty. The
        // empty range's sentinel timestamps use `Timestamp::default()` (ms),
        // so without the `is_empty()` short-circuit the unit asserts would
        // panic against a non-ms `range_meta.time_range`.
        let partition = (
            Timestamp::new_millisecond(1000),
            Timestamp::new_millisecond(2000),
        );

        let (mut ctx, part_range) = new_stream_context(
            vec![
                col("ts").gt_eq(ts_lit(1500)),
                col("ts").lt(ts_lit(1500)),
                col("k0").eq(lit("foo")),
            ],
            TimestampRange::with_unit(1500, 3000, TimeUnit::Millisecond),
            partition,
        )
        .await;

        assert!(ctx.input.implied_time_range().unwrap().is_empty());
        ctx.ranges[0].time_range = (
            Timestamp::new(1_000_000_000, TimeUnit::Nanosecond),
            Timestamp::new(2_000_000_000, TimeUnit::Nanosecond),
        );

        let key = build_range_cache_key(&ctx, &part_range).unwrap();
        // Empty implied range cannot cover, so time filters stay in the key.
        assert!(!key.scan.time_filters().is_empty());
    }

    #[tokio::test]
    async fn mixed_unit_file_range_is_compared_by_instant() {
        // After the time index unit was widened, pre-alter files keep the old
        // unit while `scan_implied_time_range` uses the current column unit.
        // The coverage decision must compare instants across units instead of
        // asserting unit equality (which panicked here for widened regions).
        let partition = (
            Timestamp::new_millisecond(1000),
            Timestamp::new_millisecond(2000),
        );

        // Covering: the implied range [0, 3000)ms contains the file range
        // [1s, 2s] expressed in microseconds; time filters are dropped from
        // the key.
        let (mut ctx, part_range) = new_stream_context(
            vec![
                col("ts").gt_eq(ts_lit(0)),
                col("ts").lt(ts_lit(3000)),
                col("k0").eq(lit("foo")),
            ],
            TimestampRange::with_unit(0, 3000, TimeUnit::Millisecond),
            partition,
        )
        .await;
        ctx.ranges[0].time_range = (
            Timestamp::new_microsecond(1_000_000),
            Timestamp::new_microsecond(2_000_000),
        );
        let key = build_range_cache_key(&ctx, &part_range).unwrap();
        assert!(key.scan.time_filters().is_empty());

        // Not covering: the implied range [1500, 3000)ms excludes the file's
        // 1s lower bound; time filters stay in the key.
        let (mut ctx, part_range) = new_stream_context(
            vec![
                col("ts").gt_eq(ts_lit(1500)),
                col("ts").lt(ts_lit(3000)),
                col("k0").eq(lit("foo")),
            ],
            TimestampRange::with_unit(1500, 3000, TimeUnit::Millisecond),
            partition,
        )
        .await;
        ctx.ranges[0].time_range = (
            Timestamp::new_microsecond(1_000_000),
            Timestamp::new_microsecond(2_000_000),
        );
        let key = build_range_cache_key(&ctx, &part_range).unwrap();
        assert!(!key.scan.time_filters().is_empty());
    }

    fn ms_ts(v: i64) -> Timestamp {
        Timestamp::new_millisecond(v)
    }

    fn implied_ms(expr: Expr) -> Option<TimestampRange> {
        implied_time_range_from_exprs("ts", TimeUnit::Millisecond, &[&expr])
    }

    #[test]
    fn implied_time_range_supported_exprs() {
        for (expr, expected) in [
            (
                col("ts").gt_eq(ts_lit(1000)),
                Some(TimestampRange::from_start(ms_ts(1000))),
            ),
            (
                col("ts").gt(ts_lit(1000)),
                Some(TimestampRange::from_start(ms_ts(1001))),
            ),
            (
                col("ts").lt_eq(ts_lit(2000)),
                Some(TimestampRange::until_end(ms_ts(2000), true)),
            ),
            (
                col("ts").lt(ts_lit(2000)),
                Some(TimestampRange::until_end(ms_ts(2000), false)),
            ),
            (
                col("ts").eq(ts_lit(1500)),
                Some(TimestampRange::single(ms_ts(1500))),
            ),
            (
                ts_lit(1000).lt_eq(col("ts")),
                Some(TimestampRange::from_start(ms_ts(1000))),
            ),
            (
                col("ts").between(ts_lit(1000), ts_lit(2000)),
                Some(TimestampRange::new_inclusive(
                    Some(ms_ts(1000)),
                    Some(ms_ts(2000)),
                )),
            ),
            (
                col("ts")
                    .gt_eq(ts_lit(1000))
                    .and(col("ts").lt(ts_lit(2000))),
                TimestampRange::with_unit(1000, 2000, TimeUnit::Millisecond),
            ),
            (
                col("ts")
                    .gt_eq(ts_lit(1000))
                    .and(col("ts").lt(ts_lit(5000)))
                    .and(col("ts").lt_eq(ts_lit(3000))),
                TimestampRange::with_unit(1000, 3001, TimeUnit::Millisecond),
            ),
        ] {
            assert_eq!(implied_ms(expr), expected);
        }

        assert_eq!(
            implied_time_range_from_exprs("ts", TimeUnit::Millisecond, &[]),
            None
        );
    }

    #[test]
    fn implied_time_range_unsupported_exprs() {
        let not_between = Expr::Between(Between {
            expr: Box::new(col("ts")),
            negated: true,
            low: Box::new(ts_lit(1000)),
            high: Box::new(ts_lit(2000)),
        });

        for expr in [
            not_between,
            col("ts").gt_eq(ts_lit(1000)).or(col("ts").lt(ts_lit(500))),
            Expr::Not(Box::new(col("ts").gt_eq(ts_lit(1000)))),
            col("ts").in_list(vec![ts_lit(1000), ts_lit(2000)], false),
            col("ts").gt_eq(col("other")),
            col("other_ts").gt_eq(ts_lit(1000)),
        ] {
            assert!(implied_ms(expr).is_none());
        }
    }

    #[test]
    fn implied_time_range_unit_conversion() {
        let second_1 = lit(ScalarValue::TimestampSecond(Some(1), None));
        let ns_1500 = lit(ScalarValue::TimestampNanosecond(Some(1_500_000_000), None));
        let ns_1500_5 = lit(ScalarValue::TimestampNanosecond(Some(1_500_500_000), None));

        for (expr, expected) in [
            (
                col("ts").gt_eq(second_1.clone()),
                Some(TimestampRange::from_start(ms_ts(1000))),
            ),
            (
                col("ts").lt_eq(second_1),
                Some(TimestampRange::until_end(ms_ts(1000), true)),
            ),
            (
                col("ts").eq(ns_1500),
                Some(TimestampRange::single(ms_ts(1500))),
            ),
            (col("ts").eq(ns_1500_5.clone()), None),
            (
                col("ts").gt_eq(ns_1500_5.clone()),
                Some(TimestampRange::from_start(ms_ts(1501))),
            ),
            (
                col("ts").lt_eq(ns_1500_5.clone()),
                Some(TimestampRange::until_end(ms_ts(1500), true)),
            ),
            (
                col("ts").gt(ns_1500_5.clone()),
                Some(TimestampRange::from_start(ms_ts(1501))),
            ),
            (
                col("ts").lt(ns_1500_5),
                Some(TimestampRange::until_end(ms_ts(1501), false)),
            ),
        ] {
            assert_eq!(implied_ms(expr), expected);
        }
    }

    #[test]
    fn selector_after_merge_changes_fingerprint() {
        let ordinary = test_scan_fingerprint(
            vec!["k0 = 'foo'".to_string()],
            vec![],
            Some(TimeSeriesRowSelector::LastRow { after_merge: false }),
            true,
            0,
        );
        let after_merge = test_scan_fingerprint(
            vec!["k0 = 'foo'".to_string()],
            vec![],
            Some(TimeSeriesRowSelector::LastRow { after_merge: true }),
            true,
            0,
        );

        assert_ne!(ordinary, after_merge);
    }

    #[test]
    fn true_selector_after_merge_is_preserved_by_fingerprint_transforms() {
        let normalized =
            test_scan_fingerprint(vec!["k0 = 'foo'".to_string()], vec![], None, true, 0);

        assert!(normalized.time_filters().is_empty());

        let fingerprint = test_scan_fingerprint(
            vec!["k0 = 'foo'".to_string()],
            vec!["ts >= 1000".to_string()],
            Some(TimeSeriesRowSelector::LastRow { after_merge: true }),
            true,
            7,
        );

        let reset = fingerprint.without_time_filters();
        let candidate = fingerprint.for_candidate_series();
        let series_data = fingerprint.for_series_data(SeriesRange::new(0, 1).unwrap());

        assert_eq!(reset.read_columns(), fingerprint.read_columns());
        assert_eq!(reset.read_column_types(), fingerprint.read_column_types());
        assert_eq!(reset.filters(), fingerprint.filters());
        assert!(reset.time_filters().is_empty());
        assert_eq!(reset.series_row_selector, fingerprint.series_row_selector);
        assert_eq!(
            fingerprint.series_row_selector,
            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
        );
        assert_eq!(
            candidate.series_row_selector,
            fingerprint.series_row_selector
        );
        assert_eq!(
            series_data.series_row_selector,
            fingerprint.series_row_selector
        );
        assert_eq!(reset.append_mode, fingerprint.append_mode);
        assert_eq!(reset.filter_deleted, fingerprint.filter_deleted);
        assert_eq!(reset.merge_mode, fingerprint.merge_mode);
        assert_eq!(
            reset.partition_expr_version,
            fingerprint.partition_expr_version
        );
    }

    fn test_schema() -> Arc<datatypes::arrow::datatypes::Schema> {
        use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema};

        Arc::new(Schema::new(vec![Field::new(
            "value",
            ArrowDataType::Int64,
            false,
        )]))
    }

    fn make_batch(values: &[i64]) -> RecordBatch {
        use datatypes::arrow::array::Int64Array;

        RecordBatch::try_new(
            test_schema(),
            vec![Arc::new(Int64Array::from(values.to_vec()))],
        )
        .unwrap()
    }

    fn make_large_binary_batch(rows: usize, bytes_per_row: usize) -> RecordBatch {
        use datatypes::arrow::array::BinaryArray;
        use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema};

        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            ArrowDataType::Binary,
            false,
        )]));
        let payload = vec![b'x'; bytes_per_row];
        let values = (0..rows).map(|_| payload.as_slice()).collect::<Vec<_>>();

        RecordBatch::try_new(schema, vec![Arc::new(BinaryArray::from_vec(values))]).unwrap()
    }

    #[test]
    fn compact_record_batches_keeps_original_boundaries() {
        let batches = vec![make_batch(&[1, 2]), make_batch(&[3]), make_batch(&[4, 5])];

        let compacted = compact_record_batches(batches).unwrap();

        assert_eq!(compacted.batch.num_rows(), 5);
        assert_eq!(compacted.slice_lengths, vec![2, 1, 2]);
    }

    #[tokio::test]
    async fn cached_flat_range_stream_replays_original_batches() {
        let value = Arc::new(RangeScanCacheValue::new(
            vec![CachedBatchSlice {
                batch: make_batch(&[1, 2, 3]),
                slice_lengths: vec![2, 1],
            }],
            make_batch(&[1, 2, 3]).get_array_memory_size(),
        ));

        let worker = RangeCacheWorker::new(1024 * 1024);
        let permit = worker
            .load_memory
            .try_acquire(value.estimated_size())
            .unwrap();
        let delivery = RangeCacheDelivery::new(value, permit);
        let replayed = cached_flat_range_stream(delivery)
            .try_collect::<Vec<_>>()
            .await
            .unwrap();

        assert_eq!(replayed.len(), 2);
        assert_eq!(replayed[0].num_rows(), 2);
        assert_eq!(replayed[1].num_rows(), 1);
    }

    #[tokio::test]
    async fn dropped_cache_buffer_releases_batches_while_waiting_for_memory() {
        let strategy = test_cache_strategy();
        let limiter = strategy.range_result_memory_limiter().unwrap();
        // Take every permit so the concat task blocks before it can compact.
        let _permit = limiter
            .acquire(limiter.available_permits() * limiter.permit_bytes())
            .await
            .unwrap();
        let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]);
        let weak = Arc::downgrade(batch.column(0));
        let mut buffer = CacheBatchBuffer::new(&strategy);
        buffer.push(batch.clone()).unwrap();
        buffer.push(batch).unwrap();
        // Both batches were handed to the task instead of staying in the buffer.
        assert!(buffer.buffered_batches.is_empty());

        // Cancel only once the task is parked on the permit, otherwise the abort
        // could land on a task that was never polled.
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            while limiter.waited_acquires() == 0 {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("concat task must wait for a memory permit");
        drop(buffer);

        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            while weak.upgrade().is_some() {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("cancelled cache work must release input without waiting for a permit");
    }

    #[tokio::test]
    async fn cache_batch_buffer_finishes_pending_batches() {
        let strategy = test_cache_strategy();
        let batch = make_batch(&[1, 2, 3]);
        let expected_size = batch.get_array_memory_size();

        let mut buffer = CacheBatchBuffer::new(&strategy);
        buffer.push(batch).unwrap();

        let value = buffer.finish().await.unwrap().unwrap();
        assert_eq!(value.cached_batches.len(), 1);
        assert_eq!(value.cached_batches[0].slice_lengths, vec![3]);
        assert_eq!(value.estimated_batches_size, expected_size);
    }

    #[tokio::test]
    async fn cache_batch_buffer_compacts_when_rows_exceed_default_batch_size() {
        let strategy = test_cache_strategy();
        let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]);

        let mut buffer = CacheBatchBuffer::new(&strategy);
        buffer.push(batch.clone()).unwrap();
        buffer.push(batch).unwrap();

        assert_eq!(buffer.buffered_rows, 0);
        assert!(buffer.buffered_batches.is_empty());

        let value = buffer.finish().await.unwrap().unwrap();
        assert_eq!(value.cached_batches.len(), 1);
        assert_eq!(
            value.cached_batches[0].slice_lengths,
            vec![
                DEFAULT_READ_BATCH_SIZE / 2 + 1,
                DEFAULT_READ_BATCH_SIZE / 2 + 1
            ]
        );
    }

    #[tokio::test]
    async fn cache_batch_buffer_compacts_when_buffered_size_exceeds_threshold() {
        let large_batch = make_large_binary_batch(DEFAULT_READ_BATCH_SIZE, 4096);
        let strategy = CacheStrategy::EnableAll(Arc::new(
            CacheManager::builder()
                .range_result_cache_size((large_batch.get_array_memory_size() * 3) as u64)
                .build(),
        ));

        let mut buffer = CacheBatchBuffer::new(&strategy);
        buffer.push(large_batch.clone()).unwrap();

        assert_eq!(buffer.buffered_rows, large_batch.num_rows());
        assert_eq!(buffer.buffered_batches.len(), 1);

        buffer.push(large_batch.clone()).unwrap();

        assert_eq!(buffer.buffered_rows, 0);
        assert!(buffer.buffered_batches.is_empty());

        let value = buffer.finish().await.unwrap().unwrap();
        assert_eq!(value.cached_batches.len(), 1);
        assert_eq!(
            value.cached_batches[0].slice_lengths,
            vec![large_batch.num_rows(), large_batch.num_rows()]
        );
    }

    #[tokio::test]
    async fn cache_batch_buffer_skips_cache_when_compacted_size_exceeds_limit() {
        let large_batch = make_large_binary_batch(DEFAULT_READ_BATCH_SIZE / 2 + 1, 4096);
        // Budget only fits two large batches.
        let budget = (large_batch.get_array_memory_size() as u64) * 2 + 1;
        let strategy = CacheStrategy::EnableAll(Arc::new(
            CacheManager::builder()
                .range_result_cache_size(budget)
                .build(),
        ));

        let mut buffer = CacheBatchBuffer::new(&strategy);
        for _ in 0..4 {
            buffer.push(large_batch.clone()).unwrap();
        }
        assert!(!matches!(buffer.finish().await, Some(Ok(_))));
    }
}