mito2/read/

scan_util.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 scanners.

use std::collections::{BinaryHeap, HashMap, VecDeque};
use std::fmt;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

use async_stream::try_stream;
use common_telemetry::tracing;
use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::timestamp::timestamp_array_to_primitive;
#[cfg(feature = "enterprise")]
use futures::TryStreamExt;
use futures::{Stream, StreamExt};
use prometheus::IntGauge;
use smallvec::SmallVec;
use snafu::ResultExt;
use store_api::storage::{RegionId, SequenceRange, TimeSeriesRowSelector};

use crate::error::{ComputeArrowSnafu, Result};
use crate::memtable::MemScanMetrics;
use crate::metrics::{
    IN_PROGRESS_SCAN, PRECISE_FILTER_ROWS_TOTAL, READ_BATCHES_RETURN, READ_ROW_GROUPS_TOTAL,
    READ_ROWS_IN_ROW_GROUP_TOTAL, READ_ROWS_RETURN, READ_STAGE_ELAPSED,
};
use crate::read::dedup::{DedupMetrics, DedupMetricsReport};
use crate::read::flat_merge::{MergeMetrics, MergeMetricsReport};
use crate::read::last_row::FlatLastTimestampSelector;
use crate::read::pruner::PartitionPruner;
use crate::read::range::{RangeMeta, RowGroupIndex};
use crate::read::scan_region::StreamContext;
use crate::read::{BoxedRecordBatchStream, ScannerMetrics};
use crate::sst::file::{FileTimeRange, RegionFileId};
use crate::sst::index::bloom_filter::applier::BloomFilterIndexApplyMetrics;
use crate::sst::index::fulltext_index::applier::FulltextIndexApplyMetrics;
use crate::sst::index::inverted_index::applier::InvertedIndexApplyMetrics;
use crate::sst::parquet::file_range::{FileRange, PreFilterMode};
use crate::sst::parquet::flat_format::{sequence_column_index, time_index_column_index};
use crate::sst::parquet::json_align;
use crate::sst::parquet::reader::{MetadataCacheMetrics, ReaderFilterMetrics, ReaderMetrics};
use crate::sst::parquet::row_group::ParquetFetchMetrics;
use crate::sst::parquet::{DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE};

/// Per-file scan metrics.
#[derive(Default, Clone)]
pub struct FileScanMetrics {
    /// Number of ranges (row groups) read from this file.
    pub num_ranges: usize,
    /// Number of rows read from this file.
    pub num_rows: usize,
    /// Time spent building file ranges/parts (file-level preparation).
    pub build_part_cost: Duration,
    /// Time spent building readers for this file (accumulated across all ranges).
    pub build_reader_cost: Duration,
    /// Time spent scanning this file (accumulated across all ranges).
    pub scan_cost: Duration,
}

impl fmt::Debug for FileScanMetrics {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{{\"build_part_cost\":\"{:?}\"", self.build_part_cost)?;

        if self.num_ranges > 0 {
            write!(f, ", \"num_ranges\":{}", self.num_ranges)?;
        }
        if self.num_rows > 0 {
            write!(f, ", \"num_rows\":{}", self.num_rows)?;
        }
        if !self.build_reader_cost.is_zero() {
            write!(
                f,
                ", \"build_reader_cost\":\"{:?}\"",
                self.build_reader_cost
            )?;
        }
        if !self.scan_cost.is_zero() {
            write!(f, ", \"scan_cost\":\"{:?}\"", self.scan_cost)?;
        }

        write!(f, "}}")
    }
}

impl FileScanMetrics {
    /// Merges another FileMetrics into this one.
    pub(crate) fn merge_from(&mut self, other: &FileScanMetrics) {
        self.num_ranges += other.num_ranges;
        self.num_rows += other.num_rows;
        self.build_part_cost += other.build_part_cost;
        self.build_reader_cost += other.build_reader_cost;
        self.scan_cost += other.scan_cost;
    }
}

/// Verbose scan metrics for a partition.
#[derive(Default)]
pub(crate) struct ScanMetricsSet {
    /// Duration to prepare the scan task.
    prepare_scan_cost: Duration,
    /// Duration to build the (merge) reader.
    build_reader_cost: Duration,
    /// Duration to scan data.
    scan_cost: Duration,
    /// Duration while waiting for `yield`.
    yield_cost: Duration,
    /// Duration to convert [`Batch`]es.
    convert_cost: Option<Time>,
    /// Duration of the scan.
    total_cost: Duration,
    /// Number of rows returned.
    num_rows: usize,
    /// Number of batches returned.
    num_batches: usize,
    /// Number of mem ranges scanned.
    num_mem_ranges: usize,
    /// Number of file ranges scanned.
    num_file_ranges: usize,

    // Memtable related metrics:
    /// Duration to scan memtables.
    mem_scan_cost: Duration,
    /// Number of rows read from memtables.
    mem_rows: usize,
    /// Number of batches read from memtables.
    mem_batches: usize,
    /// Number of series read from memtables.
    mem_series: usize,
    /// Duration of prefilter in memtable scan.
    mem_prefilter_cost: Duration,
    /// Number of rows filtered by prefilter in memtable scan.
    mem_prefilter_rows_filtered: usize,

    // SST related metrics:
    /// Duration to build file ranges.
    build_parts_cost: Duration,
    /// Duration to scan SST files.
    sst_scan_cost: Duration,
    /// Number of row groups before filtering.
    rg_total: usize,
    /// Number of row groups filtered by fulltext index.
    rg_fulltext_filtered: usize,
    /// Number of row groups filtered by inverted index.
    rg_inverted_filtered: usize,
    /// Number of row groups filtered by min-max index.
    rg_minmax_filtered: usize,
    /// Number of row groups filtered by bloom filter index.
    rg_bloom_filtered: usize,
    /// Number of rows in row group before filtering.
    rows_before_filter: usize,
    /// Number of rows in row group filtered by fulltext index.
    rows_fulltext_filtered: usize,
    /// Number of rows in row group filtered by inverted index.
    rows_inverted_filtered: usize,
    /// Number of rows in row group filtered by bloom filter index.
    rows_bloom_filtered: usize,
    /// Number of rows filtered by precise filter.
    rows_precise_filtered: usize,
    /// Number of index result cache hits for fulltext index.
    fulltext_index_cache_hit: usize,
    /// Number of index result cache misses for fulltext index.
    fulltext_index_cache_miss: usize,
    /// Number of index result cache hits for inverted index.
    inverted_index_cache_hit: usize,
    /// Number of index result cache misses for inverted index.
    inverted_index_cache_miss: usize,
    /// Number of index result cache hits for bloom filter index.
    bloom_filter_cache_hit: usize,
    /// Number of index result cache misses for bloom filter index.
    bloom_filter_cache_miss: usize,
    /// Number of index result cache hits for minmax pruning.
    minmax_cache_hit: usize,
    /// Number of index result cache misses for minmax pruning.
    minmax_cache_miss: usize,
    /// Number of pruner builder cache hits.
    pruner_cache_hit: usize,
    /// Number of pruner builder cache misses.
    pruner_cache_miss: usize,
    /// Duration spent waiting for pruner to build file ranges.
    pruner_prune_cost: Duration,
    /// Number of files filtered by manifest time-range pruning.
    files_time_range_pruned: usize,
    /// Number of record batches read from SST.
    num_sst_record_batches: usize,
    /// Number of batches decoded from SST.
    num_sst_batches: usize,
    /// Number of rows read from SST.
    num_sst_rows: usize,

    /// Elapsed time before the first poll operation.
    first_poll: Duration,

    /// Number of send full in SeriesScan.
    num_series_send_full: usize,
    /// Number of rows the series distributor scanned.
    num_distributor_rows: usize,
    /// Number of batches the series distributor scanned.
    num_distributor_batches: usize,
    /// Duration of the series distributor to scan.
    distributor_scan_cost: Duration,
    /// Duration of the series distributor to yield.
    distributor_yield_cost: Duration,
    /// Duration spent in divider operations.
    distributor_divider_cost: Duration,

    /// Merge metrics.
    merge_metrics: MergeMetrics,
    /// Dedup metrics.
    dedup_metrics: DedupMetrics,

    /// The stream reached EOF
    stream_eof: bool,

    // Optional verbose metrics:
    /// Inverted index apply metrics.
    inverted_index_apply_metrics: Option<InvertedIndexApplyMetrics>,
    /// Bloom filter index apply metrics.
    bloom_filter_apply_metrics: Option<BloomFilterIndexApplyMetrics>,
    /// Fulltext index apply metrics.
    fulltext_index_apply_metrics: Option<FulltextIndexApplyMetrics>,
    /// Parquet fetch metrics.
    fetch_metrics: Option<ParquetFetchMetrics>,
    /// Metadata cache metrics.
    metadata_cache_metrics: Option<MetadataCacheMetrics>,
    /// Per-file scan metrics, only populated when explain_verbose is true.
    per_file_metrics: Option<HashMap<RegionFileId, FileScanMetrics>>,

    /// Current memory usage for file range builders.
    build_ranges_mem_size: isize,
    /// Peak memory usage for file range builders.
    build_ranges_peak_mem_size: isize,
    /// Current number of file range builders.
    num_range_builders: isize,
    /// Peak number of file range builders.
    num_peak_range_builders: isize,
    /// Total bytes added to the range cache during this scan.
    range_cache_size: usize,
    /// Number of range cache hits during this scan.
    range_cache_hit: usize,
    /// Number of range cache misses during this scan.
    range_cache_miss: usize,
}

/// Wrapper for file metrics that compares by total cost in reverse order.
/// This allows using BinaryHeap as a min-heap for efficient top-K selection.
struct CompareCostReverse<'a> {
    total_cost: Duration,
    file_id: RegionFileId,
    metrics: &'a FileScanMetrics,
}

impl Ord for CompareCostReverse<'_> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        // Reverse comparison: smaller costs are "greater"
        other.total_cost.cmp(&self.total_cost)
    }
}

impl PartialOrd for CompareCostReverse<'_> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Eq for CompareCostReverse<'_> {}

impl PartialEq for CompareCostReverse<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.total_cost == other.total_cost
    }
}

