use std::{path::Path, sync::Arc};

use anyhow::{Result, ensure};
use log::{info, warn};
use tici_common::index::IndexKey;
use tici_grpc::FragMeta;
use tokio::task::JoinHandle;

use super::frag_data_cache::FragDataCache;
use crate::{
    FragMetaVersion, FragMetaWrapper, ShardMirrorMeta,
    compactor::frag_utils::FragsDedupStat,
    metadata::metadata_store::MetadataStore,
    reader::{
        FragReaders,
        dedup_persist::PersistedShardDedupBitsets,
        doc_cache::DocCache,
        frag_reader::FragReaderMeta,
        shard_reader::{EstimateSampleCounts, ShardReader},
    },
    utils::{
        dedup_bitset::FragAliveBitSetBytes, key::ShardingKeyRange, meta_client::MetaClient,
        schema::IndexMeta,
    },
};

#[derive(Debug)]
pub struct ShardMirrorDedupOutput {
    pub dedup_stat: FragsDedupStat,
    pub obsolete_doc_bitsets: Vec<FragAliveBitSetBytes>,
}

struct SeqFragReadersInner {
    seq: u64,
    frag_readers: Arc<FragReaders>,
}

struct SeqFragReaders {
    inner: parking_lot::RwLock<SeqFragReadersInner>,
}

impl SeqFragReaders {
    fn new(seq: u64, frag_readers: Arc<FragReaders>) -> Self {
        Self {
            inner: parking_lot::RwLock::new(SeqFragReadersInner { seq, frag_readers }),
        }
    }

    fn get_seq(&self) -> u64 {
        self.inner.read().seq
    }

    fn set_seq(&self, seq: u64) {
        self.inner.write().seq = seq;
    }

    fn get_seq_and_frag_readers(&self) -> (u64, Arc<FragReaders>) {
        let inner = self.inner.read();
        (inner.seq, inner.frag_readers.clone())
    }

    fn set_seq_and_frag_readers(&self, seq: u64, frag_readers: Arc<FragReaders>) {
        *self.inner.write() = SeqFragReadersInner { seq, frag_readers };
    }
}

// ShardMirror maintains a synchronized local replica of S3 data and monitors
// changes to the associated Shard's data on S3. Assume all data in local disk.
pub struct ShardMirror {
    index_meta: IndexMeta,
    shard_id: u64,
    index_key: IndexKey,
    bucket: String,
    shard_range: ShardingKeyRange,
    estimate_sample_counts: EstimateSampleCounts,
    seq_frag_readers: SeqFragReaders,
    doc_cache: Option<Arc<DocCache>>,
}

impl ShardMirror {
    pub async fn new(
        shard_meta: ShardMirrorMeta,
        index_meta: IndexMeta,
        frag_reader_metas: Vec<FragReaderMeta>,
        obsolete_doc_bitsets: Vec<FragAliveBitSetBytes>,
    ) -> Result<Self> {
        Self::new_with_sample_counts(
            shard_meta,
            index_meta,
            frag_reader_metas,
            EstimateSampleCounts::default(),
            obsolete_doc_bitsets,
            None,
        )
        .await
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) async fn new_with_sample_counts(
        shard_meta: ShardMirrorMeta,
        index_meta: IndexMeta,
        frag_reader_metas: Vec<FragReaderMeta>,
        estimate_sample_counts: EstimateSampleCounts,
        obsolete_doc_bitsets: Vec<FragAliveBitSetBytes>,
        doc_cache: Option<Arc<DocCache>>,
    ) -> Result<Self> {
        let shard_range = ShardingKeyRange::half_open(
            shard_meta.shard_header.start_key.clone(),
            shard_meta.shard_header.end_key.clone(),
        )?;
        let frag_readers = Self::load_frags(
            &shard_range,
            frag_reader_metas,
            obsolete_doc_bitsets,
            doc_cache.clone(),
        )?;

        let frag_readers = SeqFragReaders::new(shard_meta.shard_header.seq, Arc::new(frag_readers));
        Ok(Self {
            index_meta,
            shard_id: shard_meta.shard_header.shard_id,
            index_key: shard_meta.index_key.clone(),
            bucket: shard_meta.frag_bucket.clone(),
            shard_range,
            estimate_sample_counts,
            seq_frag_readers: frag_readers,
            doc_cache,
        })
    }

    pub fn dedup_frags(
        shard_tracing_id: String,
        shard_range: ShardingKeyRange,
        frag_reader_metas: Vec<FragReaderMeta>,
        obsolete_keys_temp_path: String,
    ) -> JoinHandle<Result<ShardMirrorDedupOutput>> {
        // dedup_frags is CPU intensive and may take an amount of time. Run in
        // spawn_blocking to avoid blocking tokio's worker threads.
        tokio::task::spawn_blocking(move || {
            let obsolete_keys_temp_dir = Path::new(&obsolete_keys_temp_path);
            let res = crate::compactor::frag_utils::dedup_frags(
                &shard_range,
                &frag_reader_metas,
                |_key| Ok(()),
                obsolete_keys_temp_dir,
            )?;
            let stat = res.dedup_stat;
            let obsolete_doc_bitsets =
                FragAliveBitSetBytes::from_doc_bitsets(res.obsolete_doc_bitsets)?;
            info!("Shard {} dedup frags done: {:?}", shard_tracing_id, stat);
            Ok(ShardMirrorDedupOutput {
                dedup_stat: stat,
                obsolete_doc_bitsets,
            })
        })
    }

    fn load_frags(
        shard_range: &ShardingKeyRange,
        frag_reader_metas: Vec<FragReaderMeta>,
        obsolete_doc_bitsets: Vec<FragAliveBitSetBytes>,
        doc_cache: Option<Arc<DocCache>>,
    ) -> Result<FragReaders> {
        let mut frag_readers = Vec::with_capacity(frag_reader_metas.len());
        ensure!(
            frag_reader_metas.len() == obsolete_doc_bitsets.len(),
            "frag reader meta count {} does not match obsolete doc bitset count {}",
            frag_reader_metas.len(),
            obsolete_doc_bitsets.len()
        );
        for (meta, obsolete_doc_bitset) in frag_reader_metas.into_iter().zip(obsolete_doc_bitsets) {
            let frag_range = meta.frag_range()?;
            let valid_range = frag_range.intersect(shard_range);
            let frag_reader = meta.load_reader_with_obsolete_doc_bitset(
                doc_cache.clone(),
                valid_range,
                obsolete_doc_bitset,
            )?;
            frag_readers.push(Arc::new(frag_reader));
        }
        Ok(frag_readers)
    }

    pub(crate) async fn refresh(
        &self,
        frag_data_cache: Arc<FragDataCache>,
        mut meta_client: Box<dyn MetaClient>,
        metadata_store: Arc<MetadataStore>,
    ) -> Result<()> {
        let shard_id = self.shard_id;
        let (curr_seq, curr_frag_readers) = self.seq_frag_readers.get_seq_and_frag_readers();
        let tracing_id = self.tracing_id();
        let key_range = tici_grpc::KeyRange {
            start_key: self.shard_range.start().raw().to_vec(),
            end_key: self.shard_range.end().raw().to_vec(),
        };
        let (new_frag_metas, new_seq) = meta_client
            .get_frags_by_keyrange(shard_id, &self.index_key, key_range)
            .await?;
        if new_seq <= curr_seq {
            info!(
                "Shard {} not changed: new_seq={}, curr_seq={}",
                tracing_id, new_seq, curr_seq
            );
            return Ok(());
        }
        let new_local_frag_paths = new_frag_metas
            .iter()
            .map(|frag_meta| frag_data_cache.local_frag_path(&frag_meta.frag_path))
            .collect::<Vec<_>>();
        let curr_local_frag_paths = curr_frag_readers
            .iter()
            .map(|reader| reader.frag_path())
            .collect::<Vec<_>>();
        if new_local_frag_paths == curr_local_frag_paths {
            info!(
                "Frags of {} not changed={} frags, new_seq={}, curr_seq={}",
                tracing_id,
                new_local_frag_paths.len(),
                new_seq,
                curr_seq,
            );
            // Update seq to avoid refresh repeatedly
            self.seq_frag_readers.set_seq(new_seq);
            let frag_meta_wrapper = FragMetaWrapper {
                version: FragMetaVersion::LATEST,
                seq: new_seq,
                frag_metas: new_frag_metas,
            };
            // Update seq asynchronously
            let _ = metadata_store
                .set_frag_meta(shard_id, &frag_meta_wrapper)
                .await;
            return Ok(());
        }

        info!(
            "Frags of {} changed: new_seq={}, curr_seq={}, curr_frag_paths({})={:?}, new_frag_paths({})={:?}",
            tracing_id,
            new_seq,
            curr_seq,
            curr_local_frag_paths.len(),
            curr_local_frag_paths,
            new_local_frag_paths.len(),
            new_local_frag_paths,
        );
        let new_frag_readers = Self::refresh_frags(
            &tracing_id,
            &self.index_key,
            shard_id,
            &self.shard_range,
            frag_data_cache.clone(),
            &self.bucket,
            &curr_frag_readers,
            new_frag_metas.clone(),
            new_local_frag_paths.clone(),
            self.doc_cache.clone(),
            new_seq,
            metadata_store,
        )
        .await?;
        self.seq_frag_readers
            .set_seq_and_frag_readers(new_seq, Arc::new(new_frag_readers));
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    async fn refresh_frags(
        tracing_id: &str,
        index_key: &IndexKey,
        shard_id: u64,
        shard_range: &ShardingKeyRange,
        frag_data_cache: Arc<FragDataCache>,
        s3_bucket: &str,
        curr_frag_readers: &FragReaders,
        new_frag_metas: Vec<FragMeta>,
        new_local_frag_paths: Vec<String>,
        doc_cache: Option<Arc<DocCache>>,
        new_seq: u64,
        metadata_store: Arc<MetadataStore>,
    ) -> Result<FragReaders> {
        let mut new_frag_reader_metas = Vec::with_capacity(new_frag_metas.len());
        for (meta, local_frag_path) in new_frag_metas.iter().zip(new_local_frag_paths) {
            match curr_frag_readers
                .iter()
                .find(|reader| reader.frag_path() == local_frag_path)
            {
                Some(frag_reader) => {
                    new_frag_reader_metas.push(frag_reader.meta().clone());
                }
                None => {
                    let frag_reader_meta = FragReaderMeta::from_frag_data_cache(
                        meta,
                        frag_data_cache.clone(),
                        s3_bucket,
                    )
                    .await?;
                    new_frag_reader_metas.push(frag_reader_meta);
                }
            }
        }

        // Update the frag metas before deduplication. If a restart happens
        // before the new bitsets are persisted, restore will validate the
        // persisted payload against this latest frag list and rebuild on miss.
        let frag_meta_wrapper = FragMetaWrapper {
            version: FragMetaVersion::LATEST,
            seq: new_seq,
            frag_metas: new_frag_metas,
        };
        metadata_store
            .set_frag_meta(shard_id, &frag_meta_wrapper)
            .await?
            .await??;

        let obsolete_keys_temp_path = frag_data_cache.get_temporary_path();
        let dedup_output = Self::dedup_frags(
            tracing_id.to_string(),
            shard_range.clone(),
            new_frag_reader_metas.clone(),
            obsolete_keys_temp_path,
        )
        .await??;

        match PersistedShardDedupBitsets::new(
            index_key,
            shard_id,
            shard_range,
            &frag_meta_wrapper.frag_metas,
            &dedup_output.obsolete_doc_bitsets,
        )
        .and_then(|persisted| persisted.encode_compressed())
        {
            Ok(persisted_bytes) => {
                if let Err(err) = async {
                    metadata_store
                        .set_shard_dedup_bitset(shard_id, &persisted_bytes)
                        .await?
                        .await??;
                    Ok::<(), anyhow::Error>(())
                }
                .await
                {
                    warn!(
                        "Persist shard dedup bitsets failed, shard={}, err={:?}",
                        tracing_id, err
                    );
                }
            }
            Err(err) => {
                warn!(
                    "Build shard dedup bitset payload failed, shard={}, err={:?}",
                    tracing_id, err
                );
            }
        }

        let frag_readers = Self::load_frags(
            shard_range,
            new_frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
            doc_cache,
        )?;

        Ok(frag_readers)
    }

    fn shard_range(&self) -> &ShardingKeyRange {
        &self.shard_range
    }

    pub fn get_reader(&self) -> ShardReader {
        let (seq, frag_readers) = self.seq_frag_readers.get_seq_and_frag_readers();
        ShardReader::new(
            self.shard_range().clone(),
            self.index_meta.clone(),
            seq,
            self.estimate_sample_counts,
            frag_readers,
        )
    }

    pub fn get_curr_seq(&self) -> u64 {
        self.seq_frag_readers.get_seq()
    }

    pub fn tracing_id(&self) -> String {
        crate::utils::shard_tracing_id(&self.index_key, self.shard_id)
    }
}

#[cfg(test)]
mod tests {
    use std::{
        collections::{BTreeMap, BTreeSet},
        sync::Arc,
    };

