mito2/sst/

index.rs

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

pub(crate) mod bloom_filter;
mod column;
#[cfg(test)]
mod column_test;
pub(crate) mod fulltext_index;
mod indexer;
pub mod intermediate;
pub(crate) mod inverted_index;
mod primary_key;
pub mod puffin_manager;
#[cfg(test)]
mod sparse_test;
mod statistics;
pub(crate) mod store;

use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};
use std::num::NonZeroUsize;
use std::sync::Arc;

use bloom_filter::creator::BloomFilterIndexer;
use common_telemetry::{debug, error, info, warn};
use datatypes::arrow::record_batch::RecordBatch;
use mito_codec::row_converter::DensePrimaryKeyCodec;
use object_store::ObjectStore;
use puffin_manager::SstPuffinManager;
use smallvec::{SmallVec, smallvec};
use snafu::ResultExt;
use statistics::{ByteCount, RowCount};
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::RegionMetadataRef;
use store_api::storage::{ColumnId, FileId, RegionId};
use strum::IntoStaticStr;
use tokio::sync::mpsc::Sender;

use crate::access_layer::{AccessLayerRef, FilePathProvider, OperationType, RegionFilePathFactory};
use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey};
use crate::cache::write_cache::{UploadOptions, UploadTracker, WriteCacheRef};
use crate::cache::{CacheManagerRef, CacheStrategy};
use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig};
use crate::error::{
    BuildIndexAsyncSnafu, Error, JoinSnafu, RegionClosedSnafu, RegionDroppedSnafu,
    RegionTruncatedSnafu, Result,
};
use crate::metrics::{
    INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL, INDEX_CREATE_MEMORY_USAGE, INDEX_PUBLICATION_STALE_TOTAL,
};
use crate::read::Batch;
use crate::region::options::IndexOptions;
use crate::region::version::VersionControlRef;
use crate::region::{
    IndexBuildSource, IndexPublication, IndexPublicationStale, ManifestContextRef,
};
use crate::request::{
    BackgroundNotify, BuildIndexRequest, IndexBuildFailed, IndexBuildFinished, IndexBuildStopped,
    WorkerRequest, WorkerRequestWithTime,
};
use crate::schedule::scheduler::{Job, SchedulerRef};
use crate::sst::file::{
    ColumnIndexMetadata, FileHandle, IndexType, IndexTypes, RegionFileId, RegionIndexId,
};
use crate::sst::file_purger::FilePurgerRef;
use crate::sst::index::fulltext_index::FULLTEXT_BLOOM_VERSION;
use crate::sst::index::fulltext_index::creator::FulltextIndexer;
use crate::sst::index::intermediate::IntermediateManager;
use crate::sst::index::inverted_index::creator::InvertedIndexer;
use crate::sst::parquet::SstInfo;
use crate::worker::WorkerListener;

pub(crate) const TYPE_INVERTED_INDEX: &str = "inverted_index";
pub(crate) const TYPE_FULLTEXT_INDEX: &str = "fulltext_index";
pub(crate) const TYPE_BLOOM_FILTER_INDEX: &str = "bloom_filter_index";

/// Evicts local state for a stale index artifact.
///
/// The remote artifact may already be referenced by another region manifest
/// because repartitioned regions share physical SST and index paths. With GC
/// enabled, remote deletion therefore stays in the reference-aware GC lifecycle.
pub(crate) async fn cleanup_stale_index_caches(
    index_id: RegionIndexId,
    access_layer: &AccessLayerRef,
    cache_manager: Option<&CacheManagerRef>,
    write_cache: Option<&WriteCacheRef>,
) {
    if let Some(cache_manager) = cache_manager {
        CacheStrategy::EnableAll(cache_manager.clone())
            .evict_puffin_cache(index_id)
            .await;
    }
    if let Some(write_cache) = write_cache {
        write_cache
            .remove(IndexKey::new(
                index_id.region_id(),
                index_id.file_id(),
                FileType::Puffin(index_id.version),
            ))
            .await;
    }
    if let Err(e) = access_layer
        .puffin_manager_factory()
        .purge_stager(index_id)
        .await
    {
        INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL.inc();
        warn!(
            e;
            "Failed to purge stale index artifact from stager, region: {}, file_id: {}, index_version: {}",
            index_id.region_id(),
            index_id.file_id(),
            index_id.version
        );
    }
}

/// Triggers background download of an index file to the local cache.
pub(crate) fn trigger_index_background_download(
    file_cache: Option<&FileCacheRef>,
    file_id: &RegionIndexId,
    file_size_hint: Option<u64>,
    path_factory: &RegionFilePathFactory,
    object_store: &ObjectStore,
) {
    if let (Some(file_cache), Some(file_size)) = (file_cache, file_size_hint) {
        let index_key = IndexKey::new(
            file_id.region_id(),
            file_id.file_id(),
            FileType::Puffin(file_id.version),
        );
        let remote_path = path_factory.build_index_file_path(file_id.file_id);
        file_cache.maybe_download_background(
            index_key,
            remote_path,
            object_store.clone(),
            file_size,
        );
    }
}

/// Output of the index creation.
#[derive(Debug, Clone, Default)]
pub struct IndexOutput {
    /// Size of the file.
    pub file_size: u64,
    /// Index version.
    pub version: u64,
    /// Inverted index output.
    pub inverted_index: InvertedIndexOutput,
    /// Fulltext index output.
    pub fulltext_index: FulltextIndexOutput,
    /// Columns of `fulltext_index` that use the bloom backend.
    pub fulltext_bloom_columns: Vec<ColumnId>,
    /// Bloom filter output.
    pub bloom_filter: BloomFilterOutput,
}

impl IndexOutput {
    pub fn build_available_indexes(&self) -> SmallVec<[IndexType; 4]> {
        let mut indexes = SmallVec::new();
        if self.inverted_index.is_available() {
            indexes.push(IndexType::InvertedIndex);
        }
        if self.fulltext_index.is_available() {
            indexes.push(IndexType::FulltextIndex);
        }
        if self.bloom_filter.is_available() {
            indexes.push(IndexType::BloomFilterIndex);
        }
        indexes
    }

    pub fn build_indexes(&self) -> Vec<ColumnIndexMetadata> {
        let mut map: HashMap<ColumnId, IndexTypes> = HashMap::new();

        if self.inverted_index.is_available() {
            for &col in &self.inverted_index.columns {
                map.entry(col).or_default().push(IndexType::InvertedIndex);
            }
        }
        if self.fulltext_index.is_available() {
            for &col in &self.fulltext_index.columns {
                map.entry(col).or_default().push(IndexType::FulltextIndex);
            }
        }
        if self.bloom_filter.is_available() {
            for &col in &self.bloom_filter.columns {
                map.entry(col)
                    .or_default()
                    .push(IndexType::BloomFilterIndex);
            }
        }

        map.into_iter()
            .map(|(column_id, created_indexes)| ColumnIndexMetadata {
                column_id,
                fulltext_bloom_version: if self.fulltext_index.is_available()
                    && self.fulltext_bloom_columns.contains(&column_id)
                {
                    FULLTEXT_BLOOM_VERSION
                } else {
                    0
                },
                created_indexes,
            })
            .collect::<Vec<_>>()
    }
}

/// Base output of the index creation.
#[derive(Debug, Clone, Default)]
pub struct IndexBaseOutput {
    /// Size of the index.
    pub index_size: ByteCount,
    /// Number of rows in the index.
    pub row_count: RowCount,
    /// Available columns in the index.
    pub columns: Vec<ColumnId>,
}

impl IndexBaseOutput {
    pub fn is_available(&self) -> bool {
        self.index_size > 0
    }
}

/// Output of the inverted index creation.
pub type InvertedIndexOutput = IndexBaseOutput;
/// Output of the fulltext index creation.
pub type FulltextIndexOutput = IndexBaseOutput;
/// Output of the bloom filter creation.
pub type BloomFilterOutput = IndexBaseOutput;

/// The index creator that hides the error handling details.
#[derive(Default)]
pub struct Indexer {
    file_id: FileId,
    /// The logical region used for metadata and intermediate index files.
    region_id: RegionId,
    /// The region that owns the physical SST and Puffin files.
    physical_region_id: RegionId,
    index_version: u64,
    puffin_manager: Option<SstPuffinManager>,
    write_cache_enabled: bool,
    inverted_indexer: Option<InvertedIndexer>,
    last_mem_inverted_index: usize,
    fulltext_indexer: Option<FulltextIndexer>,
    last_mem_fulltext_index: usize,
    bloom_filter_indexer: Option<BloomFilterIndexer>,
    last_mem_bloom_filter: usize,
    /// Present only when the active creators together need every Dense PK field.
    dense_pk_decoder: Option<DensePrimaryKeyCodec>,
    intermediate_manager: Option<IntermediateManager>,
}

impl Indexer {
    /// Wraps real creators for update-only benchmarks without a Puffin output.
    #[cfg(feature = "testing")]
    pub fn for_bench(
        metadata: &RegionMetadataRef,
        inverted_indexer: Option<InvertedIndexer>,
        bloom_filter_indexer: Option<BloomFilterIndexer>,
    ) -> Self {
        let mut indexer = Self {
            region_id: metadata.region_id,
            inverted_indexer,
            bloom_filter_indexer,
            ..Default::default()
        };
        indexer.prepare_dense_pk_decoder(metadata);
        indexer
    }

    /// Computes demand once from active creators. Field columns and duplicate
    /// requests across creators cannot turn a partial PK projection into a full one.
    fn prepare_dense_pk_decoder(&mut self, metadata: &RegionMetadataRef) {
        if metadata.primary_key_encoding != PrimaryKeyEncoding::Dense
            || metadata.primary_key.is_empty()
        {
            return;
        }
        let requested: HashSet<_> = self
            .inverted_indexer
            .iter()
            .flat_map(|indexer| indexer.column_ids())
            .chain(
                self.bloom_filter_indexer
                    .iter()
                    .flat_map(|indexer| indexer.column_ids()),
            )
            .collect();
        if metadata.primary_key.iter().all(|id| requested.contains(id)) {
            self.dense_pk_decoder = Some(DensePrimaryKeyCodec::new(metadata));
        }
    }

    /// Called before any creator consumes the batch, so the full decode is shared.
    fn prepare_primary_key(&self, batch: &mut Batch) -> Result<()> {
        if let Some(codec) = &self.dense_pk_decoder {
            batch.ensure_dense_pk_decoded(codec)?;
        }
        Ok(())
    }

    /// Updates the index with the given batch.
    pub async fn update(&mut self, batch: &mut Batch) {
        self.do_update(batch).await;

        self.flush_mem_metrics();
    }

    /// Updates the index with the given flat format RecordBatch.
    pub async fn update_flat(&mut self, batch: &RecordBatch) {
        self.do_update_flat(batch).await;

        self.flush_mem_metrics();
    }

    /// Finalizes the index creation.
    pub async fn finish(&mut self) -> IndexOutput {
        let output = self.do_finish().await;

        self.flush_mem_metrics();
        output
    }

    /// Aborts the index creation.
    pub async fn abort(&mut self) {
        self.do_abort().await;

        self.flush_mem_metrics();
    }

