use std::{cmp::Ordering, sync::Arc, time::Instant};

use anyhow::{Context, Result, anyhow, bail, ensure};
use tantivy::TantivyDocument;
use tici_common::metrics::OperationMetricsGuard;

use crate::{
    reader::{
        FragReaders,
        fast_merge::search_with_top_collector,
        frag_reader::FragRowsStream,
        merge_iter::{HeapEntry, MergeIter},
        query_info::{QueryInfo, SearchExecutionPlan},
        query_result::SearchRows,
        shard_reader::{ShardReader, ShardReaderStat},
    },
    utils::{
        cmp_doc::compare_docs,
        metrics::{begin_shard_reader_query_metrics, observe_shard_search_stream_timing},
        tantivy_doc::DocFieldAccess,
    },
};

pub const DEFAULT_SEARCH_STREAM_DRAIN_BATCH_SIZE: usize = 8192;

type MergeComparator = Box<dyn Fn(&HeapEntry, &HeapEntry) -> Ordering>;

pub struct SearchRowsBatch {
    pub rows: SearchRows,
    pub is_eof: bool,
}

enum SearchStreamPlan {
    FragScan {
        frag_streams: Vec<FragRowsStream>,
        frag_idx: usize,
    },
    StoreSortTopK {
        stream: StoreSortTopKRowsStream,
    },
    // Not truly streaming: top-k docs are collected eagerly at open time (see
    // `open_fast_field_topk_rows_stream`); `BufferedDocsRowsStream` only slices the buffer.
    FastFieldTopK {
        stream: BufferedDocsRowsStream,
    },
}

pub struct ShardSearchStream {
    query_info: Arc<QueryInfo>,
    stat: ShardReaderStat,
    metrics_guard: OperationMetricsGuard,
    plan: SearchStreamPlan,
    failed: bool,
    done: bool,
}

impl ShardSearchStream {
    pub fn next_batch(&mut self, batch_size: usize) -> Result<SearchRowsBatch> {
        ensure!(batch_size > 0, "batch_size must be greater than 0");
        if self.failed {
            bail!("search stream already failed");
        }

        let batch_start = Instant::now();
        let batch = self.next_batch_inner(batch_size);
        self.stat.timing.next_batch_us += batch_start.elapsed().as_micros() as u64;
        if batch.is_err() {
            self.failed = true;
            self.metrics_guard.mark_error();
        } else if let Ok(ref batch) = batch {
            self.stat.stream.batch_count += (batch.rows.row_count > 0) as u64;
            self.stat.docs.returned_rows += batch.rows.row_count;
        }
        batch
    }

    pub fn stat_snapshot(&self) -> ShardReaderStat {
        self.stat.clone()
    }

    pub fn query_info(&self) -> &QueryInfo {
        self.query_info.as_ref()
    }

    fn next_batch_inner(&mut self, batch_size: usize) -> Result<SearchRowsBatch> {
        if self.done {
            return self.done_batch();
        }

        match &mut self.plan {
            SearchStreamPlan::FragScan {
                frag_streams,
                frag_idx,
            } => {
                let rows = drain_frag_streams(
                    self.query_info.as_ref(),
                    &mut self.stat,
                    frag_streams,
                    frag_idx,
                    batch_size,
                )?;
                let is_eof = *frag_idx >= frag_streams.len();
                self.done = is_eof;
                Ok(SearchRowsBatch { rows, is_eof })
            }
            SearchStreamPlan::StoreSortTopK { stream } => {
                let (rows, is_eof) = stream.next_batch(batch_size, &mut self.stat)?;
                self.done = is_eof;
                Ok(SearchRowsBatch { rows, is_eof })
            }
            SearchStreamPlan::FastFieldTopK { stream } => {
                let (rows, is_eof) = stream.next_batch(batch_size, &mut self.stat)?;
                self.done = is_eof;
                Ok(SearchRowsBatch { rows, is_eof })
            }
        }
    }

    fn done_batch(&self) -> Result<SearchRowsBatch> {
        Ok(SearchRowsBatch {
            rows: SearchRows::empty(self.query_info.return_fields()?)?,
            is_eof: true,
        })
    }
}

impl Drop for ShardSearchStream {
    fn drop(&mut self) {
        observe_shard_search_stream_timing(
            self.stat.plan.as_str(),
            self.stat.timing.open_us,
            self.stat.timing.next_batch_us,
            self.stat.stream.batch_count,
        );
    }
}

