From aaa08fe39150c19cbfc6241f873123b0866d0848 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:49:01 +0800 Subject: [PATCH] feat(mito2): adapt bulk memtable encode threshold to write buffer size [Backport release/v1.2] (#9061) * feat(mito2): adapt bulk memtable encode threshold to write buffer size (#9056) * feat(mito2): adapt bulk memtable encode bytes threshold to write buffer size The default encode_bytes_threshold is now max(64MB, min(global_write_buffer_size / 32, 512MB)) instead of a fixed 64MB, so it scales with the memtable budget. GREPTIME_BULK_ENCODE_BYTES_THRESHOLD and the per-region option still override the default. Signed-off-by: Lei, HUANG * test(mito2): clarify binary units in bulk threshold test The threshold test uses powers of 1024, so label its values as MiB and GiB instead of decimal MB and GB. Signed-off-by: Lei, HUANG * refactor(mito2): resolve bulk encode threshold in memtable provider Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG (cherry picked from commit 307fe0a692848422f2eb0d1c51990160437ebd55) * test(compat): normalize CPU-dependent round-robin repartition Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(compat): scope repartition normalization to frontend stage Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: Lei, HUANG Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com> Co-authored-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/memtable.rs | 87 ++++++++++++++++++- src/mito2/src/memtable/bulk.rs | 60 +++++++++++-- src/mito2/src/region/opener.rs | 5 +- src/mito2/src/region/options.rs | 16 ++++ .../verify.result | 1 + .../verify.sql | 1 + 6 files changed, 158 insertions(+), 12 deletions(-) diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index 8c107d15f7..f6481fbcdc 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -14,7 +14,7 @@ //! Memtables are write buffers for regions. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::fmt; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -31,12 +31,12 @@ use mito_codec::row_converter::{PrimaryKeyCodec, build_primary_key_codec}; use snafu::ensure; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; -use store_api::storage::{ColumnId, SequenceNumber, SequenceRange}; +use store_api::storage::{ColumnId, RegionId, SequenceNumber, SequenceRange}; use crate::config::MitoConfig; use crate::error::{InvalidRegionOptionsSnafu, Result, UnsupportedOperationSnafu}; use crate::flush::WriteBufferManagerRef; -use crate::memtable::bulk::{BulkMemtableBuilder, CompactDispatcher}; +use crate::memtable::bulk::{BulkMemtableBuilder, BulkMemtableConfig, CompactDispatcher}; use crate::memtable::time_series::TimeSeriesMemtableBuilder; use crate::metrics::WRITE_BUFFER_BYTES; use crate::read::Batch; @@ -398,6 +398,7 @@ impl Drop for AllocTracker { pub(crate) struct MemtableBuilderProvider { write_buffer_manager: Option, config: Arc, + default_bulk_memtable_config: BulkMemtableConfig, compact_dispatcher: Arc, } @@ -428,14 +429,31 @@ impl MemtableBuilderProvider { ) -> Self { let compact_dispatcher = Arc::new(CompactDispatcher::new(config.max_background_compactions)); + let default_bulk_memtable_config = BulkMemtableConfig::default_for_write_buffer_size( + config.global_write_buffer_size.as_bytes() as usize, + ); Self { write_buffer_manager, config, + default_bulk_memtable_config, compact_dispatcher, } } + /// Parses region options with this provider's default bulk memtable config. + pub(crate) fn parse_options( + &self, + region_id: RegionId, + options: &HashMap, + ) -> Result { + RegionOptions::try_from_options_with_bulk_config( + region_id, + options, + &self.default_bulk_memtable_config, + ) + } + pub(crate) fn builder_for_options(&self, options: &RegionOptions) -> MemtableBuilderRef { let dedup = options.need_dedup(); let merge_mode = options.merge_mode(); @@ -495,6 +513,7 @@ impl MemtableBuilderProvider { !dedup, // append_mode: true if not dedup, false if dedup merge_mode, ) + .with_config(self.default_bulk_memtable_config.clone()) .with_row_group_size(options.row_group_size()) .with_compact_dispatcher(self.compact_dispatcher.clone()); @@ -790,12 +809,15 @@ impl MemtableRange { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; + use common_base::readable_size::ReadableSize; use common_error::ext::WhateverResult; use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use store_api::metadata::RegionMetadataBuilder; + use store_api::storage::RegionId; use super::*; use crate::flush::{WriteBufferManager, WriteBufferManagerImpl}; @@ -875,6 +897,65 @@ mod tests { assert_eq!(&config, builder.config()); } + #[test] + fn test_provider_uses_adaptive_config_for_implicit_bulk_builder() { + let config = MitoConfig { + global_write_buffer_size: ReadableSize::gb(8), + ..Default::default() + }; + let provider = MemtableBuilderProvider::new(None, Arc::new(config)); + let options = RegionOptions::default(); + + let builder = + provider.bulk_memtable_builder(options.need_dedup(), options.merge_mode(), &options); + + assert_eq!(256 * 1024 * 1024, builder.config().encode_bytes_threshold); + } + + #[test] + fn test_provider_parses_bulk_memtable_with_adaptive_config() { + let config = MitoConfig { + global_write_buffer_size: ReadableSize::gb(8), + ..Default::default() + }; + let provider = MemtableBuilderProvider::new(None, Arc::new(config)); + let options = HashMap::from([("memtable.type".to_string(), "bulk".to_string())]); + + let options = provider + .parse_options(RegionId::new(0, 0), &options) + .unwrap(); + + let Some(MemtableOptions::Bulk(config)) = options.memtable else { + panic!("expected bulk memtable options"); + }; + assert_eq!(256 * 1024 * 1024, config.encode_bytes_threshold); + } + + #[test] + fn test_provider_preserves_explicit_bulk_encode_bytes_threshold() { + let config = MitoConfig { + global_write_buffer_size: ReadableSize::gb(8), + ..Default::default() + }; + let provider = MemtableBuilderProvider::new(None, Arc::new(config)); + let options = HashMap::from([ + ("memtable.type".to_string(), "bulk".to_string()), + ( + "memtable.bulk.encode_bytes_threshold".to_string(), + "13".to_string(), + ), + ]); + + let options = provider + .parse_options(RegionId::new(0, 0), &options) + .unwrap(); + + let Some(MemtableOptions::Bulk(config)) = options.memtable else { + panic!("expected bulk memtable options"); + }; + assert_eq!(13, config.encode_bytes_threshold); + } + #[test] fn test_json2_requires_bulk_memtable() -> WhateverResult<()> { let mut metadata = sst_region_metadata(); diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 1b67d487a5..8e1346b992 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -92,15 +92,31 @@ pub(crate) static ENCODE_ROW_THRESHOLD: LazyLock = LazyLock::new(|| { /// Default bytes threshold for encoding. const DEFAULT_ENCODE_BYTES_THRESHOLD: usize = 64 * 1024 * 1024; -/// Bytes threshold for encoding parts. Configurable via `GREPTIME_BULK_ENCODE_BYTES_THRESHOLD`. -/// When estimated bytes exceed this threshold, parts are encoded as EncodedBulkPart. -static ENCODE_BYTES_THRESHOLD: LazyLock = LazyLock::new(|| { - env_usize( - "GREPTIME_BULK_ENCODE_BYTES_THRESHOLD", - DEFAULT_ENCODE_BYTES_THRESHOLD, - ) +/// Maximum bytes threshold for encoding when adapting to the write buffer size. +const MAX_ENCODE_BYTES_THRESHOLD: usize = 512 * 1024 * 1024; + +/// Divisor to derive the encode bytes threshold from the global write buffer size. +const ENCODE_BYTES_THRESHOLD_DIVISOR: usize = 32; + +/// Optional bytes threshold override from `GREPTIME_BULK_ENCODE_BYTES_THRESHOLD`. +static ENCODE_BYTES_THRESHOLD_OVERRIDE: LazyLock> = LazyLock::new(|| { + std::env::var("GREPTIME_BULK_ENCODE_BYTES_THRESHOLD") + .ok() + .and_then(|v| v.parse().ok()) }); +/// Computes the encode bytes threshold adapted to the global write buffer size: +/// `max(64 MiB, min(global_write_buffer_size / 32, 512 MiB))`. +fn adaptive_encode_bytes_threshold(global_write_buffer_bytes: usize) -> usize { + (global_write_buffer_bytes / ENCODE_BYTES_THRESHOLD_DIVISOR) + .clamp(DEFAULT_ENCODE_BYTES_THRESHOLD, MAX_ENCODE_BYTES_THRESHOLD) +} + +/// Returns the default bytes threshold for encoding parts. +fn default_encode_bytes_threshold() -> usize { + ENCODE_BYTES_THRESHOLD_OVERRIDE.unwrap_or(DEFAULT_ENCODE_BYTES_THRESHOLD) +} + /// Configuration for bulk memtable. #[serde_as] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -125,7 +141,7 @@ impl Default for BulkMemtableConfig { Self { merge_threshold: *MERGE_THRESHOLD, encode_row_threshold: *ENCODE_ROW_THRESHOLD, - encode_bytes_threshold: *ENCODE_BYTES_THRESHOLD, + encode_bytes_threshold: default_encode_bytes_threshold(), max_merge_groups: *MAX_MERGE_GROUPS, } .sanitize() @@ -133,6 +149,15 @@ impl Default for BulkMemtableConfig { } impl BulkMemtableConfig { + /// Returns the default config adapted to `global_write_buffer_bytes`. + pub(crate) fn default_for_write_buffer_size(global_write_buffer_bytes: usize) -> Self { + Self { + encode_bytes_threshold: ENCODE_BYTES_THRESHOLD_OVERRIDE + .unwrap_or_else(|| adaptive_encode_bytes_threshold(global_write_buffer_bytes)), + ..Default::default() + } + } + fn sanitize(mut self) -> Self { if self.merge_threshold == 0 { self.merge_threshold = DEFAULT_MERGE_THRESHOLD; @@ -1614,6 +1639,25 @@ mod tests { converter.convert() } + #[test] + fn test_adaptive_encode_bytes_threshold() { + // Below the lower bound: clamped to the default 64 MiB. + assert_eq!( + DEFAULT_ENCODE_BYTES_THRESHOLD, + adaptive_encode_bytes_threshold(1024 * 1024 * 1024) // 1 GiB / 32 = 32 MiB + ); + // In range: global_write_buffer_size / 32. + assert_eq!( + 256 * 1024 * 1024, + adaptive_encode_bytes_threshold(8 * 1024 * 1024 * 1024) // 8 GiB / 32 = 256 MiB + ); + // Above the upper bound: clamped to 512 MiB. + assert_eq!( + MAX_ENCODE_BYTES_THRESHOLD, + adaptive_encode_bytes_threshold(64 * 1024 * 1024 * 1024) // 64 GiB / 32 = 2 GiB + ); + } + #[test] fn test_bulk_memtable_sanitizes_zero_merge_threshold() { let metadata = metadata_for_test(); diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index f4e83272b7..0565e1bdce 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -237,7 +237,10 @@ impl RegionOpener { /// Parses and sets options for the region. pub(crate) fn parse_options(self, options: HashMap) -> Result { let region_id = self.region_id; - self.options(RegionOptions::try_from_options(region_id, &options)?) + let options = self + .memtable_builder_provider + .parse_options(region_id, &options)?; + self.options(options) } /// Sets the replay checkpoint for the region. diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 9023a87711..b2bb23c09a 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -45,6 +45,7 @@ const DEFAULT_INDEX_SEGMENT_ROW_COUNT: usize = 1024; const COMPACTION_TWCS_PREFIX: &str = "compaction.twcs."; const MEMTABLE_PARTITION_TREE_PREFIX: &str = "memtable.partition_tree."; const MEMTABLE_BULK_PREFIX: &str = "memtable.bulk."; +const MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD: &str = "memtable.bulk.encode_bytes_threshold"; /// Legacy memtable type identifier accepted for backward compatibility. /// The partition tree memtable has been removed; parsing this value falls @@ -291,6 +292,21 @@ impl RegionOptions { Ok(opts) } + + /// Parses region options using `default_bulk_config` for an implicit bulk threshold. + pub(crate) fn try_from_options_with_bulk_config( + region_id: RegionId, + options_map: &HashMap, + default_bulk_config: &BulkMemtableConfig, + ) -> Result { + let mut options = Self::try_from_options(region_id, options_map)?; + if !options_map.contains_key(MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD) + && let Some(MemtableOptions::Bulk(config)) = &mut options.memtable + { + config.encode_bytes_threshold = default_bulk_config.encode_bytes_threshold; + } + Ok(options) + } } /// Options for compactions diff --git a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result index 07d0b8c764..266a8a4af1 100644 --- a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result +++ b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result @@ -3,6 +3,7 @@ -- SQLNESS REPLACE (elapsed_compute.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (?m)^(\|\s0_\|\s0_\|_ProjectionExec:.*\n\|_\|_\|_AggregateExec:\smode=Final,.*\n\|_\|_\|_CoalescePartitionsExec.*\n\|_\|_\|_AggregateExec:\smode=Partial,.*\n)\|_\|_\|_RepartitionExec:\spartitioning=RoundRobinBatch\(\d+\),\sinput_partitions=2\b.*\n(\|_\|_\|_MergeScanExec:.*\n) ${1}${2} -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED diff --git a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.sql b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.sql index 81eb489e66..cb0e532b6d 100644 --- a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.sql +++ b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.sql @@ -3,6 +3,7 @@ -- SQLNESS REPLACE (elapsed_compute.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (?m)^(\|\s0_\|\s0_\|_ProjectionExec:.*\n\|_\|_\|_AggregateExec:\smode=Final,.*\n\|_\|_\|_CoalescePartitionsExec.*\n\|_\|_\|_AggregateExec:\smode=Partial,.*\n)\|_\|_\|_RepartitionExec:\spartitioning=RoundRobinBatch\(\d+\),\sinput_partitions=2\b.*\n(\|_\|_\|_MergeScanExec:.*\n) ${1}${2} -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED