From 674ba7266949c3c05b71e6e7fb99891a71cc8d2e Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 12 Oct 2026 13:16:52 +0000 Subject: [PATCH] fix(mito): preserve version ordering in bulk merges and flushes (#9517) * fix(mito): preserve delete barriers in flat last-non-null dedup Signed-off-by: evenyag * fix(mito): respect sequence ranges in bulk merges and flushes Signed-off-by: evenyag * perf(mito): avoid merge selection on the bulk write path Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/src/access_layer.rs | 2 +- src/mito2/src/engine/flush_test.rs | 313 ++++++++++++++++- src/mito2/src/flush.rs | 330 +++++++++--------- src/mito2/src/memtable.rs | 26 ++ src/mito2/src/memtable/bulk.rs | 514 ++++++++++++++++------------ src/mito2/src/memtable/bulk/part.rs | 32 +- src/mito2/src/read/flat_dedup.rs | 163 +++++---- 7 files changed, 924 insertions(+), 456 deletions(-) diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 3959e0c26a4..b1457935c16 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -579,7 +579,7 @@ pub struct SstWriteRequest { pub source: FlatSource, pub cache_manager: CacheManagerRef, /// Optional uniform row sequence for writes that do not preserve sequences. - /// Compaction passes `None` to retain the reader's effective input sequences. + /// Compaction and LastNonNull flushes pass `None` to retain input sequences. pub max_sequence: Option, pub sst_write_format: FormatType, diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index b040ae5ec54..313860c0857 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -27,6 +27,7 @@ use common_base::readable_size::ReadableSize; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; +use datafusion_expr::{col, lit}; use parquet::basic::{Encoding, Type as PhysicalType}; use rstest::rstest; use rstest_reuse::{self, apply}; @@ -35,9 +36,9 @@ use store_api::metadata::RegionMetadataRef; use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ - AlterKind, PathType, RegionAlterRequest, RegionCloseRequest, RegionFlushRequest, - RegionOpenRequest, RegionPutRequest, RegionRequest, ReplayCheckpoint, SetRegionOption, - UnsetRegionOption, + AlterKind, PathType, RegionAlterRequest, RegionCloseRequest, RegionCompactRequest, + RegionFlushRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, ReplayCheckpoint, + SetRegionOption, UnsetRegionOption, }; use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::Notify; @@ -49,9 +50,11 @@ use crate::engine::listener::{EventListener, FlushListener, StallListener}; use crate::engine::region_hook::{RegionGcInfo, RegionHook, RegionHookRef, SstFileInfo}; use crate::error::Error; use crate::manifest::action::RegionMetaActionList; +use crate::memtable::RangesOptions; use crate::test_util::{ - CreateRequestBuilder, LogStoreFactory, MockWriteBufferManager, TestEnv, build_rows, - build_rows_for_key, flush_region, kafka_log_store_factory, multiple_log_store_factories, + CreateRequestBuilder, LogStoreFactory, MockWriteBufferManager, TestEnv, + build_delete_rows_for_key, build_rows, build_rows_for_key, build_rows_with_fields, delete_rows, + delete_rows_schema, flush_region, kafka_log_store_factory, multiple_log_store_factories, prepare_test_for_kafka_log_store, put_rows, raft_engine_log_store_factory, reopen_region, rows_schema, single_kafka_log_store_factory, }; @@ -1818,3 +1821,303 @@ async fn test_region_hook() { manifest_count, ); } + +/// Explicitly finish a merge instead of sleeping for the background dispatcher. +fn compact_bulk_memtables(engine: &MitoEngine, region_id: RegionId) { + let version = engine.get_region(region_id).unwrap().version(); + for memtable in version.memtables.list_memtables() { + memtable.compact(false).unwrap(); + } +} + +async fn scan_bulk_test_keys(engine: &MitoEngine, region_id: RegionId) -> String { + let mut schema = None; + let mut batches = Vec::new(); + for key in ["deleted", "k"] { + let request = ScanRequest { + filters: vec![col("tag_0").eq(lit(key))], + ..Default::default() + }; + let stream = engine.scan_to_stream(region_id, request).await.unwrap(); + schema = Some(stream.schema()); + batches.extend(RecordBatches::try_collect(stream).await.unwrap().take()); + } + RecordBatches::try_new(schema.unwrap(), batches) + .unwrap() + .pretty_print() + .unwrap() +} + +#[rstest] +#[tokio::test] +async fn test_bulk_flush_keeps_newer_version(#[values(false, true)] delete_newer: bool) { + let mut env = TestEnv::new().await; + let config = MitoConfig { + default_flat_format: true, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + let engine = env.create_engine(config.clone()).await; + let region_id = RegionId::new(1, 1); + 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(); + let schema = rows_schema(&request); + let delete_schema = delete_rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("k", 1, 2, 1), + }, + ) + .await; + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("a", 0, 1024, 0), + }, + ) + .await; + if delete_newer { + let mut rows = build_delete_rows_for_key("b", 0, 1023); + rows.extend(build_delete_rows_for_key("k", 1, 2)); + delete_rows( + &engine, + region_id, + Rows { + schema: delete_schema, + rows, + }, + ) + .await; + } else { + let mut rows = build_rows_for_key("b", 0, 1023, 0); + rows.extend(build_rows_for_key("k", 1, 2, 2)); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + } + compact_bulk_memtables(&engine, region_id); + // Previously this advanced the old unordered row above the newer encoded row. + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows_for_key("c", 0, 1, 0), + }, + ) + .await; + let expected = if delete_newer { + "\ ++-------+---------+----+ +| tag_0 | field_0 | ts | ++-------+---------+----+ ++-------+---------+----+" + } else { + "\ ++-------+---------+---------------------+ +| tag_0 | field_0 | ts | ++-------+---------+---------------------+ +| k | 2.0 | 1970-01-01T00:00:01 | ++-------+---------+---------------------+" + }; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + flush_region(&engine, region_id, None).await; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + let engine = env.reopen_engine(engine, config).await; + reopen_region( + &engine, + region_id, + request.table_dir, + false, + request.options, + ) + .await; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); +} + +#[rstest] +#[tokio::test] +async fn test_bulk_last_non_null_merge_and_flush(#[values(false, true)] encoded: bool) { + let mut env = TestEnv::new().await; + let config = MitoConfig { + default_flat_format: true, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + let engine = env.create_engine(config.clone()).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() + .field_num(2) + .insert_option("memtable.type", "bulk") + .insert_option("merge_mode", "last_non_null") + .insert_option("memtable.bulk.merge_threshold", "4") + .insert_option( + "memtable.bulk.encode_row_threshold", + if encoded { "1" } else { "1000000" }, + ) + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "2") + .build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + // An older SST must remain hidden behind a delete retained in the memtable. + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows_with_fields("deleted", &[1], &[(Some(1), Some(1))]), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + + let mut a = build_rows_with_fields( + "a", + &(0..1023).collect::>(), + &[(Some(0), None); 1023], + ); + a.extend(build_rows_with_fields("k", &[1], &[(Some(1), None)])); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: a, + }, + ) + .await; + // Intermediate writes must be bulk parts, not an unordered merge barrier. + let mut b = build_rows_with_fields( + "b", + &(0..1023).collect::>(), + &[(Some(0), None); 1023], + ); + b.extend(build_rows_with_fields("k", &[1], &[(Some(2), None)])); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: b, + }, + ) + .await; + let mut deletes = build_delete_rows_for_key("d", 0, 1023); + deletes.extend(build_delete_rows_for_key("deleted", 1, 2)); + delete_rows( + &engine, + region_id, + Rows { + schema: delete_rows_schema(&request), + rows: deletes, + }, + ) + .await; + let before = scan_bulk_test_keys(&engine, region_id).await; + assert!(before.contains("| k | 2.0")); + assert!(!before.contains("| deleted")); + + let mut c = build_rows_with_fields( + "c", + &(0..1022).collect::>(), + &[(Some(0), None); 1022], + ); + c.extend(build_rows_with_fields("k", &[1], &[(None, Some(3))])); + c.extend(build_rows_with_fields("deleted", &[1], &[(None, Some(3))])); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: c, + }, + ) + .await; + compact_bulk_memtables(&engine, region_id); + // All four sources must have merged, with the requested output format. + let version = engine.get_region(region_id).unwrap().version(); + let mut ranges = Vec::new(); + for memtable in version.memtables.list_memtables() { + ranges.extend( + memtable + .ranges(None, RangesOptions::default()) + .unwrap() + .ranges + .into_values(), + ); + } + assert_eq!(1, ranges.len()); + assert_eq!(encoded, ranges[0].encoded().is_some()); + drop(ranges); + drop(version); + let expected = "\ ++---------+---------+---------+---------------------+ +| tag_0 | field_0 | field_1 | ts | ++---------+---------+---------+---------------------+ +| deleted | | 3.0 | 1970-01-01T00:00:01 | +| k | 2.0 | 3.0 | 1970-01-01T00:00:01 | ++---------+---------+---------+---------------------+"; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + // Leave a raw source beside merged parts so flush also exercises grouping. + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows_with_fields("z", &[0], &[(Some(0), None)]), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await + .unwrap(); + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); + let engine = env.reopen_engine(engine, config).await; + reopen_region( + &engine, + region_id, + request.table_dir, + false, + request.options, + ) + .await; + assert_eq!(expected, scan_bulk_test_keys(&engine, region_id).await); +} diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index e910c81f896..d5191bdcadb 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -46,7 +46,10 @@ use crate::error::{ }; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::memtable::bulk::ENCODE_ROW_THRESHOLD; -use crate::memtable::{BoxedRecordBatchIterator, EncodedRange, MemtableRanges, RangesOptions}; +use crate::memtable::{ + BoxedRecordBatchIterator, EncodedRange, MemtableRange, MemtableRanges, RangesOptions, + group_by_sequence, +}; use crate::metrics::{ FLUSH_BYTES_TOTAL, FLUSH_ELAPSED, FLUSH_FAILURE_TOTAL, FLUSH_FILE_TOTAL, FLUSH_REQUESTS_TOTAL, INFLIGHT_FLUSH_COUNT, @@ -850,7 +853,10 @@ impl RegionFlushTask { metadata: version.metadata.clone(), source, cache_manager: self.cache_manager.clone(), - max_sequence: Some(max_sequence), + // Sources have disjoint sequence intervals. LastNonNull can still + // contain Puts and retained delete barriers for the same row. + max_sequence: (version.options.merge_mode() != MergeMode::LastNonNull) + .then_some(max_sequence), sst_write_format: if flat_format { FormatType::Flat } else { @@ -920,7 +926,7 @@ struct FlatSources { encoded: SmallVec<[(EncodedRange, SequenceNumber); 4]>, } -/// Returns the max sequence and [FlatSource] for the given memtable. +/// Builds sequence-disjoint sources and their maximum input sequences. fn memtable_flat_sources( schema: SchemaRef, mem_ranges: MemtableRanges, @@ -928,190 +934,158 @@ fn memtable_flat_sources( options: &RegionOptions, field_column_start: usize, ) -> Result { - let MemtableRanges { ranges } = mem_ranges; + let components = group_by_sequence( + mem_ranges.ranges.into_values().collect(), + |range: &MemtableRange| { + let stats = range.stats(); + (stats.min_sequence, stats.max_sequence()) + }, + ); let mut flat_sources = FlatSources { sources: SmallVec::new(), encoded: SmallVec::new(), }; - - if ranges.len() == 1 { - debug!("Flushing single flat range"); - - let only_range = ranges.into_values().next().unwrap(); - let max_sequence = only_range.stats().max_sequence(); - if let Some(encoded) = only_range.encoded() { - flat_sources.encoded.push((encoded, max_sequence)); - } else { - let schema = only_range.record_batch_schema_hint().unwrap_or(schema); - let iter = only_range.build_record_batch_iter(None, None)?; - // Dedup according to append mode and merge mode. - // Even single range may have duplicate rows. - let iter = maybe_dedup_one( - options.append_mode, - options.merge_mode(), - field_column_start, - iter, - ); - flat_sources - .sources - .push((FlatSource::new_iter(schema, iter), max_sequence)); - }; - } else { - let min_flush_rows = *ENCODE_ROW_THRESHOLD; - // Calculate total rows from non-encoded ranges. - let total_rows: usize = ranges - .values() - .filter(|r| r.encoded().is_none()) - .map(|r| r.num_rows()) - .sum(); - debug!( - "Flushing multiple flat ranges, total_rows: {}, min_flush_rows: {}, num_ranges: {}", - total_rows, - min_flush_rows, - ranges.len() - ); - let mut rows_remaining = total_rows; - let mut last_iter_rows = 0; - let num_ranges = ranges.len(); - let mut input_iters = Vec::with_capacity(num_ranges); - let mut current_ranges = Vec::new(); - - let schemas = ranges - .values() - .filter(|range| range.encoded().is_none()) - .map(|range| { - ( - range - .record_batch_schema_hint() - .unwrap_or_else(|| schema.clone()), - range.num_rows() as u64, - ) - }) - .collect::>(); - let plans = Arc::new(collect_json2_rewrite_plans(metadata, &schemas)?); - let schema = rewrite_json2_schema( - schemas.first().map(|(schema, _)| schema).unwrap_or(&schema), - &plans, - ); - - for (_range_id, range) in ranges { - if let Some(encoded) = range.encoded() { - let max_sequence = range.stats().max_sequence(); - flat_sources.encoded.push((encoded, max_sequence)); - continue; - } - - if let Some(actual) = range.record_batch_schema_hint() { - let actual = rewrite_json2_schema(&actual, &plans); - ensure!( - actual == schema, - UnexpectedSnafu { - reason: format!( - "Different schemas found in a MemtableRanges, expected: {}, actual: {}", - schema, actual, - ), - } - ) - } - - let iter = range.build_record_batch_iter(None, None)?; - let iter: BoxedRecordBatchIterator = if plans.is_empty() { - iter - } else { - let plans = plans.clone(); - Box::new(iter.map(move |batch| rewrite_json2_batch(batch?, &plans))) - }; - input_iters.push(iter); - let range_rows = range.num_rows(); - last_iter_rows += range_rows; - rows_remaining -= range_rows; - current_ranges.push(range); - - // Flush if we have enough rows, but don't flush if the remaining rows - // would be less than DEFAULT_ROW_GROUP_SIZE (to avoid small last files). - if last_iter_rows >= min_flush_rows - && (rows_remaining == 0 || rows_remaining >= DEFAULT_ROW_GROUP_SIZE) - { - debug!( - "Flush batch ready, rows: {}, min_rows: {}, num_iters: {}, remaining: {}", - last_iter_rows, - min_flush_rows, - input_iters.len(), - rows_remaining - ); - - // Calculate max_sequence from all merged ranges - let max_sequence = current_ranges - .iter() - .map(|r| r.stats().max_sequence()) - .max() - .unwrap_or(0); - let batch_size = - crate::batch_size::estimate_batch_size(current_ranges.iter().map(|range| { - let stats = range.stats(); - (stats.num_rows() as u64, stats.bytes_allocated() as u64) - })); - - let input_iters = - std::mem::replace(&mut input_iters, Vec::with_capacity(num_ranges)); - - let maybe_dedup = merge_and_dedup_with_batch_size( + let mut rows_remaining = components + .iter() + .filter(|component| component.len() != 1 || component[0].encoded().is_none()) + .flatten() + .map(MemtableRange::num_rows) + .sum::(); + let mut pending = Vec::new(); + let mut pending_rows = 0; + for component in components { + // A standalone encoded component is an ordering barrier. Do not merge + // raw ranges on either side while uploading the intervening rows as-is. + if component.len() == 1 + && let Some(encoded) = component[0].encoded() + { + if !pending.is_empty() { + flat_sources.sources.push(build_flush_source( &schema, - options.append_mode, - options.merge_mode(), + std::mem::take(&mut pending), + metadata, + options, field_column_start, - input_iters, - batch_size, - )?; - - flat_sources.sources.push(( - FlatSource::new_iter(schema.clone(), maybe_dedup), - max_sequence, - )); - last_iter_rows = 0; - current_ranges.clear(); + )?); + pending_rows = 0; } - } - - // Handle remaining iters. - if !input_iters.is_empty() { - debug!( - "Flush remaining batch, rows: {}, min_rows: {}, num_iters: {}, remaining: {}", - last_iter_rows, - min_flush_rows, - input_iters.len(), - rows_remaining - ); - - let max_sequence = current_ranges - .iter() - .map(|r| r.stats().max_sequence()) - .max() - .unwrap_or(0); - let batch_size = - crate::batch_size::estimate_batch_size(current_ranges.iter().map(|range| { - let stats = range.stats(); - (stats.num_rows() as u64, stats.bytes_allocated() as u64) - })); - - let maybe_dedup = merge_and_dedup_with_batch_size( - &schema, - options.append_mode, - options.merge_mode(), - field_column_start, - input_iters, - batch_size, - )?; - flat_sources - .sources - .push((FlatSource::new_iter(schema, maybe_dedup), max_sequence)); + .encoded + .push((encoded, component[0].stats().max_sequence())); + continue; + } + let component_rows = component.iter().map(MemtableRange::num_rows).sum::(); + pending_rows += component_rows; + rows_remaining -= component_rows; + pending.extend(component); + // Only split between disjoint components, retaining the existing + // minimum-size rule for the final source. + if pending_rows >= *ENCODE_ROW_THRESHOLD + && (rows_remaining == 0 || rows_remaining >= DEFAULT_ROW_GROUP_SIZE) + { + flat_sources.sources.push(build_flush_source( + &schema, + std::mem::take(&mut pending), + metadata, + options, + field_column_start, + )?); + pending_rows = 0; } } - + if !pending.is_empty() { + flat_sources.sources.push(build_flush_source( + &schema, + pending, + metadata, + options, + field_column_start, + )?); + } Ok(flat_sources) } +/// Builds one sequence-contiguous source, decoding encoded ranges if their +/// intervals overlap other ranges. JSON2 layouts are reconciled for this group. +fn build_flush_source( + schema: &SchemaRef, + ranges: Vec, + metadata: &RegionMetadataRef, + options: &RegionOptions, + field_column_start: usize, +) -> Result<(FlatSource, SequenceNumber)> { + let max_sequence = ranges + .iter() + .map(|range| range.stats().max_sequence()) + .max() + .unwrap_or(0); + if let [range] = ranges.as_slice() { + // A single range needs no layout reconciliation, including legacy JSON2. + let schema = range + .record_batch_schema_hint() + .unwrap_or_else(|| schema.clone()); + let iter = maybe_dedup_one( + options.append_mode, + options.merge_mode(), + field_column_start, + range.build_record_batch_iter(None, None)?, + ); + return Ok((FlatSource::new_iter(schema, iter), max_sequence)); + } + let batch_size = crate::batch_size::estimate_batch_size(ranges.iter().map(|range| { + let stats = range.stats(); + (stats.num_rows() as u64, stats.bytes_allocated() as u64) + })); + let schemas = ranges + .iter() + .map(|range| { + ( + range + .record_batch_schema_hint() + .unwrap_or_else(|| schema.clone()), + range.num_rows() as u64, + ) + }) + .collect::>(); + let plans = Arc::new(collect_json2_rewrite_plans(metadata, &schemas)?); + let schema = rewrite_json2_schema( + schemas.first().map(|(schema, _)| schema).unwrap_or(schema), + &plans, + ); + let mut iters = Vec::with_capacity(ranges.len()); + for range in ranges { + if let Some(actual) = range.record_batch_schema_hint() { + let actual = rewrite_json2_schema(&actual, &plans); + ensure!( + actual == schema, + UnexpectedSnafu { + reason: format!( + "Different schemas found in a MemtableRanges, expected: {}, actual: {}", + schema, actual + ), + } + ); + } + let iter = range.build_record_batch_iter(None, None)?; + let iter: BoxedRecordBatchIterator = if plans.is_empty() { + iter + } else { + let plans = plans.clone(); + Box::new(iter.map(move |batch| rewrite_json2_batch(batch?, &plans))) + }; + iters.push(iter); + } + let iter = merge_and_dedup_with_batch_size( + &schema, + options.append_mode, + options.merge_mode(), + field_column_start, + iters, + batch_size, + )?; + Ok((FlatSource::new_iter(schema, iter), max_sequence)) +} + /// Merges multiple record batch iterators and applies deduplication based on the specified mode. /// /// This function is used during the flush process to combine data from multiple memtable ranges @@ -1143,7 +1117,7 @@ fn memtable_flat_sources( /// 3. If `append_mode` is false, wraps the merge iterator with a `FlatDedupIterator` that /// applies the specified merge mode: /// - `LastRow`: Removes duplicate rows, keeping only the last one -/// - `LastNonNull`: Removes duplicates but preserves the last non-null value for each field +/// - `LastNonNull`: Folds fields between delete markers, retaining every marker /// /// # Examples /// @@ -2373,7 +2347,7 @@ mod tests { 1, std::iter::once(1000), std::iter::once(Some(1.0)), - 1, + 2, ); converter.append_key_values(&kvs)?; let batch = converter.convert()?.batch; @@ -2388,6 +2362,8 @@ mod tests { )), MemtableStats { num_rows: 1, + min_sequence: id as u64 + 1, + max_sequence: id as u64 + 1, ..Default::default() }, ) @@ -2406,7 +2382,7 @@ mod tests { ranges.insert( 1, new_range( - 0, + 1, TestIterBuilder { schema: schema.clone(), batch: Some(batch), diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index 0f0a50875cc..afb546340b1 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -64,6 +64,32 @@ pub use bulk::part::{ #[cfg(any(test, feature = "test"))] pub use time_partition::filter_record_batch; +/// Groups sequence intervals into ordered, disjoint components. A component must +/// not be split during a partial merge: an excluded interval could contain an +/// intermediate row version. Bounds may conservatively include removed rows. +pub(crate) fn group_by_sequence( + mut parts: Vec, + bounds: impl Fn(&T) -> (SequenceNumber, SequenceNumber), +) -> Vec> { + parts.sort_by_key(&bounds); + let mut groups: Vec> = Vec::new(); + let mut group_max = 0; + for part in parts { + let (min, max) = bounds(&part); + match groups.last_mut() { + Some(group) if min <= group_max => { + group.push(part); + group_max = group_max.max(max); + } + _ => { + groups.push(vec![part]); + group_max = max; + } + } + } + groups +} + /// Id for memtables. /// /// Should be unique under the same region. diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index afede8904d8..7402eeb0582 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -21,7 +21,7 @@ pub mod part_reader; mod row_group_reader; use std::collections::{BTreeMap, HashSet}; -use std::sync::atomic::{AtomicI64, AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, Mutex, RwLock}; use std::time::Instant; @@ -60,6 +60,7 @@ use crate::memtable::{ AllocTracker, BoxedBatchIterator, BoxedRecordBatchIterator, EncodedBulkPart, EncodedRange, IterBuilder, KeyValues, MemScanMetrics, Memtable, MemtableBuilder, MemtableId, MemtableRange, MemtableRangeContext, MemtableRanges, MemtableRef, MemtableStats, RangesOptions, + group_by_sequence, }; use crate::read::flat_dedup::{FlatDedupIterator, FlatLastNonNull, FlatLastRow}; use crate::read::flat_merge::FlatMergeIterator; @@ -177,7 +178,7 @@ enum MergedPart { /// Result of collecting parts to merge struct CollectedParts { - /// Groups of parts ready for merging (each group has up to 16 parts) + /// Consecutive groups of parts ready for merging. groups: Vec>, } @@ -202,38 +203,18 @@ impl BulkParts { self.unordered_part.is_empty() && self.parts.is_empty() } - /// Returns true if enough parts of the same type are available to merge. - fn should_merge_parts(&self, merge_threshold: usize, encode_bytes_threshold: usize) -> bool { - let mut bulk_count = 0; - let mut encoded_count = 0; - - for wrapper in &self.parts { - if !Self::is_merge_candidate(wrapper, encode_bytes_threshold) { - continue; - } - - if wrapper.part.is_encoded() { - encoded_count += 1; - } else { - bulk_count += 1; - } - - if bulk_count >= merge_threshold || encoded_count >= merge_threshold { - return true; - } - } - - false - } - - /// Returns whether a part is small enough to benefit from merging. - fn is_merge_candidate(wrapper: &BulkPartWrapper, encode_bytes_threshold: usize) -> bool { - !wrapper.merging - && Self::is_merge_candidate_by_size( - wrapper.part.is_encoded(), - wrapper.part.estimated_size(), - encode_bytes_threshold, - ) + /// Returns a conservative scheduling hint without sorting or allocating. + /// The picker checks sequence barriers and sizes when compaction runs. + fn may_merge_parts(&self, merge_threshold: usize) -> bool { + let merge_threshold = merge_threshold.max(2); + self.parts.len() >= merge_threshold + && self + .parts + .iter() + .filter(|wrapper| !wrapper.merging) + .take(merge_threshold) + .count() + == merge_threshold } fn is_merge_candidate_by_size( @@ -250,79 +231,105 @@ impl BulkParts { || self.unordered_part.estimated_bytes() > encode_bytes_threshold } - /// Collects unmerged parts and marks them as being merged. - /// Only collects parts of types that meet the threshold. - /// Parts are grouped by row count for parallel processing. + /// Converts buffered small writes into a sorted part once a threshold is reached. + fn seal_unordered(&mut self, metadata: &RegionMetadataRef) -> Result<()> { + if let Some(part) = self.unordered_part.to_bulk_part(metadata)? { + self.parts.push(BulkPartWrapper { + part: PartToMerge::Bulk { + part, + file_id: FileId::random(), + }, + merging: false, + }); + self.unordered_part.clear(); + } + Ok(()) + } + + /// Selects consecutive sequence components without crossing unordered, busy, + /// or oversized parts. Overlapping components are indivisible, even above the size limit. + fn pick_merge_groups( + &self, + merge_threshold: usize, + max_merge_groups: usize, + encode_bytes_threshold: usize, + ) -> Vec> { + if self.parts.len() < merge_threshold.max(2) || max_merge_groups == 0 { + return Vec::new(); + } + // The unordered interval participates in overlap grouping as a barrier, + // so neither overlapping parts nor parts on opposite sides can be merged. + let intervals = self + .parts + .iter() + .enumerate() + .map(|(index, wrapper)| { + let part = &wrapper.part; + (Some(index), (part.min_sequence(), part.max_sequence())) + }) + .chain( + self.unordered_part + .sequence_bounds() + .map(|bounds| (None, bounds)), + ) + .collect(); + let components = group_by_sequence(intervals, |(_, bounds)| *bounds); + let mut groups = Vec::new(); + let mut pending = Vec::new(); + for component in components { + if groups.len() >= max_merge_groups { + break; + } + let Some(component) = component + .into_iter() + .map(|(index, _)| index) + .collect::>>() + else { + pending.clear(); + continue; + }; + let busy = component.iter().any(|&i| self.parts[i].merging); + let oversized = component.len() == 1 && { + let part = &self.parts[component[0]].part; + !Self::is_merge_candidate_by_size( + part.is_encoded(), + part.estimated_size(), + encode_bytes_threshold, + ) + }; + if busy || oversized { + pending.clear(); + continue; + } + pending.extend(component); + if pending.len() >= merge_threshold.max(2) { + groups.push(std::mem::take(&mut pending)); + } + } + groups + } + fn collect_parts_to_merge( &mut self, merge_threshold: usize, max_merge_groups: usize, encode_bytes_threshold: usize, ) -> CollectedParts { - let mut bulk_indices = Vec::new(); - let mut encoded_indices = Vec::new(); - - for (idx, wrapper) in self.parts.iter().enumerate() { - if !Self::is_merge_candidate(wrapper, encode_bytes_threshold) { - continue; - } - let num_rows = wrapper.part.num_rows(); - if wrapper.part.is_encoded() { - encoded_indices.push((idx, num_rows)); - } else { - bulk_indices.push((idx, num_rows)); - } - } - - let mut groups = Vec::new(); - - // Process bulk parts if threshold met - if bulk_indices.len() >= merge_threshold { - groups.extend(self.collect_and_group_parts( - bulk_indices, - merge_threshold, - max_merge_groups, - )); - } - - // Process encoded parts if threshold met - if encoded_indices.len() >= merge_threshold { - groups.extend(self.collect_and_group_parts( - encoded_indices, - merge_threshold, - max_merge_groups, - )); - } - - CollectedParts { groups } - } - - /// Sorts indices by row count, groups into chunks, marks as merging, and returns groups. - fn collect_and_group_parts( - &mut self, - mut indices: Vec<(usize, usize)>, - merge_threshold: usize, - max_merge_groups: usize, - ) -> Vec> { - if indices.is_empty() { - return Vec::new(); - } - - indices.sort_unstable_by_key(|(_, num_rows)| *num_rows); - indices - .chunks(merge_threshold) - .take(max_merge_groups) - .map(|chunk| { - chunk - .iter() - .map(|(idx, _)| { - let wrapper = &mut self.parts[*idx]; + let groups = self + .pick_merge_groups(merge_threshold, max_merge_groups, encode_bytes_threshold) + .into_iter() + .map(|group| { + group + .into_iter() + .map(|i| { + let wrapper = &mut self.parts[i]; wrapper.merging = true; wrapper.part.clone() }) .collect() }) - .collect() + .collect(); + CollectedParts { groups } } /// Installs merged parts and removes the original parts by file ids. @@ -431,6 +438,8 @@ pub struct BulkMemtable { compactor: Arc>, /// Dispatcher for scheduling compaction tasks compact_dispatcher: Option>, + /// Whether a write-triggered compaction is queued or running. + compacting: Arc, /// Whether the append mode is enabled append_mode: bool, /// Mode to handle duplicate rows while merging @@ -493,18 +502,8 @@ impl Memtable for BulkMemtable { bulk_parts.unordered_part.push(fragment); // Compacts unordered_part if the row or byte threshold is exceeded. - if bulk_parts.should_compact_unordered_part(self.config.encode_bytes_threshold) - && let Some(bulk_part) = - bulk_parts.unordered_part.to_bulk_part(&self.metadata)? - { - bulk_parts.parts.push(BulkPartWrapper { - part: PartToMerge::Bulk { - part: bulk_part, - file_id: FileId::random(), - }, - merging: false, - }); - bulk_parts.unordered_part.clear(); + if bulk_parts.should_compact_unordered_part(self.config.encode_bytes_threshold) { + bulk_parts.seal_unordered(&self.metadata)?; } } else { bulk_parts.parts.push(BulkPartWrapper { @@ -523,9 +522,7 @@ impl Memtable for BulkMemtable { self.update_stats(local_metrics); } - if self.should_compact() { - self.schedule_compact(); - } + self.schedule_compact(); Ok(()) } @@ -688,21 +685,12 @@ impl Memtable for BulkMemtable { return Ok(()); } - // Unified merge for all parts - let should_merge = self.parts.read().unwrap().should_merge_parts( - self.config.merge_threshold, - self.config.encode_bytes_threshold, - ); - if should_merge { - compactor.merge_parts( - &self.parts, - &self.metadata, - !self.append_mode, - self.merge_mode, - )?; - } - - Ok(()) + compactor.merge_parts( + &self.parts, + &self.metadata, + !self.append_mode, + self.merge_mode, + ) } } @@ -727,6 +715,7 @@ impl BulkMemtable { self.float_field_encoding, ))), compact_dispatcher: self.compact_dispatcher.clone(), + compacting: Arc::new(AtomicBool::new(false)), append_mode: self.append_mode, merge_mode: self.merge_mode, row_group_size: self.row_group_size, @@ -815,6 +804,7 @@ impl BulkMemtable { float_field_encoding, ))), compact_dispatcher, + compacting: Arc::new(AtomicBool::new(false)), append_mode, merge_mode, row_group_size, @@ -870,22 +860,27 @@ impl BulkMemtable { .sum() } - /// Returns whether the memtable should be compacted. - fn should_compact(&self) -> bool { - let parts = self.parts.read().unwrap(); - parts.should_merge_parts( - self.config.merge_threshold, - self.config.encode_bytes_threshold, - ) - } - - /// Schedules a compaction task using the CompactDispatcher. + /// Schedules at most one write-triggered compaction at a time. fn schedule_compact(&self) { + if self.config.max_merge_groups == 0 || self.compacting.load(Ordering::Relaxed) { + return; + } + if !self + .parts + .read() + .unwrap() + .may_merge_parts(self.config.merge_threshold) + { + return; + } + let Some(guard) = CompactTaskGuard::try_new(self.compacting.clone()) else { + return; + }; if let Some(dispatcher) = &self.compact_dispatcher { let task = MemCompactTask { metadata: self.metadata.clone(), parts: self.parts.clone(), - config: self.config.clone(), + _guard: guard, compactor: self.compactor.clone(), append_mode: self.append_mode, merge_mode: self.merge_mode, @@ -1110,6 +1105,15 @@ impl PartToMerge { } } + /// Returns the lower bound of input sequences. + fn min_sequence(&self) -> SequenceNumber { + match self { + PartToMerge::Bulk { part, .. } => part.min_sequence, + PartToMerge::Multi { part, .. } => part.min_sequence(), + PartToMerge::Encoded { part, .. } => part.metadata().min_sequence, + } + } + /// Gets the maximum sequence number of this part. fn max_sequence(&self) -> u64 { match self { @@ -1261,12 +1265,14 @@ impl MemtableCompactor { ) -> Result<()> { let start = Instant::now(); - // Collect pre-grouped parts - let collected = bulk_parts.write().unwrap().collect_parts_to_merge( - self.config.merge_threshold, - self.config.max_merge_groups, - self.config.encode_bytes_threshold, - ); + let collected = { + let mut parts = bulk_parts.write().unwrap(); + parts.collect_parts_to_merge( + self.config.merge_threshold, + self.config.max_merge_groups, + self.config.encode_bytes_threshold, + ) + }; if collected.groups.is_empty() { return Ok(()); @@ -1355,6 +1361,11 @@ impl MemtableCompactor { .map(|p| p.max_timestamp()) .max() .unwrap_or(i64::MIN); + let min_sequence = parts_to_merge + .iter() + .map(PartToMerge::min_sequence) + .min() + .unwrap_or(0); let max_sequence = parts_to_merge .iter() .map(|p| p.max_sequence()) @@ -1450,7 +1461,7 @@ impl MemtableCompactor { schema, min_timestamp, max_timestamp, - max_sequence, + min_sequence..=max_sequence, &mut metrics, )?; @@ -1476,6 +1487,7 @@ impl MemtableCompactor { batches, min_timestamp, max_timestamp, + min_sequence, max_sequence, estimated_series_count, metadata, @@ -1492,12 +1504,31 @@ impl MemtableCompactor { } } +/// Keeps the compaction flag set while queued or running and clears it on drop, +/// including failures and cancellation while waiting for a dispatcher permit. +struct CompactTaskGuard(Arc); + +impl CompactTaskGuard { + fn try_new(compacting: Arc) -> Option { + // The compactor mutex protects merging; this flag only avoids duplicate tasks. + compacting + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .ok() + .map(|_| Self(compacting)) + } +} + +impl Drop for CompactTaskGuard { + fn drop(&mut self) { + self.0.store(false, Ordering::Relaxed); + } +} + /// A memtable compact task to run in background. struct MemCompactTask { metadata: RegionMetadataRef, parts: Arc>, - /// Configuration for the bulk memtable. - config: BulkMemtableConfig, + _guard: CompactTaskGuard, /// Compactor for merging bulk parts compactor: Arc>, /// Whether the append mode is enabled @@ -1508,22 +1539,12 @@ struct MemCompactTask { impl MemCompactTask { fn compact(&self) -> Result<()> { - let mut compactor = self.compactor.lock().unwrap(); - - let should_merge = self.parts.read().unwrap().should_merge_parts( - self.config.merge_threshold, - self.config.encode_bytes_threshold, - ); - if should_merge { - compactor.merge_parts( - &self.parts, - &self.metadata, - !self.append_mode, - self.merge_mode, - )?; - } - - Ok(()) + self.compactor.lock().unwrap().merge_parts( + &self.parts, + &self.metadata, + !self.append_mode, + self.merge_mode, + ) } } @@ -1794,11 +1815,16 @@ mod tests { assert_eq!(expected_min, memtable.stats().min_sequence); assert_eq!(expected_min, memtable.min_sequence()); assert_eq!(expected_ranges, memtable.stats().num_ranges); + assert!(!memtable.compacting.load(Ordering::Relaxed)); } // The third write synchronously merges all parts without a dispatcher. // Dedup can remove the oldest rows, but their lower bound remains safe. assert_eq!(10, memtable.stats().min_sequence); assert_eq!(201, memtable.stats().max_sequence); + let ranges = memtable.ranges(None, RangesOptions::default()).unwrap(); + let merged = ranges.ranges.values().next().unwrap(); + assert!(merged.encoded().is_some()); + assert_eq!(10, merged.stats().min_sequence); let fork = memtable.fork(2, &metadata); assert!(fork.is_empty()); assert_eq!(0, fork.min_sequence()); @@ -2615,7 +2641,7 @@ mod tests { assert_eq!(5, large_part.num_rows()); memtable.write_bulk(large_part).unwrap(); - // Write another small part (2 rows) - should trigger compaction of unordered_part + // A later small write fills the unordered part to its compact threshold. let part = create_bulk_part_with_converter( "small_2", 2, @@ -2780,10 +2806,10 @@ mod tests { } #[test] - fn test_should_merge_parts_below_threshold() { + fn test_may_merge_parts_below_threshold() { let mut bulk_parts = BulkParts::default(); - // Add 7 bulk parts (below DEFAULT_MERGE_THRESHOLD of 8) + // Add one fewer part than the merge threshold. for i in 0..DEFAULT_MERGE_THRESHOLD - 1 { let part = create_bulk_part_with_converter( &format!("key_{}", i), @@ -2794,14 +2820,18 @@ mod tests { ) .unwrap(); bulk_parts.parts.push(create_bulk_part_wrapper(part)); + if i == 0 { + assert!(!bulk_parts.may_merge_parts(0)); + assert!(!bulk_parts.may_merge_parts(1)); + } } - // Should not trigger merge since we have only 7 parts - assert!(!bulk_parts.should_merge_parts(DEFAULT_MERGE_THRESHOLD, usize::MAX)); + // There are not enough parts to merge. + assert!(!bulk_parts.may_merge_parts(DEFAULT_MERGE_THRESHOLD)); } #[test] - fn test_should_merge_parts_at_threshold() { + fn test_may_merge_parts_at_threshold() { let mut bulk_parts = BulkParts::default(); let merge_threshold = 8; @@ -2819,7 +2849,7 @@ mod tests { } // Should trigger merge since we have 8 parts - assert!(bulk_parts.should_merge_parts(merge_threshold, usize::MAX)); + assert!(bulk_parts.may_merge_parts(merge_threshold)); } #[test] @@ -2839,7 +2869,7 @@ mod tests { } let max_size = bulk_parts.parts[0].part.estimated_size(); - assert!(bulk_parts.should_merge_parts(merge_threshold, max_size)); + assert!(bulk_parts.may_merge_parts(merge_threshold)); assert_eq!( 1, bulk_parts @@ -2850,7 +2880,7 @@ mod tests { } #[test] - fn test_should_merge_parts_with_merging_flag() { + fn test_may_merge_parts_with_merging_flag() { let mut bulk_parts = BulkParts::default(); let merge_threshold = 8; @@ -2868,7 +2898,7 @@ mod tests { } // Should trigger merge since we have 10 parts - assert!(bulk_parts.should_merge_parts(merge_threshold, usize::MAX)); + assert!(bulk_parts.may_merge_parts(merge_threshold)); // Mark first 3 parts as merging for wrapper in bulk_parts.parts.iter_mut().take(3) { @@ -2876,53 +2906,115 @@ mod tests { } // Now only 7 parts are available for merging, should not trigger - assert!(!bulk_parts.should_merge_parts(merge_threshold, usize::MAX)); + assert!(!bulk_parts.may_merge_parts(merge_threshold)); } #[test] - fn test_collect_parts_to_merge_grouping() { - let mut bulk_parts = BulkParts::default(); + fn test_compact_task_guard() { + let compacting = Arc::new(AtomicBool::new(false)); + let guard = CompactTaskGuard::try_new(compacting.clone()).unwrap(); + assert!(compacting.load(Ordering::Relaxed)); + assert!(CompactTaskGuard::try_new(compacting.clone()).is_none()); + drop(guard); + assert!(!compacting.load(Ordering::Relaxed)); + assert!(CompactTaskGuard::try_new(compacting.clone()).is_some()); + assert!(!compacting.load(Ordering::Relaxed)); + } - // Add 16 bulk parts with different row counts - for i in 0..16 { - let num_rows = (i % 4) + 1; // 1 to 4 rows - let timestamps: Vec = (0..num_rows) - .map(|j| 1000 + i as i64 * 100 + j as i64) - .collect(); - let values: Vec> = - (0..num_rows).map(|j| Some((i * 10 + j) as f64)).collect(); - let part = create_bulk_part_with_converter( - &format!("key_{}", i), - i as u32, - timestamps, - values, - 100 + i as u64, - ) - .unwrap(); - bulk_parts.parts.push(create_bulk_part_wrapper(part)); - } - - // Should trigger merge since we have 16 parts - assert!(bulk_parts.should_merge_parts(DEFAULT_MERGE_THRESHOLD, usize::MAX)); - - // Collect parts to merge - let collected = bulk_parts.collect_parts_to_merge( - DEFAULT_MERGE_THRESHOLD, - DEFAULT_MAX_MERGE_GROUPS, - usize::MAX, + #[test] + fn test_compact_scheduling_disabled() { + let memtable = BulkMemtable::new( + 0, + BulkMemtableConfig { + merge_threshold: 2, + max_merge_groups: 0, + ..Default::default() + }, + metadata_for_test(), + None, + Some(Arc::new(CompactDispatcher::new(0))), + false, + MergeMode::LastRow, ); + let part = + create_bulk_part_with_converter("key", 0, vec![1000], vec![Some(1.0)], 0).unwrap(); + memtable.parts.write().unwrap().parts = vec![ + create_bulk_part_wrapper(part.clone()), + create_bulk_part_wrapper(part), + ]; + memtable.schedule_compact(); + assert!(!memtable.compacting.load(Ordering::Relaxed)); + } - // Should have groups - assert!(!collected.groups.is_empty()); + #[test] + fn test_collect_parts_to_merge_sequence_components() { + let make_part = |min, max| { + let mut part = + create_bulk_part_with_converter("key", 0, vec![1000], vec![Some(1.0)], max) + .unwrap(); + part.min_sequence = min; + part + }; + let metadata = metadata_for_test(); + let encoder = BulkPartEncoder::new(metadata.clone(), DEFAULT_ROW_GROUP_SIZE).unwrap(); + // Deliberately use installation order rather than sequence order. + let mut parts = BulkParts { + parts: vec![ + create_bulk_part_wrapper(make_part(70, 70)), + create_bulk_part_wrapper(make_part(10, 30)), + BulkPartWrapper { + part: PartToMerge::Multi { + part: MultiBulkPart::from_bulk_part(make_part(35, 50), &metadata), + file_id: FileId::random(), + }, + merging: false, + }, + create_bulk_part_wrapper(make_part(60, 60)), + BulkPartWrapper { + part: PartToMerge::Encoded { + part: encoder.encode_part(&make_part(20, 40)).unwrap().unwrap(), + file_id: FileId::random(), + }, + merging: false, + }, + ], + ..Default::default() + }; + // The oversized encoded part connects three intervals transitively; + // neither the merge threshold nor its size may split that component. + assert!(parts.may_merge_parts(parts.parts.len())); + assert_eq!( + vec![vec![1, 4, 2], vec![3, 0]], + parts.pick_merge_groups(2, 2, 1) + ); + // The hint may schedule a merge even when unordered writes block all parts. + parts.unordered_part.push(make_part(0, 100)); + assert!(parts.may_merge_parts(2)); + assert!(parts.pick_merge_groups(2, 2, 1).is_empty()); + parts.unordered_part.clear(); - // All groups should have parts - for group in &collected.groups { - assert!(!group.is_empty()); - } + // Touching the unordered interval blocks the entire transitive component. + parts.unordered_part.push(make_part(50, 55)); + assert_eq!(vec![vec![3, 0]], parts.pick_merge_groups(2, 2, 1)); + parts.unordered_part.clear(); - // Total parts collected should be 16 - let total_parts: usize = collected.groups.iter().map(|g| g.len()).sum(); - assert_eq!(16, total_parts); + // Even without overlap, a group cannot straddle pending small writes. + parts.unordered_part.push(make_part(65, 65)); + assert_eq!(vec![vec![1, 4, 2]], parts.pick_merge_groups(2, 2, 1)); + parts.unordered_part.clear(); + + parts.parts[4].merging = true; + let selected = parts.collect_parts_to_merge(2, 2, 1); + assert_eq!(1, selected.groups.len()); + assert_eq!( + vec![60, 70], + selected.groups[0] + .iter() + .map(PartToMerge::max_sequence) + .collect::>() + ); + assert!(!parts.parts[1].merging); + assert!(!parts.parts[2].merging); } #[test] diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 444cc59e317..fecfceafeee 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -15,6 +15,7 @@ //! Bulk part encoder/decoder. use std::collections::{HashMap, HashSet}; +use std::ops::RangeInclusive; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -436,6 +437,11 @@ impl UnorderedPart { self.total_bytes } + /// Returns the sequence interval occupied by pending small writes. + pub(super) fn sequence_bounds(&self) -> Option<(SequenceNumber, SequenceNumber)> { + (!self.is_empty()).then_some((self.min_sequence, self.max_sequence)) + } + /// Adds a BulkPart to this unordered collection. pub fn push(&mut self, part: BulkPart) { self.total_rows += part.num_rows(); @@ -1204,7 +1210,7 @@ impl EncodedBulkPart { num_rows: meta.num_rows, num_ranges: 1, max_sequence: meta.max_sequence, - min_sequence: 0, + min_sequence: meta.min_sequence, series_count: meta.num_series as usize, } } @@ -1291,6 +1297,8 @@ pub struct BulkPartMeta { pub region_metadata: RegionMetadataRef, /// Number of series. pub num_series: u64, + /// Lower bound of input sequences, including rows removed by deduplication. + pub min_sequence: SequenceNumber, /// Maximum sequence number in part. pub max_sequence: u64, } @@ -1356,7 +1364,7 @@ impl BulkPartEncoder { arrow_schema: SchemaRef, min_timestamp: i64, max_timestamp: i64, - max_sequence: u64, + sequence_bounds: RangeInclusive, metrics: &mut BulkPartEncodeMetrics, ) -> Result> { let mut buf = Vec::with_capacity(4096); @@ -1408,7 +1416,8 @@ impl BulkPartEncoder { parquet_metadata, region_metadata: self.metadata.clone(), num_series, - max_sequence, + min_sequence: *sequence_bounds.start(), + max_sequence: *sequence_bounds.end(), }, schema: arrow_schema, })) @@ -1444,6 +1453,7 @@ impl BulkPartEncoder { region_metadata: self.metadata.clone(), num_series: part.estimated_series_count() as u64, max_sequence: part.sequence, + min_sequence: part.min_sequence, }, schema: arrow_schema, })) @@ -1653,6 +1663,8 @@ pub struct MultiBulkPart { max_timestamp: i64, /// Min timestamp in part. min_timestamp: i64, + /// Lower bound of input sequences, including rows removed by deduplication. + min_sequence: SequenceNumber, /// Max sequence number in part. max_sequence: SequenceNumber, /// Number of series. @@ -1676,6 +1688,7 @@ impl MultiBulkPart { total_rows: num_rows, max_timestamp: part.max_timestamp, min_timestamp: part.min_timestamp, + min_sequence: part.min_sequence, max_sequence: part.sequence, series_count, batch_stats, @@ -1688,6 +1701,7 @@ impl MultiBulkPart { /// * `batches` - Ordered record batches /// * `min_timestamp` - Minimum timestamp across all batches /// * `max_timestamp` - Maximum timestamp across all batches + /// * `min_sequence` - Lower bound of input sequences /// * `max_sequence` - Maximum sequence number across all batches /// * `series_count` - Number of series in the batches /// * `metadata` - Region metadata for computing batch statistics @@ -1698,6 +1712,7 @@ impl MultiBulkPart { batches: Vec, min_timestamp: i64, max_timestamp: i64, + min_sequence: SequenceNumber, max_sequence: SequenceNumber, series_count: usize, metadata: &RegionMetadata, @@ -1712,6 +1727,7 @@ impl MultiBulkPart { total_rows, max_timestamp, min_timestamp, + min_sequence, max_sequence, series_count, batch_stats, @@ -1737,6 +1753,11 @@ impl MultiBulkPart { self.max_timestamp } + /// Returns the lower bound of input sequences. + pub fn min_sequence(&self) -> SequenceNumber { + self.min_sequence + } + /// Returns the maximum sequence number. pub fn max_sequence(&self) -> SequenceNumber { self.max_sequence @@ -1827,7 +1848,7 @@ impl MultiBulkPart { num_rows: self.num_rows(), num_ranges: 1, max_sequence: self.max_sequence, - min_sequence: 0, + min_sequence: self.min_sequence, series_count: self.series_count, } } @@ -3147,12 +3168,14 @@ mod tests { let mut all_batches = Vec::new(); let mut min_ts = i64::MAX; let mut max_ts = i64::MIN; + let mut min_seq = u64::MAX; let mut max_seq = 0u64; for inputs in groups { let part = build_converted_bulk_part(inputs); min_ts = min_ts.min(part.min_timestamp); max_ts = max_ts.max(part.max_timestamp); + min_seq = min_seq.min(part.min_sequence); max_seq = max_seq.max(part.sequence); all_batches.push(part.batch); } @@ -3161,6 +3184,7 @@ mod tests { all_batches, min_ts, max_ts, + min_seq, max_seq, groups.len(), &metadata, diff --git a/src/mito2/src/read/flat_dedup.rs b/src/mito2/src/read/flat_dedup.rs index 64c78895a99..b166d4c7f32 100644 --- a/src/mito2/src/read/flat_dedup.rs +++ b/src/mito2/src/read/flat_dedup.rs @@ -160,8 +160,8 @@ impl Drop for FlatDedupReader { /// Strategy to remove duplicate rows from sorted record batches. pub trait RecordBatchDedupStrategy: Send { /// Pushes a batch to the dedup strategy. - /// Returns a batch if the strategy ensures there is no duplications based on - /// the input batch. + /// Returns rows that no longer need to be updated by subsequent input. + /// A strategy for partial merges may retain duplicates separated by Deletes. fn push_batch( &mut self, batch: RecordBatch, @@ -297,11 +297,11 @@ impl RecordBatchDedupStrategy for FlatLastRow { pub struct FlatLastNonNull { /// The start index of field columns: field_column_start: usize, - /// Filter deleted rows. + /// Filter deleted rows; otherwise preserve every delete barrier. filter_deleted: bool, /// Buffered batch to check whether the next batch have duplicated rows with this batch. /// Fields in the last row of this batch may be updated by the next batch. - /// The buffered batch should contain no duplication. + /// In preserve mode, duplicate keys may be separated by delete markers. buffer: Option, /// Whether the last row range contains a delete operation. /// If so, we don't need to update null fields. @@ -326,7 +326,7 @@ impl RecordBatchDedupStrategy for FlatLastNonNull { // If the buffer is None, dedup the batch, put the batch into the buffer and return. // There is no previous batch with the same key, we can pass contains_delete as false. let (record_batch, contains_delete) = - Self::dedup_one_batch(batch, self.field_column_start, false)?; + Self::dedup_one_batch(batch, self.field_column_start, false, self.filter_deleted)?; metrics.num_unselected_rows += row_before_dedup - record_batch.num_rows(); self.buffer = BatchLastRow::try_new(record_batch); self.contains_delete = contains_delete; @@ -341,7 +341,7 @@ impl RecordBatchDedupStrategy for FlatLastNonNull { // Dedup the batch. // There is no previous batch with the same key, we can pass contains_delete as false. let (record_batch, contains_delete) = - Self::dedup_one_batch(batch, self.field_column_start, false)?; + Self::dedup_one_batch(batch, self.field_column_start, false, self.filter_deleted)?; metrics.num_unselected_rows += row_before_dedup - record_batch.num_rows(); debug_assert!(record_batch.num_rows() > 0); self.buffer = BatchLastRow::try_new(record_batch); @@ -369,8 +369,12 @@ impl RecordBatchDedupStrategy for FlatLastNonNull { let merged = concat_batches(&schema, &[last_row, batch]).context(ComputeArrowSnafu)?; let merged_row_count = merged.num_rows(); // Dedup the merged batch and update the buffer. - let (record_batch, contains_delete) = - Self::dedup_one_batch(merged, self.field_column_start, self.contains_delete)?; + let (record_batch, contains_delete) = Self::dedup_one_batch( + merged, + self.field_column_start, + self.contains_delete, + self.filter_deleted, + )?; metrics.num_unselected_rows += merged_row_count - record_batch.num_rows(); debug_assert!(record_batch.num_rows() > 0); self.buffer = BatchLastRow::try_new(record_batch); @@ -397,7 +401,11 @@ impl RecordBatchDedupStrategy for FlatLastNonNull { } impl FlatLastNonNull { - /// Creates a new strategy with the given `filter_deleted` flag. + /// Creates a strategy with the given `filter_deleted` flag. When false, + /// every delete barrier is preserved and Puts are folded only between + /// barriers so later reads cannot fill fields from a deleted row. + /// Inputs for partial merges must form a sequence-contiguous group to avoid + /// folding fields across an intermediate version left outside the merge. pub fn new(field_column_start: usize, filter_deleted: bool) -> Self { Self { field_column_start, @@ -413,6 +421,7 @@ impl FlatLastNonNull { batch: RecordBatch, field_column_start: usize, prev_batch_contains_delete: bool, + filter_deleted: bool, ) -> Result<(RecordBatch, bool)> { // Get op type array for checking delete operations let op_type_column = batch @@ -460,6 +469,7 @@ impl FlatLastNonNull { field_column_start, op_types, prev_batch_contains_delete, + filter_deleted, ) } @@ -471,9 +481,29 @@ impl FlatLastNonNull { field_column_start: usize, op_types: &UInt8Array, first_range_contains_delete: bool, + filter_deleted: bool, ) -> Result<(RecordBatch, bool)> { - let ranges = partitions.ranges(); - let contains_delete = Self::last_range_has_delete(&ranges, op_types); + let mut ranges = partitions.ranges(); + let contains_delete = filter_deleted && Self::last_range_has_delete(&ranges, op_types); + if !filter_deleted && op_types.values().contains(&(OpType::Delete as u8)) { + let mut segments = Vec::with_capacity(ranges.len()); + for range in ranges { + let mut start = range.start; + for i in range.clone() { + if op_types.value(i) == OpType::Delete as u8 { + if start < i { + segments.push(start..i); + } + segments.push(i..i + 1); + start = i + 1; + } + } + if start < range.end { + segments.push(start..range.end); + } + } + ranges = segments; + } // Each range at least has 1 row. let num_duplications: usize = ranges.iter().map(|r| r.end - r.start - 1).sum(); @@ -498,7 +528,7 @@ impl FlatLastNonNull { &ranges, column, op_types, - first_range_contains_delete, + first_range_contains_delete && filter_deleted, ); take(column, &field_indices, take_options.clone()).context(ComputeArrowSnafu) } else { @@ -1239,52 +1269,6 @@ mod tests { check_record_batches_equal(&expected_filter_deleted, &result); assert_eq!(6, dedup_iter.metrics.num_unselected_rows); assert_eq!(2, dedup_iter.metrics.num_deleted_rows); - - // Test with filter_deleted = false - let expected_no_filter = vec![ - new_record_batch_multi_fields( - &[b"k1"], - &[1], - &[13], - &[OpType::Put], - &[(Some(11), Some(11))], - ), - new_record_batch_multi_fields( - &[b"k1", b"k1", b"k1"], - &[2, 3, 4], - &[11, 13, 13], - &[OpType::Put, OpType::Put, OpType::Delete], - &[(Some(12), Some(22)), (Some(13), None), (None, Some(14))], - ), - new_record_batch_multi_fields( - &[b"k2"], - &[1], - &[20], - &[OpType::Put], - &[(Some(101), Some(101))], - ), - new_record_batch_multi_fields( - &[b"k2"], - &[2], - &[20], - &[OpType::Delete], - &[(None, None)], - ), - new_record_batch_multi_fields( - &[b"k3"], - &[2], - &[20], - &[OpType::Put], - &[(Some(202), Some(202))], - ), - ]; - - let iter = input.clone().into_iter().map(Ok); - let mut dedup_iter = FlatDedupIterator::new(iter, FlatLastNonNull::new(1, false)); - let result = collect_iterator_results(&mut dedup_iter); - check_record_batches_equal(&expected_no_filter, &result); - assert_eq!(4, dedup_iter.metrics.num_unselected_rows); - assert_eq!(0, dedup_iter.metrics.num_deleted_rows); } #[test] @@ -1588,6 +1572,69 @@ mod tests { ); } + #[test] + fn test_last_non_null_preserves_delete_barriers() { + let input = new_record_batch_multi_fields( + &[b"k1".as_slice(); 9], + &[1; 9], + &[9, 8, 7, 6, 5, 4, 3, 2, 1], + &[ + OpType::Put, + OpType::Put, + OpType::Delete, + OpType::Delete, + OpType::Put, + OpType::Put, + OpType::Delete, + OpType::Put, + OpType::Put, + ], + &[ + (None, Some(9)), + (Some(8), None), + (None, None), + (None, None), + (None, Some(5)), + (Some(4), None), + (None, None), + (None, None), + (Some(1), Some(1)), + ], + ); + let expected = new_record_batch_multi_fields( + &[b"k1".as_slice(); 6], + &[1; 6], + &[9, 7, 6, 5, 3, 2], + &[ + OpType::Put, + OpType::Delete, + OpType::Delete, + OpType::Put, + OpType::Delete, + OpType::Put, + ], + &[ + (Some(8), Some(9)), + (None, None), + (None, None), + (Some(4), Some(5)), + (None, None), + (Some(1), Some(1)), + ], + ); + // Check both intra-batch partitioning and barriers split across batches. + for batch_size in [1, input.num_rows()] { + let batches = (0..input.num_rows()) + .step_by(batch_size) + .map(|start| Ok(input.slice(start, batch_size.min(input.num_rows() - start)))); + let output = FlatDedupIterator::new(batches, FlatLastNonNull::new(1, false)) + .collect::>>() + .unwrap(); + let actual = concat_batches(&input.schema(), &output).unwrap(); + assert_eq!(expected, actual); + } + } + #[test] fn test_flat_last_non_null_strategy_delete_middle() { let input = vec![