impl fmt::Debug for ScanMetricsSet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let ScanMetricsSet {
            prepare_scan_cost,
            build_reader_cost,
            scan_cost,
            yield_cost,
            convert_cost,
            total_cost,
            num_rows,
            num_batches,
            num_mem_ranges,
            num_file_ranges,
            build_parts_cost,
            sst_scan_cost,
            rg_total,
            rg_fulltext_filtered,
            rg_inverted_filtered,
            rg_minmax_filtered,
            rg_bloom_filtered,
            rows_before_filter,
            rows_fulltext_filtered,
            rows_inverted_filtered,
            rows_bloom_filtered,
            rows_precise_filtered,
            fulltext_index_cache_hit,
            fulltext_index_cache_miss,
            inverted_index_cache_hit,
            inverted_index_cache_miss,
            bloom_filter_cache_hit,
            bloom_filter_cache_miss,
            minmax_cache_hit,
            minmax_cache_miss,
            pruner_cache_hit,
            pruner_cache_miss,
            pruner_prune_cost,
            files_time_range_pruned,
            num_sst_record_batches,
            num_sst_batches,
            num_sst_rows,
            first_poll,
            num_series_send_full,
            num_distributor_rows,
            num_distributor_batches,
            distributor_scan_cost,
            distributor_yield_cost,
            distributor_divider_cost,
            merge_metrics,
            dedup_metrics,
            stream_eof,
            mem_scan_cost,
            mem_rows,
            mem_batches,
            mem_series,
            mem_prefilter_cost,
            mem_prefilter_rows_filtered,
            inverted_index_apply_metrics,
            bloom_filter_apply_metrics,
            fulltext_index_apply_metrics,
            fetch_metrics,
            metadata_cache_metrics,
            per_file_metrics,
            build_ranges_mem_size: _,
            build_ranges_peak_mem_size,
            num_range_builders: _,
            num_peak_range_builders,
            range_cache_size,
            range_cache_hit,
            range_cache_miss,
        } = self;

        // Write core metrics
        write!(
            f,
            "{{\"prepare_scan_cost\":\"{prepare_scan_cost:?}\", \
            \"build_reader_cost\":\"{build_reader_cost:?}\", \
            \"scan_cost\":\"{scan_cost:?}\", \
            \"yield_cost\":\"{yield_cost:?}\", \
            \"total_cost\":\"{total_cost:?}\", \
            \"num_rows\":{num_rows}, \
            \"num_batches\":{num_batches}, \
            \"num_mem_ranges\":{num_mem_ranges}, \
            \"num_file_ranges\":{num_file_ranges}, \
            \"build_parts_cost\":\"{build_parts_cost:?}\", \
            \"sst_scan_cost\":\"{sst_scan_cost:?}\", \
            \"rg_total\":{rg_total}, \
            \"rows_before_filter\":{rows_before_filter}, \
            \"num_sst_record_batches\":{num_sst_record_batches}, \
            \"num_sst_batches\":{num_sst_batches}, \
            \"num_sst_rows\":{num_sst_rows}, \
            \"first_poll\":\"{first_poll:?}\""
        )?;

        // Write convert_cost if present
        if let Some(time) = convert_cost {
            let duration = Duration::from_nanos(time.value() as u64);
            write!(f, ", \"convert_cost\":\"{duration:?}\"")?;
        }

        // Write non-zero filter counters
        if *files_time_range_pruned > 0 {
            write!(f, ", \"files_time_range_pruned\":{files_time_range_pruned}")?;
        }
        if *rg_fulltext_filtered > 0 {
            write!(f, ", \"rg_fulltext_filtered\":{rg_fulltext_filtered}")?;
        }
        if *rg_inverted_filtered > 0 {
            write!(f, ", \"rg_inverted_filtered\":{rg_inverted_filtered}")?;
        }
        if *rg_minmax_filtered > 0 {
            write!(f, ", \"rg_minmax_filtered\":{rg_minmax_filtered}")?;
        }
        if *rg_bloom_filtered > 0 {
            write!(f, ", \"rg_bloom_filtered\":{rg_bloom_filtered}")?;
        }
        if *rows_fulltext_filtered > 0 {
            write!(f, ", \"rows_fulltext_filtered\":{rows_fulltext_filtered}")?;
        }
        if *rows_inverted_filtered > 0 {
            write!(f, ", \"rows_inverted_filtered\":{rows_inverted_filtered}")?;
        }
        if *rows_bloom_filtered > 0 {
            write!(f, ", \"rows_bloom_filtered\":{rows_bloom_filtered}")?;
        }
        if *rows_precise_filtered > 0 {
            write!(f, ", \"rows_precise_filtered\":{rows_precise_filtered}")?;
        }
        if *fulltext_index_cache_hit > 0 {
            write!(
                f,
                ", \"fulltext_index_cache_hit\":{fulltext_index_cache_hit}"
            )?;
        }
        if *fulltext_index_cache_miss > 0 {
            write!(
                f,
                ", \"fulltext_index_cache_miss\":{fulltext_index_cache_miss}"
            )?;
        }
        if *inverted_index_cache_hit > 0 {
            write!(
                f,
                ", \"inverted_index_cache_hit\":{inverted_index_cache_hit}"
            )?;
        }
        if *inverted_index_cache_miss > 0 {
            write!(
                f,
                ", \"inverted_index_cache_miss\":{inverted_index_cache_miss}"
            )?;
        }
        if *bloom_filter_cache_hit > 0 {
            write!(f, ", \"bloom_filter_cache_hit\":{bloom_filter_cache_hit}")?;
        }
        if *bloom_filter_cache_miss > 0 {
            write!(f, ", \"bloom_filter_cache_miss\":{bloom_filter_cache_miss}")?;
        }
        if *minmax_cache_hit > 0 {
            write!(f, ", \"minmax_cache_hit\":{minmax_cache_hit}")?;
        }
        if *minmax_cache_miss > 0 {
            write!(f, ", \"minmax_cache_miss\":{minmax_cache_miss}")?;
        }
        if *pruner_cache_hit > 0 {
            write!(f, ", \"pruner_cache_hit\":{pruner_cache_hit}")?;
        }
        if *pruner_cache_miss > 0 {
            write!(f, ", \"pruner_cache_miss\":{pruner_cache_miss}")?;
        }
        if !pruner_prune_cost.is_zero() {
            write!(f, ", \"pruner_prune_cost\":\"{pruner_prune_cost:?}\"")?;
        }

        // Write non-zero distributor metrics
        if *num_series_send_full > 0 {
            write!(f, ", \"num_series_send_full\":{num_series_send_full}")?;
        }
        if *num_distributor_rows > 0 {
            write!(f, ", \"num_distributor_rows\":{num_distributor_rows}")?;
        }
        if *num_distributor_batches > 0 {
            write!(f, ", \"num_distributor_batches\":{num_distributor_batches}")?;
        }
        if !distributor_scan_cost.is_zero() {
            write!(
                f,
                ", \"distributor_scan_cost\":\"{distributor_scan_cost:?}\""
            )?;
        }
        if !distributor_yield_cost.is_zero() {
            write!(
                f,
                ", \"distributor_yield_cost\":\"{distributor_yield_cost:?}\""
            )?;
        }
        if !distributor_divider_cost.is_zero() {
            write!(
                f,
                ", \"distributor_divider_cost\":\"{distributor_divider_cost:?}\""
            )?;
        }

        // Write non-zero memtable metrics
        if *mem_rows > 0 {
            write!(f, ", \"mem_rows\":{mem_rows}")?;
        }
        if *mem_batches > 0 {
            write!(f, ", \"mem_batches\":{mem_batches}")?;
        }
        if *mem_series > 0 {
            write!(f, ", \"mem_series\":{mem_series}")?;
        }
        if !mem_scan_cost.is_zero() {
            write!(f, ", \"mem_scan_cost\":\"{mem_scan_cost:?}\"")?;
        }
        if !mem_prefilter_cost.is_zero() {
            write!(f, ", \"mem_prefilter_cost\":\"{mem_prefilter_cost:?}\"")?;
        }
        if *mem_prefilter_rows_filtered > 0 {
            write!(
                f,
                ", \"mem_prefilter_rows_filtered\":{mem_prefilter_rows_filtered}"
            )?;
        }

        // Write optional verbose metrics if they are not empty
        if let Some(metrics) = inverted_index_apply_metrics
            && !metrics.is_empty()
        {
            write!(f, ", \"inverted_index_apply_metrics\":{:?}", metrics)?;
        }
        if let Some(metrics) = bloom_filter_apply_metrics
            && !metrics.is_empty()
        {
            write!(f, ", \"bloom_filter_apply_metrics\":{:?}", metrics)?;
        }
        if let Some(metrics) = fulltext_index_apply_metrics
            && !metrics.is_empty()
        {
            write!(f, ", \"fulltext_index_apply_metrics\":{:?}", metrics)?;
        }
        if let Some(metrics) = fetch_metrics
            && !metrics.is_empty()
        {
            write!(f, ", \"fetch_metrics\":{:?}", metrics)?;
        }
        if let Some(metrics) = metadata_cache_metrics
            && !metrics.is_empty()
        {
            write!(f, ", \"metadata_cache_metrics\":{:?}", metrics)?;
        }

        // Write merge metrics if not empty
        if !merge_metrics.scan_cost.is_zero() {
            write!(f, ", \"merge_metrics\":{:?}", merge_metrics)?;
        }

        // Write dedup metrics if not empty
        if !dedup_metrics.dedup_cost.is_zero() {
            write!(f, ", \"dedup_metrics\":{:?}", dedup_metrics)?;
        }

        // Write top file metrics if present and non-empty
        if let Some(file_metrics) = per_file_metrics
            && !file_metrics.is_empty()
        {
            // Use min-heap (BinaryHeap with reverse comparison) to keep only top 10
            let mut heap = BinaryHeap::new();
            for (file_id, metrics) in file_metrics.iter() {
                let total_cost =
                    metrics.build_part_cost + metrics.build_reader_cost + metrics.scan_cost;

                // If the file has been pruned by a pruner, the build part cost may be zero.
                // If we didn't read any ranges from it, we don't output the file.
                if total_cost.is_zero() && metrics.num_ranges == 0 {
                    continue;
                }

                if heap.len() < 10 {
                    // Haven't reached 10 yet, just push
                    heap.push(CompareCostReverse {
                        total_cost,
                        file_id: *file_id,
                        metrics,
                    });
                } else if let Some(min_entry) = heap.peek() {
                    // If current cost is higher than the minimum in our top-10, replace it
                    if total_cost > min_entry.total_cost {
                        heap.pop();
                        heap.push(CompareCostReverse {
                            total_cost,
                            file_id: *file_id,
                            metrics,
                        });
                    }
                }
            }

            let top_files = heap.into_sorted_vec();
            write!(f, ", \"top_file_metrics\": {{")?;
            for (i, item) in top_files.iter().enumerate() {
                let CompareCostReverse {
                    total_cost: _,
                    file_id,
                    metrics,
                } = item;
                if i > 0 {
                    write!(f, ", ")?;
                }
                write!(f, "\"{}\": {:?}", file_id, metrics)?;
            }
            write!(f, "}}")?;
        }

        if *range_cache_size > 0 {
            write!(f, ", \"range_cache_size\":{range_cache_size}")?;
        }
        if *range_cache_hit > 0 {
            write!(f, ", \"range_cache_hit\":{range_cache_hit}")?;
        }
        if *range_cache_miss > 0 {
            write!(f, ", \"range_cache_miss\":{range_cache_miss}")?;
        }

        write!(
            f,
            ", \"build_ranges_peak_mem_size\":{build_ranges_peak_mem_size}, \
             \"num_peak_range_builders\":{num_peak_range_builders}, \
             \"stream_eof\":{stream_eof}}}"
        )
    }
}
impl ScanMetricsSet {
    /// Attaches the `prepare_scan_cost` to the metrics set.
    fn with_prepare_scan_cost(mut self, cost: Duration) -> Self {
        self.prepare_scan_cost += cost;
        self
    }