impl ShardReader {
    pub fn open_search_stream(&self, query_info: Arc<QueryInfo>) -> Result<ShardSearchStream> {
        let metrics_guard = begin_shard_reader_query_metrics(query_info.execution_plan().as_str());
        let mut stat = ShardReaderStat {
            plan: query_info.execution_plan(),
            ..Default::default()
        };
        let open_start = Instant::now();
        let plan = match query_info.execution_plan() {
            SearchExecutionPlan::StoreSortTopKRows => {
                open_store_sort_topk_rows_stream(&self.frag_readers, query_info.clone(), &mut stat)
                    .map(|stream| SearchStreamPlan::StoreSortTopK { stream })
            }
            SearchExecutionPlan::FastFieldTopKRows => {
                open_fast_field_topk_rows_stream(&self.frag_readers, query_info.clone(), &mut stat)
                    .map(|stream| SearchStreamPlan::FastFieldTopK { stream })
            }
            SearchExecutionPlan::FragScanRows => {
                open_frag_scan_rows_streams(&self.frag_readers, query_info.clone(), &mut stat).map(
                    |frag_streams| SearchStreamPlan::FragScan {
                        frag_streams,
                        frag_idx: 0,
                    },
                )
            }
            SearchExecutionPlan::FragScanCount => {
                bail!("SearchExecutionPlan::FragScanCount is not supported in ShardSearchStream")
            }
        };
        stat.timing.open_us += open_start.elapsed().as_micros() as u64;
        match plan {
            Ok(plan) => Ok(ShardSearchStream {
                query_info,
                stat,
                metrics_guard,
                plan,
                failed: false,
                done: false,
            }),
            Err(err) => {
                metrics_guard.mark_error();
                Err(err)
            }
        }
    }
}

struct StoreSortTopKRowsStream {
    query_info: Arc<QueryInfo>,
    merge_iter: MergeIter<MergeComparator>,
    emitted: usize,
    exhausted: bool,
    topk_limit: usize,
}

impl StoreSortTopKRowsStream {
    fn next_batch(
        &mut self,
        batch_size: usize,
        stat: &mut ShardReaderStat,
    ) -> Result<(SearchRows, bool)> {
        let topk_limit = self.topk_limit;
        ensure!(
            topk_limit >= self.emitted,
            "topk_limit ({}) must be greater than or equal to emitted ({})",
            topk_limit,
            self.emitted
        );
        let mut docs = Vec::with_capacity(batch_size.min(topk_limit - self.emitted));
        while docs.len() < batch_size && self.emitted < topk_limit {
            let Some(doc) = self.merge_iter.next_doc()? else {
                self.exhausted = true;
                break;
            };
            stat.docs.tantivy_matched += 1;
            let handle = doc
                .get_bytes(self.query_info.handle_field())?
                .ok_or_else(|| anyhow!("handle field is not set"))?;
            if !self.query_info.in_read_ranges(handle) {
                stat.docs.read_range_filtered += 1;
                continue;
            }
            docs.push(doc);
            self.emitted += 1;
        }

        let rows = materialize_docs(&docs, &self.query_info)?;
        let is_eof = self.exhausted || self.emitted >= topk_limit;
        Ok((rows, is_eof))
    }
}

/// Holds the full top-k result loaded at stream open; `next_batch` only
/// advances `offset`.
struct BufferedDocsRowsStream {
    query_info: Arc<QueryInfo>,
    docs: Vec<Arc<TantivyDocument>>,
    offset: usize,
}

impl BufferedDocsRowsStream {
    fn next_batch(
        &mut self,
        batch_size: usize,
        _stat: &mut ShardReaderStat,
    ) -> Result<(SearchRows, bool)> {
        let end = (self.offset + batch_size).min(self.docs.len());
        let rows = materialize_docs(&self.docs[self.offset..end], &self.query_info)?;
        self.offset = end;
        Ok((rows, self.offset >= self.docs.len()))
    }
}

fn open_frag_scan_rows_streams(
    frag_readers: &Arc<FragReaders>,
    query_info: Arc<QueryInfo>,
    stat: &mut ShardReaderStat,
) -> Result<Vec<FragRowsStream>> {
    let mut frag_streams = Vec::with_capacity(frag_readers.len());
    for frag_reader in frag_readers.iter() {
        frag_streams.push(
            frag_reader
                .clone()
                .open_rows_stream(query_info.clone(), stat)?,
        );
    }
    Ok(frag_streams)
}

