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

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

* refactor(mito2): resolve bulk encode threshold in memtable provider

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

---------

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
(cherry picked from commit 307fe0a692)

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

@@ -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<WriteBufferManagerRef>,
     config: Arc<MitoConfig>,
+    default_bulk_memtable_config: BulkMemtableConfig,
     compact_dispatcher: Arc<CompactDispatcher>,
 }
 
@@ -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<String, String>,
+    ) -> Result<RegionOptions> {
+        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();

src/mito2/src/memtable/bulk.rs

@@ -92,15 +92,31 @@ pub(crate) static ENCODE_ROW_THRESHOLD: LazyLock<usize> = 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<usize> = 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<Option<usize>> = 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();

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<String, String>) -> Result<Self> {
         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.

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<String, String>,
+        default_bulk_config: &BulkMemtableConfig,
+    ) -> Result<Self> {
+        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

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

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