    /// Attaches the `convert_cost` to the metrics set.
    fn with_convert_cost(mut self, time: Time) -> Self {
        self.convert_cost = Some(time);
        self
    }

    /// Merges the local scanner metrics.
    fn merge_scanner_metrics(&mut self, other: &ScannerMetrics) {
        let ScannerMetrics {
            scan_cost,
            yield_cost,
            num_batches,
            num_rows,
        } = other;

        self.scan_cost += *scan_cost;
        self.yield_cost += *yield_cost;
        self.num_rows += *num_rows;
        self.num_batches += *num_batches;
    }

    /// Merges the local reader metrics.
    fn merge_reader_metrics(&mut self, other: &ReaderMetrics) {
        let ReaderMetrics {
            build_cost,
            filter_metrics:
                ReaderFilterMetrics {
                    rg_total,
                    rg_fulltext_filtered,
                    rg_inverted_filtered,
                    rg_minmax_filtered,
                    rg_bloom_filtered,
                    rows_total,
                    rows_fulltext_filtered,
                    rows_inverted_filtered,
                    rows_bloom_filtered,
                    rows_precise_filtered,
                    fulltext_index_cache_hit,
                    fulltext_index_cache_miss,
                    inverted_index_cache_hit,
                    inverted_index_cache_miss,
                    bloom_filter_cache_hit,
                    bloom_filter_cache_miss,
                    minmax_cache_hit,
                    minmax_cache_miss,
                    pruner_cache_hit,
                    pruner_cache_miss,
                    pruner_prune_cost,
                    files_time_range_pruned,
                    inverted_index_apply_metrics,
                    bloom_filter_apply_metrics,
                    fulltext_index_apply_metrics,
                },
            num_record_batches,
            num_batches,
            num_rows,
            scan_cost,
            metadata_cache_metrics,
            fetch_metrics,
            metadata_mem_size,
            num_range_builders,
        } = other;

        self.build_parts_cost += *build_cost;
        self.sst_scan_cost += *scan_cost;

        self.files_time_range_pruned += *files_time_range_pruned;

        self.rg_total += *rg_total;
        self.rg_fulltext_filtered += *rg_fulltext_filtered;
        self.rg_inverted_filtered += *rg_inverted_filtered;
        self.rg_minmax_filtered += *rg_minmax_filtered;
        self.rg_bloom_filtered += *rg_bloom_filtered;

        self.rows_before_filter += *rows_total;
        self.rows_fulltext_filtered += *rows_fulltext_filtered;
        self.rows_inverted_filtered += *rows_inverted_filtered;
        self.rows_bloom_filtered += *rows_bloom_filtered;
        self.rows_precise_filtered += *rows_precise_filtered;

        self.fulltext_index_cache_hit += *fulltext_index_cache_hit;
        self.fulltext_index_cache_miss += *fulltext_index_cache_miss;
        self.inverted_index_cache_hit += *inverted_index_cache_hit;
        self.inverted_index_cache_miss += *inverted_index_cache_miss;
        self.bloom_filter_cache_hit += *bloom_filter_cache_hit;
        self.bloom_filter_cache_miss += *bloom_filter_cache_miss;
        self.minmax_cache_hit += *minmax_cache_hit;
        self.minmax_cache_miss += *minmax_cache_miss;
        self.pruner_cache_hit += *pruner_cache_hit;
        self.pruner_cache_miss += *pruner_cache_miss;
        self.pruner_prune_cost += *pruner_prune_cost;

        self.num_sst_record_batches += *num_record_batches;
        self.num_sst_batches += *num_batches;
        self.num_sst_rows += *num_rows;

        // Merge optional verbose metrics
        if let Some(metrics) = inverted_index_apply_metrics {
            self.inverted_index_apply_metrics
                .get_or_insert_with(InvertedIndexApplyMetrics::default)
                .merge_from(metrics);
        }
        if let Some(metrics) = bloom_filter_apply_metrics {
            self.bloom_filter_apply_metrics
                .get_or_insert_with(BloomFilterIndexApplyMetrics::default)
                .merge_from(metrics);
        }
        if let Some(metrics) = fulltext_index_apply_metrics {
            self.fulltext_index_apply_metrics
                .get_or_insert_with(FulltextIndexApplyMetrics::default)
                .merge_from(metrics);
        }
        if let Some(metrics) = fetch_metrics {
            self.fetch_metrics
                .get_or_insert_with(ParquetFetchMetrics::default)
                .merge_from(metrics);
        }
        self.metadata_cache_metrics
            .get_or_insert_with(MetadataCacheMetrics::default)
            .merge_from(metadata_cache_metrics);

        // Track memory usage and update peak.
        self.build_ranges_mem_size += *metadata_mem_size;
        if self.build_ranges_mem_size > self.build_ranges_peak_mem_size {
            self.build_ranges_peak_mem_size = self.build_ranges_mem_size;
        }

        // Track number of builders and update peak.
        self.num_range_builders += *num_range_builders;
        if self.num_range_builders > self.num_peak_range_builders {
            self.num_peak_range_builders = self.num_range_builders;
        }
    }

    /// Merges per-file metrics.
    fn merge_per_file_metrics(&mut self, other: &HashMap<RegionFileId, FileScanMetrics>) {
        let self_file_metrics = self.per_file_metrics.get_or_insert_with(HashMap::new);
        for (file_id, metrics) in other {
            self_file_metrics
                .entry(*file_id)
                .or_default()
                .merge_from(metrics);
        }
    }

    /// Sets distributor metrics.
    fn set_distributor_metrics(&mut self, distributor_metrics: &SeriesDistributorMetrics) {
        let SeriesDistributorMetrics {
            num_series_send_full,
            num_rows,
            num_batches,
            scan_cost,
            yield_cost,
            divider_cost,
        } = distributor_metrics;

        self.num_series_send_full += *num_series_send_full;
        self.num_distributor_rows += *num_rows;
        self.num_distributor_batches += *num_batches;
        self.distributor_scan_cost += *scan_cost;
        self.distributor_yield_cost += *yield_cost;
        self.distributor_divider_cost += *divider_cost;
    }

    /// Observes metrics.
    fn observe_metrics(&self) {
        READ_STAGE_ELAPSED
            .with_label_values(&["prepare_scan"])
            .observe(self.prepare_scan_cost.as_secs_f64());
        READ_STAGE_ELAPSED
            .with_label_values(&["build_reader"])
            .observe(self.build_reader_cost.as_secs_f64());
        READ_STAGE_ELAPSED
            .with_label_values(&["scan"])
            .observe(self.scan_cost.as_secs_f64());
        READ_STAGE_ELAPSED
            .with_label_values(&["yield"])
            .observe(self.yield_cost.as_secs_f64());
        if let Some(time) = &self.convert_cost {
            READ_STAGE_ELAPSED
                .with_label_values(&["convert"])
                .observe(Duration::from_nanos(time.value() as u64).as_secs_f64());
        }
        READ_STAGE_ELAPSED
            .with_label_values(&["total"])
            .observe(self.total_cost.as_secs_f64());
        READ_ROWS_RETURN.observe(self.num_rows as f64);
        READ_BATCHES_RETURN.observe(self.num_batches as f64);

        READ_STAGE_ELAPSED
            .with_label_values(&["build_parts"])
            .observe(self.build_parts_cost.as_secs_f64());

        READ_ROW_GROUPS_TOTAL
            .with_label_values(&["before_filtering"])
            .inc_by(self.rg_total as u64);
        READ_ROW_GROUPS_TOTAL
            .with_label_values(&["fulltext_index_filtered"])
            .inc_by(self.rg_fulltext_filtered as u64);
        READ_ROW_GROUPS_TOTAL
            .with_label_values(&["inverted_index_filtered"])
            .inc_by(self.rg_inverted_filtered as u64);
        READ_ROW_GROUPS_TOTAL
            .with_label_values(&["minmax_index_filtered"])
            .inc_by(self.rg_minmax_filtered as u64);
        READ_ROW_GROUPS_TOTAL
            .with_label_values(&["bloom_filter_index_filtered"])
            .inc_by(self.rg_bloom_filtered as u64);

        PRECISE_FILTER_ROWS_TOTAL
            .with_label_values(&["parquet"])
            .inc_by(self.rows_precise_filtered as u64);
        READ_ROWS_IN_ROW_GROUP_TOTAL
            .with_label_values(&["before_filtering"])
            .inc_by(self.rows_before_filter as u64);
        READ_ROWS_IN_ROW_GROUP_TOTAL
            .with_label_values(&["fulltext_index_filtered"])
            .inc_by(self.rows_fulltext_filtered as u64);
        READ_ROWS_IN_ROW_GROUP_TOTAL
            .with_label_values(&["inverted_index_filtered"])
            .inc_by(self.rows_inverted_filtered as u64);
        READ_ROWS_IN_ROW_GROUP_TOTAL
            .with_label_values(&["bloom_filter_index_filtered"])
            .inc_by(self.rows_bloom_filtered as u64);
    }
}

struct PartitionMetricsInner {
    region_id: RegionId,
    /// Index of the partition to scan.
    partition: usize,
    /// Label to distinguish different scan operation.
    scanner_type: &'static str,
    /// Query start time.
    query_start: Instant,
    /// Whether to use verbose logging.
    explain_verbose: bool,
    /// Verbose scan metrics that only log to debug logs by default.
    metrics: Mutex<ScanMetricsSet>,
    in_progress_scan: IntGauge,

