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 7e02c4945c..ca9baeef34 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 48480f3475..7bfb8d0d65 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 @@ -313,6 +314,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