feat: support per table memtable options

src/metric-engine/src/engine/create.rs

@@ -36,7 +36,9 @@ use store_api::region_request::{AffectedRows, RegionCreateRequest, RegionRequest
 use store_api::storage::consts::ReservedColumnId;
 use store_api::storage::RegionId;
 
-use crate::engine::options::set_index_options_for_data_region;
+use crate::engine::options::{
+ set_index_options_for_data_region, set_memtable_options_for_data_region,
+};
 use crate::engine::MetricEngineInner;
 use crate::error::{
  ConflictRegionOptionSnafu, CreateMitoRegionSnafu, InternalColumnOccupiedSnafu,
@@ -380,6 +382,9 @@ impl MetricEngineInner {
  // set index options
  set_index_options_for_data_region(&mut data_region_request.options);
 
+ // Set memtable options.
+ set_memtable_options_for_data_region(&mut data_region_request.options);
+
  data_region_request
  }
 

src/metric-engine/src/engine/options.rs

@@ -42,3 +42,8 @@ pub fn set_index_options_for_data_region(options: &mut HashMap<String, String>)
  SEG_ROW_COUNT_FOR_DATA_REGION.to_string(),
  );
 }
+
+/// Set memtable options for the data region.
+pub fn set_memtable_options_for_data_region(options: &mut HashMap<String, String>) {
+ options.insert("memtable.type".to_string(), "experimental".to_string());
+}

src/mito2/Cargo.toml

@@ -58,7 +58,7 @@ prost.workspace = true
 puffin.workspace = true
 rand.workspace = true
 regex = "1.5"
-serde = { version = "1.0", features = ["derive"] }
+serde.workspace = true
 serde_json.workspace = true
 serde_with.workspace = true
 smallvec.workspace = true

src/mito2/src/engine/create_test.rs

@@ -14,12 +14,15 @@
 
 use std::time::Duration;
 
+use api::v1::Rows;
+use common_recordbatch::RecordBatches;
 use store_api::region_engine::RegionEngine;
 use store_api::region_request::{RegionCloseRequest, RegionRequest};
-use store_api::storage::RegionId;
+use store_api::storage::{RegionId, ScanRequest};
 
 use crate::config::MitoConfig;
-use crate::test_util::{CreateRequestBuilder, TestEnv};
+use crate::region::options::MemtableOptions;
+use crate::test_util::{build_rows, put_rows, rows_schema, CreateRequestBuilder, TestEnv};
 
 #[tokio::test]
 async fn test_engine_create_new_region() {
@@ -198,3 +201,45 @@ async fn test_engine_create_with_custom_store() {
  .await
  .unwrap());
 }
+
+#[tokio::test]
+async fn test_engine_create_with_memtable_opts() {
+ let mut env = TestEnv::new();
+ let engine = env.create_engine(MitoConfig::default()).await;
+
+ let region_id = RegionId::new(1, 1);
+ let request = CreateRequestBuilder::new()
+ .insert_option("memtable.type", "experimental")
+ .insert_option("memtable.experimental.index_max_keys_per_shard", "2")
+ .build();
+ let column_schemas = rows_schema(&request);
+ engine
+ .handle_request(region_id, RegionRequest::Create(request))
+ .await
+ .unwrap();
+ let region = engine.get_region(region_id).unwrap();
+ let Some(MemtableOptions::Experimental(memtable_opts)) = &region.version().options.memtable
+ else {
+ unreachable!();
+ };
+ assert_eq!(2, memtable_opts.index_max_keys_per_shard);
+
+ let rows = Rows {
+ schema: column_schemas,
+ rows: build_rows(0, 3),
+ };
+ put_rows(&engine, region_id, rows).await;
+
+ let request = ScanRequest::default();
+ let stream = engine.handle_query(region_id, request).await.unwrap();
+ let batches = RecordBatches::try_collect(stream).await.unwrap();
+ let expected = "\
++-------+---------+---------------------+
+| tag_0 | field_0 | ts |
++-------+---------+---------------------+
+| 0 | 0.0 | 1970-01-01T00:00:00 |
+| 1 | 1.0 | 1970-01-01T00:00:01 |
+| 2 | 2.0 | 1970-01-01T00:00:02 |
++-------+---------+---------------------+";
+ assert_eq!(expected, batches.pretty_print().unwrap());
+}

src/mito2/src/error.rs

@@ -572,6 +572,9 @@ pub enum Error {
  #[snafu(source)]
  error: parquet::errors::ParquetError,
  },