    // Normal metrics that always report to the [ExecutionPlanMetricsSet]:
    /// Duration to build file ranges.
    build_parts_cost: Time,
    /// Duration to build the (merge) reader.
    build_reader_cost: Time,
    /// Duration to scan data.
    scan_cost: Time,
    /// Duration while waiting for `yield`.
    yield_cost: Time,
    /// Duration to convert [`Batch`]es.
    convert_cost: Time,
    /// Aggregated compute time reported to DataFusion.
    elapsed_compute: Time,
}

impl PartitionMetricsInner {
    fn on_finish(&self, stream_eof: bool) {
        let mut metrics = self.metrics.lock().unwrap();
        if metrics.total_cost.is_zero() {
            metrics.total_cost = self.query_start.elapsed();
        }
        if !metrics.stream_eof {
            metrics.stream_eof = stream_eof;
        }
    }
}

impl MergeMetricsReport for PartitionMetricsInner {
    fn report(&self, metrics: &mut MergeMetrics) {
        let mut scan_metrics = self.metrics.lock().unwrap();
        // Merge the metrics into scan_metrics
        scan_metrics.merge_metrics.merge(metrics);

        // Reset the input metrics
        *metrics = MergeMetrics::default();
    }
}

impl DedupMetricsReport for PartitionMetricsInner {
    fn report(&self, metrics: &mut DedupMetrics) {
        let mut scan_metrics = self.metrics.lock().unwrap();
        // Merge the metrics into scan_metrics
        scan_metrics.dedup_metrics.merge(metrics);

        // Reset the input metrics
        *metrics = DedupMetrics::default();
    }
}

impl Drop for PartitionMetricsInner {
    fn drop(&mut self) {
        self.on_finish(false);
        let metrics = self.metrics.lock().unwrap();
        metrics.observe_metrics();
        self.in_progress_scan.dec();

        if self.explain_verbose {
            common_telemetry::info!(
                "{} finished, region_id: {}, partition: {}, scan_metrics: {:?}",
                self.scanner_type,
                self.region_id,
                self.partition,
                metrics,
            );
        } else {
            common_telemetry::debug!(
                "{} finished, region_id: {}, partition: {}, scan_metrics: {:?}",
                self.scanner_type,
                self.region_id,
                self.partition,
                metrics,
            );
        }
    }
}

/// List of PartitionMetrics.
#[derive(Default)]
pub(crate) struct PartitionMetricsList(Mutex<Vec<Option<PartitionMetrics>>>);

impl PartitionMetricsList {
    /// Sets a new [PartitionMetrics] at the specified partition.
    pub(crate) fn set(&self, partition: usize, metrics: PartitionMetrics) {
        let mut list = self.0.lock().unwrap();
        if list.len() <= partition {
            list.resize(partition + 1, None);
        }
        list[partition] = Some(metrics);
    }

    /// Format verbose metrics for each partition for explain.
    pub(crate) fn format_verbose_metrics(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let list = self.0.lock().unwrap();
        write!(f, ", \"metrics_per_partition\": ")?;
        f.debug_list()
            .entries(list.iter().filter_map(|p| p.as_ref()))
            .finish()?;
        write!(f, "}}")
    }
}

/// Metrics while reading a partition.
#[derive(Clone)]
pub struct PartitionMetrics(Arc<PartitionMetricsInner>);

impl PartitionMetrics {
    pub(crate) fn new(
        region_id: RegionId,
        partition: usize,
        scanner_type: &'static str,
        query_start: Instant,
        explain_verbose: bool,
        metrics_set: &ExecutionPlanMetricsSet,
    ) -> Self {
        let partition_str = partition.to_string();
        let in_progress_scan = IN_PROGRESS_SCAN.with_label_values(&[scanner_type, &partition_str]);
        in_progress_scan.inc();
        let convert_cost = MetricBuilder::new(metrics_set).subset_time("convert_cost", partition);
        let metrics = ScanMetricsSet::default()
            .with_prepare_scan_cost(query_start.elapsed())
            .with_convert_cost(convert_cost.clone());
        let inner = PartitionMetricsInner {
            region_id,
            partition,
            scanner_type,
            query_start,
            explain_verbose,
            metrics: Mutex::new(metrics),
            in_progress_scan,
            build_parts_cost: MetricBuilder::new(metrics_set)
                .subset_time("build_parts_cost", partition),
            build_reader_cost: MetricBuilder::new(metrics_set)
                .subset_time("build_reader_cost", partition),
            scan_cost: MetricBuilder::new(metrics_set).subset_time("scan_cost", partition),
            yield_cost: MetricBuilder::new(metrics_set).subset_time("yield_cost", partition),
            convert_cost,
            elapsed_compute: MetricBuilder::new(metrics_set).elapsed_compute(partition),
        };
        Self(Arc::new(inner))
    }

    pub(crate) fn on_first_poll(&self) {
        let mut metrics = self.0.metrics.lock().unwrap();
        metrics.first_poll = self.0.query_start.elapsed();
    }

    pub(crate) fn inc_num_mem_ranges(&self, num: usize) {
        let mut metrics = self.0.metrics.lock().unwrap();
        metrics.num_mem_ranges += num;
    }

    pub fn inc_num_file_ranges(&self, num: usize) {
        let mut metrics = self.0.metrics.lock().unwrap();
        metrics.num_file_ranges += num;
    }

    fn record_elapsed_compute(&self, duration: Duration) {
        if duration.is_zero() {
            return;
        }
        self.0.elapsed_compute.add_duration(duration);
    }

    /// Merges `build_reader_cost`.
    pub(crate) fn inc_build_reader_cost(&self, cost: Duration) {
        self.0.build_reader_cost.add_duration(cost);

        let mut metrics = self.0.metrics.lock().unwrap();
        metrics.build_reader_cost += cost;
    }

    pub(crate) fn inc_convert_batch_cost(&self, cost: Duration) {
        self.0.convert_cost.add_duration(cost);
        self.record_elapsed_compute(cost);
    }

    /// Reports memtable scan metrics.
    pub(crate) fn report_mem_scan_metrics(&self, data: &crate::memtable::MemScanMetricsData) {
        let mut metrics = self.0.metrics.lock().unwrap();
        metrics.mem_scan_cost += data.scan_cost;
        metrics.mem_rows += data.num_rows;
        metrics.mem_batches += data.num_batches;
        metrics.mem_series += data.total_series;
        metrics.mem_prefilter_cost += data.prefilter_cost;
        metrics.mem_prefilter_rows_filtered += data.prefilter_rows_filtered;
    }

    /// Merges [ScannerMetrics], `build_reader_cost`, `scan_cost` and `yield_cost`.
    pub(crate) fn merge_metrics(&self, metrics: &ScannerMetrics) {
        self.0.scan_cost.add_duration(metrics.scan_cost);
        self.record_elapsed_compute(metrics.scan_cost);
        self.0.yield_cost.add_duration(metrics.yield_cost);
        self.record_elapsed_compute(metrics.yield_cost);

        let mut metrics_set = self.0.metrics.lock().unwrap();
        metrics_set.merge_scanner_metrics(metrics);
    }

    /// Merges [ReaderMetrics] and `build_reader_cost`.
    pub fn merge_reader_metrics(
        &self,
        metrics: &ReaderMetrics,
        per_file_metrics: Option<&HashMap<RegionFileId, FileScanMetrics>>,
    ) {
        self.0.build_parts_cost.add_duration(metrics.build_cost);

        let mut metrics_set = self.0.metrics.lock().unwrap();
        metrics_set.merge_reader_metrics(metrics);

        // Merge per-file metrics if provided
        if let Some(file_metrics) = per_file_metrics {
            metrics_set.merge_per_file_metrics(file_metrics);
        }
    }

    /// Finishes the query.
    pub(crate) fn on_finish(&self) {
        self.0.on_finish(true);
    }

    /// Sets the distributor metrics.
    pub(crate) fn set_distributor_metrics(&self, metrics: &SeriesDistributorMetrics) {
        let mut metrics_set = self.0.metrics.lock().unwrap();
        metrics_set.set_distributor_metrics(metrics);
    }

    /// Returns whether verbose explain is enabled.
    pub(crate) fn explain_verbose(&self) -> bool {
        self.0.explain_verbose
    }

    /// Returns a MergeMetricsReport trait object for reporting merge metrics.
    pub(crate) fn merge_metrics_reporter(&self) -> Arc<dyn MergeMetricsReport> {
        self.0.clone()
    }

    /// Returns a DedupMetricsReport trait object for reporting dedup metrics.
    pub(crate) fn dedup_metrics_reporter(&self) -> Arc<dyn DedupMetricsReport> {
        self.0.clone()
    }

    /// Increments the total bytes added to the range cache.
    pub(crate) fn inc_range_cache_size(&self, size: usize) {
        let mut metrics = self.0.metrics.lock().unwrap();
        metrics.range_cache_size += size;
    }

    /// Increments the range cache hit counter.
    pub(crate) fn inc_range_cache_hit(&self) {
        let mut metrics = self.0.metrics.lock().unwrap();
        metrics.range_cache_hit += 1;
    }

    /// Increments the range cache miss counter.
    pub(crate) fn inc_range_cache_miss(&self) {
        let mut metrics = self.0.metrics.lock().unwrap();
        metrics.range_cache_miss += 1;
    }
}

impl fmt::Debug for PartitionMetrics {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let metrics = self.0.metrics.lock().unwrap();
        write!(
            f,
            r#"{{"partition":{}, "metrics":{:?}}}"#,
            self.0.partition, metrics
        )
    }
}

/// Metrics for the series distributor.
#[derive(Default)]
pub(crate) struct SeriesDistributorMetrics {
    /// Number of send full in SeriesScan.
    pub(crate) num_series_send_full: usize,
    /// Number of rows the series distributor scanned.
    pub(crate) num_rows: usize,
    /// Number of batches the series distributor scanned.
    pub(crate) num_batches: usize,
    /// Duration of the series distributor to scan.
    pub(crate) scan_cost: Duration,
    /// Duration of the series distributor to yield.
    pub(crate) yield_cost: Duration,
    /// Duration spent in divider operations.
    pub(crate) divider_cost: Duration,
}

/// Aligns memtable and extension batches; SST readers align before schema compatibility.
pub(crate) fn align_query_source_batch(
    context: &StreamContext,
    batch: RecordBatch,
) -> Result<RecordBatch> {
    let query_schema = context.input.mapper.input_arrow_schema(false);
    json_align::align_query_batch(batch, &query_schema)
}

/// Scans memtable ranges at `index` using flat format that returns RecordBatch.
#[tracing::instrument(
    skip_all,
    fields(
        region_id = %stream_ctx.input.region_metadata().region_id,
        row_group_index = %index.index,
        source = "mem_flat"
    )
)]
pub(crate) fn scan_flat_mem_ranges(
    stream_ctx: Arc<StreamContext>,
    part_metrics: PartitionMetrics,
    index: RowGroupIndex,
    time_range: FileTimeRange,
) -> impl Stream<Item = Result<RecordBatch>> {
    try_stream! {
        let ranges = stream_ctx.input.build_mem_ranges(index);
        part_metrics.inc_num_mem_ranges(ranges.len());
        for range in ranges {
            let build_reader_start = Instant::now();
            let mem_scan_metrics = Some(MemScanMetrics::default());
            let mut iter = range.build_record_batch_iter(Some(time_range), mem_scan_metrics.clone())?;
            part_metrics.inc_build_reader_cost(build_reader_start.elapsed());

            while let Some(record_batch) = iter.next().transpose()? {
                yield record_batch;
            }

            // Report the memtable scan metrics to partition metrics
            if let Some(ref metrics) = mem_scan_metrics {
                let data = metrics.data();
                part_metrics.report_mem_scan_metrics(&data);
            }
        }
    }
}

/// Aligns memtable JSON2 columns to query types before merge/dedup.
///
/// For append-only queries with `TimeSeriesRowSelector::LastRow { after_merge: false }`,
/// first keeps all rows tied at each series' maximum timestamp within each sorted
/// PUT-only batch, avoiding JSON2 conversion for discarded rows. Final cross-source
/// LastRow selection still runs. Batches containing deletes are not reduced.
/// This optimization uses the query's series row selector, not its merge mode;
/// non-append tables retain all input rows for deduplication and last-non-null backfill.
pub(crate) fn prepare_flat_mem_stream(
    context: Arc<StreamContext>,
    stream: impl Stream<Item = Result<RecordBatch>> + Send + 'static,
) -> BoxedRecordBatchStream {
    if context.input.read_cols.json_target_types().is_empty() {
        return Box::pin(stream);
    }

    let select_last = context.input.append_mode
        && matches!(
            context.input.series_row_selector,
            Some(TimeSeriesRowSelector::LastRow { after_merge: false })
        );
    Box::pin(stream.map(move |batch| {
        let batch = batch?;
        // Select raw candidates before decoding JSON2; final cross-source selection remains.
        let batch = if select_last {
            FlatLastTimestampSelector::select_batch(batch)?
        } else {
            batch
        };
        align_query_source_batch(&context, batch)
    }))
}

/// Files with row count greater than this threshold can contribute to the estimation.
const SPLIT_ROW_THRESHOLD: u64 = DEFAULT_ROW_GROUP_SIZE as u64;
/// Number of series threshold for splitting batches.
const NUM_SERIES_THRESHOLD: u64 = 10240;
/// Minimum batch size after splitting. The batch size is less than 60 because a series may only have
/// 60 samples per hour.
pub(crate) const BATCH_SIZE_THRESHOLD: u64 = 50;
/// Minimum estimated rows per series for compaction to split batches. Shorter series
/// would make compaction write one small batch per series.
pub(crate) const MIN_COMPACTION_SPLIT_ROWS: usize = DEFAULT_READ_BATCH_SIZE / 2;

/// Returns the estimated rows per batch after splitting if splitting flat record batches
/// may improve merge performance. Returns `None` if splitting is not beneficial.
pub(crate) fn should_split_flat_batches_for_merge(
    stream_ctx: &Arc<StreamContext>,
    range_meta: &RangeMeta,
) -> Option<usize> {
    // Number of files to split and scan.
    let mut num_files_to_split = 0;
    let mut num_mem_rows = 0;
    let mut num_mem_series = 0;
    // Total rows and series for estimating batch size after splitting.
    let mut total_rows: u64 = 0;
    let mut total_series: u64 = 0;
    // Checks each file range, returns early if any range is not splittable.
    // For mem ranges, we collect the total number of rows and series because the number of rows in a
    // mem range may be too small.
    for index in &range_meta.row_group_indices {
        if stream_ctx.is_mem_range_index(*index) {
            let memtable = &stream_ctx.input.memtables[index.index];
            // Is mem range
            let stats = memtable.stats();
            num_mem_rows += stats.num_rows();
            num_mem_series += stats.series_count();
        } else if stream_ctx.is_file_range_index(*index) {
            // This is a file range.
            let file_index = index.index - stream_ctx.input.num_memtables();
            let file = &stream_ctx.input.files[file_index];
            let file_meta = file.meta_ref();
            if file_meta.level == 0 {
                // Short series can leave too many tiny batches retained by the merge.
                if file_meta.num_series > 0 {
                    if !can_split_series(file_meta.num_rows, file_meta.num_series) {
                        return None;
                    }
                    total_rows += file_meta.num_rows;
                    total_series += file_meta.num_series;
                }
                // Older L0 files may lack series statistics; keep their fallback behavior.
                num_files_to_split += 1;
                continue;
            } else if file_meta.num_rows < SPLIT_ROW_THRESHOLD || file_meta.num_series == 0 {
                // If the file doesn't have enough rows, or the number of series is unavailable, skips it.
                continue;
            }
            debug_assert!(file_meta.num_rows > 0);
            if !can_split_series(file_meta.num_rows, file_meta.num_series) {
                // We can't split batches in a file.
                common_telemetry::trace!(
                    "Can't split series for file {}, level: {}, num_rows: {}, num_series: {}",
                    file_meta.file_id,
                    file_meta.level,
                    file_meta.num_rows,
                    file_meta.num_series,
                );
                return None;
            } else {
                num_files_to_split += 1;
                total_rows += file.meta_ref().num_rows;
                total_series += file.meta_ref().num_series;
            }
        }
        // Skips non-file and non-mem ranges.
    }

    let should_split = if num_files_to_split > 0 {
        // We mainly consider file ranges because they have enough data for sampling.
        true
    } else if num_mem_series > 0
        && num_mem_rows > 0
        && can_split_series(num_mem_rows as u64, num_mem_series as u64)
    {
        total_rows += num_mem_rows as u64;
        total_series += num_mem_series as u64;
        true
    } else {
        false
    };

    if !should_split {
        return None;
    }

    // Estimate rows per batch after splitting.
    let estimated_batch_size = if total_series > 0 && total_rows > 0 {
        ((total_rows / total_series) as usize).clamp(1, DEFAULT_READ_BATCH_SIZE)
    } else {
        // No valid estimate available, use a conservative fallback.
        DEFAULT_READ_BATCH_SIZE / 4
    };
    Some(estimated_batch_size)
}

/// Computes the channel size for parallel scan based on the estimated rows per batch.
/// The channel should buffer approximately `2 * DEFAULT_READ_BATCH_SIZE` rows.
pub(crate) fn compute_parallel_channel_size(estimated_rows_per_batch: usize) -> usize {
    let size = 2 * DEFAULT_READ_BATCH_SIZE / estimated_rows_per_batch.max(1);
    size.clamp(2, 64)
}

/// Computes the average estimated rows per batch across multiple range readers.
pub(crate) fn compute_average_batch_size(
    estimated_rows_per_batch: impl IntoIterator<Item = usize>,
) -> usize {
    let mut total = 0usize;
    let mut count = 0usize;
    for size in estimated_rows_per_batch {
        total += size;
        count += 1;
    }

    if count == 0 {
        return DEFAULT_READ_BATCH_SIZE;
    }

    (total / count).clamp(1, DEFAULT_READ_BATCH_SIZE)
}

fn can_split_series(num_rows: u64, num_series: u64) -> bool {
    if num_rows == 0 || num_series == 0 {
        return false;
    }

    // It doesn't have too many series or it will have enough rows for each batch.
    num_series < NUM_SERIES_THRESHOLD || num_rows / num_series >= BATCH_SIZE_THRESHOLD
}

#[cfg(test)]
mod split_tests {
    use std::sync::Arc;

    use common_time::Timestamp;
    use smallvec::smallvec;
    use store_api::storage::FileId;

    use super::*;
    use crate::read::flat_projection::FlatProjectionMapper;
    use crate::read::range::{RangeMeta, RowGroupIndex, SourceIndex};
    use crate::read::scan_region::{ScanInput, StreamContext};
    use crate::sst::file::FileHandle;
    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;

    async fn new_stream_context_with_files(files: Vec<FileHandle>) -> StreamContext {
        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 input = ScanInput::builder(env.access_layer.clone(), mapper)
            .with_files(files)
            .build();

        StreamContext {
            input,
            ranges: vec![],
            query_start: std::time::Instant::now(),
        }
    }

    fn single_file_range_meta() -> RangeMeta {
        RangeMeta {
            time_range: (
                Timestamp::new_millisecond(0),
                Timestamp::new_millisecond(1000),
            ),
            indices: smallvec![SourceIndex {
                index: 0,
                num_row_groups: 1,
            }],
            row_group_indices: smallvec![RowGroupIndex {
                index: 0,
                row_group_index: 0,
            }],
            num_rows: 1024,
        }
    }

    #[tokio::test]
    async fn should_split_level_zero_file_even_when_series_stats_are_missing() {
        let mut file = sst_file_handle_with_file_id(FileId::random(), 0, 1000)
            .meta_ref()
            .clone();
        file.level = 0;
        file.num_rows = DEFAULT_ROW_GROUP_SIZE as u64;
        file.num_row_groups = 1;
        file.num_series = 0;

        let file = FileHandle::new(file, crate::test_util::new_noop_file_purger());
        let stream_ctx = Arc::new(new_stream_context_with_files(vec![file]).await);

        assert!(
            should_split_flat_batches_for_merge(&stream_ctx, &single_file_range_meta()).is_some()
        );
    }

    #[test]
    fn can_split_series_returns_false_for_zero_inputs() {
        assert!(!can_split_series(0, 1));
        assert!(!can_split_series(1, 0));
        assert!(!can_split_series(0, 0));
    }
}

/// Creates a new [ReaderFilterMetrics] with optional apply metrics initialized
/// based on the `explain_verbose` flag.
pub(crate) fn new_filter_metrics(explain_verbose: bool) -> ReaderFilterMetrics {
    if explain_verbose {
        ReaderFilterMetrics {
            inverted_index_apply_metrics: Some(InvertedIndexApplyMetrics::default()),
            bloom_filter_apply_metrics: Some(BloomFilterIndexApplyMetrics::default()),
            fulltext_index_apply_metrics: Some(FulltextIndexApplyMetrics::default()),
            ..Default::default()
        }
    } else {
        ReaderFilterMetrics::default()
    }
}

