diff --git a/config/config.md b/config/config.md
index bd9b1944e67..cc5b8135b59 100644
--- a/config/config.md
+++ b/config/config.md
@@ -222,6 +222,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. |
@@ -699,6 +700,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 cc6eb2200c3..3594bf7ebaa 100644
--- a/config/datanode.example.toml
+++ b/config/datanode.example.toml
@@ -733,6 +733,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 9f94888d2ae..84df6b4d3c8 100644
--- a/config/standalone.example.toml
+++ b/config/standalone.example.toml
@@ -932,6 +932,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 4929137b497..a691475326d 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,
@@ -391,10 +392,7 @@ impl RegionFlushTask {
};
let worker_request = match flush_result {
- Ok(FlushedEdit {
- edit,
- unindexed_files,
- }) => {
+ Ok(edit) => {
let memtables_to_remove = version_data
.version
.memtables
@@ -410,7 +408,6 @@ impl RegionFlushTask {
senders: std::mem::take(&mut self.senders),
_timer: timer,
edit,
- unindexed_files,
memtables_to_remove,
is_staging: self.is_staging,
};
@@ -448,7 +445,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;
@@ -467,7 +464,6 @@ impl RegionFlushTask {
let DoFlushMemtablesResult {
file_metas,
- unindexed_files,
flushed_bytes,
series_count,
encoded_part_count,
@@ -580,10 +576,7 @@ impl RegionFlushTask {
self.reason.as_str()
);
- Ok(FlushedEdit {
- edit,
- unindexed_files,
- })
+ Ok(edit)
}
async fn do_flush_memtables(
@@ -598,7 +591,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();
@@ -645,12 +637,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;
@@ -669,9 +656,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()
@@ -705,7 +689,6 @@ impl RegionFlushTask {
Ok(DoFlushMemtablesResult {
file_metas,
- unindexed_files,
flushed_bytes,
series_count,
encoded_part_count,
@@ -751,41 +734,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);
}
@@ -921,30 +906,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,
@@ -1954,7 +1923,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 7402eeb0582..6a09f1d228e 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 faf5c77f975..0a743b2557e 100644
--- a/src/mito2/src/request.rs
+++ b/src/mito2/src/request.rs
@@ -946,8 +946,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 28b6218cd7b..a8d8bb8282b 100644
--- a/src/mito2/src/sst/file.rs
+++ b/src/mito2/src/sst/file.rs
@@ -906,6 +906,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 ff734e3241b..3371bca587e 100644
--- a/src/mito2/src/sst/index.rs
+++ b/src/mito2/src/sst/index.rs
@@ -34,12 +34,14 @@ use std::ops::ControlFlow;
use std::sync::Arc;
use bloom_filter::creator::BloomFilterIndexer;
+use bytes::Bytes;
use common_telemetry::{debug, error, info, warn};
use datatypes::arrow::record_batch::RecordBatch;
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;
@@ -49,6 +51,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};
@@ -439,6 +442,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,
@@ -1113,7 +1190,7 @@ impl IndexBuildTask {
// Cleans index files on failure.
upload_tracker
.clean(
- &smallvec![SstInfo {
+ &[SstInfo {
file_id,
index_metadata: output,
..Default::default()
@@ -1542,6 +1619,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 0c1f200d599..58b339a80d9 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 1bca8b9c119..eea3bc5a6a1 100644
--- a/src/mito2/src/worker/handle_flush.rs
+++ b/src/mito2/src/worker/handle_flush.rs
@@ -436,21 +436,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,