    fn flush_mem_metrics(&mut self) {
        let inverted_mem = self
            .inverted_indexer
            .as_ref()
            .map_or(0, |creator| creator.memory_usage());
        INDEX_CREATE_MEMORY_USAGE
            .with_label_values(&[TYPE_INVERTED_INDEX])
            .add(inverted_mem as i64 - self.last_mem_inverted_index as i64);
        self.last_mem_inverted_index = inverted_mem;

        let fulltext_mem = self
            .fulltext_indexer
            .as_ref()
            .map_or(0, |creator| creator.memory_usage());
        INDEX_CREATE_MEMORY_USAGE
            .with_label_values(&[TYPE_FULLTEXT_INDEX])
            .add(fulltext_mem as i64 - self.last_mem_fulltext_index as i64);
        self.last_mem_fulltext_index = fulltext_mem;

        let bloom_filter_mem = self
            .bloom_filter_indexer
            .as_ref()
            .map_or(0, |creator| creator.memory_usage());
        INDEX_CREATE_MEMORY_USAGE
            .with_label_values(&[TYPE_BLOOM_FILTER_INDEX])
            .add(bloom_filter_mem as i64 - self.last_mem_bloom_filter as i64);
        self.last_mem_bloom_filter = bloom_filter_mem;
    }
}

#[async_trait::async_trait]
pub trait IndexerBuilder {
    /// Builds an indexer for the physical SST file.
    async fn build(
        &self,
        region_file_id: RegionFileId,
        index_version: u64,
        row_group_size: Option<usize>,
    ) -> Indexer;
}
#[derive(Clone)]
pub(crate) struct IndexerBuilderImpl {
    pub(crate) build_type: IndexBuildType,
    pub(crate) metadata: RegionMetadataRef,
    pub(crate) puffin_manager: SstPuffinManager,
    pub(crate) write_cache_enabled: bool,
    pub(crate) intermediate_manager: IntermediateManager,
    pub(crate) index_options: IndexOptions,
    pub(crate) inverted_index_config: InvertedIndexConfig,
    pub(crate) fulltext_index_config: FulltextIndexConfig,
    pub(crate) bloom_filter_index_config: BloomFilterConfig,
}

#[async_trait::async_trait]
impl IndexerBuilder for IndexerBuilderImpl {
    /// Sanity check for arguments and create a new [Indexer] if arguments are valid.
    async fn build(
        &self,
        region_file_id: RegionFileId,
        index_version: u64,
        row_group_size: Option<usize>,
    ) -> Indexer {
        let mut indexer = Indexer {
            file_id: region_file_id.file_id(),
            region_id: self.metadata.region_id,
            physical_region_id: region_file_id.region_id(),
            index_version,
            write_cache_enabled: self.write_cache_enabled,
            ..Default::default()
        };

        indexer.inverted_indexer =
            self.build_inverted_indexer(region_file_id.file_id(), row_group_size);
        indexer.fulltext_indexer = self.build_fulltext_indexer(region_file_id.file_id()).await;
        indexer.bloom_filter_indexer = self.build_bloom_filter_indexer(region_file_id.file_id());
        indexer.prepare_dense_pk_decoder(&self.metadata);
        indexer.intermediate_manager = Some(self.intermediate_manager.clone());

        let has_any_indexer = indexer.inverted_indexer.is_some()
            || indexer.fulltext_indexer.is_some()
            || indexer.bloom_filter_indexer.is_some();

        if !has_any_indexer {
            indexer.abort().await;
            return Indexer::default();
        }

        indexer.puffin_manager = Some(self.puffin_manager.clone());
        indexer
    }
}

impl IndexerBuilderImpl {
    fn build_inverted_indexer(
        &self,
        file_id: FileId,
        row_group_size: Option<usize>,
    ) -> Option<InvertedIndexer> {
        let create = match self.build_type {
            IndexBuildType::Flush => self.inverted_index_config.create_on_flush.auto(),
            IndexBuildType::Compact => self.inverted_index_config.create_on_compaction.auto(),
            _ => true,
        };

        if !create {
            debug!(
                "Skip creating inverted index due to config, region_id: {}, file_id: {}",
                self.metadata.region_id, file_id,
            );
            return None;
        }

        let indexed_column_ids = self.metadata.inverted_indexed_column_ids(
            self.index_options.inverted_index.ignore_column_ids.iter(),
        );
        if indexed_column_ids.is_empty() {
            debug!(
                "No columns to be indexed, skip creating inverted index, region_id: {}, file_id: {}",
                self.metadata.region_id, file_id,
            );
            return None;
        }

        let Some(mut segment_row_count) =
            NonZeroUsize::new(self.index_options.inverted_index.segment_row_count)
        else {
            warn!(
                "Segment row count is 0, skip creating index, region_id: {}, file_id: {}",
                self.metadata.region_id, file_id,
            );
            return None;
        };

        // if segment row count not aligned with row group size, adjust it to be aligned.
        if let Some(row_group_size) = row_group_size.and_then(NonZeroUsize::new)
            && row_group_size.get() % segment_row_count.get() != 0
        {
            segment_row_count = row_group_size;
        }

        let indexer = InvertedIndexer::new(
            file_id,
            &self.metadata,
            self.intermediate_manager.clone(),
            self.inverted_index_config.mem_threshold_on_create(),
            segment_row_count,
            indexed_column_ids,
        );

        Some(indexer)
    }

    async fn build_fulltext_indexer(&self, file_id: FileId) -> Option<FulltextIndexer> {
        let create = match self.build_type {
            IndexBuildType::Flush => self.fulltext_index_config.create_on_flush.auto(),
            IndexBuildType::Compact => self.fulltext_index_config.create_on_compaction.auto(),
            _ => true,
        };

        if !create {
            debug!(
                "Skip creating full-text index due to config, region_id: {}, file_id: {}",
                self.metadata.region_id, file_id,
            );
            return None;
        }

        let mem_limit = self.fulltext_index_config.mem_threshold_on_create();
        let creator = FulltextIndexer::new(
            &self.metadata.region_id,
            &file_id,
            &self.intermediate_manager,
            &self.metadata,
            self.fulltext_index_config.compress,
            mem_limit,
        )
        .await;

        let err = match creator {
            Ok(creator) => {
                if creator.is_none() {
                    debug!(
                        "Skip creating full-text index due to no columns require indexing, region_id: {}, file_id: {}",
                        self.metadata.region_id, file_id,
                    );
                }
                return creator;
            }
            Err(err) => err,
        };

        if cfg!(any(test, feature = "test")) {
            panic!(
                "Failed to create full-text indexer, region_id: {}, file_id: {}, err: {:?}",
                self.metadata.region_id, file_id, err
            );
        } else {
            warn!(
                err; "Failed to create full-text indexer, region_id: {}, file_id: {}",
                self.metadata.region_id, file_id,
            );
        }

        None
    }

    fn build_bloom_filter_indexer(&self, file_id: FileId) -> Option<BloomFilterIndexer> {
        let create = match self.build_type {
            IndexBuildType::Flush => self.bloom_filter_index_config.create_on_flush.auto(),
            IndexBuildType::Compact => self.bloom_filter_index_config.create_on_compaction.auto(),
            _ => true,
        };

        if !create {
            debug!(
                "Skip creating bloom filter due to config, region_id: {}, file_id: {}",
                self.metadata.region_id, file_id,
            );
            return None;
        }

        let mem_limit = self.bloom_filter_index_config.mem_threshold_on_create();
        let indexer = BloomFilterIndexer::new(
            file_id,
            &self.metadata,
            self.intermediate_manager.clone(),
            mem_limit,
        );

        let err = match indexer {
            Ok(indexer) => {
                if indexer.is_none() {
                    debug!(
                        "Skip creating bloom filter due to no columns require indexing, region_id: {}, file_id: {}",
                        self.metadata.region_id, file_id,
                    );
                }
                return indexer;
            }
            Err(err) => err,
        };

        if cfg!(any(test, feature = "test")) {
            panic!(
                "Failed to create bloom filter, region_id: {}, file_id: {}, err: {:?}",
                self.metadata.region_id, file_id, err
            );
        } else {
            warn!(
                err; "Failed to create bloom filter, region_id: {}, file_id: {}",
                self.metadata.region_id, file_id,
            );
        }

        None
    }
}

/// Type of an index build task.
#[derive(Debug, Clone, IntoStaticStr, PartialEq)]
pub enum IndexBuildType {
    /// Build index when schema change.
    SchemaChange,
    /// Create or update index after flush.
    Flush,
    /// Create or update index after compact.
    Compact,
    /// Manually build index.
    Manual,
}

impl IndexBuildType {
    fn as_str(&self) -> &'static str {
        self.into()
    }

    // Higher value means higher priority.
    fn priority(&self) -> u8 {
        match self {
            IndexBuildType::Manual => 3,
            IndexBuildType::SchemaChange => 2,
            IndexBuildType::Flush => 1,
            IndexBuildType::Compact => 0,
        }
    }
}

impl From<OperationType> for IndexBuildType {
    fn from(op_type: OperationType) -> Self {
        match op_type {
            OperationType::Flush => IndexBuildType::Flush,
            OperationType::Compact => IndexBuildType::Compact,
        }
    }
}

/// Outcome of an index build task.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum IndexBuildOutcome {
    Finished,
    Aborted(String),
}

/// Mpsc output result sender.
pub type ResultMpscSender = Sender<Result<IndexBuildOutcome>>;

#[derive(Clone)]
pub struct IndexBuildTask {
    /// The logical region whose manifest and in-memory version this task updates.
    pub region_id: RegionId,
    /// The SST file handle to build index for.
    pub file: FileHandle,
    /// The target region metadata used to decode rows from the SST.
    ///
    /// An SST may originate in another region while being visible in the target
    /// manifest. This metadata defines the target schema and sequence domain;
    /// applying the staging manifest only makes imported files visible. Index
    /// rebuild happens later when a flush, compaction, schema change, or manual
    /// index build request schedules it.
    pub(crate) target_region_metadata: RegionMetadataRef,
    /// The manifest state this build is based on.
    pub(crate) source: IndexBuildSource,
    pub reason: IndexBuildType,
    pub access_layer: AccessLayerRef,
    pub(crate) listener: WorkerListener,
    pub(crate) manifest_ctx: ManifestContextRef,
    pub write_cache: Option<WriteCacheRef>,
    pub cache_manager: Option<CacheManagerRef>,
    pub file_purger: FilePurgerRef,
    /// When write cache is enabled, the indexer builder should be built from the write cache.
    /// Otherwise, it should be built from the access layer.
    pub indexer_builder: Arc<dyn IndexerBuilder + Send + Sync>,
    /// Request sender to notify the region worker.
    pub(crate) request_sender: Sender<WorkerRequestWithTime>,
    /// Index build result sender.
    pub(crate) result_sender: ResultMpscSender,
}

impl std::fmt::Debug for IndexBuildTask {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("IndexBuildTask")
            .field("region_id", &self.region_id)
            .field("origin_region_id", &self.source.file_meta.region_id)
            .field("file_id", &self.source.file_meta.file_id)
            .field("schema_version", &self.source.schema_version)
            .field("reason", &self.reason)
            .finish()
    }
}

impl IndexBuildTask {
    /// Notify the caller the job is success.
    pub async fn on_success(&self, outcome: IndexBuildOutcome) {
        let _ = self.result_sender.send(Ok(outcome)).await;
    }

    /// Send index build error to waiter.
    pub async fn on_failure(&self, err: Arc<Error>) {
        let _ = self
            .result_sender
            .send(Err(err.clone()).context(BuildIndexAsyncSnafu {
                region_id: self.region_id,
            }))
            .await;
    }

    fn into_index_build_job(mut self, version_control: VersionControlRef) -> Job {
        Box::pin(async move {
            self.do_index_build(version_control).await;
        })
    }