/// Scans file ranges at `index` using flat reader that returns RecordBatch.
#[tracing::instrument(
    skip_all,
    fields(
        region_id = %stream_ctx.input.region_metadata().region_id,
        row_group_index = %index.index,
        source = read_type
    )
)]
pub(crate) async fn scan_flat_file_ranges(
    stream_ctx: Arc<StreamContext>,
    part_metrics: PartitionMetrics,
    index: RowGroupIndex,
    read_type: &'static str,
    partition_pruner: Arc<PartitionPruner>,
) -> Result<impl Stream<Item = Result<RecordBatch>>> {
    let mut reader_metrics = ReaderMetrics {
        filter_metrics: new_filter_metrics(part_metrics.explain_verbose()),
        ..Default::default()
    };
    let ranges = partition_pruner
        .build_file_ranges(index, &part_metrics, &mut reader_metrics)
        .await?;
    part_metrics.inc_num_file_ranges(ranges.len());
    part_metrics.merge_reader_metrics(&reader_metrics, None);

    // Creates initial per-file metrics with build_part_cost.
    let init_per_file_metrics = if part_metrics.explain_verbose() {
        let file = stream_ctx.input.file_from_index(index);
        let file_id = file.file_id();

        let mut map = HashMap::new();
        map.insert(
            file_id,
            FileScanMetrics {
                build_part_cost: reader_metrics.build_cost,
                ..Default::default()
            },
        );
        Some(map)
    } else {
        None
    };

    Ok(build_flat_file_range_scan_stream(
        stream_ctx,
        part_metrics,
        read_type,
        ranges,
        init_per_file_metrics,
    ))
}

/// Filters a flat-format record batch by the exact sequence range.
///
/// Returns `None` when the entire batch is filtered out. The sequence column is
/// the second-to-last internal column of the flat format
/// (`__primary_key`, `__sequence`, `__op_type`), so it must be applied before any
/// projection/compat conversion that drops internal columns.
///
/// `file_sequence_trusted` is the per-file trust decision. Foreign reader
/// batches have already been virtualized to the target-local file barrier, so
/// they are filtered using the effective batch sequence. Local untrusted files
/// pass through because exact capability excludes them before row filtering.
pub(crate) fn filter_flat_batch_by_sequence(
    record_batch: RecordBatch,
    sequence_range: Option<SequenceRange>,
    file_sequence_trusted: bool,
) -> Result<Option<RecordBatch>> {
    let Some(sequence) = sequence_range else {
        return Ok(Some(record_batch));
    };
    if !file_sequence_trusted {
        return Ok(Some(record_batch));
    }

    let num_rows = record_batch.num_rows();
    if num_rows == 0 {
        return Ok(Some(record_batch));
    }
    let sequence_column = record_batch.column(sequence_column_index(record_batch.num_columns()));
    let predicate = sequence
        .filter(sequence_column)
        .context(ComputeArrowSnafu)?;
    let select_count = predicate.true_count();
    if select_count == 0 {
        return Ok(None);
    }
    if select_count == num_rows {
        return Ok(Some(record_batch));
    }
    let filtered_batch = datatypes::arrow::compute::filter_record_batch(&record_batch, &predicate)
        .context(ComputeArrowSnafu)?;
    Ok(Some(filtered_batch))
}

/// Build the stream of scanning the input [`FileRange`]s using flat reader that returns RecordBatch.
#[tracing::instrument(
    skip_all,
    fields(read_type = read_type, range_count = ranges.len())
)]
pub fn build_flat_file_range_scan_stream(
    stream_ctx: Arc<StreamContext>,
    part_metrics: PartitionMetrics,
    read_type: &'static str,
    ranges: SmallVec<[FileRange; 2]>,
    mut per_file_metrics: Option<HashMap<RegionFileId, FileScanMetrics>>,
) -> impl Stream<Item = Result<RecordBatch>> {
    try_stream! {
        let fetch_metrics = if part_metrics.explain_verbose() {
            Some(Arc::new(ParquetFetchMetrics::default()))
        } else {
            None
        };
        let reader_metrics = &mut ReaderMetrics {
            fetch_metrics: fetch_metrics.clone(),
            ..Default::default()
        };
        for range in ranges {
            let build_reader_start = Instant::now();
            let Some(mut reader) = range
                .flat_reader(
                    match stream_ctx.input.series_row_selector {
                        Some(TimeSeriesRowSelector::LastRow { after_merge: false })
                            if stream_ctx.input.sequence_range.is_none() =>
                        {
                            stream_ctx.input.series_row_selector
                        }
                        _ => None,
                    },
                    fetch_metrics.as_deref(),
                )
                .await?
            else {
                continue;
            };
            let build_cost = build_reader_start.elapsed();
            part_metrics.inc_build_reader_cost(build_cost);

            let may_compat = range.compat_batch();
            let file_sequence_trusted = range
                .file_handle()
                .is_effective_target_sequence_trusted(stream_ctx.input.region_metadata().region_id);

            let mapper = range.compaction_projection_mapper();
            while let Some(record_batch) = reader.next_batch().await? {
                let record_batch = if let Some(mapper) = mapper {
                    mapper.project(record_batch)?
                } else {
                    record_batch
                };
                let Some(record_batch) = filter_flat_batch_by_sequence(
                    record_batch,
                    stream_ctx.input.sequence_range,
                    file_sequence_trusted,
                )? else {
                    continue;
                };

                if let Some(flat_compat) = may_compat {
                    let batch = flat_compat.compat(record_batch)?;
                    yield batch;
                } else {
                    yield record_batch;
                }
            }

            let prune_metrics = reader.metrics();

            // Update per-file metrics if tracking is enabled
            if let Some(file_metrics_map) = per_file_metrics.as_mut() {
                let file_id = range.file_handle().file_id();
                let file_metrics = file_metrics_map
                    .entry(file_id)
                    .or_insert_with(FileScanMetrics::default);

                file_metrics.num_ranges += 1;
                file_metrics.num_rows += prune_metrics.num_rows;
                file_metrics.build_reader_cost += build_cost;
                file_metrics.scan_cost += prune_metrics.scan_cost;
            }

            reader_metrics.merge_from(&prune_metrics);
        }

        // Reports metrics.
        reader_metrics.observe_rows(read_type);
        reader_metrics.filter_metrics.observe();
        part_metrics.merge_reader_metrics(reader_metrics, per_file_metrics.as_ref());
    }
}

/// Build the stream of scanning the extension range in flat format denoted by the [`RowGroupIndex`].
#[cfg(feature = "enterprise")]
pub(crate) async fn scan_flat_extension_range(
    context: Arc<StreamContext>,
    index: RowGroupIndex,
    partition_metrics: PartitionMetrics,
    options: crate::extension::ExtensionRangeReadOptions,
) -> Result<BoxedRecordBatchStream> {
    use snafu::ResultExt;

    let range = context.input.extension_range(index.index);
    let reader = range.flat_reader(context.as_ref(), options);
    let stream = reader
        .read(context.clone(), partition_metrics, index)
        .await
        .context(crate::error::ScanExternalRangeSnafu)?;
    if context.input.read_cols.json_target_types().is_empty() {
        return Ok(stream);
    }
    Ok(Box::pin(stream.and_then(move |batch| {
        std::future::ready(align_query_source_batch(&context, batch))
    })))
}

pub(crate) async fn maybe_scan_flat_other_ranges(
    context: &Arc<StreamContext>,
    index: RowGroupIndex,
    metrics: &PartitionMetrics,
    pre_filter_mode: PreFilterMode,
) -> Result<BoxedRecordBatchStream> {
    #[cfg(feature = "enterprise")]
    {
        let options = crate::extension::ExtensionRangeReadOptions { pre_filter_mode };
        scan_flat_extension_range(context.clone(), index, metrics.clone(), options).await
    }

    #[cfg(not(feature = "enterprise"))]
    {
        let _ = context;
        let _ = index;
        let _ = metrics;
        let _ = pre_filter_mode;

        crate::error::UnexpectedSnafu {
            reason: "no other ranges scannable in flat format",
        }
        .fail()
    }
}

/// A stream wrapper that splits record batches from an inner stream.
pub(crate) struct SplitRecordBatchStream<S> {
    /// The inner stream that yields record batches.
    inner: S,
    /// Buffer for split batches.
    batches: VecDeque<RecordBatch>,
}

impl<S> SplitRecordBatchStream<S> {
    /// Creates a new splitting stream wrapper.
    pub(crate) fn new(inner: S) -> Self {
        Self {
            inner,
            batches: VecDeque::new(),
        }
    }
}

impl<S> Stream for SplitRecordBatchStream<S>
where
    S: Stream<Item = Result<RecordBatch>> + Unpin,
{
    type Item = Result<RecordBatch>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            // First, check if we have buffered split batches
            if let Some(batch) = self.batches.pop_front() {
                return Poll::Ready(Some(Ok(batch)));
            }

            // Poll the inner stream for the next batch
            let record_batch = match futures::ready!(Pin::new(&mut self.inner).poll_next(cx)) {
                Some(Ok(batch)) => batch,
                Some(Err(e)) => return Poll::Ready(Some(Err(e))),
                None => return Poll::Ready(None),
            };

            // Split the batch and buffer the results
            split_record_batch(record_batch, &mut self.batches);
            // Continue the loop to return the first split batch
        }
    }
}

/// Splits the batch so each sub-batch has strictly increasing timestamps.
///
/// # Panics
/// Panics if the timestamp array is invalid.
pub(crate) fn split_record_batch(record_batch: RecordBatch, batches: &mut VecDeque<RecordBatch>) {
    let batch_rows = record_batch.num_rows();
    if batch_rows == 0 {
        return;
    }
    if batch_rows < 2 {
        batches.push_back(record_batch);
        return;
    }

    let time_index_pos = time_index_column_index(record_batch.num_columns());
    let timestamps = record_batch.column(time_index_pos);
    let (ts_values, _unit) = timestamp_array_to_primitive(timestamps).unwrap();
    let mut offsets = Vec::with_capacity(16);
    offsets.push(0);
    let values = ts_values.values();
    for (i, &value) in values.iter().take(batch_rows - 1).enumerate() {
        if value >= values[i + 1] {
            offsets.push(i + 1);
        }
    }
    offsets.push(values.len());

    // Splits the batch by offsets.
    for (i, &start) in offsets[..offsets.len() - 1].iter().enumerate() {
        let end = offsets[i + 1];
        let rows_in_batch = end - start;
        batches.push_back(record_batch.slice(start, rows_in_batch));
    }
}