    use anyhow::{Context, Result};
    use rand::Rng;
    use tempfile::TempDir;
    use tici_grpc::KeyRange;
    use tici_test_utils::test_utils::init_logger_for_ut;
    use tokio::sync::mpsc;

    use crate::{
        config::{FragWriterConfig, WorkerNodeConfig},
        reader::{
            frag_reader::FragReaderMeta,
            query_info::{QueryInfo, SearchExecutionPlan, TantivySchemaExt},
            query_result::SearchRows,
            shard_mirror::ShardMirror,
            shard_reader::ShardReader,
        },
        test_utils::{index_meta::TEST_HDFSLOGS_INDEX_META_FULLTEXT, *},
        utils::{
            key::{ShardingKey, ShardingKeyRange},
            meta_client_pool::{MetaClientPool, MockMetaClientPool},
            s3_frag,
            schema::{IndexMeta, SortField, VERSION_STORE_NAME},
            tantivy_doc::RawDoc,
        },
        writer::{cdc_log::OperationType, cdc_log_frag_writer::CDCLogFragWriter},
    };

    fn verify_rows_prefix(rows: &SearchRows, all_rows: &SearchRows) {
        assert_eq!(rows.column_names(), all_rows.column_names());
        for row_idx in 0..rows.len() {
            verify_row_eq(rows, row_idx, all_rows, row_idx);
        }
    }

    fn verify_row_eq(
        actual_rows: &SearchRows,
        actual_row_idx: usize,
        expected_rows: &SearchRows,
        expected_row_idx: usize,
    ) {
        macro_rules! verify_columns {
            ($columns:ident) => {
                assert_eq!(actual_rows.$columns.len(), expected_rows.$columns.len());
                for (actual_col, expected_col) in actual_rows
                    .$columns
                    .iter()
                    .zip(expected_rows.$columns.iter())
                {
                    assert_eq!(actual_col.col_name, expected_col.col_name);
                    assert_eq!(
                        actual_col.value_at(actual_row_idx),
                        expected_col.value_at(expected_row_idx)
                    );
                }
            };
        }

        verify_columns!(i64_columns);
        verify_columns!(u64_columns);
        verify_columns!(f64_columns);
        verify_columns!(bytes_columns);
    }

    fn test_search_with_query_info_topk(
        shard_reader: &ShardReader,
        query_info: &QueryInfo,
        sort_fields: &[SortField],
        limit: usize,
        expected_size: usize,
        all_rows: Option<&SearchRows>,
    ) -> SearchRows {
        let schema = shard_reader.get_schema().unwrap();
        let query_info = query_info
            .add_topk(&schema, limit, sort_fields.to_vec(), true)
            .unwrap();
        let (result, _stat) = shard_reader
            .search_with_query_info(Arc::new(query_info))
            .unwrap();
        let rows = result;
        assert_eq!(rows.len(), expected_size);
        verify_sort_order(&rows, sort_fields);
        if let Some(all_rows) = all_rows {
            verify_rows_prefix(&rows, all_rows)
        }
        rows
    }

    fn test_shard_search_topk(
        shard_reader: &ShardReader,
        index_meta: &IndexMeta,
        query_info: &QueryInfo,
        total_count: usize,
    ) {
        let sort_fields = index_meta.get_sort_fields();

        let all_rows = test_search_with_query_info_topk(
            shard_reader,
            query_info,
            sort_fields,
            usize::MAX,
            total_count,
            None,
        );

        let _ = test_search_with_query_info_topk(
            shard_reader,
            query_info,
            sort_fields,
            10,
            std::cmp::min(10, total_count),
            Some(&all_rows),
        );

        let _ = test_search_with_query_info_topk(
            shard_reader,
            query_info,
            sort_fields,
            0,
            0,
            Some(&all_rows),
        );
    }

    async fn test_search_with_index_meta(index_meta: IndexMeta) {
        init_logger_for_ut();
        let shard_meta = test_get_shard_writer_meta();
        let temp_dir = TempDir::new().unwrap();
        let config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());

        let backend_storage = test_create_backend_storage().await.unwrap();

        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();