fn open_store_sort_topk_rows_stream(
    frag_readers: &Arc<FragReaders>,
    query_info: Arc<QueryInfo>,
    stat: &mut ShardReaderStat,
) -> Result<StoreSortTopKRowsStream> {
    let top_k_param = query_info
        .topk_params()
        .ok_or_else(|| anyhow!("topk_params is not set"))?;
    let mut seg_scanners = Vec::new();
    for frag_reader in frag_readers.iter() {
        seg_scanners.extend(frag_reader.segment_scanners(&query_info, stat)?);
    }
    let sort_fields = top_k_param.sort_fields.clone();
    let limit = top_k_param.limit;
    let cmp: MergeComparator = Box::new(move |a: &HeapEntry, b: &HeapEntry| -> Ordering {
        compare_docs(&b.doc, &a.doc, &sort_fields)
    });
    let merge_iter = MergeIter::new(seg_scanners, cmp)?;
    Ok(StoreSortTopKRowsStream {
        query_info,
        merge_iter,
        emitted: 0,
        exhausted: false,
        topk_limit: limit,
    })
}

// Eagerly loads the entire top-k via `search_with_top_collector` at open time.
// Peak memory matches a one-shot search (large limit or many fragments);
// batching does not bound it.
fn open_fast_field_topk_rows_stream(
    frag_readers: &Arc<FragReaders>,
    query_info: Arc<QueryInfo>,
    stat: &mut ShardReaderStat,
) -> Result<BufferedDocsRowsStream> {
    let docs = search_with_top_collector(frag_readers, &query_info, stat)?;
    Ok(BufferedDocsRowsStream {
        query_info,
        docs,
        offset: 0,
    })
}

fn drain_frag_streams(
    query_info: &QueryInfo,
    stat: &mut ShardReaderStat,
    frag_streams: &mut [FragRowsStream],
    frag_idx: &mut usize,
    batch_size: usize,
) -> Result<SearchRows> {
    let return_fields = query_info.return_fields()?;
    let mut rows = SearchRows::with_capacity(return_fields, batch_size)?;
    while rows.len() < batch_size && *frag_idx < frag_streams.len() {
        let remaining = batch_size - rows.len();
        let next_rows = frag_streams[*frag_idx].next_batch(remaining, stat)?;
        match next_rows {
            Some(fragment_rows) => {
                rows.append(fragment_rows)?;
            }
            None => {
                *frag_idx += 1;
            }
        }
    }
    Ok(rows)
}