#[cfg(test)]
mod tests {
    use arrow_select::concat::concat_batches;
    use common_time::Timestamp;
    use datatypes::arrow::array::{BinaryArray, TimestampMillisecondArray, UInt64Array};
    use datatypes::types::json_type::JsonNativeType;
    use futures::TryStreamExt;
    use smallvec::{SmallVec, smallvec};
    use std::sync::Arc;
    use std::time::Instant;
    use store_api::storage::RegionId;

    use super::*;
    use crate::cache::CacheStrategy;
    use crate::memtable::bulk::context::BulkIterContext;
    use crate::memtable::bulk::part_reader::BulkPartBatchIter;
    use crate::memtable::{
        BoxedBatchIterator, BoxedRecordBatchIterator, IterBuilder, MemtableRange,
        MemtableRangeContext, MemtableStats,
    };
    use crate::read::flat_projection::FlatProjectionMapper;
    use crate::read::range::{MemRangeBuilder, SourceIndex};
    use crate::read::read_columns::ReadColumns;
    use crate::read::scan_region::ScanInput;
    use crate::sst::file::{FileHandle, FileMeta};
    use crate::sst::file_purger::NoopFilePurger;
    use crate::test_util::memtable_util::metadata_for_test;
    use crate::test_util::scheduler_util::SchedulerEnv;

    struct EmptyIterBuilder;

    impl IterBuilder for EmptyIterBuilder {
        fn build(&self, _metrics: Option<MemScanMetrics>) -> Result<BoxedBatchIterator> {
            Ok(Box::new(std::iter::empty()))
        }

        fn is_record_batch(&self) -> bool {
            true
        }

        fn build_record_batch(
            &self,
            _time_range: Option<(Timestamp, Timestamp)>,
            _metrics: Option<MemScanMetrics>,
        ) -> Result<BoxedRecordBatchIterator> {
            Ok(Box::new(std::iter::empty()))
        }
    }

    async fn new_test_stream_ctx(
        files: Vec<FileHandle>,
        memtables: Vec<MemRangeBuilder>,
    ) -> Arc<StreamContext> {
        let env = SchedulerEnv::new().await;
        let metadata = metadata_for_test();
        let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap();
        let input = ScanInput::builder(env.access_layer.clone(), mapper)
            .with_cache(CacheStrategy::Disabled)
            .with_memtables(memtables)
            .with_files(files)
            .build();

        Arc::new(StreamContext {
            input,
            ranges: Vec::new(),
            query_start: Instant::now(),
        })
    }

    fn new_test_file(num_rows: u64, num_series: u64) -> FileHandle {
        let meta = FileMeta {
            region_id: RegionId::new(123, 456),
            file_id: Default::default(),
            level: 1,
            time_range: (
                Timestamp::new_millisecond(0),
                Timestamp::new_millisecond(1000),
            ),
            num_rows,
            num_series,
            ..Default::default()
        };
        FileHandle::new(meta, Arc::new(NoopFilePurger))
    }

    fn new_test_memtable(num_rows: usize, series_count: usize) -> MemRangeBuilder {
        let context = Arc::new(MemtableRangeContext::new(
            0,
            Box::new(EmptyIterBuilder),
            Default::default(),
        ));
        let stats = MemtableStats {
            time_range: Some((
                Timestamp::new_millisecond(0),
                Timestamp::new_millisecond(1000),
            )),
            num_rows,
            num_ranges: 1,
            series_count,
            ..Default::default()
        };
        let range = MemtableRange::new(context, stats.clone());
        MemRangeBuilder::new(range, stats)
    }

    fn new_test_range_meta(row_group_indices: SmallVec<[RowGroupIndex; 2]>) -> RangeMeta {
        let indices = row_group_indices
            .iter()
            .map(|row_group_index| SourceIndex {
                index: row_group_index.index,
                num_row_groups: 1,
            })
            .collect();

        RangeMeta {
            time_range: (
                Timestamp::new_millisecond(0),
                Timestamp::new_millisecond(1000),
            ),
            indices,
            row_group_indices,
            num_rows: 0,
        }
    }

