fix(mito2): build indexes for encoded bulk parts flushed in sync mode [Backport release/v1.3] (#9745)

fix(mito2): build indexes for encoded bulk parts flushed in sync mode (#9722)

* fix(mito2): build indexes for encoded bulk parts flushed in sync mode



* fix(mito2): skip reading the SST when an index build has nothing to create



---------


(cherry picked from commit e0032534d2)

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
Co-authored-by: dennis zhuang <killme2008@gmail.com>

src/mito2/src/engine/index_build_test.rs

@@ -1914,3 +1914,167 @@ async fn test_build_series_index_rejects_non_metric_region() {
     );
     engine.stop().await.unwrap();
 }
+
+/// Counts bytes read from SST files while `COUNTING_SST_READS` is set.
+struct SstReadCounter {
+    inner: object_store::layers::mock::Reader,
+    is_sst: bool,
+    bytes: Arc<AtomicUsize>,
+}
+
+impl object_store::layers::mock::Read for SstReadCounter {
+    async fn open(
+        &self,
+        range: object_store::layers::mock::BytesRange,
+    ) -> object_store::layers::mock::Result<(
+        object_store::layers::mock::RpRead,
+        Box<dyn object_store::layers::mock::ReadStreamDyn>,
+    )> {
+        self.record(&range);
+        self.inner.open_dyn(range).await
+    }
+
+    async fn read(
+        &self,
+        range: object_store::layers::mock::BytesRange,
+    ) -> object_store::layers::mock::Result<(
+        object_store::layers::mock::RpRead,
+        object_store::layers::mock::Buffer,
+    )> {
+        self.record(&range);
+        self.inner.read_dyn(range).await
+    }
+}
+
+impl SstReadCounter {
+    fn record(&self, range: &object_store::layers::mock::BytesRange) {
+        if self.is_sst {
+            self.bytes
+                .fetch_add(range.size().unwrap_or(u64::MAX) as usize, Ordering::Relaxed);
+        }
+    }
+}
+
+#[rstest::rstest]
+#[tokio::test]
+async fn test_sync_flush_indexes_encoded_bulk_parts(
+    #[values(false, true)] disable_inverted_on_flush: bool,
+) {
+    let sst_bytes_read = Arc::new(AtomicUsize::new(0));
+    let counter = sst_bytes_read.clone();
+    let layer = object_store::layers::mock::MockLayerBuilder::default()
+        .reader_factory(Arc::new(move |path, _op, inner| {
+            Box::new(SstReadCounter {
+                inner,
+                is_sst: path.ends_with(".parquet"),
+                bytes: counter.clone(),
+            }) as object_store::layers::mock::Reader
+        }))
+        .build()
+        .unwrap();
+    let mut env = TestEnv::with_prefix("test_sync_flush_indexes_encoded_bulk_parts_")
+        .await
+        .with_mock_layer(layer);
+    let listener = Arc::new(IndexBuildListener::default());
+    let mut config = MitoConfig {
+        default_flat_format: true,
+        min_compaction_interval: Duration::from_secs(3600),
+        ..Default::default()
+    };
+    assert_eq!(config.index.build_mode, IndexBuildMode::Sync);
+    if disable_inverted_on_flush {
+        config.inverted_index.create_on_flush = Mode::Disable;
+    }
+    let engine = env
+        .create_engine_with(config, None, Some(listener.clone()), None)
+        .await;
+
+    let region_id = RegionId::new(1, 1);
+    env.get_schema_metadata_manager()
+        .register_region_table_info(
+            region_id.table_id(),
+            "test_table",
+            "test_catalog",
+            "test_schema",
+            None,
+            env.get_kv_backend(),
+        )
+        .await;
+    let mut request = CreateRequestBuilder::new()
+        .insert_option("memtable.type", "bulk")
+        .insert_option("memtable.bulk.merge_threshold", "2")
+        .insert_option("memtable.bulk.encode_row_threshold", "1")
+        .build_with_index();
+    if !disable_inverted_on_flush {
+        // `build_with_index` only adds inverted indexes on tags; also cover a skipping index.
+        let field = request
+            .column_metadatas
+            .iter_mut()
+            .find(|c| c.column_schema.name == "field_0")
+            .unwrap();
+        field.column_schema = field
+            .column_schema
+            .clone()
+            .with_skipping_options(SkippingIndexOptions::new_unchecked(
+                1024,
+                0.01,
+                SkippingIndexType::BloomFilter,
+            ))
+            .unwrap();
+    }
+    let column_schemas = rows_schema(&request);
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+
+    // Each write is large enough to become its own bulk part, and the SST is larger than the
+    // footer prefetch so reading the rows is distinguishable from reading the metadata.
+    for (start, end) in [(0, 100_000), (100_000, 200_000)] {
+        let rows = Rows {
+            schema: column_schemas.clone(),
+            rows: build_rows(start, end),
+        };
+        put_rows(&engine, region_id, rows).await;
+    }
+    // Merges the two parts into an encoded part, which flush uploads as-is.
+    let version = engine.get_region(region_id).unwrap().version();
+    for memtable in version.memtables.list_memtables() {
+        memtable.compact(false).unwrap();
+    }
+    flush_region(&engine, region_id, None).await;
+
+    let files = current_file_metas(&engine, region_id).await;
+    if disable_inverted_on_flush {
+        // Only inverted indexes are configured and they are disabled on flush, so the builds
+        // have nothing to create. They must stop before reading the rows of the SSTs.
+        let total_size: usize = files.iter().map(|f| f.file_size as usize).sum();
+        // Stopping the engine waits for scheduled index builds.
+        engine.stop().await.unwrap();
+        assert!(listener.begin_count() > 0);
+        assert!(files.iter().all(|f| f.available_indexes.is_empty()));
+        let read = sst_bytes_read.load(Ordering::Relaxed);
+        assert!(read < total_size / 2, "read {read} of {total_size} bytes");
+        return;
+    }
+
+    let metadata = engine.get_region(region_id).unwrap().metadata();
+    let indexed = async {
+        loop {
+            let files = current_file_metas(&engine, region_id).await;
+            if files
+                .iter()
+                .all(|f| f.is_index_consistent_with_region(&metadata.column_metadatas))
+            {
+                return files;
+            }
+            tokio::time::sleep(Duration::from_millis(20)).await;
+        }
+    };
+    let files = tokio::time::timeout(Duration::from_secs(10), indexed)
+        .await
+        .expect("flushed encoded parts have no index");
+    assert!(listener.begin_count() > 0);
+    assert!(files.iter().all(|f| f.inverted_index_available()));
+    assert!(files.iter().all(|f| f.bloom_filter_index_available()));
+}

src/mito2/src/flush.rs

@@ -29,7 +29,7 @@ use smallvec::{SmallVec, smallvec};
 use snafu::{ResultExt, ensure};
 use store_api::metadata::RegionMetadataRef;
 use store_api::region_request::RegionFlushReason;
-use store_api::storage::{RegionId, SequenceNumber};
+use store_api::storage::{FileId, RegionId, SequenceNumber};
 use strum::IntoStaticStr;
 use tokio::sync::{Semaphore, mpsc, watch};
 
@@ -391,7 +391,10 @@ impl RegionFlushTask {
         };
 
         let worker_request = match flush_result {
-            Ok(edit) => {
+            Ok(FlushedEdit {
+                edit,
+                unindexed_files,
+            }) => {
                 let memtables_to_remove = version_data
                     .version
                     .memtables
@@ -407,6 +410,7 @@ impl RegionFlushTask {
                     senders: std::mem::take(&mut self.senders),
                     _timer: timer,
                     edit,
+                    unindexed_files,
                     memtables_to_remove,
                     is_staging: self.is_staging,
                 };
@@ -444,7 +448,7 @@ impl RegionFlushTask {
         version_data: &VersionControlData,
         state: &CancellableTaskState,
         uncommitted: &UncommittedSsts,
-    ) -> Result<RegionEdit> {
+    ) -> Result<FlushedEdit> {
         // We must use the immutable memtables list and entry ids from the `version_data`
         // for consistency as others might already modify the version in the `version_control`.
         let version = &version_data.version;
@@ -463,6 +467,7 @@ impl RegionFlushTask {
 
         let DoFlushMemtablesResult {
             file_metas,
+            unindexed_files,
             flushed_bytes,
             series_count,
             encoded_part_count,
@@ -575,7 +580,10 @@ impl RegionFlushTask {
             self.reason.as_str()
         );
 
-        Ok(edit)
+        Ok(FlushedEdit {
+            edit,
+            unindexed_files,
+        })
     }
 
     async fn do_flush_memtables(
@@ -590,6 +598,7 @@ impl RegionFlushTask {
         let mut flushed_bytes = 0;
         let mut series_count = 0;
         let mut encoded_part_count = 0;
+        let mut unindexed_files = Vec::new();
         let mut flush_metrics = Metrics::new(WriteType::Flush);
         let partition_expr = parse_partition_expr(self.partition_expr.as_deref())?;
         let hook = self.manifest_ctx.hook();
@@ -636,7 +645,12 @@ impl RegionFlushTask {
                 .await?;
             encoded_part_count += num_encoded;
             for (source_idx, result) in results.into_iter().enumerate() {
-                let (max_sequence, ssts_written, metrics) = result?;
+                let FlushedSsts {
+                    max_sequence,
+                    ssts: ssts_written,
+                    metrics,
+                    indexed,
+                } = result?;
                 if ssts_written.is_empty() {
                     // No data written.
                     continue;
@@ -655,6 +669,9 @@ impl RegionFlushTask {
 
                 for sst_info in &ssts_written {
                     flushed_bytes += sst_info.file_size;
+                    if !indexed {
+                        unindexed_files.push(sst_info.file_id);
+                    }
                     let pk_range = sst_info
                         .file_metadata
                         .as_ref()
@@ -688,6 +705,7 @@ impl RegionFlushTask {
 
         Ok(DoFlushMemtablesResult {
             file_metas,
+            unindexed_files,
             flushed_bytes,
             series_count,
             encoded_part_count,
@@ -733,7 +751,12 @@ impl RegionFlushTask {
                     .await?;
                 uncommitted.track(&ssts);
                 FLUSH_FILE_TOTAL.inc_by(ssts.len() as u64);
-                Ok((max_sequence, ssts, metrics))
+                Ok(FlushedSsts {
+                    max_sequence,
+                    ssts,
+                    metrics,
+                    indexed: true,
+                })
             });
             tasks.push(task);
         }
@@ -756,7 +779,13 @@ impl RegionFlushTask {
                     .await?;
                 uncommitted.track(std::slice::from_ref(&encoded.sst_info));
                 FLUSH_FILE_TOTAL.inc();
-                Ok((max_sequence, smallvec![encoded.sst_info], metrics))
+                // Encoded parts are uploaded as-is, without passing through the indexer.
+                Ok(FlushedSsts {
+                    max_sequence,
+                    ssts: smallvec![encoded.sst_info],
+                    metrics,
+                    indexed: false,
+                })
             });
             tasks.push(task);
         }
@@ -892,14 +921,30 @@ impl RegionFlushTask {
     }
 }
 
+struct FlushedEdit {
+    edit: RegionEdit,
+    /// Files in `edit` written without building their indexes.
+    unindexed_files: Vec<FileId>,
+}
+
+struct FlushedSsts {
+    max_sequence: SequenceNumber,
+    ssts: SstInfoArray,
+    metrics: Metrics,
+    /// Whether the indexer ran while writing `ssts`.
+    indexed: bool,
+}
+
 struct FlushFlatMemResult {
     num_encoded: usize,
     num_sources: usize,
-    results: Vec<Result<(SequenceNumber, SstInfoArray, Metrics)>>,
+    results: Vec<Result<FlushedSsts>>,
 }
 
 struct DoFlushMemtablesResult {
     file_metas: Vec<FileMeta>,
+    /// Files written without building their indexes.
+    unindexed_files: Vec<FileId>,
     flushed_bytes: u64,
     series_count: usize,
     encoded_part_count: usize,
@@ -1909,6 +1954,7 @@ mod tests {
                     flushed_sequence: None,
                     committed_sequence: None,
                 },
+                unindexed_files: Vec::new(),
                 memtables_to_remove: smallvec![],
                 is_staging: false,
             }),

src/mito2/src/request.rs

@@ -948,6 +948,8 @@ pub(crate) struct FlushFinished {
     pub(crate) _timer: HistogramTimer,
     /// Region edit to apply.
     pub(crate) edit: RegionEdit,
+    /// Files in `edit` written without building their indexes.
+    pub(crate) unindexed_files: Vec<FileId>,
     /// Memtables to remove.
     pub(crate) memtables_to_remove: SmallVec<[MemtableId; 2]>,
     /// Whether the region is in staging mode.

src/mito2/src/sst/index.rs

@@ -266,6 +266,13 @@ pub struct Indexer {
 }
 
 impl Indexer {
+    /// Returns true if no index creator is active.
+    pub(crate) fn is_empty(&self) -> bool {
+        self.inverted_indexer.is_none()
+            && self.fulltext_indexer.is_none()
+            && self.bloom_filter_indexer.is_none()
+    }
+
     /// Wraps real creators for update-only benchmarks without a Puffin output.
     #[cfg(feature = "testing")]
     pub fn for_bench(
@@ -421,11 +428,7 @@ impl IndexerBuilder for IndexerBuilderImpl {
         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 {
+        if indexer.is_empty() {
             indexer.abort().await;
             return Indexer::default();
         }
@@ -840,6 +843,11 @@ impl IndexBuildTask {
             .indexer_builder
             .build(region_file_id, new_index_version, row_group_size)
             .await;
+        // The build type or index options may leave nothing to build. Return before reading
+        // any row; an empty indexer finishes with no output, which publishes nothing.
+        if indexer.is_empty() {
+            return Ok(IndexBuildOutcome::Finished);
+        }
 
         if let Some(mut parquet_reader) = parquet_reader.take() {
             // TODO(SNC123): optimize index batch

src/mito2/src/worker/handle_flush.rs

@@ -436,10 +436,18 @@ impl<S: LogStore> RegionWorkerLoop<S> {
 
         let flush_on_close = request.flush_reason == FlushReason::Closing;
 
-        let index_build_file_metas = std::mem::take(&mut request.edit.files_to_add);
+        let mut index_build_file_metas = std::mem::take(&mut request.edit.files_to_add);
+        if self.config.index.build_mode == IndexBuildMode::Sync {
+            // Sync flush builds indexes while writing SSTs, except for encoded bulk parts,
+            // which are uploaded as-is. Build the missing indexes for those files here.
+            let metadata = region.metadata();
+            index_build_file_metas.retain(|meta| {
+                request.unindexed_files.contains(&meta.file_id)
+                    && !meta.is_index_consistent_with_region(&metadata.column_metadatas)
+            });
+        }
 
-        // In async mode, create indexes after flush.
-        if self.config.index.build_mode == IndexBuildMode::Async {
+        if !index_build_file_metas.is_empty() {
             self.handle_rebuild_index(
                 BuildIndexRequest {
                     region_id,