    async fn do_index_build(&mut self, version_control: VersionControlRef) {
        self.listener
            .on_index_build_begin(RegionFileId::new(
                self.source.file_meta.region_id,
                self.source.file_meta.file_id,
            ))
            .await;
        let result = if self.reason == IndexBuildType::Compact {
            let mut task = self.clone();
            // Keep the scheduler slot occupied until the compact runtime finishes the build.
            match common_runtime::spawn_compact(
                async move { task.index_build(version_control).await },
            )
            .await
            {
                Ok(result) => result,
                Err(err) => Err(err).context(JoinSnafu),
            }
        } else {
            self.index_build(version_control).await
        };
        match result {
            Ok(outcome) => self.on_success(outcome).await,
            Err(e) => {
                warn!(
                    e; "Index build task failed, region: {}, file_id: {}",
                    self.region_id, self.source.file_meta.file_id,
                );
                self.on_failure(e.into()).await
            }
        }
        let worker_request = WorkerRequest::Background {
            region_id: self.region_id,
            notify: BackgroundNotify::IndexBuildStopped(IndexBuildStopped {
                file_id: self.source.file_meta.file_id,
            }),
        };
        let _ = self
            .request_sender
            .send(WorkerRequestWithTime::new(worker_request))
            .await;
    }

    // Checks if the SST file still exists in object store and version to avoid conflict with compaction.
    async fn check_sst_file_exists(&self, version_control: &VersionControlRef) -> bool {
        let file_id = self.source.file_meta.file_id;
        let level = self.source.file_meta.level;
        // We should check current version instead of the version when the job is created.
        let version = version_control.current().version;

        let Some(level_files) = version.ssts.levels().get(level as usize) else {
            warn!(
                "File id {} not found in level {} for index build, region: {}",
                file_id, level, self.region_id
            );
            return false;
        };

        match level_files.files.get(&file_id) {
            Some(handle) if !handle.is_deleted() && !handle.compacting() => {
                // If the file's metadata is present in the current version, the physical SST file
                // is guaranteed to exist on object store. The file purger removes the physical
                // file only after its metadata is removed from the version.
                true
            }
            _ => {
                warn!(
                    "File id {} not found in region version for index build, region: {}",
                    file_id, self.region_id
                );
                false
            }
        }
    }

    async fn index_build(
        &mut self,
        version_control: VersionControlRef,
    ) -> Result<IndexBuildOutcome> {
        let new_index_version = self
            .source
            .file_meta
            .index_version()
            .map_or(0, |version| version + 1);

        // Check SST file existence before building index to avoid failure of parquet reader.
        if !self.check_sst_file_exists(&version_control).await {
            self.listener
                .on_index_build_abort(RegionFileId::new(
                    self.source.file_meta.region_id,
                    self.source.file_meta.file_id,
                ))
                .await;
            return Ok(IndexBuildOutcome::Aborted(format!(
                "SST file not found during index build, region: {}, file_id: {}",
                self.region_id, self.source.file_meta.file_id
            )));
        }

        let mut parquet_reader = self
            .access_layer
            .read_sst(self.file.clone()) // use the latest file handle instead of creating a new one
            .expected_metadata(Some(self.target_region_metadata.clone()))
            .build()
            .await?;

        let row_group_size = parquet_reader.as_ref().and_then(|reader| {
            reader
                .parquet_metadata()
                .row_groups()
                .first()
                .map(|row_group| row_group.num_rows() as usize)
                .filter(|size| *size > 0)
        });

        // Use the same file_id but with new version for index file.
        let region_file_id = RegionFileId::new(
            self.source.file_meta.region_id,
            self.source.file_meta.file_id,
        );
        let index_file_id = region_file_id.file_id();
        let mut indexer = self
            .indexer_builder
            .build(region_file_id, new_index_version, row_group_size)
            .await;

        if let Some(mut parquet_reader) = parquet_reader.take() {
            // TODO(SNC123): optimize index batch
            loop {
                match parquet_reader.next_record_batch().await {
                    Ok(Some(batch)) => {
                        indexer.update_flat(&batch).await;
                    }
                    Ok(None) => break,
                    Err(e) => {
                        indexer.abort().await;
                        return Err(e);
                    }
                }
            }
        }
        let index_output = indexer.finish().await;

        if index_output.file_size > 0 {
            // Check SST file existence again after building index.
            if !self.check_sst_file_exists(&version_control).await {
                // Calls abort to clean up index files.
                indexer.abort().await;
                self.listener
                    .on_index_build_abort(RegionFileId::new(
                        self.source.file_meta.region_id,
                        self.source.file_meta.file_id,
                    ))
                    .await;
                return Ok(IndexBuildOutcome::Aborted(format!(
                    "SST file not found during index build, region: {}, file_id: {}",
                    self.region_id, self.source.file_meta.file_id
                )));
            }

            // Upload index file if write cache is enabled.
            self.maybe_upload_index_file(index_output.clone(), index_file_id, new_index_version)
                .await?;

            self.listener
                .on_index_build_before_manifest_commit(region_file_id)
                .await;

            let worker_request = match self.update_manifest(index_output, new_index_version).await {
                Ok(IndexPublication::Committed {
                    manifest_version,
                    file_meta,
                }) => {
                    self.listener
                        .on_index_build_manifest_committed(region_file_id)
                        .await;
                    let index_build_finished = IndexBuildFinished {
                        manifest_version,
                        file_meta,
                    };
                    WorkerRequest::Background {
                        region_id: self.region_id,
                        notify: BackgroundNotify::IndexBuildFinished(index_build_finished),
                    }
                }
                Ok(IndexPublication::Stale(stale)) => {
                    INDEX_PUBLICATION_STALE_TOTAL
                        .with_label_values(&["manifest_commit"])
                        .inc();
                    let index_id = RegionIndexId::new(region_file_id, new_index_version);
                    // If no successor publishes the same version, GC collects the
                    // remote artifact. Repartition requires GC, while local mode
                    // accepts this narrow orphan window when GC is disabled.
                    cleanup_stale_index_caches(
                        index_id,
                        &self.access_layer,
                        self.cache_manager.as_ref(),
                        self.write_cache.as_ref(),
                    )
                    .await;
                    self.listener.on_index_build_abort(region_file_id).await;
                    if stale == IndexPublicationStale::SchemaChanged {
                        // An index-unrelated schema change also invalidates the
                        // generation fence. Retry to avoid leaving the SST
                        // unindexed; per-SST coalescing limits this to one active
                        // and one pending task, so repeated changes cannot storm.
                        let retry = BuildIndexRequest {
                            region_id: self.region_id,
                            build_type: IndexBuildType::SchemaChange,
                            file_metas: vec![self.source.file_meta.clone()],
                        };
                        let worker_request = WorkerRequest::Background {
                            region_id: self.region_id,
                            notify: BackgroundNotify::IndexBuildRetry(retry),
                        };
                        let _ = self
                            .request_sender
                            .send(WorkerRequestWithTime::new(worker_request))
                            .await;
                    }
                    return Ok(IndexBuildOutcome::Aborted(format!(
                        "Index build source changed before publication, region: {}, file_id: {}",
                        self.region_id, self.source.file_meta.file_id
                    )));
                }
                Err(e) => {
                    let err = Arc::new(e);
                    WorkerRequest::Background {
                        region_id: self.region_id,
                        notify: BackgroundNotify::IndexBuildFailed(IndexBuildFailed { err }),
                    }
                }
            };

            let _ = self
                .request_sender
                .send(WorkerRequestWithTime::new(worker_request))
                .await;
        }
        Ok(IndexBuildOutcome::Finished)
    }

    async fn maybe_upload_index_file(
        &self,
        output: IndexOutput,
        index_file_id: FileId,
        index_version: u64,
    ) -> Result<()> {
        if let Some(write_cache) = &self.write_cache {
            let file_id = self.source.file_meta.file_id;
            let region_id = self.source.file_meta.region_id;
            let remote_store = self.access_layer.object_store();
            let mut upload_tracker = UploadTracker::new(region_id);
            let mut err = None;
            let puffin_key =
                IndexKey::new(region_id, index_file_id, FileType::Puffin(output.version));
            let index_id = RegionIndexId::new(RegionFileId::new(region_id, file_id), index_version);
            let puffin_path = RegionFilePathFactory::new(
                self.access_layer.table_dir().to_string(),
                self.access_layer.path_type(),
            )
            .build_index_file_path_with_version(index_id);
            if let Err(e) = write_cache
                // Index rebuild is background maintenance work, so its uploads
                // are reported as compaction uploads.
                .upload(
                    puffin_key,
                    &puffin_path,
                    remote_store,
                    UploadOptions {
                        op_type: OperationType::Compact,
                        write_buffer_size: crate::sst::DEFAULT_WRITE_BUFFER_SIZE,
                    },
                )
                .await
            {
                err = Some(e);
            }
            upload_tracker.push_uploaded_file(puffin_path);
            if let Some(err) = err {
                // Cleans index files on failure.
                upload_tracker
                    .clean(
                        &smallvec![SstInfo {
                            file_id,
                            index_metadata: output,
                            ..Default::default()
                        }],
                        &write_cache.file_cache(),
                        remote_store,
                    )
                    .await;
                return Err(err);
            }
        } else {
            debug!("write cache is not available, skip uploading index file");
        }
        Ok(())
    }

    async fn update_manifest(
        &self,
        output: IndexOutput,
        new_index_version: u64,
    ) -> Result<IndexPublication> {
        let mut updated = self.source.file_meta.clone();
        updated.available_indexes = output.build_available_indexes();
        updated.indexes = output.build_indexes();
        updated.index_file_size = output.file_size;
        updated.index_version = new_index_version;
        let publication = self
            .manifest_ctx
            .update_manifest_for_index(&self.source, updated)
            .await?;
        if let IndexPublication::Committed {
            manifest_version, ..
        } = &publication
        {
            info!(
                "Successfully update manifest version to {}, region: {}, reason: {}",
                manifest_version,
                self.region_id,
                self.reason.as_str()
            );
        }
        Ok(publication)
    }
}

impl PartialEq for IndexBuildTask {
    fn eq(&self, other: &Self) -> bool {
        self.reason.priority() == other.reason.priority()
    }
}

impl Eq for IndexBuildTask {}

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

impl Ord for IndexBuildTask {
    fn cmp(&self, other: &Self) -> Ordering {
        self.reason.priority().cmp(&other.reason.priority())
    }
}

#[derive(Clone)]
struct PendingIndexBuild {
    task: IndexBuildTask,
    version_control: VersionControlRef,
}

impl PartialEq for PendingIndexBuild {
    fn eq(&self, other: &Self) -> bool {
        self.task == other.task
    }
}

impl Eq for PendingIndexBuild {}

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

impl Ord for PendingIndexBuild {
    fn cmp(&self, other: &Self) -> Ordering {
        self.task.cmp(&other.task)
    }
}

impl PendingIndexBuild {
    /// Returns whether this task should replace another pending task for the
    /// same SST.
    fn supersedes(&self, other: &Self) -> bool {
        match self
            .task
            .source
            .schema_version
            .cmp(&other.task.source.schema_version)
        {
            Ordering::Greater => true,
            Ordering::Less => false,
            Ordering::Equal => match self
                .task
                .source
                .file_meta
                .index_version()
                .cmp(&other.task.source.file_meta.index_version())
            {
                Ordering::Greater => true,
                Ordering::Less => false,
                Ordering::Equal => self.task.reason.priority() > other.task.reason.priority(),
            },
        }
    }
}

#[derive(Default)]
struct PendingIndexBuilds {
    tasks: HashMap<FileId, PendingIndexBuild>,
}

impl PendingIndexBuilds {
    fn insert(&mut self, pending: PendingIndexBuild) -> Option<IndexBuildTask> {
        let file_id = pending.task.source.file_meta.file_id;
        match self.tasks.entry(file_id) {
            std::collections::hash_map::Entry::Vacant(entry) => {
                entry.insert(pending);
                None
            }
            std::collections::hash_map::Entry::Occupied(mut entry)
                if pending.supersedes(entry.get()) =>
            {
                Some(entry.insert(pending).task)
            }
            std::collections::hash_map::Entry::Occupied(_) => Some(pending.task),
        }
    }

