rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-13 07:22:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mito): preserve version ordering in bulk merges and flushes (#9517)

Browse Source 
* fix(mito): preserve delete barriers in flat last-non-null dedup

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix(mito): respect sequence ranges in bulk merges and flushes

Signed-off-by: evenyag <realevenyag@gmail.com>

* perf(mito): avoid merge selection on the bulk write path

Signed-off-by: evenyag <realevenyag@gmail.com>

---------

Signed-off-by: evenyag <realevenyag@gmail.com>
This commit is contained in:
Yingwen
2026-10-12 13:16:52 +00:00
committed by GitHub
parent 4ef312ecf2
commit 674ba72669
7 changed files with 924 additions and 456 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mito2/src/access_layer.rs
+1 -1
View File
@@ -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<SequenceNumber>,
pub sst_write_format: FormatType,
src/mito2/src/engine/flush_test.rs
+308 -5
View File
@@ -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::<Vec<_>>(),
&[(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::<Vec<_>>(),
&[(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::<Vec<_>>(),
&[(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);
}
src/mito2/src/flush.rs
+153 -177
View File
@@ -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<FlatSources> {
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::<Vec<_>>();
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::<usize>();
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::<usize>();
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<MemtableRange>,
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::<Vec<_>>();
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),
src/mito2/src/memtable.rs
+26
View File
@@ -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<T>(
mut parts: Vec<T>,
bounds: impl Fn(&T) -> (SequenceNumber, SequenceNumber),
) -> Vec<Vec<T>> {
parts.sort_by_key(&bounds);
let mut groups: Vec<Vec<T>> = 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.
src/mito2/src/memtable/bulk.rs
+303 -211
View File
@@ -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<Vec<PartToMerge>>,
}
@@ -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<Vec<usize>> {
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::<Option<Vec<_>>>()
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<Vec<PartToMerge>> {
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<Mutex<MemtableCompactor>>,
/// Dispatcher for scheduling compaction tasks
compact_dispatcher: Option<Arc<CompactDispatcher>>,
/// Whether a write-triggered compaction is queued or running.
compacting: Arc<AtomicBool>,
/// 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<AtomicBool>);
impl CompactTaskGuard {
fn try_new(compacting: Arc<AtomicBool>) -> Option<Self> {
// 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<RwLock<BulkParts>>,
/// Configuration for the bulk memtable.
config: BulkMemtableConfig,
_guard: CompactTaskGuard,
/// Compactor for merging bulk parts
compactor: Arc<Mutex<MemtableCompactor>>,
/// 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<i64> = (0..num_rows)
.map(|j| 1000 + i as i64 * 100 + j as i64)
.collect();
let values: Vec<Option<f64>> =
(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::<Vec<_>>()
);
assert!(!parts.parts[1].merging);
assert!(!parts.parts[2].merging);
}
#[test]
src/mito2/src/memtable/bulk/part.rs
+28 -4
View File
@@ -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<SequenceNumber>,
metrics: &mut BulkPartEncodeMetrics,
) -> Result<Option<EncodedBulkPart>> {
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<RecordBatch>,
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,
src/mito2/src/read/flat_dedup.rs
+105 -58
View File
@@ -160,8 +160,8 @@ impl<I, S> Drop for FlatDedupReader<I, S> {
/// 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<BatchLastRow>,
/// 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::<Result<Vec<_>>>()
.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![