        let mut mock_client = meta_client_pool.get_meta_client().await.unwrap();
        let (frag_metas, _seq) = mock_client
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &shard_meta.index_key,
                KeyRange::default(),
            )
            .await
            .unwrap();
        let mut frag_reader_metas = Vec::with_capacity(frag_metas.len());
        for frag_meta in frag_metas {
            let property = frag_meta.property.as_ref().unwrap();
            let local_frag_path = frag_meta.frag_path.replace(
                &shard_meta.frag_location.prefix,
                &config.storage.frag_base_path(),
            );
            assert!(!std::fs::exists(local_frag_path).unwrap());

            let local_frag_path = s3_frag::download_frag(
                backend_storage.as_ref(),
                &shard_meta.frag_location.bucket,
                &frag_meta,
                &config.storage.frag_base_path(),
            )
            .await
            .unwrap();
            assert!(is_sorted_frag(&local_frag_path, &index_meta));
            frag_reader_metas.push(
                FragReaderMeta::new(
                    &local_frag_path,
                    property.min_handle.clone(),
                    property.max_handle.clone(),
                )
                .unwrap(),
            );
        }

        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            config.storage.temporary_path(),
        )
        .await
        .unwrap()
        .unwrap();
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await
        .unwrap();
        let reader = shard_mirror.get_reader();
        let total_count = 150000;
        let query_info = Arc::new(
            QueryInfo::query_all_for_test(
                &index_meta,
                &reader.get_schema().unwrap(),
                vec![test_get_inf_range()],
            )
            .unwrap(),
        );
        let (all_docs, _search_stat) = reader.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(all_docs.len(), total_count);
        test_shard_search_topk(&reader, &index_meta, &query_info, total_count);

        // Test returning only primary key columns (id and tenant_id)
        let pk_query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &reader.get_schema().unwrap(),
            vec![test_get_inf_range()],
            vec!["id", "tenant_id"],
        )
        .unwrap();

        let (pk_docs, _search_stat) = reader.search_with_query_info(pk_query_info).unwrap();

        // Verify the result - we expect the same number of documents to be returned
        assert_eq!(pk_docs.len(), total_count);

        // Compare values of first row from all_docs and pk_docs
        if !all_docs.is_empty() && !pk_docs.is_empty() {
            assert!(
                all_docs.column_count() >= 4,
                "all_docs first row should have at least 4 columns"
            );
            assert_eq!(
                pk_docs.i64_columns.len(),
                2,
                "pk_docs first row should have 2 columns"
            );
            let pk_first_col_name = pk_docs.i64_columns[0].col_name.clone();
            let pk_second_col_name = pk_docs.i64_columns[1].col_name.clone();
            let all_id_value = rows_i64_value(&all_docs, 0, &pk_first_col_name).unwrap();
            let all_tenant_id_value = rows_i64_value(&all_docs, 0, &pk_second_col_name).unwrap();
            let pk_id_value = rows_i64_value(&pk_docs, 0, &pk_first_col_name).unwrap();
            let pk_tenant_id_value = rows_i64_value(&pk_docs, 0, &pk_second_col_name).unwrap();

            assert_eq!(
                all_id_value, pk_id_value,
                "first pk column in all_docs should match pk_docs"
            );
            assert_eq!(
                all_tenant_id_value, pk_tenant_id_value,
                "second pk column in all_docs should match pk_docs"
            );
        }
    }

    #[tokio::test]
    async fn test_search_with_mixed_index_meta() {
        test_search_with_index_meta(TEST_HDFSLOGS_INDEX_META.clone()).await;
    }

    #[tokio::test]
    async fn test_search_with_fulltext_index_meta() {
        test_search_with_index_meta(TEST_HDFSLOGS_INDEX_META_FULLTEXT.clone()).await;
    }

    #[tokio::test]
    async fn test_update_delete() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let shard_meta = test_get_shard_writer_meta();
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let (raw_docs, expected_raw_docs) = test_gen_raw_docs(10000, 3000, 1000);
        let frag_count = 10;
        let docs_per_frag = raw_docs.len() / frag_count;
        let mut frag_reader_metas = Vec::with_capacity(frag_count);
        let mut current_docs = 0;
        for raw_doc in raw_docs {
            cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();
            current_docs += 1;
            if current_docs >= docs_per_frag {
                let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap();
                let frag_res = frag_writer.finalize().unwrap();
                frag_reader_metas.push(
                    FragReaderMeta::new(
                        &frag_writer.frag_path(),
                        frag_res.handle_stat.min_handle,
                        frag_res.handle_stat.max_handle,
                    )
                    .unwrap(),
                );
                current_docs = 0;
            }
        }
        assert_eq!(frag_reader_metas.len(), frag_count);
        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await
        .unwrap()
        .unwrap();
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await
        .unwrap();
        let reader = shard_mirror.get_reader();
        let query_info = Arc::new(
            QueryInfo::query_all_for_test(
                &index_meta,
                &reader.get_schema().unwrap(),
                vec![test_get_inf_range()],
            )
            .unwrap(),
        );
        let (all_docs, _search_stat1) = reader.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(all_docs.len(), expected_raw_docs.len());

        test_shard_search_topk(&reader, &index_meta, &query_info, expected_raw_docs.len());
    }

    #[tokio::test]
    async fn test_update_delete_binary_pk() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let index_meta = TEST_INDEX_META_BINARY_PK.clone();
        let shard_meta = test_get_shard_writer_meta();
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let (raw_docs, mut expected_raw_docs) = test_gen_raw_docs_binary_pk(10000, 3000, 1000);
        let frag_count = 10;
        let docs_per_frag = raw_docs.len() / frag_count;
        let mut frag_reader_metas = Vec::with_capacity(frag_count);
        let mut current_docs = 0;
        for raw_doc in raw_docs {
            cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();
            current_docs += 1;
            if current_docs >= docs_per_frag {
                let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap();
                let frag_res = frag_writer.finalize().unwrap();
                frag_reader_metas.push(
                    FragReaderMeta::new(
                        &frag_writer.frag_path(),
                        frag_res.handle_stat.min_handle,
                        frag_res.handle_stat.max_handle,
                    )
                    .unwrap(),
                );
                current_docs = 0;
            }
        }
        assert_eq!(frag_reader_metas.len(), frag_count);
        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await
        .unwrap()
        .unwrap();
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await
        .unwrap();

        let reader = shard_mirror.get_reader();
        let mut return_cols = index_meta
            .get_column_infos()
            .iter()
            .map(|column| column.column_name.as_str())
            .collect::<Vec<_>>();
        return_cols.push(VERSION_STORE_NAME);
        let query_info = QueryInfo::query_cols_for_test(
            &index_meta,
            &reader.get_schema().unwrap(),
            vec![test_get_inf_range()],
            return_cols,
        )
        .unwrap();
        let (all_docs, _search_stat) = reader.search_with_query_info(query_info).unwrap();
        let rows = all_docs;

        assert_eq!(rows.len(), expected_raw_docs.len());

        for row_idx in 0..rows.len() {
            let description_col = rows_bytes_value(&rows, row_idx, "column_1").unwrap();
            let binary_pk = rows_bytes_value(&rows, row_idx, "column_2").unwrap();
            let handle = ShardingKey::from_raw(binary_pk.to_vec());
            let expected_doc = expected_raw_docs.remove(&handle).unwrap();
            let expected_description = expected_doc
                .cols
                .iter()
                .find(|(col_id, _)| *col_id == 1)
                .and_then(|(_, value)| value.as_deref())
                .unwrap();

            assert_eq!(Some(description_col), Some(expected_description.as_bytes()));
            assert_eq!(Some(binary_pk), Some(expected_doc.handle.raw()));
            assert_eq!(
                rows_u64_value(&rows, row_idx, VERSION_STORE_NAME).unwrap(),
                expected_doc.version
            );
        }

        assert!(expected_raw_docs.is_empty());
    }

    #[tokio::test]
    async fn test_insert_only_duplication() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let shard_meta = test_get_shard_writer_meta();
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        // Insert only
        let (raw_docs, expected_raw_docs) = test_gen_raw_docs(100, 0, 0);
        let mut frag_reader_metas = Vec::new();
        let mut write_one_frag = async || {
            for raw_doc in &raw_docs {
                cdc_log_frag_writer.add_raw_doc(raw_doc.clone()).unwrap();
            }
            let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap();
            let frag_res = frag_writer.finalize().unwrap();
            frag_reader_metas.push(
                FragReaderMeta::new(
                    &frag_writer.frag_path(),
                    frag_res.handle_stat.min_handle,
                    frag_res.handle_stat.max_handle,
                )
                .unwrap(),
            );
        };
        write_one_frag().await;
        write_one_frag().await;
        assert_eq!(frag_reader_metas.len(), 2);

        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await
        .unwrap()
        .unwrap();
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await
        .unwrap();
        let reader = shard_mirror.get_reader();
        let query_info = Arc::new(
            QueryInfo::query_all_for_test(
                &index_meta,
                &reader.get_schema().unwrap(),
                vec![test_get_inf_range()],
            )
            .unwrap(),
        );
        let (all_docs, _search_stat) = reader.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(all_docs.len(), expected_raw_docs.len());

        test_shard_search_topk(&reader, &index_meta, &query_info, expected_raw_docs.len());
    }

    fn random_range(
        map: &BTreeMap<ShardingKey, RawDoc>,
    ) -> (ShardingKeyRange, BTreeMap<ShardingKey, RawDoc>) {
        let mut rng = rand::rng();
        let len = map.len();
        assert!(len > 0);
        let start_idx = rng.random_range(0..len / 2);
        let end_idx = rng.random_range(start_idx + 1..len); // Make sure end_idx > start_idx
        let keys: Vec<_> = map.keys().cloned().collect();
        let start = keys[start_idx].clone();
        let end = keys[end_idx].clone();
        // [start, end)
        let range_map = map
            .range(start.clone()..end.clone())
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect::<BTreeMap<_, _>>();
        let range = ShardingKeyRange::half_open(start.into_vec(), end.into_vec()).unwrap();
        (range, range_map)
    }

    #[tokio::test]
    async fn test_read_by_range() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let shard_meta = test_get_shard_writer_meta();
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let (raw_docs, expected_raw_docs) = test_gen_raw_docs(10000, 3000, 1000);
        let frag_count = 10;
        let docs_per_frag = raw_docs.len() / frag_count;
        let mut frag_reader_metas = Vec::with_capacity(frag_count);
        let mut current_docs = 0;
        for raw_doc in raw_docs {
            cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();
            current_docs += 1;
            if current_docs >= docs_per_frag {
                let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap();
                let frag_res = frag_writer.finalize().unwrap();
                frag_reader_metas.push(
                    FragReaderMeta::new(
                        &frag_writer.frag_path(),
                        frag_res.handle_stat.min_handle,
                        frag_res.handle_stat.max_handle,
                    )
                    .unwrap(),
                );
                current_docs = 0;
            }
        }
        assert_eq!(frag_reader_metas.len(), frag_count);
        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await
        .unwrap()
        .unwrap();
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await
        .unwrap();
        let reader = shard_mirror.get_reader();

        for _ in 0..10 {
            let (read_range, expected_docs) = random_range(&expected_raw_docs);
            let query_info = Arc::new(
                QueryInfo::query_all_for_test(
                    &index_meta,
                    &reader.get_schema().unwrap(),
                    vec![read_range],
                )
                .unwrap(),
            );
            let (all_docs, _search_stat1) =
                reader.search_with_query_info(query_info.clone()).unwrap();
            assert_eq!(all_docs.len(), expected_docs.len());

            test_shard_search_topk(&reader, &index_meta, &query_info, expected_docs.len());
        }
    }

    #[tokio::test]
    async fn test_delete_only() {
        let backend_storage = test_create_backend_storage().await.unwrap();
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let shard_meta = test_get_shard_writer_meta();
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        // Write delete-only frag
        let (raw_docs, _expected_raw_docs) = test_gen_raw_docs(100, 0, 0);
        for mut raw_doc in raw_docs.into_iter() {
            raw_doc.op_type = OperationType::Delete;
            cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();
        }
        let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap();
        let frag_res = frag_writer.finalize().unwrap();
        assert!(frag_res.seg_ids.is_empty());

        // Upload delete-only frag
        let frag_batch_uploader = test_get_frag_batch_uploader().await.unwrap();
        let frag_meta = s3_frag::upload_frag(
            &frag_batch_uploader,
            TEST_S3_BUCKET,
            &TEST_S3_PREFIX,
            local_data_path,
            &frag_writer.frag_path(),
            frag_res.handle_stat.clone(),
            &frag_res.seg_ids,
        )
        .await
        .unwrap();

        // Download delete-only frag
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let local_frag_path = s3_frag::download_frag(
            backend_storage.as_ref(),
            TEST_S3_BUCKET,
            &frag_meta,
            local_data_path,
        )
        .await
        .unwrap();

        // Read delete-only frag
        let frag_reader_metas = vec![
            FragReaderMeta::new(
                &local_frag_path,
                frag_res.handle_stat.min_handle,
                frag_res.handle_stat.max_handle,
            )
            .unwrap(),
        ];
        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await
        .unwrap()
        .unwrap();
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await
        .unwrap();
        let reader = shard_mirror.get_reader();
        let query_info = Arc::new(
            QueryInfo::query_all_for_test(
                &index_meta,
                &reader.get_schema().unwrap(),
                vec![test_get_inf_range()],
            )
            .unwrap(),
        );
        let (all_docs, _search_stat) = reader.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(all_docs.len(), 0);
        test_shard_search_topk(&reader, &index_meta, &query_info, 0);
    }

    #[tokio::test]
    async fn test_mock_read_after_split_shard() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let (raw_docs, expected_raw_docs) = test_gen_raw_docs(10000, 3000, 1000);
        let frag_count = 10;
        let docs_per_frag = raw_docs.len() / frag_count;
        let mut frag_reader_metas = Vec::with_capacity(frag_count);
        let mut current_docs = 0;
        for raw_doc in raw_docs {
            cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();
            current_docs += 1;
            if current_docs >= docs_per_frag {
                let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap();
                let frag_res = frag_writer.finalize().unwrap();
                frag_reader_metas.push(
                    FragReaderMeta::new(
                        &frag_writer.frag_path(),
                        frag_res.handle_stat.min_handle,
                        frag_res.handle_stat.max_handle,
                    )
                    .unwrap(),
                );
                current_docs = 0;
            }
        }
        assert_eq!(frag_reader_metas.len(), frag_count);

        // Shrinking range to mock split shard
        let start = expected_raw_docs.keys().next().unwrap();
        let end = expected_raw_docs.keys().nth(1000).unwrap();
        let mut expected_query_docs_count = 0;
        for (_key, value) in expected_raw_docs.iter().take(1000) {
            let col = value.cols.iter().find(|(id, _)| *id == 1).unwrap();
            if col.1.as_ref().unwrap().contains(" rwp09000") {
                expected_query_docs_count += 1;
            }
        }

        let mut shard_meta = test_get_shard_writer_meta();
        shard_meta.shard_header.start_key = start.raw().to_vec();
        shard_meta.shard_header.end_key = end.raw().to_vec();
        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await
        .unwrap()
        .unwrap();
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await
        .unwrap();

        let reader = shard_mirror.get_reader();
        let read_ranges = vec![test_get_inf_range()];
        let query_info = Arc::new(
            QueryInfo::query_word(
                &index_meta,
                &reader.get_schema().unwrap(),
                &read_ranges,
                "column_1".to_string(),
                "rwp09000".to_string(),
            )
            .unwrap(),
        );
        let (docs, search_stat1) = reader.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(docs.len(), expected_query_docs_count);
        // Read ranges is infinite, so all frags are subset of read ranges, no overlap
        // frag.
        assert_eq!(search_stat1.frags.overlap, 0);
        test_shard_search_topk(&reader, &index_meta, &query_info, expected_query_docs_count);
    }

    #[tokio::test]
    async fn test_search_count_matches_rows_for_same_range() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let local_data_path = temp_dir
            .path()
            .to_str()
            .context("Failed to convert path to string")?;
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let shard_meta = test_get_shard_writer_meta();
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await?;

        let (raw_docs, expected_raw_docs) = test_gen_raw_docs(2_000, 0, 0);
        let frag_count = 10;
        let docs_per_frag = raw_docs.len() / frag_count;
        let mut frag_reader_metas = Vec::with_capacity(frag_count);
        let mut current_docs = 0;
        for raw_doc in raw_docs {
            cdc_log_frag_writer.add_raw_doc(raw_doc)?;
            current_docs += 1;
            if current_docs >= docs_per_frag {
                let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await?;
                let frag_res = frag_writer.finalize()?;
                frag_reader_metas.push(FragReaderMeta::new(
                    &frag_writer.frag_path(),
                    frag_res.handle_stat.min_handle,
                    frag_res.handle_stat.max_handle,
                )?);
                current_docs = 0;
            }
        }
        assert_eq!(frag_reader_metas.len(), frag_count);

        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta
                .shard_range()
                .context("Missing shard range in shard mirror meta")?,
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await??;
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await?;

        let handles = expected_raw_docs.keys().collect::<Vec<_>>();
        let start_handle = handles
            .get(200)
            .context("Missing start handle for read range")?;
        let end_handle = handles
            .get(800)
            .context("Missing end handle for read range")?;
        let read_range =
            ShardingKeyRange::half_open(start_handle.raw().to_vec(), end_handle.raw().to_vec())?;
        let read_ranges = vec![read_range];
        let reader = shard_mirror.get_reader();
        let schema = reader.get_schema()?;
        let rows_query = QueryInfo::query_word(
            &index_meta,
            &schema,
            &read_ranges,
            "column_1".to_string(),
            "rwp09000".to_string(),
        )?;
        let count_query = QueryInfo::query_word_count(
            &index_meta,
            &schema,
            &read_ranges,
            "column_1".to_string(),
            "rwp09000".to_string(),
        )?;

        let (rows_result, _rows_stat) = reader.search_with_query_info(rows_query)?;
        let (count, _count_stat) = reader.count_with_query_info(count_query)?;
        let rows = rows_result;
        assert_eq!(rows.len() as u64, count);
        Ok(())
    }

    #[tokio::test]
    async fn test_null_value() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let local_data_path = temp_dir
            .path()
            .to_str()
            .context("Failed to convert path to string")?;
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let shard_meta = test_get_shard_writer_meta();
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await?;

        let h1 = "with_null_value".to_string();
        let raw_doc1 = RawDoc {
            handle: ShardingKey::from_raw(h1.as_bytes().to_vec()),
            version: 1,
            op_type: OperationType::Insert,
            cols: vec![
                (1, None),
                (2, None),
                (3, Some(h1.clone())),
                (4, Some(h1.clone())),
            ],
            is_local_timezone: true,
        };

        let h2 = "no_null_value".to_string();
        let raw_doc2 = RawDoc {
            handle: ShardingKey::from_raw(h2.as_bytes().to_vec()),
            version: 1,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("This is a sentence".to_string())),
                (2, Some("This is a paragraph".to_string())),
                (3, Some(h2.clone())),
                (4, Some(h2.clone())),
            ],
            is_local_timezone: true,
        };

        cdc_log_frag_writer.add_raw_doc(raw_doc1)?;
        cdc_log_frag_writer.add_raw_doc(raw_doc2)?;

        let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await?;
        let frag_res = frag_writer.finalize()?;
        let frag_reader_meta = FragReaderMeta::new(
            &frag_writer.frag_path(),
            frag_res.handle_stat.min_handle,
            frag_res.handle_stat.max_handle,
        )?;
        let frag_reader_metas = vec![frag_reader_meta];

        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await??;
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await?;
        let reader = shard_mirror.get_reader();
        let query_info = QueryInfo::query_all_for_test(
            &index_meta,
            &reader.get_schema()?,
            vec![test_get_inf_range()],
        )?;
        let (all_docs, _search_stat) = reader.search_with_query_info(query_info)?;
        let rows = all_docs;

        assert_eq!(rows.len(), 2);
        assert!(rows_is_null(&rows, 0, "column_1")?);
        assert!(rows_is_null(&rows, 0, "column_2")?);
        assert_eq!(rows_bytes_value(&rows, 0, "column_3")?, h1.as_bytes());
        assert_eq!(rows_bytes_value(&rows, 0, "column_4")?, h1.as_bytes());
        assert_eq!(
            rows_bytes_value(&rows, 1, "column_1")?,
            b"This is a sentence"
        );
        assert_eq!(
            rows_bytes_value(&rows, 1, "column_2")?,
            b"This is a paragraph"
        );
        assert_eq!(rows_bytes_value(&rows, 1, "column_3")?, h2.as_bytes());
        assert_eq!(rows_bytes_value(&rows, 1, "column_4")?, h2.as_bytes());

        Ok(())
    }

    #[tokio::test]
    async fn test_search_with_top_collector_sort() {
        init_logger_for_ut();
        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_ALL_TYPES_META.clone();
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();

        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);

        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let raw_docs = test_gen_raw_docs_for_all_types(1000);
        let frag_count = 20;
        let mut docs_per_frag = vec![Vec::new(); frag_count];
        for (idx, raw_doc) in raw_docs.into_iter().enumerate() {
            let target_frag_idx = idx % frag_count;
            docs_per_frag[target_frag_idx].push(raw_doc);
        }

        let mut frag_reader_metas = Vec::with_capacity(frag_count);
        for docs in docs_per_frag {
            for raw_doc in docs {
                cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();
            }
            let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap();
            let frag_res = frag_writer.finalize().unwrap();
            frag_reader_metas.push(
                FragReaderMeta::new(
                    &frag_writer.frag_path(),
                    frag_res.handle_stat.min_handle,
                    frag_res.handle_stat.max_handle,
                )
                .unwrap(),
            );
        }
        assert_eq!(frag_reader_metas.len(), frag_count);

        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await
        .unwrap()
        .unwrap();
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await
        .unwrap();

        let reader = shard_mirror.get_reader();
        let base_query_info = QueryInfo::query_all_for_test(
            &index_meta,
            &reader.get_schema().unwrap(),
            vec![test_get_inf_range()],
        )
        .unwrap();

        let schema = reader.get_schema().unwrap();
        let column_count = index_meta.get_column_infos().len();
        let sort_fields_all = (0..column_count)
            .flat_map(|i| {
                let df = schema.doc_field(&format!("column_{}", i + 1)).unwrap();
                // Tantivy doesn't support ordering by bytes type, so we skip it
                if df.value_type == tantivy::schema::Type::Bytes {
                    Vec::new()
                } else {
                    vec![
                        SortField {
                            field: df.clone(),
                            is_asc: true,
                        },
                        SortField {
                            field: df,
                            is_asc: false,
                        },
                    ]
                }
            })
            .collect::<Vec<_>>();

        for (i, sf) in sort_fields_all.iter().enumerate() {
            let sort_fields = vec![sf.clone()];
            let q0 = base_query_info
                .add_topk(&schema, 100, sort_fields.clone(), false)
                .unwrap();
            let (res0, _stat0) = reader.search_with_query_info(q0).unwrap();
            let rows0 = res0;
            if i < 2 {
                assert_eq!(rows0.len(), 1000);
            } else {
                assert_eq!(rows0.len(), 100);
                verify_sort_order(&rows0, &sort_fields);
            }
        }

        let two_fields = vec![sort_fields_all[4].clone(), sort_fields_all[2].clone()];
        let q2 = base_query_info
            .add_topk(&schema, 100, two_fields, false)
            .unwrap();
        let (res2, _stat2) = reader.search_with_query_info(q2).unwrap();
        assert_eq!(res2.len(), 1000);

        let start_key = format!("all_type_{:07}", 990).as_bytes().to_vec();
        let end_key = format!("all_type_{:07}", 1000).as_bytes().to_vec();
        let specific_range = ShardingKeyRange::half_open(start_key, end_key).unwrap();

        let mut all_expr: tipb::Expr = tipb::Expr::default();
        all_expr.set_tp(tipb::ExprType::ScalarFunc);
        all_expr.set_sig(tipb::ScalarFuncSig::UnaryNotInt);
        let mut lt_expr: tipb::Expr = tipb::Expr::default();
        lt_expr.set_tp(tipb::ExprType::ScalarFunc);
        lt_expr.set_sig(tipb::ScalarFuncSig::LtInt);
        let mut col_ref: tipb::Expr = tipb::Expr::default();
        col_ref.set_tp(tipb::ExprType::ColumnRef);
        col_ref.set_sig(tipb::ScalarFuncSig::Unspecified);
        col_ref.set_val(b"column_1".to_vec());
        let mut zero: tipb::Expr = tipb::Expr::default();
        zero.set_tp(tipb::ExprType::Int64);
        zero.set_sig(tipb::ScalarFuncSig::Unspecified);
        zero.set_val(0i64.to_le_bytes().to_vec());
        lt_expr.mut_children().push(col_ref);
        lt_expr.mut_children().push(zero);
        all_expr.mut_children().push(lt_expr);

        let ranged_q = reader
            .parse_expr(
                &all_expr,
                &["column_2".to_string()],
                &[specific_range],
                20,
                vec!["column_2".to_string()],
                vec![true],
            )
            .unwrap();
        let (ranged_res, _ranged_stat) = reader.search_with_query_info(ranged_q).unwrap();
        let ranged_rows = ranged_res;
        assert_eq!(ranged_rows.len(), 10);
        verify_sort_order(
            &ranged_rows,
            &[SortField {
                field: schema.doc_field("column_2").unwrap(),
                is_asc: true,
            }],
        );
    }

    #[tokio::test]
    async fn test_fast_field_topk_rows_respects_partial_read_range() -> Result<()> {
        init_logger_for_ut();
        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_ALL_TYPES_META.clone();
        let temp_dir = TempDir::new()?;
        let local_data_path = temp_dir
            .path()
            .to_str()
            .context("TempDir path is not valid UTF-8")?;

        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);

        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await?;

        let raw_docs = test_gen_raw_docs_for_all_types(240);
        let frag_size = 40;
        assert_eq!(raw_docs.len() % frag_size, 0);

        let mut frag_reader_metas = Vec::with_capacity(raw_docs.len() / frag_size);
        for docs in raw_docs.chunks(frag_size) {
            for raw_doc in docs {
                cdc_log_frag_writer.add_raw_doc(raw_doc.clone())?;
            }
            let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await?;
            let frag_res = frag_writer.finalize()?;
            frag_reader_metas.push(FragReaderMeta::new(
                &frag_writer.frag_path(),
                frag_res.handle_stat.min_handle,
                frag_res.handle_stat.max_handle,
            )?);
        }

        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta
                .shard_range()
                .context("Missing shard range in shard mirror meta")?,
            frag_reader_metas.clone(),
            format!("{}/temp", local_data_path),
        )
        .await??;
        let shard_mirror = ShardMirror::new(
            shard_mirror_meta,
            index_meta.clone(),
            frag_reader_metas,
            dedup_output.obsolete_doc_bitsets,
        )
        .await?;

        let reader = shard_mirror.get_reader();
        let schema = reader.get_schema()?;
        let start_idx = 30usize;
        let end_idx = 170usize;
        let read_range = ShardingKeyRange::half_open(
            raw_docs[start_idx].handle.raw().to_vec(),
            raw_docs[end_idx].handle.raw().to_vec(),
        )?;
        let base_query =
            QueryInfo::query_cols_for_test(&index_meta, &schema, vec![read_range], vec!["id"])?;
        let sort_field = SortField {
            field: schema.doc_field("column_2")?,
            is_asc: false,
        };
        let query_info = base_query.add_topk(&schema, 10, vec![sort_field.clone()], false)?;

        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FastFieldTopKRows
        );

        let (result, stat) = reader.search_with_query_info(query_info)?;
        let rows = result;

        assert_eq!(rows.len(), 10);
        assert_eq!(stat.frags.overlap, 2);
        assert_eq!(stat.frags.subset, 3);
        assert_eq!(stat.frags.disjoint, 1);
        assert_eq!(
            stat.frags.subset + stat.frags.overlap + stat.frags.disjoint,
            6
        );
        assert_eq!(stat.docs.read_range_filtered, 0);

        verify_sort_order(&rows, &[sort_field]);

        let returned_pairs = (0..rows.len())
            .map(|row_idx| {
                Ok((
                    rows_i64_value(&rows, row_idx, "column_1")?,
                    rows_i64_value(&rows, row_idx, "column_2")?,
                ))
            })
            .collect::<Result<Vec<_>>>()?;

        for (id, _) in returned_pairs.iter() {
            let doc_idx = (*id - 1000) as usize;
            assert!(
                (start_idx..end_idx).contains(&doc_idx),
                "returned id={} maps to doc_idx={} outside read range [{}, {})",
                id,
                doc_idx,
                start_idx,
                end_idx
            );
        }

        let mut expected_pairs = raw_docs[start_idx..end_idx]
            .iter()
            .map(|raw_doc| Ok((raw_doc_i64_col(raw_doc, 1)?, raw_doc_i64_col(raw_doc, 2)?)))
            .collect::<Result<Vec<_>>>()?;
        expected_pairs.sort_by(|a, b| b.1.cmp(&a.1));
        let expected_top10 = expected_pairs.into_iter().take(10).collect::<Vec<_>>();

        assert_eq!(
            returned_pairs.iter().copied().collect::<BTreeSet<_>>(),
            expected_top10.iter().copied().collect::<BTreeSet<_>>()
        );
        assert_eq!(returned_pairs, expected_top10);

        Ok(())
    }

    fn rows_i64_value(rows: &SearchRows, row_idx: usize, field_name: &str) -> Result<i64> {
        let col = rows
            .i64_column(field_name)
            .with_context(|| format!("Field '{field_name}' not found in rows"))?;
        col.value_at(row_idx)
            .with_context(|| format!("Field '{field_name}' has no integer value at row {row_idx}"))
    }

    fn rows_u64_value(rows: &SearchRows, row_idx: usize, field_name: &str) -> Result<u64> {
        let col = rows
            .u64_column(field_name)
            .with_context(|| format!("Field '{field_name}' not found in rows"))?;
        col.value_at(row_idx).with_context(|| {
            format!("Field '{field_name}' has no unsigned integer value at row {row_idx}")
        })
    }

    /// Float column value at `row_idx`, or `None` when NULL (matches
    /// [`crate::utils::cmp_doc::cmp_opt_f64`]).
    fn rows_opt_f64_value(
        rows: &SearchRows,
        row_idx: usize,
        field_name: &str,
    ) -> Result<Option<f64>> {
        let col = rows
            .f64_column(field_name)
            .with_context(|| format!("Field '{field_name}' not found in rows"))?;
        Ok(col.value_at(row_idx))
    }

    fn rows_bytes_value<'a>(
        rows: &'a SearchRows,
        row_idx: usize,
        field_name: &str,
    ) -> Result<&'a [u8]> {
        let col = rows
            .bytes_column(field_name)
            .with_context(|| format!("Field '{field_name}' not found in rows"))?;
        col.value_at(row_idx)
            .with_context(|| format!("Field '{field_name}' has no bytes value at row {row_idx}"))
    }

    fn rows_is_null(rows: &SearchRows, row_idx: usize, field_name: &str) -> Result<bool> {
        if let Some(col) = rows.i64_column(field_name) {
            return Ok(col.is_null(row_idx));
        }
        if let Some(col) = rows.u64_column(field_name) {
            return Ok(col.is_null(row_idx));
        }
        if let Some(col) = rows.f64_column(field_name) {
            return Ok(col.is_null(row_idx));
        }
        if let Some(col) = rows.bytes_column(field_name) {
            return Ok(col.is_null(row_idx));
        }
        anyhow::bail!("Field '{field_name}' not found in rows");
    }

    fn raw_doc_i64_col(raw_doc: &RawDoc, col_id: i64) -> Result<i64> {
        let value = raw_doc
            .cols
            .iter()
            .find(|(id, _)| *id == col_id)
            .with_context(|| format!("Column {col_id} not found in raw doc: {raw_doc:?}"))?
            .1
            .as_deref()
            .with_context(|| format!("Column {col_id} is NULL in raw doc: {raw_doc:?}"))?;
        value.parse::<i64>().with_context(|| {
            format!(
                "Failed to parse column {col_id} value '{value}' as i64 in raw doc: {raw_doc:?}"
            )
        })
    }

    fn verify_sort_order(rows: &SearchRows, sort_fields: &[SortField]) {
        use std::cmp::Ordering as CmpOrdering;

        use tantivy::schema::Type as TvType;
        for row_idx in 0..rows.len().saturating_sub(1) {
            for sf in sort_fields {
                let fname = &sf.field.name;
                let ord = match sf.field.value_type {
                    TvType::I64 => rows_i64_value(rows, row_idx, fname)
                        .unwrap()
                        .cmp(&rows_i64_value(rows, row_idx + 1, fname).unwrap()),
                    TvType::U64 => rows_u64_value(rows, row_idx, fname)
                        .unwrap()
                        .cmp(&rows_u64_value(rows, row_idx + 1, fname).unwrap()),
                    TvType::F64 => crate::utils::cmp_doc::cmp_opt_f64(
                        rows_opt_f64_value(rows, row_idx, fname).unwrap(),
                        rows_opt_f64_value(rows, row_idx + 1, fname).unwrap(),
                    ),
                    TvType::Str | TvType::Bytes => rows_bytes_value(rows, row_idx, fname)
                        .unwrap()
                        .cmp(rows_bytes_value(rows, row_idx + 1, fname).unwrap()),
                    _ => CmpOrdering::Equal,
                };
                let actual_ord = if sf.is_asc { ord } else { ord.reverse() };
                if actual_ord != CmpOrdering::Equal {
                    assert!(
                        actual_ord != CmpOrdering::Greater,
                        "rows not sorted by sort_field"
                    );
                    break;
                }
            }
        }
    }
}