    fn highest_ready(&self, building_files: &HashSet<FileId>) -> Option<PendingIndexBuild> {
        self.tasks
            .values()
            .filter(|pending| !building_files.contains(&pending.task.source.file_meta.file_id))
            .max()
            .cloned()
    }

    fn remove(&mut self, file_id: FileId) {
        self.tasks.remove(&file_id);
    }

    fn drain(&mut self) -> impl Iterator<Item = PendingIndexBuild> {
        std::mem::take(&mut self.tasks).into_values()
    }

    #[cfg(test)]
    fn len(&self) -> usize {
        self.tasks.len()
    }

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

/// Tracks the index build status of a region scheduled by the [IndexBuildScheduler].
struct IndexBuildStatus {
    building_files: HashSet<FileId>,
    /// At most one coalesced pending task is kept for each SST.
    pending_tasks: PendingIndexBuilds,
    /// Whether the active builds belong to a region incarnation that has
    /// stopped accepting index publications.
    retiring: bool,
}

impl IndexBuildStatus {
    fn new() -> Self {
        IndexBuildStatus {
            building_files: HashSet::new(),
            pending_tasks: PendingIndexBuilds::default(),
            retiring: false,
        }
    }

    async fn fail_pending(&mut self, err: Arc<Error>) {
        for pending in self.pending_tasks.drain() {
            pending.task.on_failure(err.clone()).await;
        }
    }

    async fn retire(&mut self, err: Arc<Error>) {
        self.retiring = !self.building_files.is_empty();
        self.fail_pending(err).await;
    }
}

pub struct IndexBuildScheduler {
    /// Background job scheduler.
    scheduler: SchedulerRef,
    /// Tracks regions need to build index.
    region_status: HashMap<RegionId, IndexBuildStatus>,
    /// Limit of files allowed to build index concurrently on this worker.
    files_limit: usize,
}

/// Manager background index build tasks of a worker.
impl IndexBuildScheduler {
    pub fn new(scheduler: SchedulerRef, files_limit: usize) -> Self {
        IndexBuildScheduler {
            scheduler,
            region_status: HashMap::new(),
            files_limit,
        }
    }

    pub(crate) async fn schedule_build(
        &mut self,
        version_control: &VersionControlRef,
        task: IndexBuildTask,
    ) -> Result<()> {
        let status = self
            .region_status
            .entry(task.region_id)
            .or_insert_with(IndexBuildStatus::new);

        let file_id = task.source.file_meta.file_id;
        let region_file_id = RegionFileId::new(
            task.source.file_meta.region_id,
            task.source.file_meta.file_id,
        );
        let can_wait_for_active = status.retiring || task.reason == IndexBuildType::SchemaChange;
        let (rejected, coalesced) =
            if status.building_files.contains(&file_id) && !can_wait_for_active {
                (Some(task), None)
            } else {
                let pending = PendingIndexBuild {
                    task,
                    version_control: version_control.clone(),
                };
                (None, status.pending_tasks.insert(pending))
            };
        let should_schedule = !status.retiring;

        if let Some(rejected) = rejected {
            debug!(
                "Rejecting index build because region file {:?} is already being built",
                region_file_id
            );
            rejected
                .on_success(IndexBuildOutcome::Aborted(format!(
                    "Index is already being built for region file {:?}",
                    region_file_id
                )))
                .await;
            rejected.listener.on_index_build_abort(region_file_id).await;
        }

        if let Some(coalesced) = coalesced {
            debug!(
                "Coalescing redundant index build task for region file {:?}",
                region_file_id
            );
            coalesced
                .on_success(IndexBuildOutcome::Aborted(format!(
                    "Index build was coalesced for region file {:?}",
                    region_file_id
                )))
                .await;
            coalesced
                .listener
                .on_index_build_abort(region_file_id)
                .await;
        }

        if should_schedule {
            self.schedule_next_build_batch();
        }
        Ok(())
    }

    /// Schedule tasks until reaching the files limit or no more tasks.
    fn schedule_next_build_batch(&mut self) {
        let mut building_count = self
            .region_status
            .values()
            .map(|status| status.building_files.len())
            .sum::<usize>();

        while building_count < self.files_limit {
            let Some(pending) = self.find_next_task() else {
                break;
            };

            let task = pending.task;
            let region_id = task.region_id;
            let file_id = task.source.file_meta.file_id;
            let job = task.clone().into_index_build_job(pending.version_control);
            match self.scheduler.schedule(job) {
                Ok(()) => {
                    if let Some(status) = self.region_status.get_mut(&region_id) {
                        status.pending_tasks.remove(file_id);
                        status.building_files.insert(file_id);
                        building_count += 1;
                    } else {
                        error!(
                            "Region status not found when scheduling index build task, region: {}",
                            region_id
                        );
                    }
                }
                Err(err) => {
                    error!(
                        err;
                        "Failed to schedule index build job, region: {}, file_id: {}",
                        region_id,
                        file_id
                    );
                    if let Some(status) = self.region_status.get_mut(&region_id) {
                        status.pending_tasks.remove(file_id);
                    }
                    common_runtime::spawn_global(async move {
                        task.on_failure(Arc::new(err)).await;
                    });
                }
            }
        }

        self.region_status.retain(|_, status| {
            !status.building_files.is_empty() || !status.pending_tasks.is_empty()
        });
    }

    /// Find the ready task with the highest priority.
    fn find_next_task(&self) -> Option<PendingIndexBuild> {
        self.region_status
            .values()
            .filter(|status| !status.retiring)
            .filter_map(|status| status.pending_tasks.highest_ready(&status.building_files))
            .max()
    }

    pub(crate) fn on_task_stopped(&mut self, region_id: RegionId, file_id: FileId) {
        if let Some(status) = self.region_status.get_mut(&region_id) {
            if !status.building_files.remove(&file_id) {
                debug!(
                    "Index build task is not tracked as building, region: {}, file: {}",
                    region_id, file_id
                );
                return;
            }
            if status.building_files.is_empty() && status.pending_tasks.is_empty() {
                // No more tasks for this region, remove it.
                self.region_status.remove(&region_id);
            } else if status.building_files.is_empty() {
                // All builds from the previous region incarnation have stopped.
                status.retiring = false;
            }
        }

        self.schedule_next_build_batch();
    }

    pub(crate) async fn on_failure(&mut self, region_id: RegionId, err: Arc<Error>) {
        let Some(status) = self.region_status.get_mut(&region_id) else {
            error!(err; "Index build failed after scheduler state was removed, region: {}", region_id);
            return;
        };
        if status.retiring {
            error!(err; "Index build from a retiring region incarnation failed, region: {}", region_id);
            return;
        }
        error!(
            err; "Index build scheduler encountered failure for region {}, failing pending tasks.",
            region_id
        );
        status.fail_pending(err).await;
        if status.building_files.is_empty() {
            self.region_status.remove(&region_id);
        }
    }

    /// Notifies the scheduler that the region is dropped.
    pub(crate) async fn on_region_dropped(&mut self, region_id: RegionId) {
        self.retire_region(
            region_id,
            Arc::new(RegionDroppedSnafu { region_id }.build()),
        )
        .await;
    }

    /// Notifies the scheduler that the region is closed.
    pub(crate) async fn on_region_closed(&mut self, region_id: RegionId) {
        self.retire_region(region_id, Arc::new(RegionClosedSnafu { region_id }.build()))
            .await;
    }

    /// Notifies the scheduler that the region is truncated.
    pub(crate) async fn on_region_truncated(&mut self, region_id: RegionId) {
        self.retire_region(
            region_id,
            Arc::new(RegionTruncatedSnafu { region_id }.build()),
        )
        .await;
    }

    async fn retire_region(&mut self, region_id: RegionId, err: Arc<Error>) {
        let Some(status) = self.region_status.get_mut(&region_id) else {
            return;
        };
        status.retire(err).await;
        if status.building_files.is_empty() {
            self.region_status.remove(&region_id);
        }
    }
}

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

    use api::v1::SemanticType;
    use common_base::readable_size::ReadableSize;
    use datafusion_common::HashMap;
    use datatypes::data_type::ConcreteDataType;
    use datatypes::schema::{
        ColumnSchema, FulltextOptions, SkippingIndexOptions, SkippingIndexType,
    };
    use datatypes::value::Value;
    use index::inverted_index::format::reader::InvertedIndexReader;
    use object_store::ObjectStore;
    use object_store::services::Memory;
    use partition::expr::col;
    use puffin::puffin_manager::{PuffinManager, PuffinReader};
    use puffin_manager::PuffinManagerFactory;
    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
    use tokio::sync::mpsc;

