chore: write to mito2 (#2250)

* chore: write to mito2

* fix: clippy

* feat: brdige memtable

* chore: rebase develop

src/mito2/src/engine.rs

@@ -65,8 +65,7 @@ impl MitoEngine {
         region_id: RegionId,
         request: RegionRequest,
     ) -> Result<Output> {
-        self.inner.handle_request(region_id, request).await?;
-        Ok(Output::AffectedRows(0))
+        self.inner.handle_request(region_id, request).await
     }
 
     /// Returns true if the specific region exists.
@@ -98,9 +97,8 @@ impl EngineInner {
         self.workers.stop().await
     }
 
-    // TODO(yingwen): return `Output` instead of `Result<()>`.
     /// Handles [RequestBody] and return its executed result.
-    async fn handle_request(&self, region_id: RegionId, request: RegionRequest) -> Result<()> {
+    async fn handle_request(&self, region_id: RegionId, request: RegionRequest) -> Result<Output> {
         let body = RequestBody::try_from_region_request(region_id, request)?;
         let (request, receiver) = RegionTask::from_request(region_id, body);
         self.workers.submit_to_worker(request).await?;

src/mito2/src/engine/tests.rs

@@ -16,7 +16,11 @@
 
 use std::collections::HashMap;
 
-use store_api::region_request::{RegionCloseRequest, RegionOpenRequest};
+use api::helper::ColumnDataTypeWrapper;
+use api::v1::value::ValueData;
+use api::v1::{Row, Rows};
+use store_api::metadata::ColumnMetadata;
+use store_api::region_request::{RegionCloseRequest, RegionOpenRequest, RegionPutRequest};
 use store_api::storage::RegionId;
 
 use super::*;
@@ -220,3 +224,64 @@ async fn test_engine_reopen_region() {
         .unwrap();
     assert!(engine.is_region_exists(region_id));
 }
+
+fn column_metadata_to_column_schema(metadata: &ColumnMetadata) -> api::v1::ColumnSchema {
+    api::v1::ColumnSchema {
+        column_name: metadata.column_schema.name.clone(),
+        datatype: ColumnDataTypeWrapper::try_from(metadata.column_schema.data_type.clone())
+            .unwrap()
+            .datatype() as i32,
+        semantic_type: metadata.semantic_type as i32,
+    }
+}
+
+fn build_rows(num_rows: usize) -> Vec<Row> {
+    (0..num_rows)
+        .map(|i| api::v1::Row {
+            values: vec![
+                api::v1::Value {
+                    value_data: Some(ValueData::StringValue(i.to_string())),
+                },
+                api::v1::Value {
+                    value_data: Some(ValueData::F64Value(i as f64)),
+                },
+                api::v1::Value {
+                    value_data: Some(ValueData::TsMillisecondValue(i as i64)),
+                },
+            ],
+        })
+        .collect()
+}
+
+#[tokio::test]
+async fn test_write_to_region() {
+    let env = TestEnv::with_prefix("write-to-region");
+    let engine = env.create_engine(MitoConfig::default()).await;
+
+    let region_id = RegionId::new(1, 1);
+    let request = CreateRequestBuilder::new().build();
+
+    let column_schemas = request
+        .column_metadatas
+        .iter()
+        .map(column_metadata_to_column_schema)
+        .collect::<Vec<_>>();
+    engine
+        .handle_request(region_id, RegionRequest::Create(request))
+        .await
+        .unwrap();
+
+    let num_rows = 42;
+    let rows = Rows {
+        schema: column_schemas,
+        rows: build_rows(num_rows),
+    };
+    let output = engine
+        .handle_request(region_id, RegionRequest::Put(RegionPutRequest { rows }))
+        .await
+        .unwrap();
+    let Output::AffectedRows(rows_inserted) = output else {
+        unreachable!()
+    };
+    assert_eq!(num_rows, rows_inserted);
+}

src/mito2/src/memtable/time_series.rs

@@ -15,6 +15,7 @@
 use std::collections::btree_map::Entry;
 use std::collections::{BTreeMap, Bound, HashSet};
 use std::fmt::{Debug, Formatter};
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Arc, RwLock};
 
 use api::v1::OpType;
@@ -31,13 +32,27 @@ use store_api::metadata::RegionMetadataRef;
 use store_api::storage::{ColumnId, ScanRequest};
 
 use crate::error::{ComputeArrowSnafu, ConvertVectorSnafu, PrimaryKeyLengthMismatchSnafu, Result};
-use crate::memtable::{BoxedBatchIterator, KeyValues, Memtable, MemtableId};
+use crate::memtable::{
+    BoxedBatchIterator, KeyValues, Memtable, MemtableBuilder, MemtableId, MemtableRef,
+};
 use crate::read::{Batch, BatchBuilder, BatchColumn};
 use crate::row_converter::{McmpRowCodec, RowCodec, SortField};
 
 /// Initial vector builder capacity.
 const INITIAL_BUILDER_CAPACITY: usize = 32;
 
+#[derive(Debug, Default)]
+pub struct TimeSeriesMemtableBuilder {
+    id: AtomicU32,
+}
+
+impl MemtableBuilder for TimeSeriesMemtableBuilder {
+    fn build(&self, metadata: &RegionMetadataRef) -> MemtableRef {
+        let id = self.id.fetch_add(1, Ordering::Relaxed);
+        Arc::new(TimeSeriesMemtable::new(metadata.clone(), id))
+    }
+}
+
 /// Memtable implementation that groups rows by their primary key.
 pub struct TimeSeriesMemtable {
     id: MemtableId,
@@ -47,7 +62,7 @@ pub struct TimeSeriesMemtable {
 }
 
 impl TimeSeriesMemtable {
-    pub fn new(region_metadata: RegionMetadataRef, id: MemtableId) -> Result<Self> {
+    pub fn new(region_metadata: RegionMetadataRef, id: MemtableId) -> Self {
         let row_codec = McmpRowCodec::new(
             region_metadata
                 .primary_key_columns()
@@ -55,12 +70,12 @@ impl TimeSeriesMemtable {
                 .collect(),
         );
         let series_set = SeriesSet::new(region_metadata.clone());
-        Ok(Self {
+        Self {
             id,
             region_metadata,
             series_set,
             row_codec,
-        })
+        }
     }
 }
 
@@ -749,7 +764,7 @@ mod tests {
         common_telemetry::init_default_ut_logging();
         let schema = schema_for_test();
         let kvs = build_key_values(&schema, "hello".to_string(), 42, 100);
-        let memtable = TimeSeriesMemtable::new(schema, 42).unwrap();
+        let memtable = TimeSeriesMemtable::new(schema, 42);
         memtable.write(&kvs).unwrap();
 
         let expected_ts = kvs
@@ -780,7 +795,7 @@ mod tests {
         common_telemetry::init_default_ut_logging();
         let schema = schema_for_test();
         let kvs = build_key_values(&schema, "hello".to_string(), 42, 100);
-        let memtable = TimeSeriesMemtable::new(schema, 42).unwrap();
+        let memtable = TimeSeriesMemtable::new(schema, 42);
         memtable.write(&kvs).unwrap();
 
         let iter = memtable.iter(ScanRequest {

src/mito2/src/request.rs

@@ -23,6 +23,7 @@ use api::helper::{
 };
 use api::v1::{ColumnDataType, ColumnSchema, OpType, Rows, Value};
 use common_base::readable_size::ReadableSize;
+use common_query::Output;
 use snafu::{ensure, OptionExt, ResultExt};
 use store_api::metadata::{ColumnMetadata, RegionMetadata};
 use store_api::region_request::{
@@ -313,7 +314,7 @@ pub(crate) fn validate_proto_value(
 /// Sender and write request.
 pub(crate) struct SenderWriteRequest {
     /// Result sender.
-    pub(crate) sender: Option<Sender<Result<()>>>,
+    pub(crate) sender: Option<Sender<Result<Output>>>,
     pub(crate) request: WriteRequest,
 }
 
@@ -333,7 +334,7 @@ pub(crate) struct RegionTask {
     ///
     /// Now the result is a `Result<()>`, but we could replace the empty tuple
     /// with an enum if we need to carry more information.
-    pub(crate) sender: Option<Sender<Result<()>>>,
+    pub(crate) sender: Option<Sender<Result<Output>>>,
     /// Request body.
     pub(crate) body: RequestBody,
     /// Region identifier.
@@ -345,7 +346,7 @@ impl RegionTask {
     pub(crate) fn from_request(
         region_id: RegionId,
         body: RequestBody,
-    ) -> (RegionTask, Receiver<Result<()>>) {
+    ) -> (RegionTask, Receiver<Result<Output>>) {
         let (sender, receiver) = oneshot::channel();
         (
             RegionTask {

src/mito2/src/worker.rs

@@ -24,6 +24,7 @@ use std::hash::{Hash, Hasher};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
+use common_query::Output;
 use common_runtime::JoinHandle;
 use common_telemetry::{error, info, warn};
 use futures::future::try_join_all;
@@ -36,7 +37,8 @@ use tokio::sync::{mpsc, Mutex};
 
 use crate::config::MitoConfig;
 use crate::error::{JoinSnafu, Result, WorkerStoppedSnafu};
-use crate::memtable::{DefaultMemtableBuilder, MemtableBuilderRef};
+use crate::memtable::time_series::TimeSeriesMemtableBuilder;
+use crate::memtable::MemtableBuilderRef;
 use crate::region::{MitoRegionRef, RegionMap, RegionMapRef};
 use crate::request::{RegionTask, RequestBody, SenderWriteRequest, WorkerRequest};
 use crate::wal::Wal;
@@ -188,7 +190,7 @@ impl RegionWorker {
             wal: Wal::new(log_store),
             object_store,
             running: running.clone(),
-            memtable_builder: Arc::new(DefaultMemtableBuilder::default()),
+            memtable_builder: Arc::new(TimeSeriesMemtableBuilder::default()),
         };
         let handle = common_runtime::spawn_write(async move {
             worker_thread.run().await;
@@ -373,7 +375,7 @@ impl<S> RegionWorkerLoop<S> {
         }
 
         for task in ddl_tasks {
-            let res: std::result::Result<(), crate::error::Error> = match task.body {
+            let res: std::result::Result<Output, crate::error::Error> = match task.body {
                 RequestBody::Create(req) => self.handle_create_request(task.region_id, req).await,
                 RequestBody::Open(req) => self.handle_open_request(task.region_id, req).await,
                 RequestBody::Close(_) => self.handle_close_request(task.region_id).await,

src/mito2/src/worker/handle_close.rs

@@ -14,6 +14,7 @@
 
 //! Handling close request.
 
+use common_query::Output;
 use common_telemetry::info;
 use store_api::storage::RegionId;
 
@@ -21,9 +22,9 @@ use crate::error::Result;
 use crate::worker::RegionWorkerLoop;
 
 impl<S> RegionWorkerLoop<S> {
-    pub(crate) async fn handle_close_request(&mut self, region_id: RegionId) -> Result<()> {
+    pub(crate) async fn handle_close_request(&mut self, region_id: RegionId) -> Result<Output> {
         let Some(region) = self.regions.get_region(region_id) else {
-            return Ok(());
+            return Ok(Output::AffectedRows(0));
         };
 
         info!("Try to close region {}", region_id);
@@ -33,6 +34,6 @@ impl<S> RegionWorkerLoop<S> {
 
         info!("Region {} closed", region_id);
 
-        Ok(())
+        Ok(Output::AffectedRows(0))
     }
 }

src/mito2/src/worker/handle_create.rs

@@ -16,6 +16,7 @@
 
 use std::sync::Arc;
 
+use common_query::Output;
 use common_telemetry::info;
 use snafu::{ensure, ResultExt};
 use store_api::metadata::RegionMetadataBuilder;
@@ -31,7 +32,7 @@ impl<S> RegionWorkerLoop<S> {
         &mut self,
         region_id: RegionId,
         request: RegionCreateRequest,
-    ) -> Result<()> {
+    ) -> Result<Output> {
         // Checks whether the table exists.
         if self.regions.is_region_exists(region_id) {
             ensure!(
@@ -40,7 +41,7 @@ impl<S> RegionWorkerLoop<S> {
             );
 
             // Region already exists.
-            return Ok(());
+            return Ok(Output::AffectedRows(0));
         }
 
         // Convert the request into a RegionMetadata and validate it.
@@ -70,6 +71,6 @@ impl<S> RegionWorkerLoop<S> {
         // Insert the MitoRegion into the RegionMap.
         self.regions.insert_region(Arc::new(region));
 
-        Ok(())
+        Ok(Output::AffectedRows(0))
     }
 }

src/mito2/src/worker/handle_open.rs

@@ -16,6 +16,7 @@
 
 use std::sync::Arc;
 
+use common_query::Output;
 use common_telemetry::info;
 use store_api::region_request::RegionOpenRequest;
 use store_api::storage::RegionId;
@@ -29,9 +30,9 @@ impl<S> RegionWorkerLoop<S> {
         &mut self,
         region_id: RegionId,
         request: RegionOpenRequest,
-    ) -> Result<()> {
+    ) -> Result<Output> {
         if self.regions.is_region_exists(region_id) {
-            return Ok(());
+            return Ok(Output::AffectedRows(0));
         }
 
         info!("Try to open region {}", region_id);
@@ -51,6 +52,6 @@ impl<S> RegionWorkerLoop<S> {
         // Insert the MitoRegion into the RegionMap.
         self.regions.insert_region(Arc::new(region));
 
-        Ok(())
+        Ok(Output::AffectedRows(0))
     }
 }

src/mito2/src/worker/handle_write.rs

@@ -19,6 +19,7 @@ use std::mem;
 use std::sync::Arc;
 
 use api::v1::{Mutation, WalEntry};
+use common_query::Output;
 use snafu::ResultExt;
 use store_api::logstore::LogStore;
 use store_api::metadata::RegionMetadata;
@@ -123,7 +124,7 @@ fn maybe_fill_missing_columns(request: &mut WriteRequest, metadata: &RegionMetad
 }
 
 /// Send result to the request.
-fn send_result(sender: Option<Sender<Result<()>>>, res: Result<()>) {
+fn send_result(sender: Option<Sender<Result<Output>>>, res: Result<Output>) {
     if let Some(sender) = sender {
         // Ignore send result.
         let _ = sender.send(res);
@@ -135,13 +136,19 @@ struct WriteNotify {
     /// Error to send to the waiter.
     err: Option<Arc<Error>>,
     /// Sender to send write result to the waiter for this mutation.
-    sender: Option<Sender<Result<()>>>,
+    sender: Option<Sender<Result<Output>>>,
+    /// Number of rows to be written.
+    num_rows: usize,
 }
 
 impl WriteNotify {
     /// Creates a new notify from the `sender`.
-    fn new(sender: Option<Sender<Result<()>>>) -> WriteNotify {
-        WriteNotify { err: None, sender }
+    fn new(sender: Option<Sender<Result<Output>>>, num_rows: usize) -> WriteNotify {
+        WriteNotify {
+            err: None,
+            sender,
+            num_rows,
+        }
     }
 
     /// Send result to the waiter.
@@ -154,7 +161,7 @@ impl WriteNotify {
             let _ = sender.send(Err(err.clone()).context(WriteGroupSnafu));
         } else {
             // Send success result.
-            let _ = sender.send(Ok(()));
+            let _ = sender.send(Ok(Output::AffectedRows(self.num_rows)));
         }
     }
 }
@@ -208,7 +215,7 @@ impl RegionWriteCtx {
 
     /// Push [SenderWriteRequest] to the context.
     fn push_sender_request(&mut self, sender_req: SenderWriteRequest) {
-        let num_rows = sender_req.request.rows.rows.len() as u64;
+        let num_rows = sender_req.request.rows.rows.len();
 
         self.wal_entry.mutations.push(Mutation {
             op_type: sender_req.request.op_type as i32,
@@ -216,10 +223,11 @@ impl RegionWriteCtx {
             rows: Some(sender_req.request.rows),
         });
         // Notifiers are 1:1 map to mutations.
-        self.notifiers.push(WriteNotify::new(sender_req.sender));
+        self.notifiers
+            .push(WriteNotify::new(sender_req.sender, num_rows));
 
         // Increase sequence number.
-        self.next_sequence += num_rows;
+        self.next_sequence += num_rows as u64;
     }
 
     /// Encode and add WAL entry to the writer.