From 5b4efd908ca03c8f5243d308d96146365600c73c Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 27 Oct 2026 10:49:16 +0000 Subject: [PATCH] perf(mito2): index encoded bulk parts from memory during flush (#9913) * perf(mito2): build encoded bulk part indexes from in-memory bytes during flush Sync flush uploaded encoded bulk parts without indexes and rebuilt them after the flush, which read every encoded SST back from the object store and needed a second manifest commit. Build the indexes from the encoded bytes still in memory before the flush commits, the same way the regular SST writer does, and drop the post-flush rebuild for these files. Signed-off-by: Dennis Zhuang * fix(mito2): keep flush going when indexing encoded bytes fails Match ParquetWriter: a failure while building indexes from the encoded bytes logs and leaves the SST unindexed instead of failing the flush after the SST was uploaded. Signed-off-by: Dennis Zhuang * fix(mito2): own encoded SSTs for cleanup before uploading them Index creation now runs after an encoded SST is uploaded, so a truncate or drop could cancel the flush after the upload but before UncommittedSsts tracked the file, leaving the SST and its index behind. Track the file before the upload, and fill the index metadata on the caller's SstInfo instead of a partial one. Signed-off-by: Dennis Zhuang * refactor(mito2): name encoded SST index build params explicitly Signed-off-by: Dennis Zhuang * fix(mito2): drop index metadata when the puffin file fails and size index decode batches Indexer::do_finish kept the per-index metadata when finishing the puffin file failed, so callers could publish indexes that were not written. Return an empty output instead. Indexing encoded bulk parts now decodes with the batch size estimated from the largest row group, the same estimate bulk merges use, instead of a fixed 8192 rows. Signed-off-by: Dennis Zhuang * docs(config): document the mito index build_mode option Signed-off-by: Dennis Zhuang * fix(mito2): decode encoded SSTs for indexing on the blocking pool Indexing encoded bulk parts decoded Parquet batches inside the flush task on the global async runtime. Decode each batch with spawn_blocking_global and keep the indexer update on the async task, with one batch in flight. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- config/config.md | 2 + config/datanode.example.toml | 6 + config/standalone.example.toml | 6 + src/mito2/src/access_layer.rs | 69 ++++++- src/mito2/src/batch_size.rs | 19 ++ src/mito2/src/cache/write_cache.rs | 79 ++++++-- src/mito2/src/engine/index_build_test.rs | 214 +++++++++++++++++++--- src/mito2/src/flush.rs | 86 +++------ src/mito2/src/memtable/bulk.rs | 17 +- src/mito2/src/request.rs | 2 - src/mito2/src/sst/file.rs | 5 + src/mito2/src/sst/index.rs | 82 ++++++++- src/mito2/src/sst/index/indexer/finish.rs | 5 + src/mito2/src/worker/handle_flush.rs | 17 +- 14 files changed, 468 insertions(+), 141 deletions(-) diff --git a/config/config.md b/config/config.md index f04689d97f6..15cd4915f9d 100644 --- a/config/config.md +++ b/config/config.md @@ -225,6 +225,7 @@ | `region_engine.mito.index.aux_path` | String | `""` | Auxiliary directory path for the index in filesystem, used to store intermediate files for
creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.
The default name for this directory is `index_intermediate` for backward compatibility.

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | | `region_engine.mito.index.staging_ttl` | String | `7d` | The TTL of the staging directory.
Defaults to 7 days.
Setting it to "0s" to disable TTL. | +| `region_engine.mito.index.build_mode` | String | `sync` | When to build indexes for new SST files.
- `sync`: build indexes while writing SST files during flush and compaction.
- `async`: build indexes in the background after the SST files are written. Queries read
these files without indexes until the build finishes. | | `region_engine.mito.index.metadata_cache_size` | String | `64MiB` | Cache size for inverted index metadata. | | `region_engine.mito.index.content_cache_size` | String | `128MiB` | Cache size for inverted index content. | | `region_engine.mito.index.content_cache_page_size` | String | `64KiB` | Page size for inverted index content cache. | @@ -703,6 +704,7 @@ | `region_engine.mito.index.aux_path` | String | `""` | Auxiliary directory path for the index in filesystem, used to store intermediate files for
creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.
The default name for this directory is `index_intermediate` for backward compatibility.

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | | `region_engine.mito.index.staging_ttl` | String | `7d` | The TTL of the staging directory.
Defaults to 7 days.
Setting it to "0s" to disable TTL. | +| `region_engine.mito.index.build_mode` | String | `sync` | When to build indexes for new SST files.
- `sync`: build indexes while writing SST files during flush and compaction.
- `async`: build indexes in the background after the SST files are written. Queries read
these files without indexes until the build finishes. | | `region_engine.mito.index.metadata_cache_size` | String | `64MiB` | Cache size for inverted index metadata. | | `region_engine.mito.index.content_cache_size` | String | `128MiB` | Cache size for inverted index content. | | `region_engine.mito.index.content_cache_page_size` | String | `64KiB` | Page size for inverted index content cache. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index f6e9a2571f8..acc65582e81 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -734,6 +734,12 @@ staging_size = "2GB" ## Setting it to "0s" to disable TTL. staging_ttl = "7d" +## When to build indexes for new SST files. +## - `sync`: build indexes while writing SST files during flush and compaction. +## - `async`: build indexes in the background after the SST files are written. Queries read +## these files without indexes until the build finishes. +build_mode = "sync" + ## Cache size for inverted index metadata. metadata_cache_size = "64MiB" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index f569ff1b952..5cfbb95729e 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -949,6 +949,12 @@ staging_size = "2GB" ## Setting it to "0s" to disable TTL. staging_ttl = "7d" +## When to build indexes for new SST files. +## - `sync`: build indexes while writing SST files during flush and compaction. +## - `async`: build indexes in the background after the SST files are written. Queries read +## these files without indexes until the build finishes. +build_mode = "sync" + ## Cache size for inverted index metadata. metadata_cache_size = "64MiB" diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index b1457935c16..3067c6098a1 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -44,9 +44,9 @@ use crate::metrics::{COMPACTION_STAGE_ELAPSED, FLUSH_ELAPSED}; use crate::read::FlatSource; use crate::region::options::IndexOptions; use crate::sst::file::{FileHandle, RegionFileId, RegionIndexId}; -use crate::sst::index::IndexerBuilderImpl; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager}; +use crate::sst::index::{IndexerBuilderImpl, build_index_from_parquet_bytes}; use crate::sst::location::{self, region_dir_from_table_dir}; use crate::sst::parquet::reader::ParquetReaderBuilder; use crate::sst::parquet::writer::ParquetWriter; @@ -477,13 +477,19 @@ impl AccessLayer { } /// Puts encoded SST bytes to the write cache (if enabled) and uploads it to the object store. + /// + /// Builds the SST's indexes from `data` into `sst_info.index_metadata` when + /// `index_build_params` is set. pub(crate) async fn put_sst( &self, data: &bytes::Bytes, - region_file_id: RegionFileId, + region_id: RegionId, + sst_info: &mut SstInfo, cache_manager: &CacheManagerRef, write_buffer_size: ReadableSize, + index_build_params: Option, ) -> Result { + let region_file_id = RegionFileId::new(region_id, sst_info.file_id); if let Some(write_cache) = cache_manager.write_cache() { // Write to cache and upload to remote store let upload_request = SstUploadRequest { @@ -494,7 +500,14 @@ impl AccessLayer { remote_store: self.object_store.clone(), }; write_cache - .put_and_upload_sst(data, region_file_id, upload_request, write_buffer_size) + .put_and_upload_sst( + data, + region_id, + sst_info, + upload_request, + write_buffer_size, + index_build_params, + ) .await } else { let start = Instant::now(); @@ -519,6 +532,19 @@ impl AccessLayer { } let mut metrics = Metrics::new(WriteType::Flush); metrics.write_batch = start.elapsed(); + + if let Some(params) = index_build_params { + let start = Instant::now(); + let builder = params.indexer_builder( + self.puffin_manager_factory + .build(self.object_store.clone(), path_provider), + false, + self.intermediate_manager.clone(), + ); + sst_info.index_metadata = + build_index_from_parquet_bytes(&builder, region_file_id, data).await; + metrics.update_index = start.elapsed(); + } Ok(metrics) } } @@ -593,6 +619,36 @@ pub struct SstWriteRequest { pub bloom_filter_index_config: BloomFilterConfig, } +/// Index inputs for SST bytes that were encoded outside [ParquetWriter]. +pub(crate) struct EncodedSstIndexBuildParams { + pub(crate) metadata: RegionMetadataRef, + pub(crate) index_options: IndexOptions, + pub(crate) inverted_index_config: InvertedIndexConfig, + pub(crate) fulltext_index_config: FulltextIndexConfig, + pub(crate) bloom_filter_index_config: BloomFilterConfig, +} + +impl EncodedSstIndexBuildParams { + pub(crate) fn indexer_builder( + self, + puffin_manager: SstPuffinManager, + write_cache_enabled: bool, + intermediate_manager: IntermediateManager, + ) -> IndexerBuilderImpl { + IndexerBuilderImpl { + build_type: OperationType::Flush.into(), + metadata: self.metadata, + puffin_manager, + write_cache_enabled, + intermediate_manager, + index_options: self.index_options, + inverted_index_config: self.inverted_index_config, + fulltext_index_config: self.fulltext_index_config, + bloom_filter_index_config: self.bloom_filter_index_config, + } + } +} + /// Cleaner to remove temp files on the atomic write dir. pub(crate) struct TempFileCleaner { region_id: RegionId, @@ -830,9 +886,14 @@ mod tests { access_layer .put_sst( &encoded, - file_id, + file_id.region_id(), + &mut SstInfo { + file_id: file_id.file_id(), + ..Default::default() + }, &cache_manager, ReadableSize(chunk_size as u64), + None, ) .await .unwrap(); diff --git a/src/mito2/src/batch_size.rs b/src/mito2/src/batch_size.rs index 92fdd75673f..b71265ca7ee 100644 --- a/src/mito2/src/batch_size.rs +++ b/src/mito2/src/batch_size.rs @@ -14,11 +14,30 @@ //! Utilities for choosing an execution batch size from source statistics. +use parquet::file::metadata::ParquetMetaData; + use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; /// Approximate decoded bytes to produce in one batch. pub(crate) const TARGET_BATCH_BYTES: usize = 64 * 1024 * 1024; +/// Returns `(num_rows, uncompressed_bytes)` of the largest row group in a Parquet file, which +/// bounds the decoded size of a batch read from it. +pub(crate) fn largest_row_group_statistic(metadata: &ParquetMetaData) -> Option<(u64, u64)> { + metadata + .row_groups() + .iter() + .map(|row_group| { + let uncompressed_bytes = row_group + .columns() + .iter() + .map(|column| column.uncompressed_size() as u64) + .sum(); + (row_group.num_rows() as u64, uncompressed_bytes) + }) + .max_by_key(|(_, uncompressed_bytes)| *uncompressed_bytes) +} + /// Estimates a batch size from `(num_rows, estimated_bytes)` pairs. /// /// Uses the widest source with valid statistics and targets approximately 64 MiB of decoded data. diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index afe943eef50..24e2e957219 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -26,8 +26,9 @@ use store_api::storage::RegionId; use tokio::sync::mpsc::{UnboundedSender, unbounded_channel}; use crate::access_layer::{ - FilePathProvider, Metrics, OperationType, RegionFilePathFactory, SstInfoArray, SstWriteRequest, - TempFileCleaner, WriteCachePathProvider, WriteType, new_fs_cache_store, + EncodedSstIndexBuildParams, FilePathProvider, Metrics, OperationType, RegionFilePathFactory, + SstInfoArray, SstWriteRequest, TempFileCleaner, WriteCachePathProvider, WriteType, + new_fs_cache_store, }; use crate::cache::file_cache::{FileCache, FileCacheRef, FileType, IndexKey, IndexValue}; use crate::cache::manifest_cache::ManifestCache; @@ -35,11 +36,11 @@ use crate::error::{self, Result}; use crate::metrics::UPLOAD_BYTES_TOTAL; use crate::region::opener::RegionLoadCacheTask; use crate::sst::file::RegionFileId; -use crate::sst::index::IndexerBuilderImpl; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager}; -use crate::sst::parquet::WriteOptions; +use crate::sst::index::{IndexerBuilderImpl, build_index_from_parquet_bytes}; use crate::sst::parquet::writer::ParquetWriter; +use crate::sst::parquet::{SstInfo, WriteOptions}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; /// Wraps the remote object store used by write cache uploads. @@ -173,15 +174,20 @@ impl WriteCache { } /// Put encoded SST data to the cache and upload to the remote object store. + /// + /// Builds the SST's indexes from `data` when `index_build_params` is set and uploads them + /// with the SST. pub(crate) async fn put_and_upload_sst( &self, data: &bytes::Bytes, - region_file_id: RegionFileId, + region_id: RegionId, + sst_info: &mut SstInfo, upload_request: SstUploadRequest, write_buffer_size: ReadableSize, + index_build_params: Option, ) -> Result { - let region_id = region_file_id.region_id(); - let file_id = region_file_id.file_id(); + let file_id = sst_info.file_id; + let region_file_id = RegionFileId::new(region_id, file_id); let mut metrics = Metrics::new(WriteType::Flush); // Create index key for the SST file @@ -204,31 +210,70 @@ impl WriteCache { metrics.write_batch = cache_start.elapsed(); + if let Some(params) = index_build_params { + let start = Instant::now(); + let builder = params.indexer_builder( + self.build_puffin_manager(), + true, + self.intermediate_manager.clone(), + ); + sst_info.index_metadata = + build_index_from_parquet_bytes(&builder, region_file_id, data).await; + metrics.update_index = start.elapsed(); + } + // Upload to remote store + let remote_store = &upload_request.remote_store; + let mut upload_tracker = UploadTracker::new(region_id); let upload_start = Instant::now(); let remote_path = upload_request .dest_path_provider .build_sst_file_path(region_file_id); - - if let Err(e) = self + let mut result = self .upload( parquet_key, &remote_path, - &upload_request.remote_store, + remote_store, UploadOptions { // `put_and_upload_sst` is only used by the flush path. op_type: OperationType::Flush, write_buffer_size, }, ) - .await - { - // Clean up cache on failure - self.remove(parquet_key).await; - return Err(e); + .await; + metrics.upload_parquet = upload_start.elapsed(); + + if result.is_ok() && sst_info.index_metadata.file_size > 0 { + upload_tracker.push_uploaded_file(remote_path); + let puffin_key = IndexKey::new(region_id, file_id, FileType::Puffin(0)); + let puffin_path = upload_request + .dest_path_provider + .build_index_file_path(region_file_id); + let start = Instant::now(); + result = self + .upload( + puffin_key, + &puffin_path, + remote_store, + UploadOptions { + op_type: OperationType::Flush, + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + }, + ) + .await; + metrics.upload_puffin = start.elapsed(); } - metrics.upload_parquet = upload_start.elapsed(); + if let Err(e) = result { + upload_tracker + .clean( + std::slice::from_ref(sst_info), + &self.file_cache, + remote_store, + ) + .await; + return Err(e); + } Ok(metrics) } @@ -537,7 +582,7 @@ impl UploadTracker { /// Cleans uploaded files and files in the file cache at best effort. pub(crate) async fn clean( &self, - sst_info: &SstInfoArray, + sst_info: &[SstInfo], file_cache: &FileCacheRef, remote_store: &ObjectStore, ) { diff --git a/src/mito2/src/engine/index_build_test.rs b/src/mito2/src/engine/index_build_test.rs index 672c5e472d5..2877ac109d9 100644 --- a/src/mito2/src/engine/index_build_test.rs +++ b/src/mito2/src/engine/index_build_test.rs @@ -23,6 +23,8 @@ use api::v1::region::{StrictWindow, compact_request}; use api::v1::{Rows, SemanticType}; use async_trait::async_trait; use common_base::readable_size::ReadableSize; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; use datafusion_expr::expr::InList; use datafusion_expr::{Expr, col, lit}; @@ -34,7 +36,7 @@ use store_api::metadata::ColumnMetadata; use store_api::region_engine::RegionEngine; use store_api::region_request::{ AddColumn, AlterKind, RegionAlterRequest, RegionBuildIndexRequest, RegionCloseRequest, - RegionCompactRequest, RegionRequest, SetIndexOption, + RegionCompactRequest, RegionFlushRequest, RegionRequest, RegionTruncateRequest, SetIndexOption, }; use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::{Notify, Semaphore}; @@ -2297,8 +2299,9 @@ impl SstReadCounter { #[rstest::rstest] #[tokio::test] -async fn test_sync_flush_indexes_encoded_bulk_parts( - #[values(false, true)] disable_inverted_on_flush: bool, +async fn test_flush_indexes_encoded_bulk_parts( + #[values(false, true)] async_without_flush_indexes: bool, + #[values(false, true)] write_cache: bool, ) { let sst_bytes_read = Arc::new(AtomicUsize::new(0)); let counter = sst_bytes_read.clone(); @@ -2312,19 +2315,27 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( })) .build() .unwrap(); - let mut env = TestEnv::with_prefix("test_sync_flush_indexes_encoded_bulk_parts_") + let mut env = TestEnv::with_prefix("test_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 { + if async_without_flush_indexes { + config.index.build_mode = IndexBuildMode::Async; config.inverted_index.create_on_flush = Mode::Disable; } + if write_cache { + config = config.enable_write_cache( + env.data_home().join("write_cache").display().to_string(), + ReadableSize::mb(32), + None, + ); + } + let listener = Arc::new(IndexBuildListener::default()); let engine = env .create_engine_with(config, None, Some(listener.clone()), None) .await; @@ -2345,7 +2356,7 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( .insert_option("memtable.bulk.merge_threshold", "2") .insert_option("memtable.bulk.encode_row_threshold", "1") .build_with_index(); - if !disable_inverted_on_flush { + if !async_without_flush_indexes { // `build_with_index` only adds inverted indexes on tags; also cover a skipping index. let field = request .column_metadatas @@ -2385,9 +2396,10 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( 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. + assert!(!files.is_empty()); + if async_without_flush_indexes { + // Async mode builds indexes after flush. Only inverted indexes are configured and they + // are disabled on flush, so the builds 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(); @@ -2398,25 +2410,169 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( return; } - let version = engine.get_region(region_id).unwrap().version(); - let indexed = async { - loop { - let files = current_file_metas(&engine, region_id).await; - if files.iter().all(|f| { - f.is_index_consistent_with_region( - &version.metadata.column_metadatas, - &version.options.index_options, - ) - }) { - 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); + // Sync mode builds the indexes from the encoded bytes in memory during flush. + assert_eq!(0, sst_bytes_read.load(Ordering::Relaxed)); + assert_eq!(0, listener.begin_count()); assert!(files.iter().all(|f| f.inverted_index_available())); assert!(files.iter().all(|f| f.bloom_filter_index_available())); + let access_layer = engine.get_region(region_id).unwrap().access_layer.clone(); + for file in &files { + let index_path = location::index_file_path( + access_layer.table_dir(), + RegionIndexId::new(RegionFileId::new(region_id, file.file_id), 0), + access_layer.path_type(), + ); + assert!( + access_layer + .object_store() + .exists(&index_path) + .await + .unwrap() + ); + } +} + +/// Blocks closing a puffin file whose SST is already closed, which only happens for encoded +/// bulk parts: the SST writer closes the puffin file before the SST. +struct BlockingPuffinWriter { + file_name: String, + inner: object_store::layers::mock::Writer, + closed_ssts: Arc>>, + entered: Arc, + release: Arc, +} + +impl object_store::layers::mock::Write for BlockingPuffinWriter { + async fn write( + &mut self, + bs: object_store::layers::mock::Buffer, + ) -> object_store::layers::mock::Result<()> { + self.inner.write(bs).await + } + + async fn close( + &mut self, + ) -> object_store::layers::mock::Result { + if let Some(file_id) = self.file_name.strip_suffix(".puffin") + && self.closed_ssts.lock().unwrap().contains(file_id) + { + self.entered.notify_one(); + self.release.notified().await; + } + let metadata = self.inner.close().await?; + if let Some(file_id) = self.file_name.strip_suffix(".parquet") { + self.closed_ssts.lock().unwrap().insert(file_id.to_string()); + } + Ok(metadata) + } + + async fn abort(&mut self) -> object_store::layers::mock::Result<()> { + self.inner.abort().await + } +} + +#[tokio::test] +async fn test_truncate_while_indexing_encoded_bulk_part_removes_uploaded_sst() { + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let closed_ssts = Arc::new(std::sync::Mutex::new(HashSet::new())); + let (factory_entered, factory_release) = (entered.clone(), release.clone()); + let layer = object_store::layers::mock::MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _op, inner| { + Box::new(BlockingPuffinWriter { + file_name: path.rsplit('/').next().unwrap_or_default().to_string(), + inner, + closed_ssts: closed_ssts.clone(), + entered: factory_entered.clone(), + release: factory_release.clone(), + }) as object_store::layers::mock::Writer + })) + .build() + .unwrap(); + let mut env = TestEnv::with_prefix("test_truncate_while_indexing_encoded_bulk_part_") + .await + .with_mock_layer(layer); + let config = MitoConfig { + default_flat_format: true, + ..Default::default() + }; + let engine = env.create_engine_with(config, None, None, 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 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(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + 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; + } + let version = engine.get_region(region_id).unwrap().version(); + for memtable in version.memtables.list_memtables() { + memtable.compact(false).unwrap(); + } + + let flush_engine = engine.clone(); + let flush_task = tokio::spawn(async move { + flush_engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + }); + // The encoded SST is uploaded and its index is being written. + tokio::time::timeout(Duration::from_secs(5), entered.notified()) + .await + .expect("flush did not reach index creation"); + + tokio::time::timeout( + Duration::from_secs(5), + engine.handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::All), + ), + ) + .await + .expect("truncate did not finish after cancelling flush") + .unwrap(); + release.notify_one(); + let flush_err = tokio::time::timeout(Duration::from_secs(5), flush_task) + .await + .expect("cancelled flush did not finish") + .expect("flush task panicked") + .unwrap_err(); + assert_eq!(flush_err.status_code(), StatusCode::Cancelled); + + let access_layer = engine.get_region(region_id).unwrap().access_layer.clone(); + let leftovers = access_layer + .object_store() + .list_with(access_layer.table_dir()) + .recursive(true) + .await + .unwrap() + .into_iter() + .map(|entry| entry.path().to_string()) + .filter(|path| path.ends_with(".parquet") || path.ends_with(".puffin")) + .collect::>(); + assert!(leftovers.is_empty(), "{leftovers:?}"); } diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index fba0607235b..0a4f1765f39 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -29,16 +29,17 @@ use smallvec::{SmallVec, smallvec}; use snafu::{ResultExt, ensure}; use store_api::metadata::RegionMetadataRef; use store_api::region_request::RegionFlushReason; -use store_api::storage::{FileId, RegionId, SequenceNumber}; +use store_api::storage::{RegionId, SequenceNumber}; use strum::IntoStaticStr; use tokio::sync::{Semaphore, mpsc, watch}; use crate::access_layer::{ - AccessLayerRef, Metrics, OperationType, SstInfoArray, SstWriteRequest, WriteType, + AccessLayerRef, EncodedSstIndexBuildParams, Metrics, OperationType, SstInfoArray, + SstWriteRequest, WriteType, }; use crate::cache::CacheManagerRef; use crate::compaction::{collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema}; -use crate::config::MitoConfig; +use crate::config::{IndexBuildMode, MitoConfig}; use crate::engine::region_hook::SstFileInfo; use crate::error::{ Error, FlushCancelledSnafu, FlushRegionSnafu, JoinSnafu, RegionBusySnafu, RegionClosedSnafu, @@ -66,7 +67,7 @@ use crate::request::{ }; use crate::schedule::CancellableTaskState; use crate::schedule::scheduler::{Job, SchedulerRef}; -use crate::sst::file::{FileMeta, RegionFileId, UncommittedSsts}; +use crate::sst::file::{FileMeta, UncommittedSsts}; use crate::sst::parquet::metadata::extract_primary_key_range; use crate::sst::parquet::{ DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE, SstInfo, WriteOptions, flat_format, @@ -394,10 +395,7 @@ impl RegionFlushTask { }; let worker_request = match flush_result { - Ok(FlushedEdit { - edit, - unindexed_files, - }) => { + Ok(edit) => { let memtables_to_remove = version_data .version .memtables @@ -413,7 +411,6 @@ impl RegionFlushTask { senders: std::mem::take(&mut self.senders), _timer: timer, edit, - unindexed_files, memtables_to_remove, is_staging: self.is_staging, }; @@ -451,7 +448,7 @@ impl RegionFlushTask { version_data: &VersionControlData, state: &CancellableTaskState, uncommitted: &UncommittedSsts, - ) -> Result { + ) -> Result { // 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; @@ -470,7 +467,6 @@ impl RegionFlushTask { let DoFlushMemtablesResult { file_metas, - unindexed_files, flushed_bytes, series_count, encoded_part_count, @@ -594,10 +590,7 @@ impl RegionFlushTask { self.reason.as_str() ); - Ok(FlushedEdit { - edit, - unindexed_files, - }) + Ok(edit) } async fn do_flush_memtables( @@ -612,7 +605,6 @@ 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(); @@ -659,12 +651,7 @@ impl RegionFlushTask { .await?; encoded_part_count += num_encoded; for (source_idx, result) in results.into_iter().enumerate() { - let FlushedSsts { - max_sequence, - ssts: ssts_written, - metrics, - indexed, - } = result?; + let (max_sequence, ssts_written, metrics) = result?; if ssts_written.is_empty() { // No data written. continue; @@ -683,9 +670,6 @@ 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() @@ -719,7 +703,6 @@ impl RegionFlushTask { Ok(DoFlushMemtablesResult { file_metas, - unindexed_files, flushed_bytes, series_count, encoded_part_count, @@ -765,41 +748,43 @@ impl RegionFlushTask { .await?; uncommitted.track(&ssts); FLUSH_FILE_TOTAL.inc_by(ssts.len() as u64); - Ok(FlushedSsts { - max_sequence, - ssts, - metrics, - indexed: true, - }) + Ok((max_sequence, ssts, metrics)) }); tasks.push(task); } - for (encoded, max_sequence) in flat_sources.encoded { + let build_index = self.engine_config.index.build_mode == IndexBuildMode::Sync; + for (mut encoded, max_sequence) in flat_sources.encoded { let access_layer = self.access_layer.clone(); let cache_manager = self.cache_manager.clone(); let region_id = version.metadata.region_id; let write_buffer_size = write_opts.write_buffer_size; let semaphore = self.flush_semaphore.clone(); let uncommitted = uncommitted.clone(); + let index_build_params = build_index.then(|| EncodedSstIndexBuildParams { + metadata: version.metadata.clone(), + index_options: self.index_options.clone(), + inverted_index_config: self.engine_config.inverted_index.clone(), + fulltext_index_config: self.engine_config.fulltext_index.clone(), + bloom_filter_index_config: self.engine_config.bloom_filter_index.clone(), + }); + // Index creation runs after the SST is uploaded, so cancellation can land after the + // upload. Own the files before uploading them. + uncommitted.track_before_upload(encoded.sst_info.file_id, index_build_params.is_some()); let task = common_runtime::spawn_global(async move { let _permit = semaphore.acquire().await.unwrap(); let metrics = access_layer .put_sst( &encoded.data, - RegionFileId::new(region_id, encoded.sst_info.file_id), + region_id, + &mut encoded.sst_info, &cache_manager, write_buffer_size, + index_build_params, ) .await?; uncommitted.track(std::slice::from_ref(&encoded.sst_info)); FLUSH_FILE_TOTAL.inc(); - // Encoded parts are uploaded as-is, without passing through the indexer. - Ok(FlushedSsts { - max_sequence, - ssts: smallvec![encoded.sst_info], - metrics, - indexed: false, - }) + Ok((max_sequence, smallvec![encoded.sst_info], metrics)) }); tasks.push(task); } @@ -935,30 +920,14 @@ impl RegionFlushTask { } } -struct FlushedEdit { - edit: RegionEdit, - /// Files in `edit` written without building their indexes. - unindexed_files: Vec, -} - -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>, + results: Vec>, } struct DoFlushMemtablesResult { file_metas: Vec, - /// Files written without building their indexes. - unindexed_files: Vec, flushed_bytes: u64, series_count: usize, encoded_part_count: usize, @@ -1970,7 +1939,6 @@ mod tests { flushed_sequence: None, committed_sequence: None, }, - unindexed_files: Vec::new(), memtables_to_remove: smallvec![], is_staging: false, }), diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 847a12db5ae..e323fb6b340 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1155,20 +1155,9 @@ impl PartToMerge { PartToMerge::Multi { part, .. } => { Some((part.num_rows() as u64, part.estimated_size() as u64)) } - PartToMerge::Encoded { part, .. } => part - .metadata() - .parquet_metadata - .row_groups() - .iter() - .map(|row_group| { - let uncompressed_bytes = row_group - .columns() - .iter() - .map(|column| column.uncompressed_size() as u64) - .sum(); - (row_group.num_rows() as u64, uncompressed_bytes) - }) - .max_by_key(|(_, uncompressed_bytes)| *uncompressed_bytes), + PartToMerge::Encoded { part, .. } => { + crate::batch_size::largest_row_group_statistic(&part.metadata().parquet_metadata) + } } } diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 2e81fbc4396..d7e12a5da60 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -1148,8 +1148,6 @@ 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, /// Memtables to remove. pub(crate) memtables_to_remove: SmallVec<[MemtableId; 2]>, /// Whether the region is in staging mode. diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 1d4caaf9127..4c6613dee40 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -944,6 +944,11 @@ impl UncommittedSsts { } } + /// Tracks an SST before it is uploaded, so a cancelled upload is still removed. + pub(crate) fn track_before_upload(&self, file_id: FileId, with_index: bool) { + self.files.lock().unwrap().insert(file_id, (0, with_index)); + } + /// Disarms cleanup after the manifest edit has committed or may have committed. /// /// A manifest update error does not guarantee that the edit was not persisted. Once an edit diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 4dc7002f147..fe39af344d2 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -35,14 +35,16 @@ use std::sync::Arc; use arrow_schema::extension::ExtensionType; use bloom_filter::creator::BloomFilterIndexer; +use bytes::Bytes; use common_telemetry::{debug, error, info, warn}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::Json2ExtensionType; use index::target::IndexTarget; use mito_codec::row_converter::DensePrimaryKeyCodec; use object_store::ObjectStore; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use puffin_manager::SstPuffinManager; -use smallvec::{SmallVec, smallvec}; +use smallvec::SmallVec; use snafu::ResultExt; use statistics::{ByteCount, RowCount}; use store_api::codec::PrimaryKeyEncoding; @@ -52,6 +54,7 @@ use strum::IntoStaticStr; use tokio::sync::mpsc::Sender; use crate::access_layer::{AccessLayerRef, FilePathProvider, OperationType, RegionFilePathFactory}; +use crate::batch_size::{estimate_batch_size, largest_row_group_statistic}; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; use crate::cache::write_cache::{UploadOptions, UploadTracker, WriteCacheRef}; use crate::cache::{CacheManagerRef, CacheStrategy}; @@ -442,6 +445,80 @@ impl IndexerBuilder for IndexerBuilderImpl { } } +/// Builds the indexes of an SST from its encoded parquet bytes. +/// +/// Encoded bulk parts are flushed without passing through [ParquetWriter], so their indexes +/// are built from the bytes still in memory instead of reading the uploaded SST back. Like +/// the writer, a failure only drops the index and leaves the SST unindexed. +/// +/// [ParquetWriter]: crate::sst::parquet::writer::ParquetWriter +pub(crate) async fn build_index_from_parquet_bytes( + builder: &IndexerBuilderImpl, + region_file_id: RegionFileId, + data: &Bytes, +) -> IndexOutput { + let reader_builder = match ParquetRecordBatchReaderBuilder::try_new(data.clone()) { + Ok(reader_builder) => reader_builder, + Err(e) => { + warn!(e; "Failed to read encoded SST {:?} for index creation", region_file_id); + return IndexOutput::default(); + } + }; + let row_group_size = reader_builder + .metadata() + .row_groups() + .first() + .map(|row_group| row_group.num_rows() as usize) + .filter(|size| *size > 0); + let mut indexer = builder.build(region_file_id, 0, row_group_size).await; + if indexer.is_empty() { + return IndexOutput::default(); + } + // Follow the batch size the flush reads use, so wide rows are not decoded 8192 at a time. + let batch_size = estimate_batch_size(largest_row_group_statistic(reader_builder.metadata())); + + // TODO(dennis): decode only the columns the enabled indexers need, including the encoded + // primary key for tag indexes. This needs the indexer to stop locating columns by position. + let mut reader = match reader_builder.with_batch_size(batch_size).build() { + Ok(reader) => reader, + Err(e) => { + warn!(e; "Failed to read encoded SST {:?} for index creation", region_file_id); + indexer.abort().await; + return IndexOutput::default(); + } + }; + loop { + // Decoding is CPU-bound, so each batch is decoded on the blocking pool. Only one batch + // is in flight, and an aborted flush stops after the batch being decoded. + let decoded = common_runtime::spawn_blocking_global(move || { + let batch = reader.next(); + (reader, batch) + }) + .await; + let batch = match decoded { + Ok((next_reader, batch)) => { + reader = next_reader; + batch + } + Err(e) => { + warn!(e; "Failed to decode encoded SST {:?} for index creation", region_file_id); + indexer.abort().await; + return IndexOutput::default(); + } + }; + match batch { + Some(Ok(batch)) => indexer.update_flat(&batch).await, + Some(Err(e)) => { + warn!(e; "Failed to decode encoded SST {:?} for index creation", region_file_id); + indexer.abort().await; + return IndexOutput::default(); + } + None => break, + } + } + indexer.finish().await +} + impl IndexerBuilderImpl { fn build_inverted_indexer( &self, @@ -1164,7 +1241,7 @@ impl IndexBuildTask { // Cleans index files on failure. upload_tracker .clean( - &smallvec![SstInfo { + &[SstInfo { file_id, index_metadata: output, ..Default::default() @@ -1593,6 +1670,7 @@ mod tests { use partition::expr::col; use puffin::puffin_manager::{PuffinManager, PuffinReader}; use puffin_manager::PuffinManagerFactory; + use smallvec::smallvec; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use tokio::sync::mpsc; diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs index 4e46f886ebd..1a6d90ac4f5 100644 --- a/src/mito2/src/sst/index/indexer/finish.rs +++ b/src/mito2/src/sst/index/indexer/finish.rs @@ -57,6 +57,11 @@ impl Indexer { self.do_prune_intm_sst_dir().await; output.file_size = self.do_finish_puffin_writer(writer).await; + // A failed puffin file must not advertise the indexes written into it. Not covered by + // tests: `do_finish_puffin_writer` panics on failure in test builds. + if output.file_size == 0 { + return IndexOutput::default(); + } output.version = self.index_version; output } diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index f09efe81e47..4f26911077f 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -439,21 +439,10 @@ impl RegionWorkerLoop { let flush_on_close = request.flush_reason == FlushReason::Closing; - 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 version = region.version(); - index_build_file_metas.retain(|meta| { - request.unindexed_files.contains(&meta.file_id) - && !meta.is_index_consistent_with_region( - &version.metadata.column_metadatas, - &version.options.index_options, - ) - }); - } + let index_build_file_metas = std::mem::take(&mut request.edit.files_to_add); - if !index_build_file_metas.is_empty() { + // In async mode, create indexes after flush. + if self.config.index.build_mode == IndexBuildMode::Async { self.handle_rebuild_index( BuildIndexRequest { region_id,