+
+ #[snafu(display("Invalid region options, {}", reason))]
+ InvalidRegionOptions { reason: String, location: Location },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -621,7 +624,8 @@ impl ErrorExt for Error {
  | FillDefault { .. }
  | ConvertColumnDataType { .. }
  | ColumnNotFound { .. }
- | InvalidMetadata { .. } => StatusCode::InvalidArguments,
+ | InvalidMetadata { .. }
+ | InvalidRegionOptions { .. } => StatusCode::InvalidArguments,
 
  InvalidRegionRequestSchemaVersion { .. } => StatusCode::RequestOutdated,
 

src/mito2/src/memtable.rs

@@ -28,9 +28,11 @@ use crate::error::Result;
 use crate::flush::WriteBufferManagerRef;
 use crate::memtable::key_values::KeyValue;
 pub use crate::memtable::key_values::KeyValues;
-use crate::memtable::merge_tree::MergeTreeConfig;
+use crate::memtable::merge_tree::{MergeTreeConfig, MergeTreeMemtableBuilder};
+use crate::memtable::time_series::TimeSeriesMemtableBuilder;
 use crate::metrics::WRITE_BUFFER_BYTES;
 use crate::read::Batch;
+use crate::region::options::MemtableOptions;
 
 pub mod key_values;
 pub mod merge_tree;
@@ -53,7 +55,7 @@ pub enum MemtableConfig {
 
 impl Default for MemtableConfig {
  fn default() -> Self {
- Self::Experimental(MergeTreeConfig::default())
+ Self::TimeSeries
  }
 }
 
@@ -206,6 +208,46 @@ impl Drop for AllocTracker {
  }
 }
 
+/// Provider of memtable builders for regions.
+#[derive(Clone)]
+pub(crate) struct MemtableBuilderProvider {
+ write_buffer_manager: Option<WriteBufferManagerRef>,
+ default_memtable_builder: MemtableBuilderRef,
+}
+
+impl MemtableBuilderProvider {
+ pub(crate) fn new(
+ write_buffer_manager: Option<WriteBufferManagerRef>,
+ default_memtable_builder: MemtableBuilderRef,
+ ) -> Self {
+ Self {
+ write_buffer_manager,
+ default_memtable_builder,
+ }
+ }
+
+ pub(crate) fn builder_for_options(
+ &self,
+ options: Option<&MemtableOptions>,
+ ) -> MemtableBuilderRef {
+ match options {
+ Some(MemtableOptions::TimeSeries) => Arc::new(TimeSeriesMemtableBuilder::new(
+ self.write_buffer_manager.clone(),
+ )),
+ Some(MemtableOptions::Experimental(opts)) => Arc::new(MergeTreeMemtableBuilder::new(
+ MergeTreeConfig {
+ index_max_keys_per_shard: opts.index_max_keys_per_shard,
+ data_freeze_threshold: opts.data_freeze_threshold,
+ fork_dictionary_bytes: opts.fork_dictionary_bytes,
+ ..Default::default()
+ },
+ self.write_buffer_manager.clone(),
+ )),
+ None => self.default_memtable_builder.clone(),
+ }
+ }
+}
+
 #[cfg(test)]
 mod tests {
  use common_base::readable_size::ReadableSize;

src/mito2/src/memtable/merge_tree.rs

@@ -46,6 +46,8 @@ use crate::memtable::{
 
 /// Use `1/DICTIONARY_SIZE_FACTOR` of OS memory as dictionary size.
 const DICTIONARY_SIZE_FACTOR: u64 = 8;
+pub(crate) const DEFAULT_MAX_KEYS_PER_SHARD: usize = 8192;
+pub(crate) const DEFAULT_FREEZE_THRESHOLD: usize = 131072;
 
 /// Id of a shard, only unique inside a partition.
 type ShardId = u32;
@@ -59,6 +61,9 @@ struct PkId {
  pk_index: PkIndex,
 }
 
+// TODO(yingwen): `fork_dictionary_bytes` is per region option, if we have multiple merge
+// tree memtable then we will use a lot memory. We should find a better way to control the
+// dictionary size.
 /// Config for the merge tree memtable.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 #[serde(default)]
@@ -68,6 +73,10 @@ pub struct MergeTreeConfig {
  /// Number of rows to freeze a data part.
  pub data_freeze_threshold: usize,
  /// Whether to delete duplicates rows.
+ ///
+ /// Skips deserializing as it should be determined by whether the
+ /// table is append only.
+ #[serde(skip_deserializing)]
  pub dedup: bool,
  /// Total bytes of dictionary to keep in fork.
  pub fork_dictionary_bytes: ReadableSize,
@@ -539,4 +548,17 @@ mod tests {
  assert_eq!(timestamps, read);
  }
  }
+
+ #[test]
+ fn test_deserialize_config() {
+ let config = MergeTreeConfig {
+ dedup: false,
+ ..Default::default()
+ };
+ // Creates a json with dedup = false.
+ let json = serde_json::to_string(&config).unwrap();
+ let config: MergeTreeConfig = serde_json::from_str(&json).unwrap();
+ assert!(config.dedup);
+ assert_eq!(MergeTreeConfig::default(), config);
+ }
 }

src/mito2/src/region.rs

@@ -32,7 +32,7 @@ use crate::access_layer::AccessLayerRef;
 use crate::error::{RegionNotFoundSnafu, RegionReadonlySnafu, Result};
 use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList};
 use crate::manifest::manager::RegionManifestManager;
-use crate::memtable::MemtableId;
+use crate::memtable::{MemtableBuilderRef, MemtableId};
 use crate::region::version::{VersionControlRef, VersionRef};
 use crate::request::OnFailure;
 use crate::sst::file_purger::FilePurgerRef;
@@ -83,9 +83,10 @@ pub(crate) struct MitoRegion {
  last_flush_millis: AtomicI64,
  /// Whether the region is writable.
  writable: AtomicBool,
-
  /// Provider to get current time.
  time_provider: TimeProviderRef,
+ /// Memtable builder for the region.
+ pub(crate) memtable_builder: MemtableBuilderRef,
 }
 
 pub(crate) type MitoRegionRef = Arc<MitoRegion>;

src/mito2/src/region/opener.rs

@@ -37,7 +37,7 @@ use crate::error::{
 use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions};
 use crate::manifest::storage::manifest_compress_type;
 use crate::memtable::time_partition::TimePartitions;
-use crate::memtable::MemtableBuilderRef;
+use crate::memtable::MemtableBuilderProvider;
 use crate::region::options::RegionOptions;
 use crate::region::version::{VersionBuilder, VersionControl, VersionControlRef};
 use crate::region::MitoRegion;
@@ -53,7 +53,7 @@ use crate::wal::{EntryId, Wal};
 pub(crate) struct RegionOpener {
  region_id: RegionId,
  metadata: Option<RegionMetadata>,
- memtable_builder: MemtableBuilderRef,
+ memtable_builder_provider: MemtableBuilderProvider,
  object_store_manager: ObjectStoreManagerRef,
  region_dir: String,
  scheduler: SchedulerRef,
@@ -69,15 +69,15 @@ impl RegionOpener {
  pub(crate) fn new(
  region_id: RegionId,
  region_dir: &str,
- memtable_builder: MemtableBuilderRef,
+ memtable_builder_provider: MemtableBuilderProvider,
  object_store_manager: ObjectStoreManagerRef,
  scheduler: SchedulerRef,
  intermediate_manager: IntermediateManager,
  ) -> RegionOpener {
  RegionOpener {
  region_id,
  metadata: None,
- memtable_builder,
+ memtable_builder_provider,
  object_store_manager,
  region_dir: normalize_dir(region_dir),
  scheduler,
@@ -171,11 +171,14 @@ impl RegionOpener {
  let manifest_manager =
  RegionManifestManager::new(metadata.clone(), region_manifest_options).await?;
 
+ let memtable_builder = self
+ .memtable_builder_provider
+ .builder_for_options(options.memtable.as_ref());
  // Initial memtable id is 0.
  let part_duration = options.compaction.time_window();
  let mutable = Arc::new(TimePartitions::new(
  metadata.clone(),
- self.memtable_builder,
+ memtable_builder.clone(),
  0,
  part_duration,
  ));
@@ -210,6 +213,7 @@ impl RegionOpener {
  // Region is writable after it is created.
  writable: AtomicBool::new(true),
  time_provider,
+ memtable_builder,
  })
  }
 
@@ -277,11 +281,14 @@ impl RegionOpener {
  access_layer.clone(),
  self.cache_manager.clone(),
  ));
+ let memtable_builder = self
+ .memtable_builder_provider
+ .builder_for_options(region_options.memtable.as_ref());
  // Initial memtable id is 0.
  let part_duration = region_options.compaction.time_window();
  let mutable = Arc::new(TimePartitions::new(
  metadata.clone(),
- self.memtable_builder.clone(),
+ memtable_builder.clone(),
  0,
  part_duration,
  ));
@@ -329,6 +336,7 @@ impl RegionOpener {
  // Region is always opened in read only mode.
  writable: AtomicBool::new(false),
  time_provider,
+ memtable_builder,
  };
  Ok(Some(region))
  }