fn materialize_docs(docs: &[Arc<TantivyDocument>], query_info: &QueryInfo) -> Result<SearchRows> {
    SearchRows::from_doc_refs(
        docs.iter().map(Arc::as_ref),
        docs.len(),
        query_info.return_fields()?,
    )
    .context("failed to materialize search rows from documents")
}

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

    use anyhow::Result;
    use tempfile::TempDir;

    use super::*;
    use crate::{
        reader::{
            frag_reader::FragReader,
            query_info::{SearchExecutionPlan, TantivySchemaExt},
            shard_reader::EstimateSampleCounts,
        },
        test_utils::{
            TEST_HDFSLOGS_INDEX_META, TEST_INDEX_META_NORMAL, TestDoc, build_test_frag,
            load_frag_reader, row_versions, test_get_inf_range,
        },
        utils::{
            key::ShardingKeyRange,
            schema::{IndexMeta, SortField, VERSION_STORE_NAME},
        },
    };

    fn build_loaded_frag(
        index_meta: &IndexMeta,
        segments: Vec<Vec<TestDoc>>,
    ) -> Result<(TempDir, Arc<FragReader>)> {
        let (temp_dir, meta) = build_test_frag(index_meta, segments)?;
        let frag_reader = load_frag_reader(meta, &[])?;
        Ok((temp_dir, frag_reader))
    }

    fn build_shard_reader(
        index_meta: IndexMeta,
        frag_readers: Vec<Arc<FragReader>>,
    ) -> ShardReader {
        ShardReader::new(
            test_get_inf_range(),
            index_meta,
            1,
            EstimateSampleCounts::default(),
            Arc::new(frag_readers),
        )
    }

    fn expect_rows(batch: SearchRowsBatch) -> Result<(Vec<u64>, bool)> {
        Ok((row_versions(&batch.rows)?, batch.is_eof))
    }

    fn expect_next_batch_error(
        stream: &mut ShardSearchStream,
        batch_size: usize,
        expected_message: &str,
    ) {
        let err = match stream.next_batch(batch_size) {
            Ok(_) => panic!("expected next_batch error containing {expected_message}"),
            Err(err) => err,
        };
        assert!(
            err.to_string().contains(expected_message),
            "expected error containing {expected_message}, got {err:#}"
        );
    }

    #[test]
    fn test_frag_scan_rows_stream_drains_fragments_in_batches() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td1, frag_reader_1) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 20, "beta testword"),
                TestDoc::normal(b"doc_02", 30, "gamma testword"),
            ]],
        )?;
        let (_td2, frag_reader_2) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::normal(b"doc_10", 40, "delta testword"),
                TestDoc::normal(b"doc_11", 50, "epsilon testword"),
            ]],
        )?;
        let shard_reader =
            build_shard_reader(index_meta.clone(), vec![frag_reader_1, frag_reader_2]);
        let query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &shard_reader.get_schema()?,
            vec![test_get_inf_range()],
            vec![VERSION_STORE_NAME],
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FragScanRows
        );

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;
        let mut batch_lengths = Vec::new();
        let mut eof_flags = Vec::new();
        let mut versions = Vec::new();
        loop {
            let (batch_versions, is_eof) = expect_rows(stream.next_batch(2)?)?;
            batch_lengths.push(batch_versions.len());
            eof_flags.push(is_eof);
            versions.extend(batch_versions);
            if is_eof {
                break;
            }
        }

        versions.sort_unstable();
        assert_eq!(batch_lengths, vec![2, 2, 1]);
        assert_eq!(eof_flags, vec![false, false, true]);
        assert_eq!(versions, vec![10, 20, 30, 40, 50]);
        let (done_versions, done_eof) = expect_rows(stream.next_batch(2)?)?;
        assert!(done_versions.is_empty());
        assert!(done_eof);
        Ok(())
    }

    #[test]
    fn test_count_query_cannot_open_search_stream() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td, frag_reader) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 20, "beta testword"),
            ]],
        )?;
        let shard_reader = build_shard_reader(index_meta.clone(), vec![frag_reader]);
        let schema = shard_reader.get_schema()?;
        let query_info = QueryInfo::query_word_count(
            &index_meta,
            &schema,
            &[test_get_inf_range()],
            "column_1".to_string(),
            "testword".to_string(),
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FragScanCount
        );

        let err = match shard_reader.open_search_stream(Arc::new(query_info)) {
            Ok(_) => panic!("count query unexpectedly opened a search stream"),
            Err(err) => err,
        };

        assert!(
            err.to_string().contains(
                "SearchExecutionPlan::FragScanCount is not supported in ShardSearchStream"
            ),
            "expected count stream error, got {err:#}"
        );
        Ok(())
    }

    #[test]
    fn test_count_with_query_info_counts_across_fragments() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td1, frag_reader_1) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 20, "beta testword"),
            ]],
        )?;
        let (_td2, frag_reader_2) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::normal(b"doc_10", 30, "gamma testword"),
                TestDoc::normal(b"doc_11", 40, "delta no_match"),
            ]],
        )?;
        let shard_reader =
            build_shard_reader(index_meta.clone(), vec![frag_reader_1, frag_reader_2]);
        let schema = shard_reader.get_schema()?;
        let query_info = QueryInfo::query_word_count(
            &index_meta,
            &schema,
            &[test_get_inf_range()],
            "column_1".to_string(),
            "testword".to_string(),
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FragScanCount
        );

        let (count, stat) = shard_reader.count_with_query_info(query_info)?;

        assert_eq!(count, 3);
        assert_eq!(stat.docs.tantivy_matched, count);
        assert_eq!(stat.frags.disjoint, 0);
        Ok(())
    }

    #[test]
    fn test_search_with_query_info_rejects_count_query() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td, frag_reader) = build_loaded_frag(
            &index_meta,
            vec![vec![TestDoc::normal(b"doc_00", 10, "alpha testword")]],
        )?;
        let shard_reader = build_shard_reader(index_meta.clone(), vec![frag_reader]);
        let schema = shard_reader.get_schema()?;
        let query_info = QueryInfo::query_word_count(
            &index_meta,
            &schema,
            &[test_get_inf_range()],
            "column_1".to_string(),
            "testword".to_string(),
        )?;

        let err = shard_reader.search_with_query_info(query_info).unwrap_err();

        assert!(
            err.to_string()
                .contains("ResultType is invalid, expected Rows, got Count"),
            "expected count search error, got {err:#}"
        );
        Ok(())
    }

    #[test]
    fn test_shard_stream_rejects_zero_batch_size_without_poisoning_stream() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td, frag_reader) = build_loaded_frag(
            &index_meta,
            vec![vec![TestDoc::normal(b"doc_00", 10, "alpha testword")]],
        )?;
        let shard_reader = build_shard_reader(index_meta.clone(), vec![frag_reader]);
        let query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &shard_reader.get_schema()?,
            vec![test_get_inf_range()],
            vec![VERSION_STORE_NAME],
        )?;

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;

        expect_next_batch_error(&mut stream, 0, "batch_size must be greater than 0");
        let (versions, is_eof) = expect_rows(stream.next_batch(1)?)?;
        assert_eq!(versions, vec![10]);
        assert!(!is_eof);
        let (versions, is_eof) = expect_rows(stream.next_batch(1)?)?;
        assert!(versions.is_empty());
        assert!(is_eof);
        Ok(())
    }

    #[test]
    fn test_frag_scan_rows_stream_returns_empty_batch_for_disjoint_fragment() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td, frag_reader) = build_loaded_frag(
            &index_meta,
            vec![vec![TestDoc::normal(b"doc_00", 10, "alpha testword")]],
        )?;
        let shard_reader = build_shard_reader(index_meta.clone(), vec![frag_reader]);
        let read_range = ShardingKeyRange::half_open(b"zzz".to_vec(), b"zzzz".to_vec())?;
        let query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &shard_reader.get_schema()?,
            vec![read_range],
            vec![VERSION_STORE_NAME],
        )?;

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;

        let (versions, is_eof) = expect_rows(stream.next_batch(2)?)?;
        assert!(versions.is_empty());
        assert!(is_eof);
        assert_eq!(stream.stat_snapshot().frags.disjoint, 1);
        Ok(())
    }

    #[test]
    fn test_frag_scan_rows_stream_returns_empty_batch_for_no_matches() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td, frag_reader) = build_loaded_frag(
            &index_meta,
            vec![vec![TestDoc::normal(b"doc_00", 10, "alpha testword")]],
        )?;
        let shard_reader = build_shard_reader(index_meta.clone(), vec![frag_reader]);
        let query_info = QueryInfo::query_word(
            &index_meta,
            &shard_reader.get_schema()?,
            &[test_get_inf_range()],
            "column_1".to_string(),
            "missing".to_string(),
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FragScanRows
        );

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;

        let SearchRowsBatch { rows, is_eof } = stream.next_batch(2)?;
        assert_eq!(rows.len(), 0);
        assert!(is_eof);
        assert_eq!(stream.stat_snapshot().docs.tantivy_matched, 0);
        Ok(())
    }

    #[test]
    fn test_store_sort_topk_rows_stream_batches_results() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td1, frag_reader_1) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 30, "gamma testword"),
            ]],
        )?;
        let (_td2, frag_reader_2) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::normal(b"doc_10", 20, "beta testword"),
                TestDoc::normal(b"doc_11", 40, "delta testword"),
            ]],
        )?;
        let shard_reader =
            build_shard_reader(index_meta.clone(), vec![frag_reader_1, frag_reader_2]);
        let schema = shard_reader.get_schema()?;
        let base_query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &schema,
            vec![test_get_inf_range()],
            vec![VERSION_STORE_NAME],
        )?;
        let query_info = base_query_info.add_topk(
            &schema,
            3,
            vec![SortField {
                field: schema.doc_field("column_1")?,
                is_asc: true,
            }],
            true,
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::StoreSortTopKRows
        );

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;
        let (first_versions, first_eof) = expect_rows(stream.next_batch(2)?)?;
        let (second_versions, second_eof) = expect_rows(stream.next_batch(2)?)?;

        assert_eq!(first_versions, vec![10, 20]);
        assert!(!first_eof);
        assert_eq!(second_versions, vec![40]);
        assert!(second_eof);
        Ok(())
    }

    #[test]
    fn test_store_sort_topk_rows_stream_filters_read_ranges() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td, frag_reader) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 20, "beta testword"),
                TestDoc::normal(b"doc_02", 30, "gamma testword"),
                TestDoc::normal(b"doc_03", 40, "delta testword"),
            ]],
        )?;
        let shard_reader = build_shard_reader(index_meta.clone(), vec![frag_reader]);
        let schema = shard_reader.get_schema()?;
        let read_range = ShardingKeyRange::half_open(b"doc_01".to_vec(), b"doc_03".to_vec())?;
        let base_query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &schema,
            vec![read_range],
            vec![VERSION_STORE_NAME],
        )?;
        let query_info = base_query_info.add_topk(
            &schema,
            10,
            vec![SortField {
                field: schema.doc_field("column_1")?,
                is_asc: true,
            }],
            true,
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::StoreSortTopKRows
        );

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;
        let (versions, is_eof) = expect_rows(stream.next_batch(10)?)?;

        assert_eq!(versions, vec![20, 30]);
        assert!(is_eof);
        let stat = stream.stat_snapshot();
        assert_eq!(stat.docs.tantivy_matched, 4);
        assert_eq!(stat.docs.read_range_filtered, 2);
        Ok(())
    }

    #[test]
    fn test_store_sort_topk_rows_stream_handles_zero_limit() -> Result<()> {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let (_td, frag_reader) = build_loaded_frag(
            &index_meta,
            vec![vec![TestDoc::normal(b"doc_00", 10, "alpha testword")]],
        )?;
        let shard_reader = build_shard_reader(index_meta.clone(), vec![frag_reader]);
        let schema = shard_reader.get_schema()?;
        let base_query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &schema,
            vec![test_get_inf_range()],
            vec![VERSION_STORE_NAME],
        )?;
        let query_info = base_query_info.add_topk(
            &schema,
            0,
            vec![SortField {
                field: schema.doc_field("column_1")?,
                is_asc: true,
            }],
            true,
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::StoreSortTopKRows
        );

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;
        let (versions, is_eof) = expect_rows(stream.next_batch(10)?)?;

        assert!(versions.is_empty());
        assert!(is_eof);
        assert_eq!(stream.stat_snapshot().docs.tantivy_matched, 0);
        Ok(())
    }

    #[test]
    fn test_fast_field_topk_rows_stream_batches_results() -> Result<()> {
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let (_td1, frag_reader_1) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::numeric(b"doc_00", 10, 101, 1001),
                TestDoc::numeric(b"doc_01", 30, 103, 1003),
            ]],
        )?;
        let (_td2, frag_reader_2) = build_loaded_frag(
            &index_meta,
            vec![vec![
                TestDoc::numeric(b"doc_10", 20, 102, 1002),
                TestDoc::numeric(b"doc_11", 40, 104, 1004),
            ]],
        )?;
        let shard_reader =
            build_shard_reader(index_meta.clone(), vec![frag_reader_1, frag_reader_2]);
        let schema = shard_reader.get_schema()?;
        let base_query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &schema,
            vec![test_get_inf_range()],
            vec![VERSION_STORE_NAME],
        )?;
        let query_info = base_query_info.add_topk(
            &schema,
            3,
            vec![SortField {
                field: schema.doc_field("column_3")?,
                is_asc: false,
            }],
            false,
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FastFieldTopKRows
        );

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;
        let (first_versions, first_eof) = expect_rows(stream.next_batch(2)?)?;
        let (second_versions, second_eof) = expect_rows(stream.next_batch(2)?)?;

        assert_eq!(first_versions, vec![40, 30]);
        assert!(!first_eof);
        assert_eq!(second_versions, vec![20]);
        assert!(second_eof);
        Ok(())
    }

    #[test]
    fn test_fast_field_topk_rows_stream_handles_zero_limit() -> Result<()> {
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let (_td, frag_reader) = build_loaded_frag(
            &index_meta,
            vec![vec![TestDoc::numeric(b"doc_00", 10, 101, 1001)]],
        )?;
        let shard_reader = build_shard_reader(index_meta.clone(), vec![frag_reader]);
        let schema = shard_reader.get_schema()?;
        let base_query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &schema,
            vec![test_get_inf_range()],
            vec![VERSION_STORE_NAME],
        )?;
        let query_info = base_query_info.add_topk(
            &schema,
            0,
            vec![SortField {
                field: schema.doc_field("column_3")?,
                is_asc: false,
            }],
            false,
        )?;
        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FastFieldTopKRows
        );

        let mut stream = shard_reader.open_search_stream(Arc::new(query_info))?;
        let (versions, is_eof) = expect_rows(stream.next_batch(10)?)?;

        assert!(versions.is_empty());
        assert!(is_eof);
        assert_eq!(stream.stat_snapshot().docs.tantivy_matched, 0);
        Ok(())
    }
}