From 07de65d2ac5338d1c53038115cb22131f9d5d2c5 Mon Sep 17 00:00:00 2001 From: Wei <15172118655@163.com> Date: Mon, 15 Jan 2024 18:02:53 +0800 Subject: [PATCH] test: engine with write cache (#3163) * feat: write cache test for engine * chore: unused * chore: comment * refactor: super to crate * chore: cr comment Co-authored-by: Yingwen * chore: clippy --------- Co-authored-by: Yingwen --- src/datatypes/src/types/duration_type.rs | 2 +- src/datatypes/src/types/primitive_type.rs | 2 +- src/datatypes/src/vectors/decimal.rs | 2 +- src/datatypes/src/vectors/helper.rs | 10 ++--- src/datatypes/src/vectors/interval.rs | 2 +- src/mito2/src/cache.rs | 4 +- src/mito2/src/cache/file_cache.rs | 3 +- src/mito2/src/cache/write_cache.rs | 30 ++++--------- src/mito2/src/config.rs | 9 ++++ src/mito2/src/engine/basic_test.rs | 43 +++++++++++++++++++ src/mito2/src/sst/parquet.rs | 2 +- .../src/req_convert/insert/stmt_to_region.rs | 2 +- 12 files changed, 74 insertions(+), 37 deletions(-) diff --git a/src/datatypes/src/types/duration_type.rs b/src/datatypes/src/types/duration_type.rs index db668451d32..e5e5c66be5c 100644 --- a/src/datatypes/src/types/duration_type.rs +++ b/src/datatypes/src/types/duration_type.rs @@ -25,7 +25,6 @@ use paste::paste; use serde::{Deserialize, Serialize}; use snafu::OptionExt; -use super::LogicalPrimitiveType; use crate::data_type::DataType; use crate::duration::{ DurationMicrosecond, DurationMillisecond, DurationNanosecond, DurationSecond, @@ -34,6 +33,7 @@ use crate::error; use crate::prelude::{ ConcreteDataType, LogicalTypeId, MutableVector, ScalarVectorBuilder, Value, ValueRef, Vector, }; +use crate::types::LogicalPrimitiveType; use crate::vectors::{ DurationMicrosecondVector, DurationMicrosecondVectorBuilder, DurationMillisecondVector, DurationMillisecondVectorBuilder, DurationNanosecondVector, DurationNanosecondVectorBuilder, diff --git a/src/datatypes/src/types/primitive_type.rs b/src/datatypes/src/types/primitive_type.rs index 1893b8ac27e..800be3d11be 100644 --- a/src/datatypes/src/types/primitive_type.rs +++ b/src/datatypes/src/types/primitive_type.rs @@ -22,11 +22,11 @@ use num::NumCast; use serde::{Deserialize, Serialize}; use snafu::OptionExt; -use super::boolean_type::bool_to_numeric; use crate::data_type::{ConcreteDataType, DataType}; use crate::error::{self, Result}; use crate::scalars::{Scalar, ScalarRef, ScalarVectorBuilder}; use crate::type_id::LogicalTypeId; +use crate::types::boolean_type::bool_to_numeric; use crate::types::{DateTimeType, DateType}; use crate::value::{Value, ValueRef}; use crate::vectors::{MutableVector, PrimitiveVector, PrimitiveVectorBuilder, Vector}; diff --git a/src/datatypes/src/vectors/decimal.rs b/src/datatypes/src/vectors/decimal.rs index 1daf689c9c9..2aef6863987 100644 --- a/src/datatypes/src/vectors/decimal.rs +++ b/src/datatypes/src/vectors/decimal.rs @@ -23,7 +23,6 @@ use common_decimal::decimal128::{DECIMAL128_DEFAULT_SCALE, DECIMAL128_MAX_PRECIS use common_decimal::Decimal128; use snafu::{OptionExt, ResultExt}; -use super::{MutableVector, Validity, Vector, VectorRef}; use crate::arrow::datatypes::DataType as ArrowDataType; use crate::data_type::ConcreteDataType; use crate::error::{ @@ -33,6 +32,7 @@ use crate::prelude::{ScalarVector, ScalarVectorBuilder}; use crate::serialize::Serializable; use crate::value::{Value, ValueRef}; use crate::vectors; +use crate::vectors::{MutableVector, Validity, Vector, VectorRef}; /// Decimal128Vector is a vector keep i128 values with precision and scale. #[derive(Debug, PartialEq)] diff --git a/src/datatypes/src/vectors/helper.rs b/src/datatypes/src/vectors/helper.rs index 1ec9aec45cc..ef7dab842fb 100644 --- a/src/datatypes/src/vectors/helper.rs +++ b/src/datatypes/src/vectors/helper.rs @@ -25,17 +25,15 @@ use arrow_schema::IntervalUnit; use datafusion_common::ScalarValue; use snafu::{OptionExt, ResultExt}; -use super::{ - Decimal128Vector, DurationMicrosecondVector, DurationMillisecondVector, - DurationNanosecondVector, DurationSecondVector, IntervalDayTimeVector, IntervalYearMonthVector, -}; use crate::data_type::ConcreteDataType; use crate::error::{self, Result}; use crate::scalars::{Scalar, ScalarVectorBuilder}; use crate::value::{ListValue, ListValueRef}; use crate::vectors::{ - BinaryVector, BooleanVector, ConstantVector, DateTimeVector, DateVector, Float32Vector, - Float64Vector, Int16Vector, Int32Vector, Int64Vector, Int8Vector, IntervalMonthDayNanoVector, + BinaryVector, BooleanVector, ConstantVector, DateTimeVector, DateVector, Decimal128Vector, + DurationMicrosecondVector, DurationMillisecondVector, DurationNanosecondVector, + DurationSecondVector, Float32Vector, Float64Vector, Int16Vector, Int32Vector, Int64Vector, + Int8Vector, IntervalDayTimeVector, IntervalMonthDayNanoVector, IntervalYearMonthVector, ListVector, ListVectorBuilder, MutableVector, NullVector, StringVector, TimeMicrosecondVector, TimeMillisecondVector, TimeNanosecondVector, TimeSecondVector, TimestampMicrosecondVector, TimestampMillisecondVector, TimestampNanosecondVector, TimestampSecondVector, UInt16Vector, diff --git a/src/datatypes/src/vectors/interval.rs b/src/datatypes/src/vectors/interval.rs index cf619cfc921..de81c050a92 100644 --- a/src/datatypes/src/vectors/interval.rs +++ b/src/datatypes/src/vectors/interval.rs @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::{PrimitiveVector, PrimitiveVectorBuilder}; use crate::types::{IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType}; +use crate::vectors::{PrimitiveVector, PrimitiveVectorBuilder}; pub type IntervalYearMonthVector = PrimitiveVector; pub type IntervalYearMonthVectorBuilder = PrimitiveVectorBuilder; diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 9372d57b3a8..ffafbadd0cd 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -15,12 +15,10 @@ //! Cache for the engine. mod cache_size; -// TODO(yingwen): Remove this after the write cache is ready. -#[allow(unused)] + pub(crate) mod file_cache; #[cfg(test)] pub(crate) mod test_util; -#[allow(unused)] pub(crate) mod write_cache; use std::mem; diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs index 5fa836e15c1..6e92db6e722 100644 --- a/src/mito2/src/cache/file_cache.rs +++ b/src/mito2/src/cache/file_cache.rs @@ -14,7 +14,7 @@ //! A cache for files. -use std::ops::{Range, RangeBounds}; +use std::ops::Range; use std::sync::Arc; use std::time::Instant; @@ -165,6 +165,7 @@ impl FileCache { } } + #[allow(unused)] /// Removes a file from the cache explicitly. pub(crate) async fn remove(&self, key: IndexKey) { let file_path = self.cache_file_path(key); diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 7a23cda4718..b02ce8f7f3e 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -14,27 +14,20 @@ //! A write-through cache for remote object stores. -use std::ops::Range; use std::sync::Arc; -use api::v1::region; -use bytes::Bytes; use common_base::readable_size::ReadableSize; use common_telemetry::{debug, info}; use object_store::manager::ObjectStoreManagerRef; use object_store::ObjectStore; use snafu::ResultExt; -use store_api::metadata::RegionMetadataRef; -use store_api::storage::RegionId; use crate::access_layer::{new_fs_object_store, SstWriteRequest}; use crate::cache::file_cache::{FileCache, FileCacheRef, FileType, IndexKey, IndexValue}; use crate::error::{self, Result}; use crate::metrics::{FLUSH_ELAPSED, UPLOAD_BYTES_TOTAL}; -use crate::read::Source; -use crate::sst::file::FileId; use crate::sst::index::intermediate::IntermediateManager; -use crate::sst::index::{Indexer, IndexerBuilder}; +use crate::sst::index::IndexerBuilder; use crate::sst::parquet::writer::ParquetWriter; use crate::sst::parquet::{SstInfo, WriteOptions}; use crate::sst::DEFAULT_WRITE_BUFFER_SIZE; @@ -46,6 +39,8 @@ pub struct WriteCache { /// Local file cache. file_cache: FileCacheRef, /// Object store manager. + #[allow(unused)] + /// TODO: Remove unused after implementing async write cache object_store_manager: ObjectStoreManagerRef, /// Intermediate manager for inverted index. intermediate_manager: IntermediateManager, @@ -231,24 +226,17 @@ pub struct SstUploadRequest { #[cfg(test)] mod tests { - use api::v1::OpType; + use common_base::readable_size::ReadableSize; use common_test_util::temp_dir::create_temp_dir; - use object_store::manager::ObjectStoreManager; - use object_store::services::Fs; use object_store::util::join_dir; - use object_store::ObjectStore; - use store_api::storage::RegionId; use super::*; - use crate::cache::file_cache::{self, FileCache}; use crate::cache::test_util::new_fs_store; use crate::sst::file::FileId; use crate::sst::location::{index_file_path, sst_file_path}; - use crate::test_util::sst_util::{ - new_batch_by_range, new_source, sst_file_handle, sst_region_metadata, - }; - use crate::test_util::{build_rows, new_batch_builder, CreateRequestBuilder, TestEnv}; + use crate::test_util::sst_util::{new_batch_by_range, new_source, sst_region_metadata}; + use crate::test_util::TestEnv; #[tokio::test] async fn test_write_and_upload_sst() { @@ -296,7 +284,7 @@ mod tests { cache_manager: Default::default(), }; - let request = SstUploadRequest { + let upload_request = SstUploadRequest { upload_path: upload_path.clone(), index_upload_path: index_upload_path.clone(), remote_store: mock_store.clone(), @@ -308,8 +296,8 @@ mod tests { }; // Write to cache and upload sst to mock remote store - let sst_info = write_cache - .write_and_upload_sst(write_request, request, &write_opts) + write_cache + .write_and_upload_sst(write_request, upload_request, &write_opts) .await .unwrap() .unwrap(); diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index dce77a1233f..748401b622d 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -181,6 +181,15 @@ impl MitoConfig { Ok(()) } + + /// Enable experimental write cache. + #[cfg(test)] + pub fn enable_write_cache(mut self, path: String, size: ReadableSize) -> Self { + self.enable_experimental_write_cache = true; + self.experimental_write_cache_path = path; + self.experimental_write_cache_size = size; + self + } } /// Operational mode for certain actions. diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 47aed723cf3..e4d0afc60eb 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -18,6 +18,7 @@ use std::collections::HashMap; use api::v1::value::ValueData; use api::v1::Rows; +use common_base::readable_size::ReadableSize; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; @@ -555,3 +556,45 @@ async fn test_region_usage() { // region total usage assert_eq!(region_stat.disk_usage(), 4072); } + +#[tokio::test] +async fn test_engine_with_write_cache() { + common_telemetry::init_default_ut_logging(); + + let mut env = TestEnv::new(); + let path = env.data_home().to_str().unwrap().to_string(); + let mito_config = MitoConfig::default().enable_write_cache(path, ReadableSize::mb(512)); + let engine = env.create_engine(mito_config).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows_for_key("a", 0, 3, 0), + }; + put_rows(&engine, region_id, rows).await; + + flush_region(&engine, region_id, None).await; + + let request = ScanRequest::default(); + let scanner = engine.scanner(region_id, request).unwrap(); + + 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 | 2.0 | 1970-01-01T00:00:02 | ++-------+---------+---------------------+"; + assert_eq!(expected, batches.pretty_print().unwrap()); +} diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 924df76e1dd..42c8eff87fb 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -27,8 +27,8 @@ use std::sync::Arc; use common_base::readable_size::ReadableSize; use parquet::file::metadata::ParquetMetaData; -use super::DEFAULT_WRITE_BUFFER_SIZE; use crate::sst::file::FileTimeRange; +use crate::sst::DEFAULT_WRITE_BUFFER_SIZE; /// Key of metadata in parquet SST. pub const PARQUET_METADATA_KEY: &str = "greptime:metadata"; diff --git a/src/operator/src/req_convert/insert/stmt_to_region.rs b/src/operator/src/req_convert/insert/stmt_to_region.rs index 1297adf6a5d..eb3b1b2be97 100644 --- a/src/operator/src/req_convert/insert/stmt_to_region.rs +++ b/src/operator/src/req_convert/insert/stmt_to_region.rs @@ -25,13 +25,13 @@ use sql::statements::insert::Insert; use sqlparser::ast::{ObjectName, Value as SqlValue}; use table::TableRef; -use super::semantic_type; use crate::error::{ CatalogSnafu, ColumnDataTypeSnafu, ColumnDefaultValueSnafu, ColumnNoneDefaultValueSnafu, ColumnNotFoundSnafu, InvalidSqlSnafu, MissingInsertBodySnafu, ParseSqlSnafu, Result, TableNotFoundSnafu, }; use crate::req_convert::common::partitioner::Partitioner; +use crate::req_convert::insert::semantic_type; const DEFAULT_PLACEHOLDER_VALUE: &str = "default";