From 6bf260a05c15c9ed1a5338fe44ddb30767b94953 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 25 Aug 2023 14:18:42 +0800 Subject: [PATCH] chore: write to mito2 (#2250) * chore: write to mito2 * fix: clippy * feat: brdige memtable * chore: rebase develop --- src/mito2/src/engine.rs | 6 +-- src/mito2/src/engine/tests.rs | 67 ++++++++++++++++++++++++++- src/mito2/src/memtable/time_series.rs | 27 ++++++++--- src/mito2/src/request.rs | 7 +-- src/mito2/src/worker.rs | 8 ++-- src/mito2/src/worker/handle_close.rs | 7 +-- src/mito2/src/worker/handle_create.rs | 7 +-- src/mito2/src/worker/handle_open.rs | 7 +-- src/mito2/src/worker/handle_write.rs | 24 ++++++---- 9 files changed, 126 insertions(+), 34 deletions(-) diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index dec53bac08b..219cfd29cda 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -65,8 +65,7 @@ impl MitoEngine { region_id: RegionId, request: RegionRequest, ) -> Result { - 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 { 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?; diff --git a/src/mito2/src/engine/tests.rs b/src/mito2/src/engine/tests.rs index 5e7b96bee47..ff3fc590c95 100644 --- a/src/mito2/src/engine/tests.rs +++ b/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 { + (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::>(); + 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); +} diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index 940028c085f..109832b973c 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/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 { + 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 { diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 03d7d8c17cc..7b139ff4d84 100644 --- a/src/mito2/src/request.rs +++ b/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>>, + pub(crate) sender: Option>>, 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>>, + pub(crate) sender: Option>>, /// 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>) { + ) -> (RegionTask, Receiver>) { let (sender, receiver) = oneshot::channel(); ( RegionTask { diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index ff15536aa1d..cb3b988dbe5 100644 --- a/src/mito2/src/worker.rs +++ b/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 RegionWorkerLoop { } for task in ddl_tasks { - let res: std::result::Result<(), crate::error::Error> = match task.body { + let res: std::result::Result = 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, diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index b3aba7a767e..451a0d934eb 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/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 RegionWorkerLoop { - 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 { 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 RegionWorkerLoop { info!("Region {} closed", region_id); - Ok(()) + Ok(Output::AffectedRows(0)) } } diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index 55d9124fbf2..d46f0138b67 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/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 RegionWorkerLoop { &mut self, region_id: RegionId, request: RegionCreateRequest, - ) -> Result<()> { + ) -> Result { // Checks whether the table exists. if self.regions.is_region_exists(region_id) { ensure!( @@ -40,7 +41,7 @@ impl RegionWorkerLoop { ); // Region already exists. - return Ok(()); + return Ok(Output::AffectedRows(0)); } // Convert the request into a RegionMetadata and validate it. @@ -70,6 +71,6 @@ impl RegionWorkerLoop { // Insert the MitoRegion into the RegionMap. self.regions.insert_region(Arc::new(region)); - Ok(()) + Ok(Output::AffectedRows(0)) } } diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index eac4a9173f9..7848cabcee6 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/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 RegionWorkerLoop { &mut self, region_id: RegionId, request: RegionOpenRequest, - ) -> Result<()> { + ) -> Result { 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 RegionWorkerLoop { // Insert the MitoRegion into the RegionMap. self.regions.insert_region(Arc::new(region)); - Ok(()) + Ok(Output::AffectedRows(0)) } } diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 5d3becb53fd..766f25dca52 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/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>>, res: Result<()>) { +fn send_result(sender: Option>>, res: Result) { 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>, /// Sender to send write result to the waiter for this mutation. - sender: Option>>, + sender: Option>>, + /// Number of rows to be written. + num_rows: usize, } impl WriteNotify { /// Creates a new notify from the `sender`. - fn new(sender: Option>>) -> WriteNotify { - WriteNotify { err: None, sender } + fn new(sender: Option>>, 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.