feat: Implement append mode for a region

src/mito2/src/compaction/twcs.rs

@@ -188,6 +188,7 @@ impl Picker for TwcsPicker {
             cache_manager,
             storage: current_version.options.storage.clone(),
             index_options: current_version.options.index_options.clone(),
+            append_mode: current_version.options.append_mode,
         };
         Some(Box::new(task))
     }
@@ -255,6 +256,8 @@ pub(crate) struct TwcsCompactionTask {
     pub(crate) storage: Option<String>,
     /// Index options of the region.
     pub(crate) index_options: IndexOptions,
+    /// The region is using append mode.
+    pub(crate) append_mode: bool,
 }
 
 impl Debug for TwcsCompactionTask {
@@ -264,6 +267,7 @@ impl Debug for TwcsCompactionTask {
             .field("outputs", &self.outputs)
             .field("expired_ssts", &self.expired_ssts)
             .field("compaction_time_window", &self.compaction_time_window)
+            .field("append_mode", &self.append_mode)
             .finish()
     }
 }
@@ -332,9 +336,15 @@ impl TwcsCompactionTask {
             let cache_manager = self.cache_manager.clone();
             let storage = self.storage.clone();
             let index_options = self.index_options.clone();
+            let append_mode = self.append_mode;
             futs.push(async move {
-                let reader =
-                    build_sst_reader(metadata.clone(), sst_layer.clone(), &output.inputs).await?;
+                let reader = build_sst_reader(
+                    metadata.clone(),
+                    sst_layer.clone(),
+                    &output.inputs,
+                    append_mode,
+                )
+                .await?;
                 let file_meta_opt = sst_layer
                     .write_sst(
                         SstWriteRequest {
@@ -565,9 +575,11 @@ async fn build_sst_reader(
     metadata: RegionMetadataRef,
     sst_layer: AccessLayerRef,
     inputs: &[FileHandle],
+    append_mode: bool,
 ) -> error::Result<BoxedBatchReader> {
     SeqScan::new(sst_layer, ProjectionMapper::all(&metadata)?)
         .with_files(inputs.to_vec())
+        .with_append_mode(append_mode)
         // We ignore file not found error during compaction.
         .with_ignore_file_not_found(true)
         .build_reader()

src/mito2/src/engine.rs

@@ -17,6 +17,8 @@
 #[cfg(test)]
 mod alter_test;
 #[cfg(test)]
+mod append_mode_test;
+#[cfg(test)]
 mod basic_test;
 #[cfg(test)]
 mod catchup_test;

src/mito2/src/engine/append_mode_test.rs

@@ -0,0 +1,153 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Tests for append mode.
+
+use api::v1::Rows;
+use common_recordbatch::RecordBatches;
+use store_api::region_engine::RegionEngine;
+use store_api::region_request::{RegionCompactRequest, RegionRequest};
+use store_api::storage::{RegionId, ScanRequest};
+
+use crate::config::MitoConfig;
+use crate::test_util::{
+    build_rows, build_rows_for_key, flush_region, put_rows, rows_schema, CreateRequestBuilder,
+    TestEnv,
+};
+
+#[tokio::test]
+async fn test_append_mode_write_query() {
+    common_telemetry::init_default_ut_logging();
+
+    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("append_mode", "true")
+        .build();
+
+    let column_schemas = rows_schema(&request);
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+
+    // rows 1, 2
+    let rows = build_rows(1, 3);
+    let rows = Rows {
+        schema: column_schemas.clone(),
+        rows,
+    };
+    put_rows(&engine, region_id, rows).await;
+
+    let mut rows = build_rows(0, 2);
+    rows.append(&mut build_rows(1, 2));
+    // rows 0, 1, 1
+    let rows = Rows {
+        schema: column_schemas,
+        rows,
+    };
+    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 |
+| 1     | 1.0     | 1970-01-01T00:00:01 |
+| 1     | 1.0     | 1970-01-01T00:00:01 |
+| 2     | 2.0     | 1970-01-01T00:00:02 |
++-------+---------+---------------------+";
+    assert_eq!(expected, batches.pretty_print().unwrap());
+}
+
+#[tokio::test]
+async fn test_append_mode_compaction() {
+    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("compaction.type", "twcs")
+        .insert_option("compaction.twcs.max_active_window_files", "2")
+        .insert_option("compaction.twcs.max_inactive_window_files", "2")
+        .insert_option("append_mode", "true")
+        .build();
+
+    let column_schemas = rows_schema(&request);
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+
+    // Flush 2 SSTs for compaction.
+    // a, field 1, 2
+    let rows = Rows {
+        schema: column_schemas.clone(),
+        rows: build_rows_for_key("a", 1, 3, 1),
+    };
+    put_rows(&engine, region_id, rows).await;
+    flush_region(&engine, region_id, None).await;
+    // a, field 0, 1
+    let rows = Rows {
+        schema: column_schemas.clone(),
+        rows: build_rows_for_key("a", 0, 2, 0),
+    };
+    put_rows(&engine, region_id, rows).await;
+    flush_region(&engine, region_id, None).await;
+    // b, field 0, 1
+    let rows = Rows {
+        schema: column_schemas.clone(),
+        rows: build_rows_for_key("b", 0, 2, 0),
+    };
+    put_rows(&engine, region_id, rows).await;
+    flush_region(&engine, region_id, None).await;
+
+    let output = engine
+        .handle_request(region_id, RegionRequest::Compact(RegionCompactRequest {}))
+        .await
+        .unwrap();
+    assert_eq!(output.affected_rows, 0);
+
+    // a, field 2, 3
+    let rows = Rows {
+        schema: column_schemas,
+        rows: build_rows_for_key("a", 2, 4, 2),
+    };
+    put_rows(&engine, region_id, rows).await;
+
+    let scanner = engine.scanner(region_id, ScanRequest::default()).unwrap();
+    assert_eq!(1, scanner.num_files());
+    let stream = scanner.scan().await.unwrap();
+    let batches = RecordBatches::try_collect(stream).await.unwrap();
+    let expected = "\
++-------+---------+---------------------+
+| tag_0 | field_0 | ts                  |
++-------+---------+---------------------+
+| a     | 0.0     | 1970-01-01T00:00:00 |
+| a     | 1.0     | 1970-01-01T00:00:01 |
+| a     | 1.0     | 1970-01-01T00:00:01 |
+| a     | 2.0     | 1970-01-01T00:00:02 |
+| a     | 2.0     | 1970-01-01T00:00:02 |
+| a     | 3.0     | 1970-01-01T00:00:03 |
+| b     | 0.0     | 1970-01-01T00:00:00 |
+| b     | 1.0     | 1970-01-01T00:00:01 |
++-------+---------+---------------------+";
+    assert_eq!(expected, batches.pretty_print().unwrap());
+}

src/mito2/src/memtable.rs

@@ -24,6 +24,7 @@ use store_api::metadata::RegionMetadataRef;
 use store_api::storage::ColumnId;
 use table::predicate::Predicate;
 
+use crate::config::MitoConfig;
 use crate::error::Result;
 use crate::flush::WriteBufferManagerRef;
 use crate::memtable::key_values::KeyValue;
@@ -212,40 +213,59 @@ impl Drop for AllocTracker {
 #[derive(Clone)]
 pub(crate) struct MemtableBuilderProvider {
     write_buffer_manager: Option<WriteBufferManagerRef>,
-    default_memtable_builder: MemtableBuilderRef,
+    config: Arc<MitoConfig>,
 }
 
 impl MemtableBuilderProvider {
     pub(crate) fn new(
         write_buffer_manager: Option<WriteBufferManagerRef>,
-        default_memtable_builder: MemtableBuilderRef,
+        config: Arc<MitoConfig>,
     ) -> Self {
         Self {
             write_buffer_manager,
-            default_memtable_builder,
+            config,
         }
     }
 
     pub(crate) fn builder_for_options(
         &self,
         options: Option<&MemtableOptions>,
+        dedup: bool,
     ) -> MemtableBuilderRef {
         match options {
             Some(MemtableOptions::TimeSeries) => Arc::new(TimeSeriesMemtableBuilder::new(
                 self.write_buffer_manager.clone(),
+                dedup,
             )),
             Some(MemtableOptions::PartitionTree(opts)) => {
                 Arc::new(PartitionTreeMemtableBuilder::new(
                     PartitionTreeConfig {
                         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()
+                        dedup,
                     },
                     self.write_buffer_manager.clone(),
                 ))
             }
-            None => self.default_memtable_builder.clone(),
+            None => self.default_memtable_builder(dedup),
+        }
+    }
+
+    fn default_memtable_builder(&self, dedup: bool) -> MemtableBuilderRef {
+        match &self.config.memtable {
+            MemtableConfig::PartitionTree(config) => {
+                let mut config = config.clone();
+                config.dedup = dedup;
+                Arc::new(PartitionTreeMemtableBuilder::new(
+                    config,
+                    self.write_buffer_manager.clone(),
+                ))
+            }
+            MemtableConfig::TimeSeries => Arc::new(TimeSeriesMemtableBuilder::new(
+                self.write_buffer_manager.clone(),
+                dedup,
+            )),
         }
     }
 }

src/mito2/src/memtable/time_series.rs

@@ -54,13 +54,15 @@ const INITIAL_BUILDER_CAPACITY: usize = 0;
 #[derive(Debug, Default)]
 pub struct TimeSeriesMemtableBuilder {
     write_buffer_manager: Option<WriteBufferManagerRef>,
+    dedup: bool,
 }
 
 impl TimeSeriesMemtableBuilder {
     /// Creates a new builder with specific `write_buffer_manager`.
-    pub fn new(write_buffer_manager: Option<WriteBufferManagerRef>) -> Self {
+    pub fn new(write_buffer_manager: Option<WriteBufferManagerRef>, dedup: bool) -> Self {
         Self {
             write_buffer_manager,
+            dedup,
         }
     }
 }
@@ -71,7 +73,7 @@ impl MemtableBuilder for TimeSeriesMemtableBuilder {
             metadata.clone(),
             id,
             self.write_buffer_manager.clone(),
-            true, // todo(hl): set according to region option
+            self.dedup,
         ))
     }
 }