rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-29 15:06:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(mito2): index encoded bulk parts from memory during flush (#9913)

Browse Source 
* 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 <killme2008@gmail.com>

* 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 <killme2008@gmail.com>

* 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 <killme2008@gmail.com>

* refactor(mito2): name encoded SST index build params explicitly

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* 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 <killme2008@gmail.com>

* docs(config): document the mito index build_mode option

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* 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 <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
dennis zhuang
2026-10-27 10:49:16 +00:00
committed by GitHub
parent 17ed988f49
commit 5b4efd908c
14 changed files with 468 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mito2/src/access_layer.rs
+65 -4
View File
@@ -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<EncodedSstIndexBuildParams>,
) -> Result<Metrics> {
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();
src/mito2/src/batch_size.rs
+19
View File
@@ -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.
src/mito2/src/cache/write_cache.rs
Vendored
+62 -17
View File
@@ -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<EncodedSstIndexBuildParams>,
) -> Result<Metrics> {
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,
) {
src/mito2/src/engine/index_build_test.rs
+185 -29
View File
@@ -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<std::sync::Mutex<HashSet<String>>>,
entered: Arc<Notify>,
release: Arc<Notify>,
}
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<object_store::layers::mock::Metadata> {
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::<Vec<_>>();
assert!(leftovers.is_empty(), "{leftovers:?}");
}
src/mito2/src/flush.rs
+27 -59
View File
@@ -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<FlushedEdit> {
) -> Result<RegionEdit> {
// 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<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<FlushedSsts>>,
results: Vec<Result<(SequenceNumber, SstInfoArray, Metrics)>>,
}
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,
@@ -1970,7 +1939,6 @@ mod tests {
flushed_sequence: None,
committed_sequence: None,
},
unindexed_files: Vec::new(),
memtables_to_remove: smallvec![],
is_staging: false,
}),
src/mito2/src/memtable/bulk.rs
+3 -14
View File
@@ -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)
}
}
}
src/mito2/src/request.rs
-2
View File
@@ -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<FileId>,
/// Memtables to remove.
pub(crate) memtables_to_remove: SmallVec<[MemtableId; 2]>,
/// Whether the region is in staging mode.
src/mito2/src/sst/file.rs
+5
View File
@@ -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
src/mito2/src/sst/index.rs
+80 -2
View File
@@ -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;
src/mito2/src/sst/index/indexer/finish.rs
+5
View File
@@ -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
}
src/mito2/src/worker/handle_flush.rs
+3 -14
View File
@@ -439,21 +439,10 @@ impl<S: LogStore> RegionWorkerLoop<S> {
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,