    #[tokio::test]
    async fn test_memtable_last_row_selects_before_json2_alignment()
    -> std::result::Result<(), Box<dyn std::error::Error>> {
        struct BatchIter(RecordBatch, Arc<BulkIterContext>, Option<SequenceRange>);
        impl IterBuilder for BatchIter {
            fn build(&self, _metrics: Option<MemScanMetrics>) -> Result<BoxedBatchIterator> {
                Ok(Box::new(std::iter::empty()))
            }
            fn is_record_batch(&self) -> bool {
                true
            }
            fn build_record_batch(
                &self,
                _time_range: Option<FileTimeRange>,
                metrics: Option<MemScanMetrics>,
            ) -> Result<BoxedRecordBatchIterator> {
                Ok(Box::new(BulkPartBatchIter::from_single(
                    self.0.clone(),
                    self.1.clone(),
                    self.2,
                    1,
                    metrics,
                )))
            }
        }
        let (metadata, batch) = crate::test_util::sst_util::new_json2_record_batch()?;
        let batch = concat_batches(&batch.schema(), &[batch.clone(), batch.slice(0, 1)])?;
        let mut columns = batch.columns().to_vec();
        columns[1] = Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3]));
        let seq_index = sequence_column_index(batch.num_columns());
        columns[seq_index] = Arc::new(UInt64Array::from(vec![1, 2, 3]));
        let batch = RecordBatch::try_new(batch.schema(), columns)?;
        let reader_context = Arc::new(BulkIterContext::new(
            metadata.clone(),
            None,
            None,
            false,
            DEFAULT_READ_BATCH_SIZE,
        )?);
        let time_range = (Timestamp::new_millisecond(1), Timestamp::new_millisecond(3));
        let env = SchedulerEnv::new().await;
        for (append, selector, sequence, count) in [
            (
                true,
                Some(TimeSeriesRowSelector::LastRow { after_merge: false }),
                None,
                1,
            ),
            (
                true,
                Some(TimeSeriesRowSelector::LastRow { after_merge: true }),
                None,
                3,
            ),
            (
                false,
                Some(TimeSeriesRowSelector::LastRow { after_merge: false }),
                None,
                3,
            ),
            (true, None, None, 3),
            (
                true,
                Some(TimeSeriesRowSelector::LastRow { after_merge: false }),
                Some(SequenceRange::GtLtEq { min: 0, max: 2 }),
                1,
            ),
        ] {
            let range_context = Arc::new(MemtableRangeContext::new(
                0,
                Box::new(BatchIter(batch.clone(), reader_context.clone(), sequence)),
                Default::default(),
            ));
            let stats = MemtableStats {
                time_range: Some(time_range),
                num_rows: 3,
                num_ranges: 1,
                series_count: 1,
                ..Default::default()
            };
            let range = MemtableRange::new(range_context, stats.clone());
            let read_cols = ReadColumns::new([0, 1])
                .with_json_target_types([(0, JsonNativeType::Variant)].into_iter().collect());
            let mapper =
                FlatProjectionMapper::new_with_read_columns(&metadata, vec![0, 1], read_cols)?;
            let input = ScanInput::builder(env.access_layer.clone(), mapper)
                .with_append_mode(append)
                .with_series_row_selector(selector)
                .with_sequence_range(sequence)
                .with_memtables(vec![MemRangeBuilder::new(range, stats)])
                .build();
            let context = Arc::new(StreamContext {
                input,
                ranges: vec![],
                query_start: Instant::now(),
            });
            let metrics = PartitionMetrics::new(
                metadata.region_id,
                0,
                "test_json2",
                context.query_start,
                false,
                &ExecutionPlanMetricsSet::new(),
            );
            let mut raw_stream = Box::pin(scan_flat_mem_ranges(
                context.clone(),
                metrics,
                RowGroupIndex {
                    index: 0,
                    row_group_index: -1,
                },
                time_range,
            ));
            let raw = raw_stream.try_next().await?.ok_or("expected a raw batch")?;
            let expected_raw = if sequence.is_some() {
                batch.slice(0, 2)
            } else {
                batch.clone()
            };
            assert_eq!(expected_raw, raw);
            assert!(raw_stream.try_next().await?.is_none());
            let mut stream = prepare_flat_mem_stream(context, futures::stream::iter([Ok(raw)]));
            let output = stream.try_next().await?.ok_or("expected a JSON2 batch")?;
            assert_eq!(count, output.num_rows());
            let values = output
                .column(0)
                .as_any()
                .downcast_ref::<BinaryArray>()
                .ok_or("expected aligned JSONB")?;
            assert_eq!(
                datatypes::types::parse_string_to_jsonb(if sequence.is_some() {
                    r#"{"a":2}"#
                } else {
                    r#"{"a":1}"#
                })?,
                values.value(count - 1)
            );
            assert!(stream.try_next().await?.is_none());
        }
        Ok(())
    }

    #[tokio::test]
    async fn test_should_split_flat_batches_for_merge_uses_splittable_file_rows_per_series() {
        let num_rows = SPLIT_ROW_THRESHOLD * 2;
        let num_series = (num_rows / 100).max(1);
        let stream_ctx =
            new_test_stream_ctx(vec![new_test_file(num_rows, num_series)], vec![]).await;
        let range_meta = new_test_range_meta(smallvec![RowGroupIndex {
            index: 0,
            row_group_index: 0,
        }]);

        assert_eq!(
            Some((num_rows / num_series) as usize),
            should_split_flat_batches_for_merge(&stream_ctx, &range_meta)
        );
    }

    #[tokio::test]
    async fn test_should_split_flat_batches_for_merge_skips_small_or_unknown_series_files() {
        let stream_ctx = new_test_stream_ctx(
            vec![
                new_test_file(SPLIT_ROW_THRESHOLD.saturating_sub(1), 1),
                new_test_file(SPLIT_ROW_THRESHOLD * 2, 0),
            ],
            vec![],
        )
        .await;
        let range_meta = new_test_range_meta(smallvec![
            RowGroupIndex {
                index: 0,
                row_group_index: 0,
            },
            RowGroupIndex {
                index: 1,
                row_group_index: 0,
            }
        ]);

        assert_eq!(
            None,
            should_split_flat_batches_for_merge(&stream_ctx, &range_meta)
        );
    }

    #[tokio::test]
    async fn test_should_split_flat_batches_for_merge_returns_none_for_unsplittable_file() {
        let num_series =
            (SPLIT_ROW_THRESHOLD / (BATCH_SIZE_THRESHOLD - 1)).max(NUM_SERIES_THRESHOLD) + 1;
        let stream_ctx =
            new_test_stream_ctx(vec![new_test_file(SPLIT_ROW_THRESHOLD, num_series)], vec![]).await;
        let range_meta = new_test_range_meta(smallvec![RowGroupIndex {
            index: 0,
            row_group_index: 0,
        }]);

        assert_eq!(
            None,
            should_split_flat_batches_for_merge(&stream_ctx, &range_meta)
        );
    }

    #[tokio::test]
    async fn test_should_split_flat_batches_for_merge_uses_level_zero_series_stats() {
        for (num_rows, num_series, expected) in [
            (70_144, 35_083, None),
            (100_000, 1_000, Some(100)),
            (100_000, 0, Some(DEFAULT_READ_BATCH_SIZE / 4)),
        ] {
            let mut meta = new_test_file(num_rows, num_series).meta_ref().clone();
            meta.level = 0;
            let file = FileHandle::new(meta, Arc::new(NoopFilePurger));
            let stream_ctx = new_test_stream_ctx(vec![file], vec![]).await;
            let range_meta = new_test_range_meta(smallvec![RowGroupIndex {
                index: 0,
                row_group_index: 0,
            }]);

            assert_eq!(
                expected,
                should_split_flat_batches_for_merge(&stream_ctx, &range_meta),
                "rows={num_rows}, series={num_series}",
            );
        }
    }

    #[tokio::test]
    async fn test_should_split_flat_batches_for_merge_falls_back_to_memtables() {
        let stream_ctx = new_test_stream_ctx(vec![], vec![new_test_memtable(5_000, 100)]).await;
        let range_meta = new_test_range_meta(smallvec![RowGroupIndex {
            index: 0,
            row_group_index: 0,
        }]);

        assert_eq!(
            Some(50),
            should_split_flat_batches_for_merge(&stream_ctx, &range_meta)
        );
    }

    #[tokio::test]
    async fn test_should_split_flat_batches_for_merge_clamps_estimate() {
        let stream_ctx =
            new_test_stream_ctx(vec![new_test_file(SPLIT_ROW_THRESHOLD * 2, 1)], vec![]).await;
        let range_meta = new_test_range_meta(smallvec![RowGroupIndex {
            index: 0,
            row_group_index: 0,
        }]);

        assert_eq!(
            Some(DEFAULT_READ_BATCH_SIZE),
            should_split_flat_batches_for_merge(&stream_ctx, &range_meta)
        );
    }

    #[test]
    fn test_compute_parallel_channel_size_clamps_to_max_for_small_batches() {
        assert_eq!(64, compute_parallel_channel_size(0));
        assert_eq!(64, compute_parallel_channel_size(1));
    }

    #[test]
    fn test_compute_parallel_channel_size_returns_expected_mid_range_size() {
        assert_eq!(
            4,
            compute_parallel_channel_size(DEFAULT_READ_BATCH_SIZE / 2)
        );
    }

    #[test]
    fn test_compute_parallel_channel_size_clamps_to_min_for_large_batches() {
        assert_eq!(2, compute_parallel_channel_size(DEFAULT_READ_BATCH_SIZE));
        assert_eq!(
            2,
            compute_parallel_channel_size(DEFAULT_READ_BATCH_SIZE * 2)
        );
    }

    #[test]
    fn test_compute_average_batch_size_uses_arithmetic_mean() {
        assert_eq!(24, compute_average_batch_size([16, 24, 32]));
    }

    #[test]
    fn test_compute_average_batch_size_clamps_values() {
        assert_eq!(
            DEFAULT_READ_BATCH_SIZE,
            compute_average_batch_size([DEFAULT_READ_BATCH_SIZE, DEFAULT_READ_BATCH_SIZE * 2])
        );
        assert_eq!(1, compute_average_batch_size([0, 1]));
    }

    #[test]
    fn test_compute_average_batch_size_falls_back_when_empty() {
        assert_eq!(
            DEFAULT_READ_BATCH_SIZE,
            compute_average_batch_size(std::iter::empty())
        );
    }

    /// Builds a flat-format record batch whose time index column holds `timestamps`.
    fn flat_ts_batch(timestamps: &[i64]) -> RecordBatch {
        use datatypes::arrow::array::{TimestampMillisecondArray, UInt8Array, UInt64Array};
        use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit};

        let num_rows = timestamps.len();
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "ts",
                DataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            Field::new("pk", DataType::UInt64, false),
            Field::new("seq", DataType::UInt64, false),
            Field::new("op", DataType::UInt8, false),
        ]));
        RecordBatch::try_new(
            schema,
            vec![
                Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())),
                Arc::new(UInt64Array::from(vec![0u64; num_rows])),
                Arc::new(UInt64Array::from(vec![0u64; num_rows])),
                Arc::new(UInt8Array::from(vec![0u8; num_rows])),
            ],
        )
        .unwrap()
    }

    /// Splits `timestamps` and returns the time index values of each sub-batch.
    fn split_ts(timestamps: &[i64]) -> Vec<Vec<i64>> {
        let mut batches = VecDeque::new();
        split_record_batch(flat_ts_batch(timestamps), &mut batches);
        batches
            .iter()
            .map(|batch| {
                let pos = time_index_column_index(batch.num_columns());
                let (values, _) = timestamp_array_to_primitive(batch.column(pos)).unwrap();
                values.values().to_vec()
            })
            .collect()
    }

    #[test]
    fn test_split_record_batch_on_equal_timestamps() {
        // Splits on both decreasing and equal timestamps.
        assert_eq!(
            split_ts(&[1, 2, 2, 3, 1]),
            vec![vec![1, 2], vec![2, 3], vec![1]]
        );
        // A run of equal timestamps yields single-row sub-batches.
        assert_eq!(split_ts(&[5, 5, 5]), vec![vec![5], vec![5], vec![5]]);
        // Equal-ts run at the leading edge of the batch.
        assert_eq!(split_ts(&[5, 5, 1, 2]), vec![vec![5], vec![5], vec![1, 2]]);
        // Equal-ts run at the trailing edge of the batch.
        assert_eq!(split_ts(&[1, 2, 5, 5]), vec![vec![1, 2, 5], vec![5]]);
    }

    #[test]
    fn test_split_record_batch_on_decreasing_timestamps() {
        assert_eq!(split_ts(&[1, 2, 3]), vec![vec![1, 2, 3]]);
        assert_eq!(split_ts(&[1, 3, 2, 4]), vec![vec![1, 3], vec![2, 4]]);
    }

    #[test]
    fn test_split_record_batch_empty_and_single_row() {
        let mut batches = VecDeque::new();
        split_record_batch(flat_ts_batch(&[]), &mut batches);
        assert!(batches.is_empty());

        assert_eq!(split_ts(&[42]), vec![vec![42]]);
    }
}

#[cfg(test)]
mod sequence_filter_tests {
    use std::sync::Arc;

    use datatypes::arrow::array::{Int64Array, StringArray, UInt8Array, UInt64Array};
    use datatypes::arrow::datatypes::{DataType, Field, Schema};
    use datatypes::arrow::record_batch::RecordBatch;
    use store_api::storage::SequenceRange;

    use super::filter_flat_batch_by_sequence;

    /// Builds a flat-format record batch: `(tag, field, ts, __primary_key, __sequence, __op_type)`.
    fn batch(sequences: &[u64]) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![
            Field::new("tag_0", DataType::Utf8, false),
            Field::new("field_0", DataType::Int64, false),
            Field::new(
                "ts",
                DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
            Field::new("__primary_key", DataType::UInt8, false),
            Field::new("__sequence", DataType::UInt64, false),
            Field::new("__op_type", DataType::UInt8, false),
        ]));
        let tags = StringArray::from_iter_values((0..sequences.len()).map(|i| i.to_string()));
        let fields = Int64Array::from_iter_values(0..sequences.len() as i64);
        let ts = datatypes::arrow::array::TimestampMillisecondArray::from_iter_values(
            (0..sequences.len()).map(|i| i as i64 * 1000),
        );
        let pk = UInt8Array::from(vec![0u8; sequences.len()]);
        let seq = UInt64Array::from_iter_values(sequences.iter().copied());
        let op = UInt8Array::from(vec![0u8; sequences.len()]);
        RecordBatch::try_new(
            schema,
            vec![
                Arc::new(tags),
                Arc::new(fields),
                Arc::new(ts),
                Arc::new(pk),
                Arc::new(seq),
                Arc::new(op),
            ],
        )
        .unwrap()
    }

    fn remaining_tags(batch: &RecordBatch) -> Vec<String> {
        batch
            .column(0)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap()
            .iter()
            .map(|v| v.unwrap().to_string())
            .collect()
    }

    #[test]
    fn test_filter_flat_batch_by_sequence_no_range_or_legacy_file() {
        let b = batch(&[1, 2, 3, 4]);

        let out = filter_flat_batch_by_sequence(b.clone(), None, true).unwrap();
        assert_eq!(remaining_tags(&out.unwrap()), vec!["0", "1", "2", "3"]);

        let out = filter_flat_batch_by_sequence(
            b.clone(),
            Some(SequenceRange::GtLtEq { min: 2, max: 3 }),
            false,
        )
        .unwrap();
        assert_eq!(remaining_tags(&out.unwrap()), vec!["0", "1", "2", "3"]);
    }

    #[test]
    fn test_filter_flat_batch_by_sequence_exact_range() {
        let b = batch(&[1, 2, 3, 4]);
        let out = filter_flat_batch_by_sequence(
            b.clone(),
            Some(SequenceRange::GtLtEq { min: 2, max: 3 }),
            true,
        )
        .unwrap();
        assert_eq!(remaining_tags(&out.unwrap()), vec!["2"]);

        let out = filter_flat_batch_by_sequence(
            b.clone(),
            Some(SequenceRange::GtLtEq { min: 10, max: 20 }),
            true,
        )
        .unwrap();
        assert!(out.is_none());

        let empty = batch(&[]);
        let out = filter_flat_batch_by_sequence(
            empty,
            Some(SequenceRange::GtLtEq { min: 0, max: 10 }),
            true,
        )
        .unwrap();
        assert_eq!(out.unwrap().num_rows(), 0);

        // Foreign batches are already virtualized by the reader. Their source
        // marker is irrelevant, and filtering uses the effective batch values.
        let out = filter_flat_batch_by_sequence(
            batch(&[1, 5, 9]),
            Some(SequenceRange::GtLtEq { min: 2, max: 8 }),
            true,
        )
        .unwrap();
        assert_eq!(remaining_tags(&out.unwrap()), vec!["1"]);
    }
}