    use super::*;
    use crate::access_layer::{FilePathProvider, Metrics, SstWriteRequest, WriteType};
    use crate::cache::write_cache::WriteCache;
    use crate::config::{FulltextIndexConfig, IndexBuildMode, MitoConfig, Mode};
    use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList};
    use crate::memtable::time_partition::TimePartitions;
    use crate::region::RegionLeaderState;
    use crate::region::version::{VersionBuilder, VersionControl};
    use crate::schedule::scheduler::{LocalScheduler, Scheduler};
    use crate::sst::file::{FileMeta, RegionFileId};
    use crate::sst::file_purger::NoopFilePurger;
    use crate::sst::location;
    use crate::sst::parquet::WriteOptions;
    use crate::test_util::memtable_util::EmptyMemtableBuilder;
    use crate::test_util::scheduler_util::{SchedulerEnv, VecScheduler};
    use crate::test_util::sst_util::{
        new_flat_source_from_record_batches, new_record_batch_by_range, sst_region_metadata,
    };

    struct MetaConfig {
        with_inverted: bool,
        with_fulltext: bool,
        with_skipping_bloom: bool,
    }

    async fn seed_manifest_file(manifest_ctx: &ManifestContextRef, file_meta: &FileMeta) {
        manifest_ctx
            .update_manifest(
                RegionLeaderState::Writable,
                RegionMetaActionList::with_action(RegionMetaAction::Edit(RegionEdit {
                    files_to_add: vec![file_meta.clone()],
                    files_to_remove: Vec::new(),
                    timestamp_ms: None,
                    flushed_sequence: None,
                    flushed_entry_id: None,
                    committed_sequence: None,
                    compaction_time_window: None,
                })),
                false,
            )
            .await
            .unwrap();
    }

    fn mock_region_metadata(
        MetaConfig {
            with_inverted,
            with_fulltext,
            with_skipping_bloom,
        }: MetaConfig,
    ) -> RegionMetadataRef {
        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
        let mut column_schema = ColumnSchema::new("a", ConcreteDataType::int64_datatype(), false);
        if with_inverted {
            column_schema = column_schema.with_inverted_index(true);
        }
        builder
            .push_column_metadata(ColumnMetadata {
                column_schema,
                semantic_type: SemanticType::Field,
                column_id: 1,
            })
            .push_column_metadata(ColumnMetadata {
                column_schema: ColumnSchema::new("b", ConcreteDataType::float64_datatype(), false),
                semantic_type: SemanticType::Field,
                column_id: 2,
            })
            .push_column_metadata(ColumnMetadata {
                column_schema: ColumnSchema::new(
                    "c",
                    ConcreteDataType::timestamp_millisecond_datatype(),
                    false,
                ),
                semantic_type: SemanticType::Timestamp,
                column_id: 3,
            });

        if with_fulltext {
            let column_schema =
                ColumnSchema::new("text", ConcreteDataType::string_datatype(), true)
                    .with_fulltext_options(FulltextOptions {
                        enable: true,
                        ..Default::default()
                    })
                    .unwrap();

            let column = ColumnMetadata {
                column_schema,
                semantic_type: SemanticType::Field,
                column_id: 4,
            };

            builder.push_column_metadata(column);
        }

        if with_skipping_bloom {
            let column_schema =
                ColumnSchema::new("bloom", ConcreteDataType::string_datatype(), false)
                    .with_skipping_options(SkippingIndexOptions::new_unchecked(
                        42,
                        0.01,
                        SkippingIndexType::BloomFilter,
                    ))
                    .unwrap();

            let column = ColumnMetadata {
                column_schema,
                semantic_type: SemanticType::Field,
                column_id: 5,
            };

            builder.push_column_metadata(column);
        }

        Arc::new(builder.build().unwrap())
    }

    fn mock_object_store() -> ObjectStore {
        ObjectStore::new(Memory::default()).unwrap()
    }

    async fn mock_intm_mgr(path: impl AsRef<str>) -> IntermediateManager {
        IntermediateManager::init_fs(path).await.unwrap()
    }

    fn random_region_file_id() -> RegionFileId {
        RegionFileId::new(RegionId::new(1, 2), FileId::random())
    }

    struct NoopPathProvider;

    impl FilePathProvider for NoopPathProvider {
        fn build_index_file_path(&self, _file_id: RegionFileId) -> String {
            unreachable!()
        }

        fn build_index_file_path_with_version(&self, _index_id: RegionIndexId) -> String {
            unreachable!()
        }

        fn build_sst_file_path(&self, _file_id: RegionFileId) -> String {
            unreachable!()
        }
    }

    async fn mock_sst_file(
        metadata: RegionMetadataRef,
        env: &SchedulerEnv,
        build_mode: IndexBuildMode,
    ) -> SstInfo {
        let source = new_flat_source_from_record_batches(vec![
            new_record_batch_by_range(&["a", "d"], 0, 60),
            new_record_batch_by_range(&["b", "f"], 0, 40),
            new_record_batch_by_range(&["b", "h"], 100, 200),
        ]);
        let mut index_config = MitoConfig::default().index;
        index_config.build_mode = build_mode;
        let write_request = SstWriteRequest {
            op_type: OperationType::Flush,
            metadata: metadata.clone(),
            source,
            max_sequence: None,
            sst_write_format: Default::default(),
            cache_manager: Default::default(),
            preserve_row_sequence: false,
            index_options: IndexOptions::default(),
            index_config,
            inverted_index_config: Default::default(),
            fulltext_index_config: Default::default(),
            bloom_filter_index_config: Default::default(),
        };
        let mut metrics = Metrics::new(WriteType::Flush);
        env.access_layer
            .write_sst(write_request, &WriteOptions::default(), &mut metrics)
            .await
            .unwrap()
            .remove(0)
    }

    async fn mock_version_control(
        metadata: RegionMetadataRef,
        file_purger: FilePurgerRef,
        files: HashMap<FileId, FileMeta>,
    ) -> VersionControlRef {
        let mutable = TimePartitions::new(
            metadata.clone(),
            Arc::new(EmptyMemtableBuilder::default()),
            0,
            None,
        );
        let version_builder = VersionBuilder::new(metadata, mutable)
            .add_files(file_purger, files.values().cloned())
            .build();
        Arc::new(VersionControl::new(version_builder))
    }

    async fn mock_indexer_builder(
        metadata: RegionMetadataRef,
        env: &SchedulerEnv,
    ) -> Arc<dyn IndexerBuilder + Send + Sync> {
        let (dir, factory) = PuffinManagerFactory::new_for_test_async("mock_indexer_builder").await;
        let intm_manager = mock_intm_mgr(dir.path().to_string_lossy()).await;
        let puffin_manager = factory.build(
            env.access_layer.object_store().clone(),
            RegionFilePathFactory::new(
                env.access_layer.table_dir().to_string(),
                env.access_layer.path_type(),
            ),
        );
        Arc::new(IndexerBuilderImpl {
            build_type: IndexBuildType::Flush,
            metadata,
            puffin_manager,
            write_cache_enabled: false,
            intermediate_manager: intm_manager,
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig::default(),
        })
    }

    #[tokio::test]
    async fn test_build_indexer_basic() {
        let (dir, factory) =
            PuffinManagerFactory::new_for_test_async("test_build_indexer_basic_").await;
        let intm_manager = mock_intm_mgr(dir.path().to_string_lossy()).await;

        let metadata = mock_region_metadata(MetaConfig {
            with_inverted: true,
            with_fulltext: true,
            with_skipping_bloom: true,
        });
        let indexer = IndexerBuilderImpl {
            build_type: IndexBuildType::Flush,
            metadata,
            puffin_manager: factory.build(mock_object_store(), NoopPathProvider),
            write_cache_enabled: false,
            intermediate_manager: intm_manager,
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig::default(),
        }
        .build(random_region_file_id(), 0, Some(1024))
        .await;

        assert!(indexer.inverted_indexer.is_some());
        assert!(indexer.fulltext_indexer.is_some());
        assert!(indexer.bloom_filter_indexer.is_some());
    }

    #[tokio::test]
    async fn test_build_indexer_disable_create() {
        let (dir, factory) =
            PuffinManagerFactory::new_for_test_async("test_build_indexer_disable_create_").await;
        let intm_manager = mock_intm_mgr(dir.path().to_string_lossy()).await;

        let metadata = mock_region_metadata(MetaConfig {
            with_inverted: true,
            with_fulltext: true,
            with_skipping_bloom: true,
        });
        let indexer = IndexerBuilderImpl {
            build_type: IndexBuildType::Flush,
            metadata: metadata.clone(),
            puffin_manager: factory.build(mock_object_store(), NoopPathProvider),
            write_cache_enabled: false,
            intermediate_manager: intm_manager.clone(),
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig {
                create_on_flush: Mode::Disable,
                ..Default::default()
            },
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig::default(),
        }
        .build(random_region_file_id(), 0, Some(1024))
        .await;

        assert!(indexer.inverted_indexer.is_none());
        assert!(indexer.fulltext_indexer.is_some());
        assert!(indexer.bloom_filter_indexer.is_some());

        let indexer = IndexerBuilderImpl {
            build_type: IndexBuildType::Compact,
            metadata: metadata.clone(),
            puffin_manager: factory.build(mock_object_store(), NoopPathProvider),
            write_cache_enabled: false,
            intermediate_manager: intm_manager.clone(),
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig {
                create_on_compaction: Mode::Disable,
                ..Default::default()
            },
            bloom_filter_index_config: BloomFilterConfig::default(),
        }
        .build(random_region_file_id(), 0, Some(1024))
        .await;

        assert!(indexer.inverted_indexer.is_some());
        assert!(indexer.fulltext_indexer.is_none());
        assert!(indexer.bloom_filter_indexer.is_some());

        let indexer = IndexerBuilderImpl {
            build_type: IndexBuildType::Compact,
            metadata,
            puffin_manager: factory.build(mock_object_store(), NoopPathProvider),
            write_cache_enabled: false,
            intermediate_manager: intm_manager,
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig {
                create_on_compaction: Mode::Disable,
                ..Default::default()
            },
        }
        .build(random_region_file_id(), 0, Some(1024))
        .await;

        assert!(indexer.inverted_indexer.is_some());
        assert!(indexer.fulltext_indexer.is_some());
        assert!(indexer.bloom_filter_indexer.is_none());
    }

    #[tokio::test]
    async fn test_build_indexer_no_required() {
        let (dir, factory) =
            PuffinManagerFactory::new_for_test_async("test_build_indexer_no_required_").await;
        let intm_manager = mock_intm_mgr(dir.path().to_string_lossy()).await;

        let metadata = mock_region_metadata(MetaConfig {
            with_inverted: false,
            with_fulltext: true,
            with_skipping_bloom: true,
        });
        let indexer = IndexerBuilderImpl {
            build_type: IndexBuildType::Flush,
            metadata: metadata.clone(),
            puffin_manager: factory.build(mock_object_store(), NoopPathProvider),
            write_cache_enabled: false,
            intermediate_manager: intm_manager.clone(),
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig::default(),
        }
        .build(random_region_file_id(), 0, Some(1024))
        .await;

        assert!(indexer.inverted_indexer.is_none());
        assert!(indexer.fulltext_indexer.is_some());
        assert!(indexer.bloom_filter_indexer.is_some());

        let metadata = mock_region_metadata(MetaConfig {
            with_inverted: true,
            with_fulltext: false,
            with_skipping_bloom: true,
        });
        let indexer = IndexerBuilderImpl {
            build_type: IndexBuildType::Flush,
            metadata: metadata.clone(),
            puffin_manager: factory.build(mock_object_store(), NoopPathProvider),
            write_cache_enabled: false,
            intermediate_manager: intm_manager.clone(),
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig::default(),
        }
        .build(random_region_file_id(), 0, Some(1024))
        .await;

        assert!(indexer.inverted_indexer.is_some());
        assert!(indexer.fulltext_indexer.is_none());
        assert!(indexer.bloom_filter_indexer.is_some());

        let metadata = mock_region_metadata(MetaConfig {
            with_inverted: true,
            with_fulltext: true,
            with_skipping_bloom: false,
        });
        let indexer = IndexerBuilderImpl {
            build_type: IndexBuildType::Flush,
            metadata: metadata.clone(),
            puffin_manager: factory.build(mock_object_store(), NoopPathProvider),
            write_cache_enabled: false,
            intermediate_manager: intm_manager,
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig::default(),
        }
        .build(random_region_file_id(), 0, Some(1024))
        .await;

        assert!(indexer.inverted_indexer.is_some());
        assert!(indexer.fulltext_indexer.is_some());
        assert!(indexer.bloom_filter_indexer.is_none());
    }

    #[tokio::test]
    async fn test_build_indexer_zero_row_group_hint() {
        let (dir, factory) =
            PuffinManagerFactory::new_for_test_async("test_build_indexer_zero_row_group_").await;
        let intm_manager = mock_intm_mgr(dir.path().to_string_lossy()).await;

        let metadata = mock_region_metadata(MetaConfig {
            with_inverted: true,
            with_fulltext: true,
            with_skipping_bloom: true,
        });
        let indexer = IndexerBuilderImpl {
            build_type: IndexBuildType::Flush,
            metadata,
            puffin_manager: factory.build(mock_object_store(), NoopPathProvider),
            write_cache_enabled: false,
            intermediate_manager: intm_manager,
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig::default(),
        }
        .build(random_region_file_id(), 0, Some(0))
        .await;

        assert!(indexer.inverted_indexer.is_some());
    }

    #[tokio::test]
    async fn test_index_build_task_sst_not_exist() {
        let env = SchedulerEnv::new().await;
        let (tx, _rx) = mpsc::channel(4);
        let (result_tx, mut result_rx) = mpsc::channel::<Result<IndexBuildOutcome>>(4);
        let mut scheduler = env.mock_index_build_scheduler(4);
        let metadata = Arc::new(sst_region_metadata());
        let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
        let file_purger = Arc::new(NoopFilePurger {});
        let files = HashMap::new();
        let version_control =
            mock_version_control(metadata.clone(), file_purger.clone(), files).await;
        let region_id = metadata.region_id;
        let indexer_builder = mock_indexer_builder(metadata, &env).await;

        let file_meta = FileMeta {
            region_id,
            file_id: FileId::random(),
            file_size: 100,
            ..Default::default()
        };

        let file = FileHandle::new(file_meta.clone(), file_purger.clone());

        // Create mock task.
        let task = IndexBuildTask {
            region_id,
            file,
            target_region_metadata: version_control.current().version.metadata.clone(),
            source: IndexBuildSource::new(
                file_meta,
                version_control.current().version.metadata.schema_version,
            ),
            reason: IndexBuildType::Flush,
            access_layer: env.access_layer.clone(),
            listener: WorkerListener::default(),
            manifest_ctx,
            write_cache: None,
            cache_manager: None,
            file_purger,
            indexer_builder,
            request_sender: tx,
            result_sender: result_tx,
        };

        // Schedule the build task and check result.
        scheduler
            .schedule_build(&version_control, task)
            .await
            .unwrap();
        match result_rx.recv().await.unwrap() {
            Ok(outcome) => {
                if outcome == IndexBuildOutcome::Finished {
                    panic!("Expect aborted result due to missing SST file")
                }
            }
            _ => panic!("Expect aborted result due to missing SST file"),
        }
    }

    #[tokio::test]
    async fn test_index_build_task_foreign_file_uses_target_metadata() {
        let env = SchedulerEnv::new().await;
        let mut scheduler = env.mock_index_build_scheduler(4);
        let source_metadata = Arc::new(sst_region_metadata());
        let mut target_metadata = (*source_metadata).clone();
        target_metadata.region_id = RegionId::new(1, 3);
        let mut target_builder = RegionMetadataBuilder::new(target_metadata.region_id);
        for mut column_metadata in target_metadata.column_metadatas.clone() {
            if column_metadata.column_id == 2 {
                column_metadata.column_schema =
                    column_metadata.column_schema.with_inverted_index(true);
            }
            target_builder.push_column_metadata(column_metadata);
        }
        let partition_expr = col("field_0")
            .gt_eq(Value::UInt64(100))
            .and(col("field_0").lt(Value::UInt64(200)));
        target_builder
            .primary_key(target_metadata.primary_key.clone())
            .partition_expr_json(Some(partition_expr.as_json_str().unwrap()))
            .bump_version();
        let target_metadata = Arc::new(target_builder.build().unwrap());
        let manifest_ctx = env.mock_manifest_context(target_metadata.clone()).await;
        let region_id = target_metadata.region_id;
        let file_purger = Arc::new(NoopFilePurger {});
        let sst_info = mock_sst_file(source_metadata.clone(), &env, IndexBuildMode::Async).await;
        let file_meta = FileMeta {
            region_id: source_metadata.region_id,
            file_id: sst_info.file_id,
            file_size: sst_info.file_size,
            max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size,
            available_indexes: smallvec![IndexType::InvertedIndex],
            // Old manifests may publish an index without recording its size.
            index_file_size: 0,
            index_version: 0,
            num_rows: sst_info.num_rows as u64,
            num_row_groups: sst_info.num_row_groups,
            ..Default::default()
        };
        seed_manifest_file(&manifest_ctx, &file_meta).await;
        let files = HashMap::from([(file_meta.file_id, file_meta.clone())]);
        let version_control =
            mock_version_control(target_metadata.clone(), file_purger.clone(), files).await;
        let indexer_builder = mock_indexer_builder(target_metadata.clone(), &env).await;

        let file = FileHandle::new(file_meta.clone(), file_purger.clone());

        // Create mock task.
        let (tx, mut rx) = mpsc::channel(4);
        let (result_tx, mut result_rx) = mpsc::channel::<Result<IndexBuildOutcome>>(4);
        let task = IndexBuildTask {
            region_id,
            file,
            target_region_metadata: version_control.current().version.metadata.clone(),
            source: IndexBuildSource::new(
                file_meta.clone(),
                version_control.current().version.metadata.schema_version,
            ),
            reason: IndexBuildType::Flush,
            access_layer: env.access_layer.clone(),
            listener: WorkerListener::default(),
            manifest_ctx,
            write_cache: None,
            cache_manager: None,
            file_purger,
            indexer_builder,
            request_sender: tx,
            result_sender: result_tx,
        };

        scheduler
            .schedule_build(&version_control, task)
            .await
            .unwrap();

        // The task should finish successfully.
        match result_rx.recv().await.unwrap() {
            Ok(outcome) => {
                assert_eq!(outcome, IndexBuildOutcome::Finished);
            }
            _ => panic!("Expect finished result"),
        }

        // A notification should be sent to the worker to update the manifest.
        let worker_req = rx.recv().await.unwrap().request;
        match worker_req {
            WorkerRequest::Background {
                region_id: req_region_id,
                notify: BackgroundNotify::IndexBuildFinished(finished),
            } => {
                assert_eq!(req_region_id, region_id);
                let updated_meta = &finished.file_meta;

                // The mock indexer builder creates all index types.
                assert!(!updated_meta.available_indexes.is_empty());
                assert!(updated_meta.index_file_size > 0);
                assert_eq!(updated_meta.file_id, file_meta.file_id);
                assert_eq!(updated_meta.index_version, 1);
                let field_0_index = updated_meta
                    .indexes
                    .iter()
                    .find(|index| index.column_id == 2)
                    .expect("field_0 should have an inverted index");
                assert_eq!(
                    field_0_index.created_indexes.as_slice(),
                    [IndexType::InvertedIndex]
                );
            }
            _ => panic!("Unexpected worker request: {:?}", worker_req),
        }

        let puffin_reader = env
            .access_layer
            .build_puffin_manager()
            .reader(&RegionIndexId::new(
                RegionFileId::new(source_metadata.region_id, file_meta.file_id),
                1,
            ))
            .await
            .unwrap();
        let blob = puffin_reader
            .blob(inverted_index::INDEX_BLOB_TYPE)
            .await
            .unwrap();
        let blob_reader = blob.reader().await.unwrap();
        let index_metadata =
            index::inverted_index::format::reader::InvertedIndexBlobReader::new(blob_reader)
                .metadata(None)
                .await
                .unwrap();
        assert!(index_metadata.metas.contains_key("2"));
        assert_eq!(index_metadata.total_row_count, 100);
    }

    async fn schedule_index_build_task_with_mode(build_mode: IndexBuildMode) {
        let env = SchedulerEnv::new().await;
        let mut scheduler = env.mock_index_build_scheduler(4);
        let metadata = Arc::new(sst_region_metadata());
        let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
        let file_purger = Arc::new(NoopFilePurger {});
        let region_id = metadata.region_id;
        let sst_info = mock_sst_file(metadata.clone(), &env, build_mode.clone()).await;
        let file_meta = FileMeta {
            region_id,
            file_id: sst_info.file_id,
            file_size: sst_info.file_size,
            max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size,
            index_file_size: sst_info.index_metadata.file_size,
            num_rows: sst_info.num_rows as u64,
            num_row_groups: sst_info.num_row_groups,
            ..Default::default()
        };
        seed_manifest_file(&manifest_ctx, &file_meta).await;
        let files = HashMap::from([(file_meta.file_id, file_meta.clone())]);
        let version_control =
            mock_version_control(metadata.clone(), file_purger.clone(), files).await;
        let indexer_builder = mock_indexer_builder(metadata.clone(), &env).await;

        let file = FileHandle::new(file_meta.clone(), file_purger.clone());

        // Create mock task.
        let (tx, _rx) = mpsc::channel(4);
        let (result_tx, mut result_rx) = mpsc::channel::<Result<IndexBuildOutcome>>(4);
        let task = IndexBuildTask {
            region_id,
            file,
            target_region_metadata: version_control.current().version.metadata.clone(),
            source: IndexBuildSource::new(
                file_meta.clone(),
                version_control.current().version.metadata.schema_version,
            ),
            reason: IndexBuildType::Flush,
            access_layer: env.access_layer.clone(),
            listener: WorkerListener::default(),
            manifest_ctx,
            write_cache: None,
            cache_manager: None,
            file_purger,
            indexer_builder,
            request_sender: tx,
            result_sender: result_tx,
        };

        scheduler
            .schedule_build(&version_control, task)
            .await
            .unwrap();

        let puffin_path = location::index_file_path(
            env.access_layer.table_dir(),
            RegionIndexId::new(RegionFileId::new(region_id, file_meta.file_id), 0),
            env.access_layer.path_type(),
        );

        if build_mode == IndexBuildMode::Async {
            // The index file should not exist before the task finishes.
            assert!(
                !env.access_layer
                    .object_store()
                    .exists(&puffin_path)
                    .await
                    .unwrap()
            );
        } else {
            // The index file should exist before the task finishes.
            assert!(
                env.access_layer
                    .object_store()
                    .exists(&puffin_path)
                    .await
                    .unwrap()
            );
        }

        // The task should finish successfully.
        match result_rx.recv().await.unwrap() {
            Ok(outcome) => {
                assert_eq!(outcome, IndexBuildOutcome::Finished);
            }
            _ => panic!("Expect finished result"),
        }

        // The index file should exist after the task finishes.
        assert!(
            env.access_layer
                .object_store()
                .exists(&puffin_path)
                .await
                .unwrap()
        );
    }

    #[tokio::test]
    async fn test_index_build_task_build_mode() {
        schedule_index_build_task_with_mode(IndexBuildMode::Async).await;
        schedule_index_build_task_with_mode(IndexBuildMode::Sync).await;
    }

    #[tokio::test]
    async fn test_index_build_task_no_index() {
        let env = SchedulerEnv::new().await;
        let mut scheduler = env.mock_index_build_scheduler(4);
        let mut metadata = sst_region_metadata();
        // Unset indexes in metadata to simulate no index scenario.
        metadata.column_metadatas.iter_mut().for_each(|col| {
            col.column_schema.set_inverted_index(false);
            let _ = col.column_schema.unset_skipping_options();
        });
        let region_id = metadata.region_id;
        let metadata = Arc::new(metadata);
        let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
        let file_purger = Arc::new(NoopFilePurger {});
        let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await;
        let file_meta = FileMeta {
            region_id,
            file_id: sst_info.file_id,
            file_size: sst_info.file_size,
            max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size,
            index_file_size: sst_info.index_metadata.file_size,
            num_rows: sst_info.num_rows as u64,
            num_row_groups: sst_info.num_row_groups,
            ..Default::default()
        };
        seed_manifest_file(&manifest_ctx, &file_meta).await;
        let files = HashMap::from([(file_meta.file_id, file_meta.clone())]);
        let version_control =
            mock_version_control(metadata.clone(), file_purger.clone(), files).await;
        let indexer_builder = mock_indexer_builder(metadata.clone(), &env).await;

        let file = FileHandle::new(file_meta.clone(), file_purger.clone());

        // Create mock task.
        let (tx, mut rx) = mpsc::channel(4);
        let (result_tx, mut result_rx) = mpsc::channel::<Result<IndexBuildOutcome>>(4);
        let task = IndexBuildTask {
            region_id,
            file,
            target_region_metadata: version_control.current().version.metadata.clone(),
            source: IndexBuildSource::new(
                file_meta.clone(),
                version_control.current().version.metadata.schema_version,
            ),
            reason: IndexBuildType::Flush,
            access_layer: env.access_layer.clone(),
            listener: WorkerListener::default(),
            manifest_ctx,
            write_cache: None,
            cache_manager: None,
            file_purger,
            indexer_builder,
            request_sender: tx,
            result_sender: result_tx,
        };

        scheduler
            .schedule_build(&version_control, task)
            .await
            .unwrap();

        // The task should finish successfully.
        match result_rx.recv().await.unwrap() {
            Ok(outcome) => {
                assert_eq!(outcome, IndexBuildOutcome::Finished);
            }
            _ => panic!("Expect finished result"),
        }

        // No index is built, so no notification should be sent to the worker.
        let _ = rx.recv().await.is_none();
    }

    #[tokio::test]
    async fn test_index_build_task_with_write_cache() {
        let env = SchedulerEnv::new().await;
        let mut scheduler = env.mock_index_build_scheduler(4);
        let metadata = Arc::new(sst_region_metadata());
        let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
        let file_purger = Arc::new(NoopFilePurger {});
        let region_id = metadata.region_id;

        let (dir, factory) = PuffinManagerFactory::new_for_test_async("test_write_cache").await;
        let intm_manager = mock_intm_mgr(dir.path().to_string_lossy()).await;

        // Create mock write cache
        let write_cache = Arc::new(
            WriteCache::new_fs(
                dir.path().to_str().unwrap(),
                ReadableSize::mb(10),
                None,
                None,
                true, // enable_background_worker
                factory,
                intm_manager,
                ReadableSize::mb(10),
            )
            .await
            .unwrap(),
        );
        // Indexer builder built from write cache.
        let indexer_builder = Arc::new(IndexerBuilderImpl {
            build_type: IndexBuildType::Flush,
            metadata: metadata.clone(),
            puffin_manager: write_cache.build_puffin_manager().clone(),
            write_cache_enabled: true,
            intermediate_manager: write_cache.intermediate_manager().clone(),
            index_options: IndexOptions::default(),
            inverted_index_config: InvertedIndexConfig::default(),
            fulltext_index_config: FulltextIndexConfig::default(),
            bloom_filter_index_config: BloomFilterConfig::default(),
        });

        let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await;
        let file_meta = FileMeta {
            region_id,
            file_id: sst_info.file_id,
            file_size: sst_info.file_size,
            index_file_size: sst_info.index_metadata.file_size,
            num_rows: sst_info.num_rows as u64,
            num_row_groups: sst_info.num_row_groups,
            ..Default::default()
        };
        seed_manifest_file(&manifest_ctx, &file_meta).await;
        let files = HashMap::from([(file_meta.file_id, file_meta.clone())]);
        let version_control =
            mock_version_control(metadata.clone(), file_purger.clone(), files).await;

        let file = FileHandle::new(file_meta.clone(), file_purger.clone());

        // Create mock task.
        let (tx, mut _rx) = mpsc::channel(4);
        let (result_tx, mut result_rx) = mpsc::channel::<Result<IndexBuildOutcome>>(4);
        let task = IndexBuildTask {
            region_id,
            file,
            target_region_metadata: version_control.current().version.metadata.clone(),
            source: IndexBuildSource::new(
                file_meta.clone(),
                version_control.current().version.metadata.schema_version,
            ),
            reason: IndexBuildType::Flush,
            access_layer: env.access_layer.clone(),
            listener: WorkerListener::default(),
            manifest_ctx,
            write_cache: Some(write_cache.clone()),
            cache_manager: None,
            file_purger,
            indexer_builder,
            request_sender: tx,
            result_sender: result_tx,
        };

        scheduler
            .schedule_build(&version_control, task)
            .await
            .unwrap();

        // The task should finish successfully.
        match result_rx.recv().await.unwrap() {
            Ok(outcome) => {
                assert_eq!(outcome, IndexBuildOutcome::Finished);
            }
            _ => panic!("Expect finished result"),
        }

        // The write cache should contain the uploaded index file.
        let index_key = IndexKey::new(
            region_id,
            file_meta.file_id,
            FileType::Puffin(sst_info.index_metadata.version),
        );
        assert!(write_cache.file_cache().contains_key(&index_key));
    }

    async fn create_mock_task_for_schedule(
        env: &SchedulerEnv,
        file_id: FileId,
        region_id: RegionId,
        reason: IndexBuildType,
    ) -> IndexBuildTask {
        create_mock_task_for_schedule_with_result(env, file_id, region_id, reason)
            .await
            .0
    }

    /// Like [`create_mock_task_for_schedule`] but also returns the result receiver
    /// so tests can verify pending task cancellation errors.
    async fn create_mock_task_for_schedule_with_result(
        env: &SchedulerEnv,
        file_id: FileId,
        region_id: RegionId,
        reason: IndexBuildType,
    ) -> (IndexBuildTask, mpsc::Receiver<Result<IndexBuildOutcome>>) {
        let metadata = Arc::new(sst_region_metadata());
        let schema_version = metadata.schema_version;
        let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
        let file_purger = Arc::new(NoopFilePurger {});
        let indexer_builder = mock_indexer_builder(metadata.clone(), env).await;
        let (tx, _rx) = mpsc::channel(4);
        let (result_tx, result_rx) = mpsc::channel::<Result<IndexBuildOutcome>>(4);

        let file_meta = FileMeta {
            region_id,
            file_id,
            file_size: 100,
            ..Default::default()
        };

        let file = FileHandle::new(file_meta.clone(), file_purger.clone());

        let task = IndexBuildTask {
            region_id,
            file,
            target_region_metadata: metadata,
            source: IndexBuildSource::new(file_meta, schema_version),
            reason,
            access_layer: env.access_layer.clone(),
            listener: WorkerListener::default(),
            manifest_ctx,
            write_cache: None,
            cache_manager: None,
            file_purger,
            indexer_builder,
            request_sender: tx,
            result_sender: result_tx,
        };
        (task, result_rx)
    }

    #[tokio::test]
    async fn test_scheduler_coalesces_latest_schema_generation_per_sst() {
        let job_scheduler = Arc::new(VecScheduler::default());
        let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone());
        let mut scheduler = env.mock_index_build_scheduler(2);
        let metadata = Arc::new(sst_region_metadata());
        let region_id = metadata.region_id;
        let file_id = FileId::random();
        let file_purger = Arc::new(NoopFilePurger {});
        let files = HashMap::from([(
            file_id,
            FileMeta {
                region_id,
                file_id,
                file_size: 100,
                ..Default::default()
            },
        )]);
        let version_control = mock_version_control(metadata, file_purger, files).await;

        let (active, _active_rx) = create_mock_task_for_schedule_with_result(
            &env,
            file_id,
            region_id,
            IndexBuildType::Manual,
        )
        .await;
        scheduler
            .schedule_build(&version_control, active)
            .await
            .unwrap();
        assert_eq!(job_scheduler.num_jobs(), 1);

        let (mut older, mut older_rx) = create_mock_task_for_schedule_with_result(
            &env,
            file_id,
            region_id,
            IndexBuildType::SchemaChange,
        )
        .await;
        older.source.schema_version = 1;
        scheduler
            .schedule_build(&version_control, older)
            .await
            .unwrap();

        let (mut latest, mut latest_rx) = create_mock_task_for_schedule_with_result(
            &env,
            file_id,
            region_id,
            IndexBuildType::SchemaChange,
        )
        .await;
        latest.source.schema_version = 2;
        scheduler
            .schedule_build(&version_control, latest)
            .await
            .unwrap();

        let replaced = tokio::time::timeout(std::time::Duration::from_secs(5), older_rx.recv())
            .await
            .expect("replaced pending task result sender was not completed")
            .expect("replaced pending task result channel closed");
        assert!(matches!(
            replaced,
            Ok(IndexBuildOutcome::Aborted(reason)) if reason.contains("coalesced")
        ));
        assert!(matches!(
            latest_rx.try_recv(),
            Err(mpsc::error::TryRecvError::Empty)
        ));

        let status = &scheduler.region_status[&region_id];
        assert_eq!(status.building_files.len(), 1);
        assert_eq!(status.pending_tasks.len(), 1);
        assert_eq!(job_scheduler.num_jobs(), 1);

        scheduler.on_task_stopped(region_id, file_id);
        let status = &scheduler.region_status[&region_id];
        assert_eq!(status.building_files.len(), 1);
        assert!(status.pending_tasks.is_empty());
        assert_eq!(job_scheduler.num_jobs(), 2);
    }

    #[tokio::test]
    async fn test_scheduler_completes_sender_when_job_is_rejected() {
        let job_scheduler = Arc::new(LocalScheduler::new(1));
        job_scheduler.stop(false).await.unwrap();
        let env = SchedulerEnv::new().await.scheduler(job_scheduler);
        let mut scheduler = env.mock_index_build_scheduler(1);
        let metadata = Arc::new(sst_region_metadata());
        let region_id = metadata.region_id;
        let file_id = FileId::random();
        let file_purger = Arc::new(NoopFilePurger {});
        let files = HashMap::from([(
            file_id,
            FileMeta {
                region_id,
                file_id,
                file_size: 100,
                ..Default::default()
            },
        )]);
        let version_control = mock_version_control(metadata, file_purger, files).await;
        let (task, mut result_rx) = create_mock_task_for_schedule_with_result(
            &env,
            file_id,
            region_id,
            IndexBuildType::Flush,
        )
        .await;

        scheduler
            .schedule_build(&version_control, task)
            .await
            .unwrap();

        let result = tokio::time::timeout(std::time::Duration::from_secs(5), result_rx.recv())
            .await
            .expect("scheduler rejection did not complete the result sender")
            .expect("result channel closed without a result");
        assert!(result.is_err());
        assert!(!scheduler.region_status.contains_key(&region_id));
    }

    #[tokio::test]
    async fn test_scheduler_comprehensive() {
        let env = SchedulerEnv::new().await;
        let mut scheduler = env.mock_index_build_scheduler(2);
        let metadata = Arc::new(sst_region_metadata());
        let region_id = metadata.region_id;
        let file_purger = Arc::new(NoopFilePurger {});

        // Prepare multiple files for testing
        let file_id1 = FileId::random();
        let file_id2 = FileId::random();
        let file_id3 = FileId::random();
        let file_id4 = FileId::random();
        let file_id5 = FileId::random();

        let mut files = HashMap::new();
        for file_id in [file_id1, file_id2, file_id3, file_id4, file_id5] {
            files.insert(
                file_id,
                FileMeta {
                    region_id,
                    file_id,
                    file_size: 100,
                    ..Default::default()
                },
            );
        }

        let version_control = mock_version_control(metadata, file_purger, files).await;

        // Test 1: Basic scheduling
        let task1 =
            create_mock_task_for_schedule(&env, file_id1, region_id, IndexBuildType::Flush).await;
        assert!(
            scheduler
                .schedule_build(&version_control, task1)
                .await
                .is_ok()
        );
        assert!(scheduler.region_status.contains_key(&region_id));
        let status = scheduler.region_status.get(&region_id).unwrap();
        assert_eq!(status.building_files.len(), 1);
        assert!(status.building_files.contains(&file_id1));

        // Test 2: Duplicate file scheduling (should be skipped)
        let task1_dup =
            create_mock_task_for_schedule(&env, file_id1, region_id, IndexBuildType::Flush).await;
        scheduler
            .schedule_build(&version_control, task1_dup)
            .await
            .unwrap();
        let status = scheduler.region_status.get(&region_id).unwrap();
        assert_eq!(status.building_files.len(), 1); // Still only one

        // Test 3: Fill up to limit (2 building tasks)
        let task2 =
            create_mock_task_for_schedule(&env, file_id2, region_id, IndexBuildType::Flush).await;
        scheduler
            .schedule_build(&version_control, task2)
            .await
            .unwrap();
        let status = scheduler.region_status.get(&region_id).unwrap();
        assert_eq!(status.building_files.len(), 2); // Reached limit
        assert_eq!(status.pending_tasks.len(), 0);

        // Test 4: Add tasks with different priorities to pending queue
        // Now all new tasks will be pending since we reached the limit
        let task3 =
            create_mock_task_for_schedule(&env, file_id3, region_id, IndexBuildType::Compact).await;
        let task4 =
            create_mock_task_for_schedule(&env, file_id4, region_id, IndexBuildType::SchemaChange)
                .await;
        let task5 =
            create_mock_task_for_schedule(&env, file_id5, region_id, IndexBuildType::Manual).await;

        scheduler
            .schedule_build(&version_control, task3)
            .await
            .unwrap();
        scheduler
            .schedule_build(&version_control, task4)
            .await
            .unwrap();
        scheduler
            .schedule_build(&version_control, task5)
            .await
            .unwrap();

        let status = scheduler.region_status.get(&region_id).unwrap();
        assert_eq!(status.building_files.len(), 2); // Still at limit
        assert_eq!(status.pending_tasks.len(), 3); // Three pending

        // Test 5: Task completion triggers scheduling next highest priority task (Manual)
        scheduler.on_task_stopped(region_id, file_id1);
        let status = scheduler.region_status.get(&region_id).unwrap();
        assert!(!status.building_files.contains(&file_id1));
        assert_eq!(status.building_files.len(), 2); // Should schedule next task
        assert_eq!(status.pending_tasks.len(), 2); // One less pending
        // The highest priority task (Manual) should now be building
        assert!(status.building_files.contains(&file_id5));

        // Test 6: Complete another task, should schedule SchemaChange (second highest priority)
        scheduler.on_task_stopped(region_id, file_id2);
        let status = scheduler.region_status.get(&region_id).unwrap();
        assert_eq!(status.building_files.len(), 2);
        assert_eq!(status.pending_tasks.len(), 1); // One less pending
        assert!(status.building_files.contains(&file_id4)); // SchemaChange should be building

        // Test 7: Complete remaining tasks and cleanup
        scheduler.on_task_stopped(region_id, file_id5);
        scheduler.on_task_stopped(region_id, file_id4);

        let status = scheduler.region_status.get(&region_id).unwrap();
        assert_eq!(status.building_files.len(), 1); // Last task (Compact) should be building
        assert_eq!(status.pending_tasks.len(), 0);
        assert!(status.building_files.contains(&file_id3));

        scheduler.on_task_stopped(region_id, file_id3);

        // Region should be removed when all tasks complete
        assert!(!scheduler.region_status.contains_key(&region_id));

        // Test 8: A build failure keeps active leases without retiring the region
        let task6 =
            create_mock_task_for_schedule(&env, file_id1, region_id, IndexBuildType::Flush).await;
        let task7 =
            create_mock_task_for_schedule(&env, file_id2, region_id, IndexBuildType::Flush).await;
        let task8 =
            create_mock_task_for_schedule(&env, file_id3, region_id, IndexBuildType::Manual).await;

        scheduler
            .schedule_build(&version_control, task6)
            .await
            .unwrap();
        scheduler
            .schedule_build(&version_control, task7)
            .await
            .unwrap();
        scheduler
            .schedule_build(&version_control, task8)
            .await
            .unwrap();

        assert!(scheduler.region_status.contains_key(&region_id));
        let status = scheduler.region_status.get(&region_id).unwrap();
        assert_eq!(status.building_files.len(), 2);
        assert_eq!(status.pending_tasks.len(), 1);

        scheduler
            .on_failure(
                region_id,
                Arc::new(
                    crate::error::UnexpectedSnafu {
                        reason: "index build failed".to_string(),
                    }
                    .build(),
                ),
            )
            .await;
        let status = scheduler.region_status.get(&region_id).unwrap();
        assert!(!status.retiring);
        assert_eq!(status.building_files.len(), 2);
        assert!(status.pending_tasks.is_empty());

        scheduler.on_task_stopped(region_id, file_id1);
        assert!(scheduler.region_status.contains_key(&region_id));
        scheduler.on_task_stopped(region_id, file_id2);
        assert!(!scheduler.region_status.contains_key(&region_id));
    }

    /// Helper to set up a scheduler with files_limit=1 and 3 scheduled tasks,
    /// returning the scheduler, the two pending-task result receivers, and the
    /// version control.
    async fn setup_scheduler_with_pending_tasks(
        env: &SchedulerEnv,
    ) -> (
        IndexBuildScheduler,
        mpsc::Receiver<Result<IndexBuildOutcome>>,
        mpsc::Receiver<Result<IndexBuildOutcome>>,
        VersionControlRef,
        RegionId,
        FileId, // building file_id for no-op assertion
    ) {
        let metadata = Arc::new(sst_region_metadata());
        let region_id = metadata.region_id;
        let file_purger = Arc::new(NoopFilePurger {});

        let file_id1 = FileId::random();
        let file_id2 = FileId::random();
        let file_id3 = FileId::random();

        let files = HashMap::from([
            (
                file_id1,
                FileMeta {
                    region_id,
                    file_id: file_id1,
                    file_size: 100,
                    ..Default::default()
                },
            ),
            (
                file_id2,
                FileMeta {
                    region_id,
                    file_id: file_id2,
                    file_size: 100,
                    ..Default::default()
                },
            ),
            (
                file_id3,
                FileMeta {
                    region_id,
                    file_id: file_id3,
                    file_size: 100,
                    ..Default::default()
                },
            ),
        ]);
        let version_control =
            mock_version_control(metadata.clone(), file_purger.clone(), files).await;

        let mut scheduler = env.mock_index_build_scheduler(1);

        // task1 becomes the "building" task (files_limit=1).
        // We intentionally drop its result receiver: the building task's late-stop
        // behavior is covered by the manual on_task_stopped no-op assertion below.
        let (task1, _rx1) = create_mock_task_for_schedule_with_result(
            env,
            file_id1,
            region_id,
            IndexBuildType::Flush,
        )
        .await;
        let (task2, rx2) = create_mock_task_for_schedule_with_result(
            env,
            file_id2,
            region_id,
            IndexBuildType::Flush,
        )
        .await;
        let (task3, rx3) = create_mock_task_for_schedule_with_result(
            env,
            file_id3,
            region_id,
            IndexBuildType::Flush,
        )
        .await;

        scheduler
            .schedule_build(&version_control, task1)
            .await
            .unwrap();
        scheduler
            .schedule_build(&version_control, task2)
            .await
            .unwrap();
        scheduler
            .schedule_build(&version_control, task3)
            .await
            .unwrap();

        // Verify: 1 building + 2 pending.
        assert!(scheduler.region_status.contains_key(&region_id));
        let status = scheduler.region_status.get(&region_id).unwrap();
        assert_eq!(status.building_files.len(), 1);
        assert_eq!(status.pending_tasks.len(), 2);

        (scheduler, rx2, rx3, version_control, region_id, file_id1)
    }

    /// Pattern‑matches a pending‑task cancellation error: outer **must** be
    /// [`crate::error::Error::BuildIndexAsync`] and the inner source **must** be
    /// the lifecycle variant named by `expected_source` (`"dropped"`, `"closed"`,
    /// or `"truncated"`).  Panics with a descriptive message on mismatch.
    fn assert_lifecycle_error(
        err: crate::error::Error,
        expected_source: &str,
        lifecycle_name: &str,
    ) {
        let crate::error::Error::BuildIndexAsync { source, .. } = err else {
            panic!("[{lifecycle_name}] Expected BuildIndexAsync outer error, got: {err:?}");
        };
        let actual = match &*source {
            crate::error::Error::RegionDropped { .. } => "dropped",
            crate::error::Error::RegionClosed { .. } => "closed",
            crate::error::Error::RegionTruncated { .. } => "truncated",
            other => panic!(
                "[{lifecycle_name}] Expected lifecycle source variant (RegionDropped/RegionClosed/RegionTruncated), got: {other:?}"
            ),
        };
        assert_eq!(
            actual, expected_source,
            "[{lifecycle_name}] Source variant mismatch"
        );
    }

    /// Receives a pending‑task error from `rx` with a 5‑second timeout, then
    /// delegates to [`assert_lifecycle_error`].
    async fn recv_lifecycle_error(
        rx: &mut mpsc::Receiver<Result<IndexBuildOutcome>>,
        expected_source: &str,
        lifecycle_name: &str,
    ) {
        let result = tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
            .await
            .unwrap_or_else(|_| {
                panic!(
                    "[{lifecycle_name}] Timeout (5s) waiting for lifecycle error from pending task"
                )
            });
        let result = result.unwrap_or_else(|| {
            panic!("[{lifecycle_name}] Channel closed without receiving lifecycle error")
        });
        let err = result.unwrap_err();
        assert_lifecycle_error(err, expected_source, lifecycle_name);
    }

    #[tokio::test]
    async fn test_scheduler_lifecycle_cleanup() {
        let env = SchedulerEnv::new().await;

        // --- on_region_dropped ---
        {
            let (mut scheduler, mut rx2, mut rx3, _version_control, region_id, building_file_id) =
                setup_scheduler_with_pending_tasks(&env).await;

            scheduler.on_region_dropped(region_id).await;

            let status = scheduler.region_status.get(&region_id).unwrap();
            assert!(status.retiring);
            assert_eq!(status.building_files.len(), 1);
            assert!(status.pending_tasks.is_empty());

            // Pending-task receivers get lifecycle errors (with timeout).
            recv_lifecycle_error(&mut rx2, "dropped", "on_region_dropped").await;
            recv_lifecycle_error(&mut rx3, "dropped", "on_region_dropped").await;

            // The active build keeps its lease until it stops.
            scheduler.on_task_stopped(region_id, building_file_id);
            assert!(!scheduler.region_status.contains_key(&region_id));
        }

        // --- on_region_closed ---
        {
            let (mut scheduler, mut rx2, mut rx3, version_control, region_id, building_file_id) =
                setup_scheduler_with_pending_tasks(&env).await;

            scheduler.on_region_closed(region_id).await;

            let status = scheduler.region_status.get(&region_id).unwrap();
            assert!(status.retiring);
            assert_eq!(status.building_files.len(), 1);
            assert!(status.pending_tasks.is_empty());

            recv_lifecycle_error(&mut rx2, "closed", "on_region_closed").await;
            recv_lifecycle_error(&mut rx3, "closed", "on_region_closed").await;

            // A build scheduled by the reopened region waits behind the old
            // incarnation's active lease.
            let (reopened_task, _reopened_rx) = create_mock_task_for_schedule_with_result(
                &env,
                building_file_id,
                region_id,
                IndexBuildType::Manual,
            )
            .await;
            scheduler
                .schedule_build(&version_control, reopened_task)
                .await
                .unwrap();
            let status = scheduler.region_status.get(&region_id).unwrap();
            assert!(status.retiring);
            assert_eq!(status.building_files.len(), 1);
            assert_eq!(status.pending_tasks.len(), 1);

            // A manifest error from the old incarnation must not discard the
            // reopened task before the old build reports stopped.
            scheduler
                .on_failure(region_id, Arc::new(RegionClosedSnafu { region_id }.build()))
                .await;
            assert_eq!(scheduler.region_status[&region_id].pending_tasks.len(), 1);

            scheduler.on_task_stopped(region_id, building_file_id);
            let status = scheduler.region_status.get(&region_id).unwrap();
            assert!(!status.retiring);
            assert_eq!(status.building_files.len(), 1);
            assert!(status.pending_tasks.is_empty());

            scheduler.on_task_stopped(region_id, building_file_id);
            assert!(!scheduler.region_status.contains_key(&region_id));
        }

        // --- on_region_truncated ---
        {
            let (mut scheduler, mut rx2, mut rx3, _version_control, region_id, building_file_id) =
                setup_scheduler_with_pending_tasks(&env).await;

            scheduler.on_region_truncated(region_id).await;

            let status = scheduler.region_status.get(&region_id).unwrap();
            assert!(status.retiring);
            assert_eq!(status.building_files.len(), 1);
            assert!(status.pending_tasks.is_empty());

            recv_lifecycle_error(&mut rx2, "truncated", "on_region_truncated").await;
            recv_lifecycle_error(&mut rx3, "truncated", "on_region_truncated").await;

            scheduler.on_task_stopped(region_id, building_file_id);
            assert!(!scheduler.region_status.contains_key(&region_id));
        }
    }
}