From 707a0d56261d911a3089931809a2dee62592c426 Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Tue, 28 Nov 2023 17:35:03 +0800 Subject: [PATCH 01/17] fix: urldecode when influxdb auth (#2831) * fix: add url decode when influxdb auth * chore: fmt toml --- Cargo.lock | 17 +++++++++-------- src/servers/Cargo.toml | 1 + src/servers/src/error.rs | 8 ++++++++ src/servers/src/http/authorize.rs | 6 ++++-- 4 files changed, 22 insertions(+), 10 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ba35273b8166..19e122bd18ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3190,9 +3190,9 @@ checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" [[package]] name = "form_urlencoded" -version = "1.2.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a62bc1cf6f830c2ec14a513a9fb124d0a213a629668a4186f329db21fe045652" +checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456" dependencies = [ "percent-encoding", ] @@ -3875,9 +3875,9 @@ checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" [[package]] name = "idna" -version = "0.4.0" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d20d6b07bfbc108882d88ed8e37d39636dcc260e15e30c45e6ba089610b917c" +checksum = "634d9b1461af396cad843f47fdba5597a4f9e6ddd4bfb6ff5d85028c25cb12f6" dependencies = [ "unicode-bidi", "unicode-normalization", @@ -5876,9 +5876,9 @@ dependencies = [ [[package]] name = "percent-encoding" -version = "2.3.0" +version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94" +checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "pest" @@ -8277,6 +8277,7 @@ dependencies = [ "tonic-reflection", "tower", "tower-http", + "urlencoding", ] [[package]] @@ -10248,9 +10249,9 @@ checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" [[package]] name = "url" -version = "2.4.1" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "143b538f18257fac9cad154828a57c6bf5157e1aa604d4816b5995bf6de87ae5" +checksum = "31e6302e3bb753d46e83516cae55ae196fc0c309407cf11ab35cc51a4c2a4633" dependencies = [ "form_urlencoded", "idna", diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 0ea4641cd5ee..4d4078c5d1a8 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -92,6 +92,7 @@ tonic-reflection = "0.10" tonic.workspace = true tower = { version = "0.4", features = ["full"] } tower-http = { version = "0.3", features = ["full"] } +urlencoding = "2.1" [target.'cfg(not(windows))'.dependencies] tikv-jemalloc-ctl = { version = "0.5", features = ["use_std"] } diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index be6b202859d3..e105b6885ec6 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -394,6 +394,13 @@ pub enum Error { error: serde_json::error::Error, location: Location, }, + + #[snafu(display("Failed to decode url"))] + UrlDecode { + #[snafu(source)] + error: FromUtf8Error, + location: Location, + }, } pub type Result = std::result::Result; @@ -444,6 +451,7 @@ impl ErrorExt for Error { | DataFrame { .. } | PreparedStmtTypeMismatch { .. } | TimePrecision { .. } + | UrlDecode { .. } | IncompatibleSchema { .. } => StatusCode::InvalidArguments, InfluxdbLinesWrite { source, .. } diff --git a/src/servers/src/http/authorize.rs b/src/servers/src/http/authorize.rs index e4217bb69eab..9225b6ec52dd 100644 --- a/src/servers/src/http/authorize.rs +++ b/src/servers/src/http/authorize.rs @@ -35,7 +35,7 @@ use super::header::GreptimeDbName; use super::PUBLIC_APIS; use crate::error::{ self, InvalidAuthorizationHeaderSnafu, InvalidParameterSnafu, InvisibleASCIISnafu, - NotFoundInfluxAuthSnafu, Result, UnsupportedAuthSchemeSnafu, + NotFoundInfluxAuthSnafu, Result, UnsupportedAuthSchemeSnafu, UrlDecodeSnafu, }; use crate::http::HTTP_API_PREFIX; @@ -166,7 +166,9 @@ fn get_influxdb_credentials( return Ok(None); }; - match extract_influxdb_user_from_query(query_str) { + let query_str = urlencoding::decode(query_str).context(UrlDecodeSnafu)?; + + match extract_influxdb_user_from_query(&query_str) { (None, None) => Ok(None), (Some(username), Some(password)) => { Ok(Some((username.to_string(), password.to_string().into()))) From d0d0f091f01f8d7bdde7603c1e868a2a032b73a2 Mon Sep 17 00:00:00 2001 From: ZonaHe Date: Tue, 28 Nov 2023 18:38:37 +0800 Subject: [PATCH 02/17] feat: update dashboard to v0.4.0 (#2832) Co-authored-by: ZonaHex --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index 600e6fd31c08..fb7a04cff86a 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.3.3 +v0.4.0 From abbac46c05ea0cd5742aac6589fc7ca20eb58446 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 29 Nov 2023 13:49:40 +0800 Subject: [PATCH 03/17] fix: do not expose manifest compression algorithm (#2835) * fix: don't allow to set manifest compression algorithm * docs: update config examples --- config/datanode.example.toml | 4 ++-- config/standalone.example.toml | 4 ++-- src/mito2/src/config.rs | 7 +++---- src/mito2/src/manifest/storage.rs | 1 + src/mito2/src/region/opener.rs | 5 ++++- tests-integration/tests/http.rs | 2 +- 6 files changed, 13 insertions(+), 10 deletions(-) diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 256970863184..2b4b7b4b1b9a 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -64,8 +64,8 @@ worker_channel_size = 128 worker_request_batch_size = 64 # Number of meta action updated to trigger a new checkpoint for the manifest manifest_checkpoint_distance = 10 -# Manifest compression type -manifest_compress_type = "uncompressed" +# Whether to compress manifest and checkpoint file by gzip (default false). +compress_manifest = false # Max number of running background jobs max_background_jobs = 4 # Interval to auto flush a region if it has not flushed yet. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 9055b7eca088..bb8f5c2d1cbb 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -133,8 +133,8 @@ worker_channel_size = 128 worker_request_batch_size = 64 # Number of meta action updated to trigger a new checkpoint for the manifest manifest_checkpoint_distance = 10 -# Manifest compression type -manifest_compress_type = "uncompressed" +# Whether to compress manifest and checkpoint file by gzip (default false). +compress_manifest = false # Max number of running background jobs max_background_jobs = 4 # Interval to auto flush a region if it has not flushed yet. diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index e6061ffa7a2e..edb7d3dd0d65 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -17,7 +17,6 @@ use std::time::Duration; use common_base::readable_size::ReadableSize; -use common_datasource::compression::CompressionType; use common_telemetry::warn; use serde::{Deserialize, Serialize}; @@ -42,8 +41,8 @@ pub struct MitoConfig { /// Number of meta action updated to trigger a new checkpoint /// for the manifest (default 10). pub manifest_checkpoint_distance: u64, - /// Manifest compression type (default uncompressed). - pub manifest_compress_type: CompressionType, + /// Whether to compress manifest and checkpoint file by gzip (default false). + pub compress_manifest: bool, // Background job configs: /// Max number of running background jobs (default 4). @@ -78,7 +77,7 @@ impl Default for MitoConfig { worker_channel_size: 128, worker_request_batch_size: 64, manifest_checkpoint_distance: 10, - manifest_compress_type: CompressionType::Uncompressed, + compress_manifest: false, max_background_jobs: DEFAULT_MAX_BG_JOB, auto_flush_interval: Duration::from_secs(30 * 60), global_write_buffer_size: ReadableSize::gb(1), diff --git a/src/mito2/src/manifest/storage.rs b/src/mito2/src/manifest/storage.rs index edd63ac52162..3e8860155b6c 100644 --- a/src/mito2/src/manifest/storage.rs +++ b/src/mito2/src/manifest/storage.rs @@ -42,6 +42,7 @@ const DEFAULT_MANIFEST_COMPRESSION_TYPE: CompressionType = CompressionType::Gzip /// So when we encounter problems, we need to fall back to `FALL_BACK_COMPRESS_TYPE` for processing. const FALL_BACK_COMPRESS_TYPE: CompressionType = CompressionType::Uncompressed; +/// Returns the [CompressionType] according to whether to compress manifest files. #[inline] pub const fn manifest_compress_type(compress: bool) -> CompressionType { if compress { diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 0baa8dac50b1..4c3a1ad577bd 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -33,6 +33,7 @@ use crate::cache::CacheManagerRef; use crate::config::MitoConfig; use crate::error::{EmptyRegionDirSnafu, ObjectStoreNotFoundSnafu, RegionCorruptedSnafu, Result}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; +use crate::manifest::storage::manifest_compress_type; use crate::memtable::MemtableBuilderRef; use crate::region::options::RegionOptions; use crate::region::version::{VersionBuilder, VersionControl, VersionControlRef}; @@ -259,7 +260,9 @@ impl RegionOpener { Ok(RegionManifestOptions { manifest_dir: new_manifest_dir(&self.region_dir), object_store, - compress_type: config.manifest_compress_type, + // We don't allow users to set the compression algorithm as we use it as a file suffix. + // Currently, the manifest storage doesn't have good support for changing compression algorithms. + compress_type: manifest_compress_type(config.compress_manifest), checkpoint_distance: config.manifest_checkpoint_distance, }) } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index b6caafdad641..68b0f01d1412 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -720,7 +720,7 @@ num_workers = {} worker_channel_size = 128 worker_request_batch_size = 64 manifest_checkpoint_distance = 10 -manifest_compress_type = "uncompressed" +compress_manifest = false max_background_jobs = 4 auto_flush_interval = "30m" global_write_buffer_size = "1GiB" From 92a9802343324a9c0707234456e25a39e6f9bfa0 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Wed, 29 Nov 2023 14:40:10 +0800 Subject: [PATCH 04/17] feat: canonicalize all unquoted identifier to lowercase (#2828) * feat: canonicalize all unquoted identifier to lowercase Signed-off-by: Ruihang Xia * add more tests Signed-off-by: Ruihang Xia * test altering table Signed-off-by: Ruihang Xia * primary key declare Signed-off-by: Ruihang Xia * fix primary key declare Signed-off-by: Ruihang Xia * partition by and time index Signed-off-by: Ruihang Xia * remove redundent call to canonicalize Signed-off-by: Ruihang Xia --------- Signed-off-by: Ruihang Xia --- src/sql/src/parser.rs | 24 ++++++ src/sql/src/parsers/alter_parser.rs | 13 ++-- src/sql/src/parsers/copy_parser.rs | 3 +- src/sql/src/parsers/create_parser.rs | 29 +++++-- src/sql/src/parsers/describe_parser.rs | 3 +- src/sql/src/parsers/drop_parser.rs | 3 +- src/sql/src/parsers/show_parser.rs | 3 +- src/sql/src/parsers/truncate_parser.rs | 3 +- .../standalone/common/alter/add_col.result | 17 +++++ .../cases/standalone/common/alter/add_col.sql | 4 + .../standalone/common/alter/drop_col.result | 6 +- .../standalone/common/alter/drop_col.sql | 4 +- .../common/alter/rename_table.result | 43 +++++++++++ .../standalone/common/alter/rename_table.sql | 17 +++++ .../create/upper_case_table_name.result | 75 ++++++++++++++++++- .../common/create/upper_case_table_name.sql | 36 ++++++++- .../standalone/common/delete/delete.result | 24 ++++++ .../cases/standalone/common/delete/delete.sql | 13 ++++ .../order/order_variable_size_payload.result | 12 +-- .../order/order_variable_size_payload.sql | 12 +-- .../common/truncate/truncate.result | 24 ++++++ .../standalone/common/truncate/truncate.sql | 13 ++++ 22 files changed, 344 insertions(+), 37 deletions(-) diff --git a/src/sql/src/parser.rs b/src/sql/src/parser.rs index e048384dc265..91e9b191627c 100644 --- a/src/sql/src/parser.rs +++ b/src/sql/src/parser.rs @@ -13,6 +13,7 @@ // limitations under the License. use snafu::ResultExt; +use sqlparser::ast::Ident; use sqlparser::dialect::Dialect; use sqlparser::keywords::Keyword; use sqlparser::parser::{Parser, ParserError}; @@ -166,6 +167,29 @@ impl<'a> ParserContext<'a> { pub(crate) fn peek_token_as_string(&self) -> String { self.parser.peek_token().to_string() } + + /// Canonicalize the identifier to lowercase if it's not quoted. + pub fn canonicalize_identifier(ident: Ident) -> Ident { + if ident.quote_style.is_some() { + ident + } else { + Ident { + value: ident.value.to_lowercase(), + quote_style: None, + } + } + } + + /// Like [canonicalize_identifier] but for [ObjectName]. + pub fn canonicalize_object_name(object_name: ObjectName) -> ObjectName { + ObjectName( + object_name + .0 + .into_iter() + .map(Self::canonicalize_identifier) + .collect(), + ) + } } #[cfg(test)] diff --git a/src/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index a5436d8f53b3..7f137d52c946 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -33,20 +33,22 @@ impl<'a> ParserContext<'a> { let parser = &mut self.parser; parser.expect_keywords(&[Keyword::ALTER, Keyword::TABLE])?; - let table_name = parser.parse_object_name()?; + let raw_table_name = parser.parse_object_name()?; + let table_name = Self::canonicalize_object_name(raw_table_name); let alter_operation = if parser.parse_keyword(Keyword::ADD) { if let Some(constraint) = parser.parse_optional_table_constraint()? { AlterTableOperation::AddConstraint(constraint) } else { let _ = parser.parse_keyword(Keyword::COLUMN); - let column_def = parser.parse_column_def()?; + let mut column_def = parser.parse_column_def()?; + column_def.name = Self::canonicalize_identifier(column_def.name); let location = if parser.parse_keyword(Keyword::FIRST) { Some(AddColumnLocation::First) } else if let Token::Word(word) = parser.peek_token().token { if word.value.to_ascii_uppercase() == "AFTER" { let _ = parser.next_token(); - let name = parser.parse_identifier()?; + let name = Self::canonicalize_identifier(parser.parse_identifier()?); Some(AddColumnLocation::After { column_name: name.value, }) @@ -63,7 +65,7 @@ impl<'a> ParserContext<'a> { } } else if parser.parse_keyword(Keyword::DROP) { if parser.parse_keyword(Keyword::COLUMN) { - let name = self.parser.parse_identifier()?; + let name = Self::canonicalize_identifier(self.parser.parse_identifier()?); AlterTableOperation::DropColumn { name } } else { return Err(ParserError::ParserError(format!( @@ -72,7 +74,8 @@ impl<'a> ParserContext<'a> { ))); } } else if parser.parse_keyword(Keyword::RENAME) { - let new_table_name_obj = parser.parse_object_name()?; + let new_table_name_obj_raw = parser.parse_object_name()?; + let new_table_name_obj = Self::canonicalize_object_name(new_table_name_obj_raw); let new_table_name = match &new_table_name_obj.0[..] { [table] => table.value.clone(), _ => { diff --git a/src/sql/src/parsers/copy_parser.rs b/src/sql/src/parsers/copy_parser.rs index 7b47c79783cd..17ee3df152c9 100644 --- a/src/sql/src/parsers/copy_parser.rs +++ b/src/sql/src/parsers/copy_parser.rs @@ -69,7 +69,7 @@ impl<'a> ParserContext<'a> { } fn parse_copy_table(&mut self) -> Result { - let table_name = + let raw_table_name = self.parser .parse_object_name() .with_context(|_| error::UnexpectedSnafu { @@ -77,6 +77,7 @@ impl<'a> ParserContext<'a> { expected: "a table name", actual: self.peek_token_as_string(), })?; + let table_name = Self::canonicalize_object_name(raw_table_name); if self.parser.parse_keyword(Keyword::TO) { let (with, connection, location) = self.parse_copy_to()?; diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index 6d01f8bacc32..d81c3eb6c402 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -73,7 +73,7 @@ impl<'a> ParserContext<'a> { let if_not_exists = self.parser .parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); - let table_name = self + let raw_table_name = self .parser .parse_object_name() .context(error::UnexpectedSnafu { @@ -81,6 +81,7 @@ impl<'a> ParserContext<'a> { expected: "a table name", actual: self.peek_token_as_string(), })?; + let table_name = Self::canonicalize_object_name(raw_table_name); let (columns, constraints) = self.parse_columns()?; let engine = self.parse_table_engine(common_catalog::consts::FILE_ENGINE)?; let options = self @@ -142,7 +143,7 @@ impl<'a> ParserContext<'a> { self.parser .parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); - let table_name = self + let raw_table_name = self .parser .parse_object_name() .context(error::UnexpectedSnafu { @@ -150,6 +151,7 @@ impl<'a> ParserContext<'a> { expected: "a table name", actual: self.peek_token_as_string(), })?; + let table_name = Self::canonicalize_object_name(raw_table_name); let (columns, constraints) = self.parse_columns()?; @@ -197,10 +199,14 @@ impl<'a> ParserContext<'a> { actual: self.peek_token_as_string(), })?; - let column_list = self + let raw_column_list = self .parser .parse_parenthesized_column_list(Mandatory, false) .context(error::SyntaxSnafu)?; + let column_list = raw_column_list + .into_iter() + .map(Self::canonicalize_identifier) + .collect(); let entries = self.parse_comma_separated(Self::parse_partition_entry)?; @@ -435,7 +441,7 @@ impl<'a> ParserContext<'a> { }; } Ok(ColumnDef { - name, + name: Self::canonicalize_identifier(name), data_type, collation, options, @@ -488,7 +494,8 @@ impl<'a> ParserContext<'a> { fn parse_optional_table_constraint(&mut self) -> Result> { let name = if self.parser.parse_keyword(Keyword::CONSTRAINT) { - Some(self.parser.parse_identifier().context(error::SyntaxSnafu)?) + let raw_name = self.parser.parse_identifier().context(error::SyntaxSnafu)?; + Some(Self::canonicalize_identifier(raw_name)) } else { None }; @@ -504,10 +511,14 @@ impl<'a> ParserContext<'a> { expected: "KEY", actual: self.peek_token_as_string(), })?; - let columns = self + let raw_columns = self .parser .parse_parenthesized_column_list(Mandatory, false) .context(error::SyntaxSnafu)?; + let columns = raw_columns + .into_iter() + .map(Self::canonicalize_identifier) + .collect(); Ok(Some(TableConstraint::Unique { name, columns, @@ -526,10 +537,14 @@ impl<'a> ParserContext<'a> { actual: self.peek_token_as_string(), })?; - let columns = self + let raw_columns = self .parser .parse_parenthesized_column_list(Mandatory, false) .context(error::SyntaxSnafu)?; + let columns = raw_columns + .into_iter() + .map(Self::canonicalize_identifier) + .collect::>(); ensure!( columns.len() == 1, diff --git a/src/sql/src/parsers/describe_parser.rs b/src/sql/src/parsers/describe_parser.rs index a7c6b5c21e88..3348137f9d9d 100644 --- a/src/sql/src/parsers/describe_parser.rs +++ b/src/sql/src/parsers/describe_parser.rs @@ -30,7 +30,7 @@ impl<'a> ParserContext<'a> { } fn parse_describe_table(&mut self) -> Result { - let table_idents = + let raw_table_idents = self.parser .parse_object_name() .with_context(|_| error::UnexpectedSnafu { @@ -38,6 +38,7 @@ impl<'a> ParserContext<'a> { expected: "a table name", actual: self.peek_token_as_string(), })?; + let table_idents = Self::canonicalize_object_name(raw_table_idents); ensure!( !table_idents.0.is_empty(), InvalidTableNameSnafu { diff --git a/src/sql/src/parsers/drop_parser.rs b/src/sql/src/parsers/drop_parser.rs index add3810385bb..92cddf728491 100644 --- a/src/sql/src/parsers/drop_parser.rs +++ b/src/sql/src/parsers/drop_parser.rs @@ -29,7 +29,7 @@ impl<'a> ParserContext<'a> { } let _ = self.parser.next_token(); - let table_ident = + let raw_table_ident = self.parser .parse_object_name() .with_context(|_| error::UnexpectedSnafu { @@ -37,6 +37,7 @@ impl<'a> ParserContext<'a> { expected: "a table name", actual: self.peek_token_as_string(), })?; + let table_ident = Self::canonicalize_object_name(raw_table_ident); ensure!( !table_ident.0.is_empty(), InvalidTableNameSnafu { diff --git a/src/sql/src/parsers/show_parser.rs b/src/sql/src/parsers/show_parser.rs index 48a64e86ec5f..1278ea55bf86 100644 --- a/src/sql/src/parsers/show_parser.rs +++ b/src/sql/src/parsers/show_parser.rs @@ -50,7 +50,7 @@ impl<'a> ParserContext<'a> { /// Parse SHOW CREATE TABLE statement fn parse_show_create_table(&mut self) -> Result { - let table_name = + let raw_table_name = self.parser .parse_object_name() .with_context(|_| error::UnexpectedSnafu { @@ -58,6 +58,7 @@ impl<'a> ParserContext<'a> { expected: "a table name", actual: self.peek_token_as_string(), })?; + let table_name = Self::canonicalize_object_name(raw_table_name); ensure!( !table_name.0.is_empty(), InvalidTableNameSnafu { diff --git a/src/sql/src/parsers/truncate_parser.rs b/src/sql/src/parsers/truncate_parser.rs index e3a31f2db2e0..fa5253bcf03e 100644 --- a/src/sql/src/parsers/truncate_parser.rs +++ b/src/sql/src/parsers/truncate_parser.rs @@ -26,7 +26,7 @@ impl<'a> ParserContext<'a> { let _ = self.parser.next_token(); let _ = self.parser.parse_keyword(Keyword::TABLE); - let table_ident = + let raw_table_ident = self.parser .parse_object_name() .with_context(|_| error::UnexpectedSnafu { @@ -34,6 +34,7 @@ impl<'a> ParserContext<'a> { expected: "a table name", actual: self.peek_token_as_string(), })?; + let table_ident = Self::canonicalize_object_name(raw_table_ident); ensure!( !table_ident.0.is_empty(), diff --git a/tests/cases/standalone/common/alter/add_col.result b/tests/cases/standalone/common/alter/add_col.result index 7b3c6d4fb4a1..c01420c641d8 100644 --- a/tests/cases/standalone/common/alter/add_col.result +++ b/tests/cases/standalone/common/alter/add_col.result @@ -87,6 +87,23 @@ DESC TABLE test; | idc | String | | YES | idc | FIELD | +--------+----------------------+-----+------+---------+---------------+ +ALTER TABLE test ADD COLUMN "IdC" STRING default 'idc' PRIMARY KEY; + +Affected Rows: 0 + +DESC TABLE test; + ++--------+----------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++--------+----------------------+-----+------+---------+---------------+ +| i | Int32 | | YES | | FIELD | +| j | TimestampMillisecond | PRI | NO | | TIMESTAMP | +| k | Int32 | | YES | | FIELD | +| host | String | | YES | | FIELD | +| idc | String | | YES | idc | FIELD | +| IdC | String | | YES | idc | FIELD | ++--------+----------------------+-----+------+---------+---------------+ + DROP TABLE test; Affected Rows: 0 diff --git a/tests/cases/standalone/common/alter/add_col.sql b/tests/cases/standalone/common/alter/add_col.sql index 8d506959ef2d..7cd4159a4871 100644 --- a/tests/cases/standalone/common/alter/add_col.sql +++ b/tests/cases/standalone/common/alter/add_col.sql @@ -22,4 +22,8 @@ SELECT * FROM test; DESC TABLE test; +ALTER TABLE test ADD COLUMN "IdC" STRING default 'idc' PRIMARY KEY; + +DESC TABLE test; + DROP TABLE test; diff --git a/tests/cases/standalone/common/alter/drop_col.result b/tests/cases/standalone/common/alter/drop_col.result index cf0c7251efa8..81a9d46f060f 100644 --- a/tests/cases/standalone/common/alter/drop_col.result +++ b/tests/cases/standalone/common/alter/drop_col.result @@ -6,7 +6,11 @@ INSERT INTO test VALUES (1, 1), (2, 2); Affected Rows: 2 -ALTER TABLE test DROP COLUMN i; +ALTER TABLE test DROP COLUMN "I"; + +Error: 4002(TableColumnNotFound), Column I not exists in table test + +ALTER TABLE test DROP COLUMN I; Affected Rows: 0 diff --git a/tests/cases/standalone/common/alter/drop_col.sql b/tests/cases/standalone/common/alter/drop_col.sql index 6bda8278335d..778613f0367d 100644 --- a/tests/cases/standalone/common/alter/drop_col.sql +++ b/tests/cases/standalone/common/alter/drop_col.sql @@ -2,7 +2,9 @@ CREATE TABLE test(i INTEGER, j TIMESTAMP TIME INDEX); INSERT INTO test VALUES (1, 1), (2, 2); -ALTER TABLE test DROP COLUMN i; +ALTER TABLE test DROP COLUMN "I"; + +ALTER TABLE test DROP COLUMN I; SELECT * FROM test; diff --git a/tests/cases/standalone/common/alter/rename_table.result b/tests/cases/standalone/common/alter/rename_table.result index 40982666dc05..b635d89e25bf 100644 --- a/tests/cases/standalone/common/alter/rename_table.result +++ b/tests/cases/standalone/common/alter/rename_table.result @@ -81,3 +81,46 @@ DROP TABLE new_table; Affected Rows: 0 +CREATE TABLE "AbCdE"("CoLa" INTEGER, "cOlB" TIMESTAMP TIME INDEX); + +Affected Rows: 0 + +ALTER TABLE "AbCdE" RENAME "fGhI"; + +Affected Rows: 0 + +DESC TABLE "fGhI"; + ++--------+----------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++--------+----------------------+-----+------+---------+---------------+ +| CoLa | Int32 | | YES | | FIELD | +| cOlB | TimestampMillisecond | PRI | NO | | TIMESTAMP | ++--------+----------------------+-----+------+---------+---------------+ + +SELECT * FROM "fGhI"; + +++ +++ + +ALTER TABLE "fGhI" RENAME JkLmN; + +Affected Rows: 0 + +DESC TABLE "JkLmN"; + +Error: 4001(TableNotFound), Table not found: JkLmN + +DESC TABLE JkLmN; + ++--------+----------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++--------+----------------------+-----+------+---------+---------------+ +| CoLa | Int32 | | YES | | FIELD | +| cOlB | TimestampMillisecond | PRI | NO | | TIMESTAMP | ++--------+----------------------+-----+------+---------+---------------+ + +DROP TABLE jklmn; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/alter/rename_table.sql b/tests/cases/standalone/common/alter/rename_table.sql index 85c3e8b4b9ff..377161166b84 100644 --- a/tests/cases/standalone/common/alter/rename_table.sql +++ b/tests/cases/standalone/common/alter/rename_table.sql @@ -28,3 +28,20 @@ ALTER TABLE new_table RENAME t; DROP TABLE t; DROP TABLE new_table; + + +CREATE TABLE "AbCdE"("CoLa" INTEGER, "cOlB" TIMESTAMP TIME INDEX); + +ALTER TABLE "AbCdE" RENAME "fGhI"; + +DESC TABLE "fGhI"; + +SELECT * FROM "fGhI"; + +ALTER TABLE "fGhI" RENAME JkLmN; + +DESC TABLE "JkLmN"; + +DESC TABLE JkLmN; + +DROP TABLE jklmn; diff --git a/tests/cases/standalone/common/create/upper_case_table_name.result b/tests/cases/standalone/common/create/upper_case_table_name.result index c2957f0496cc..39562efac4e6 100644 --- a/tests/cases/standalone/common/create/upper_case_table_name.result +++ b/tests/cases/standalone/common/create/upper_case_table_name.result @@ -6,11 +6,11 @@ use upper_case_table_name; Affected Rows: 0 -create table system_Metric(ts timestamp time index); +create table "system_Metric"(ts timestamp time index); Affected Rows: 0 -insert into system_Metric values (0), (1); +insert into "system_Metric" values (0), (1); Affected Rows: 2 @@ -27,7 +27,76 @@ select * from "system_Metric"; | 1970-01-01T00:00:00.001 | +-------------------------+ -drop table system_Metric; +drop table "system_Metric"; + +Affected Rows: 0 + +create table "AbCdEfG"("CoLA" string, "cOlB" string, "tS" timestamp time index, primary key ("CoLA")); + +Affected Rows: 0 + +desc table "AbCdEfG"; + ++--------+----------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++--------+----------------------+-----+------+---------+---------------+ +| CoLA | String | PRI | YES | | TAG | +| cOlB | String | | YES | | FIELD | +| tS | TimestampMillisecond | PRI | NO | | TIMESTAMP | ++--------+----------------------+-----+------+---------+---------------+ + +-- unquoted table name and column name. +create table AbCdEfGe(CoLA string, cOlB string, tS timestamp time index, primary key (cOlA)); + +Affected Rows: 0 + +desc table aBcDeFgE; + ++--------+----------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++--------+----------------------+-----+------+---------+---------------+ +| cola | String | PRI | YES | | TAG | +| colb | String | | YES | | FIELD | +| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP | ++--------+----------------------+-----+------+---------+---------------+ + +drop table "AbCdEfG"; + +Affected Rows: 0 + +drop table aBcDeFgE; + +Affected Rows: 0 + +-- unquoted column name in partition +create table AbCdEfGe( + CoLA string PRIMARY KEY, + tS timestamp time index +) PARTITION BY RANGE COLUMNS (cOlA) ( + PARTITION p0 VALUES LESS THAN (MAXVALUE) +); + +Affected Rows: 0 + +drop table abcdefge; + +Affected Rows: 0 + +-- unquoted column name in TIME INDEX +create table AbCdEfGe(CoLA string, tS timestamp, TIME INDEX (Ts)); + +Affected Rows: 0 + +desc table abcdefge; + ++--------+----------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++--------+----------------------+-----+------+---------+---------------+ +| cola | String | | YES | | FIELD | +| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP | ++--------+----------------------+-----+------+---------+---------------+ + +drop table abcdefge; Affected Rows: 0 diff --git a/tests/cases/standalone/common/create/upper_case_table_name.sql b/tests/cases/standalone/common/create/upper_case_table_name.sql index 32bb1fc3b342..90cedca0057b 100644 --- a/tests/cases/standalone/common/create/upper_case_table_name.sql +++ b/tests/cases/standalone/common/create/upper_case_table_name.sql @@ -2,14 +2,44 @@ create database upper_case_table_name; use upper_case_table_name; -create table system_Metric(ts timestamp time index); +create table "system_Metric"(ts timestamp time index); -insert into system_Metric values (0), (1); +insert into "system_Metric" values (0), (1); select * from system_Metric; select * from "system_Metric"; -drop table system_Metric; +drop table "system_Metric"; + +create table "AbCdEfG"("CoLA" string, "cOlB" string, "tS" timestamp time index, primary key ("CoLA")); + +desc table "AbCdEfG"; + +-- unquoted table name and column name. +create table AbCdEfGe(CoLA string, cOlB string, tS timestamp time index, primary key (cOlA)); + +desc table aBcDeFgE; + +drop table "AbCdEfG"; + +drop table aBcDeFgE; + +-- unquoted column name in partition +create table AbCdEfGe( + CoLA string PRIMARY KEY, + tS timestamp time index +) PARTITION BY RANGE COLUMNS (cOlA) ( + PARTITION p0 VALUES LESS THAN (MAXVALUE) +); + +drop table abcdefge; + +-- unquoted column name in TIME INDEX +create table AbCdEfGe(CoLA string, tS timestamp, TIME INDEX (Ts)); + +desc table abcdefge; + +drop table abcdefge; use public; diff --git a/tests/cases/standalone/common/delete/delete.result b/tests/cases/standalone/common/delete/delete.result index 7885c4b0cfb6..4df80c05a2f8 100644 --- a/tests/cases/standalone/common/delete/delete.result +++ b/tests/cases/standalone/common/delete/delete.result @@ -60,3 +60,27 @@ DROP TABLE monitor; Affected Rows: 0 +CREATE TABLE "MoNiToR" ("hOsT" STRING PRIMARY KEY, "tS" TIMESTAMP TIME INDEX, "cPu" DOUBLE DEFAULT 0); + +Affected Rows: 0 + +DELETE FROM "MoNiToR" WHERE "hOsT" = 'host2'; + +Affected Rows: 0 + +DROP TABLE "MoNiToR"; + +Affected Rows: 0 + +CREATE TABLE MoNiToR (hOsT STRING PRIMARY KEY, tS TIMESTAMP TIME INDEX, cPu DOUBLE DEFAULT 0); + +Affected Rows: 0 + +DELETE FROM MoNiToR WHERE hOsT = 'host2'; + +Affected Rows: 0 + +DROP TABLE MoNiToR; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/delete/delete.sql b/tests/cases/standalone/common/delete/delete.sql index f59eb73df81b..9786aa2d7bed 100644 --- a/tests/cases/standalone/common/delete/delete.sql +++ b/tests/cases/standalone/common/delete/delete.sql @@ -24,3 +24,16 @@ DELETE FROM monitor WHERE memory > 2048; SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; DROP TABLE monitor; + + +CREATE TABLE "MoNiToR" ("hOsT" STRING PRIMARY KEY, "tS" TIMESTAMP TIME INDEX, "cPu" DOUBLE DEFAULT 0); + +DELETE FROM "MoNiToR" WHERE "hOsT" = 'host2'; + +DROP TABLE "MoNiToR"; + +CREATE TABLE MoNiToR (hOsT STRING PRIMARY KEY, tS TIMESTAMP TIME INDEX, cPu DOUBLE DEFAULT 0); + +DELETE FROM MoNiToR WHERE hOsT = 'host2'; + +DROP TABLE MoNiToR; diff --git a/tests/cases/standalone/common/order/order_variable_size_payload.result b/tests/cases/standalone/common/order/order_variable_size_payload.result index becf9aa5c14f..050026cea7b8 100644 --- a/tests/cases/standalone/common/order/order_variable_size_payload.result +++ b/tests/cases/standalone/common/order/order_variable_size_payload.result @@ -347,19 +347,19 @@ select i, split_part(s, 'b', 1) from test8 order by i; | | d | +---+----------------------------------------+ -CREATE TABLE DirectReports +CREATE TABLE "DirectReports" ( - EmployeeID smallint, + "EmployeeID" smallint, "Name" varchar NOT NULL, - Title varchar NOT NULL, - EmployeeLevel int NOT NULL, + "Title" varchar NOT NULL, + "EmployeeLevel" int NOT NULL, "Sort" varchar NOT NULL, "Timestamp" TIMESTAMP TIME INDEX, ); Affected Rows: 0 -INSERT INTO DirectReports VALUES +INSERT INTO "DirectReports" VALUES (1, 'Ken Sánchez', 'Chief Executive Officer', 1, 'Ken Sánchez', 1), (273, '>Brian Welcker', 'Vice President of Sales', 2, 'Ken Sánchez>Brian Welcker', 2), (274, '>>Stephen Jiang', 'North American Sales Manager', 3, 'Ken Sánchez>Brian Welcker>Stephen Jiang', 3), @@ -426,7 +426,7 @@ DROP table test8; Affected Rows: 0 -DROP TABLE DirectReports; +DROP TABLE "DirectReports"; Affected Rows: 0 diff --git a/tests/cases/standalone/common/order/order_variable_size_payload.sql b/tests/cases/standalone/common/order/order_variable_size_payload.sql index b71ba7405e4b..0fbf6a067ebc 100644 --- a/tests/cases/standalone/common/order/order_variable_size_payload.sql +++ b/tests/cases/standalone/common/order/order_variable_size_payload.sql @@ -82,17 +82,17 @@ insert into test8 values (3, 'aba', 1), (1, 'ccbcc', 2), (NULL, 'dbdbd', 3), (2, select i, split_part(s, 'b', 1) from test8 order by i; -CREATE TABLE DirectReports +CREATE TABLE "DirectReports" ( - EmployeeID smallint, + "EmployeeID" smallint, "Name" varchar NOT NULL, - Title varchar NOT NULL, - EmployeeLevel int NOT NULL, + "Title" varchar NOT NULL, + "EmployeeLevel" int NOT NULL, "Sort" varchar NOT NULL, "Timestamp" TIMESTAMP TIME INDEX, ); -INSERT INTO DirectReports VALUES +INSERT INTO "DirectReports" VALUES (1, 'Ken Sánchez', 'Chief Executive Officer', 1, 'Ken Sánchez', 1), (273, '>Brian Welcker', 'Vice President of Sales', 2, 'Ken Sánchez>Brian Welcker', 2), (274, '>>Stephen Jiang', 'North American Sales Manager', 3, 'Ken Sánchez>Brian Welcker>Stephen Jiang', 3), @@ -125,4 +125,4 @@ DROP table test7; DROP table test8; -DROP TABLE DirectReports; +DROP TABLE "DirectReports"; diff --git a/tests/cases/standalone/common/truncate/truncate.result b/tests/cases/standalone/common/truncate/truncate.result index 42358d542a34..c0e7560f8145 100644 --- a/tests/cases/standalone/common/truncate/truncate.result +++ b/tests/cases/standalone/common/truncate/truncate.result @@ -74,3 +74,27 @@ DROP TABLE monitor; Affected Rows: 0 +CREATE TABLE "MoNiToR" ("hOsT" STRING PRIMARY KEY, "tS" TIMESTAMP TIME INDEX, "cPu" DOUBLE DEFAULT 0); + +Affected Rows: 0 + +TRUNCATE "MoNiToR"; + +Affected Rows: 0 + +DROP TABLE "MoNiToR"; + +Affected Rows: 0 + +CREATE TABLE MoNiToR (hOsT STRING PRIMARY KEY, tS TIMESTAMP TIME INDEX, cPu DOUBLE DEFAULT 0); + +Affected Rows: 0 + +TRUNCATE MoNiToR; + +Affected Rows: 0 + +DROP TABLE MoNiToR; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/truncate/truncate.sql b/tests/cases/standalone/common/truncate/truncate.sql index b5fd24985702..66cfd813cc54 100644 --- a/tests/cases/standalone/common/truncate/truncate.sql +++ b/tests/cases/standalone/common/truncate/truncate.sql @@ -31,3 +31,16 @@ TRUNCATE monitor; SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; DROP TABLE monitor; + + +CREATE TABLE "MoNiToR" ("hOsT" STRING PRIMARY KEY, "tS" TIMESTAMP TIME INDEX, "cPu" DOUBLE DEFAULT 0); + +TRUNCATE "MoNiToR"; + +DROP TABLE "MoNiToR"; + +CREATE TABLE MoNiToR (hOsT STRING PRIMARY KEY, tS TIMESTAMP TIME INDEX, cPu DOUBLE DEFAULT 0); + +TRUNCATE MoNiToR; + +DROP TABLE MoNiToR; From 445bd92c7af8d1b5e5f1f5eae64ad8a92bb469f5 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Wed, 29 Nov 2023 15:44:43 +0800 Subject: [PATCH 05/17] feat: add arg to standalone start command (#2837) * feat: add arg to standalone start command Signed-off-by: Ruihang Xia * add this arg to metasrv Signed-off-by: Ruihang Xia * remove short arg name Signed-off-by: Ruihang Xia --------- Signed-off-by: Ruihang Xia --- src/cmd/src/metasrv.rs | 7 +++++++ src/cmd/src/standalone.rs | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/src/cmd/src/metasrv.rs b/src/cmd/src/metasrv.rs index 0eb5b5270197..1de1e60099d7 100644 --- a/src/cmd/src/metasrv.rs +++ b/src/cmd/src/metasrv.rs @@ -100,6 +100,9 @@ struct StartCommand { http_timeout: Option, #[clap(long, default_value = "GREPTIMEDB_METASRV")] env_prefix: String, + /// The working home directory of this metasrv instance. + #[clap(long)] + data_home: Option, } impl StartCommand { @@ -152,6 +155,10 @@ impl StartCommand { opts.http.timeout = Duration::from_secs(http_timeout); } + if let Some(data_home) = &self.data_home { + opts.data_home = data_home.clone(); + } + // Disable dashboard in metasrv. opts.http.disable_dashboard = true; diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 5483c6c74724..8e7a56855125 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -227,6 +227,9 @@ struct StartCommand { user_provider: Option, #[clap(long, default_value = "GREPTIMEDB_STANDALONE")] env_prefix: String, + /// The working home directory of this standalone instance. + #[clap(long)] + data_home: Option, } impl StartCommand { @@ -257,6 +260,10 @@ impl StartCommand { opts.http.addr = addr.clone() } + if let Some(data_home) = &self.data_home { + opts.storage.data_home = data_home.clone(); + } + if let Some(addr) = &self.rpc_addr { // frontend grpc addr conflict with datanode default grpc addr let datanode_grpc_addr = DatanodeOptions::default().rpc_addr; From 7c53f92e4bcd74bb21435e7a183a9ebfeeb12c24 Mon Sep 17 00:00:00 2001 From: ZonaHe Date: Wed, 29 Nov 2023 16:50:25 +0800 Subject: [PATCH 06/17] feat: update dashboard to v0.4.1 (#2841) Co-authored-by: ZonaHex --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index fb7a04cff86a..5aff472ddfbf 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.4.0 +v0.4.1 From 616eb0491452f6a5eaa71916eae330db71f99896 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 29 Nov 2023 17:59:42 +0900 Subject: [PATCH 07/17] chore: bump version to 0.4.4 (#2840) chore: bump to v0.4.4 --- Cargo.lock | 122 ++++++++++++++++++++++++++--------------------------- Cargo.toml | 2 +- 2 files changed, 62 insertions(+), 62 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 19e122bd18ea..099d0ead5708 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -196,7 +196,7 @@ checksum = "8f1f8f5a6f3d50d89e3797d7593a50f96bb2aaa20ca0cc7be1fb673232c91d72" [[package]] name = "api" -version = "0.4.3" +version = "0.4.4" dependencies = [ "common-base", "common-decimal", @@ -655,7 +655,7 @@ dependencies = [ [[package]] name = "auth" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-trait", @@ -828,7 +828,7 @@ dependencies = [ [[package]] name = "benchmarks" -version = "0.4.3" +version = "0.4.4" dependencies = [ "arrow", "chrono", @@ -1160,7 +1160,7 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "catalog" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "arc-swap", @@ -1429,7 +1429,7 @@ checksum = "702fc72eb24e5a1e48ce58027a675bc24edd52096d5397d4aea7c6dd9eca0bd1" [[package]] name = "client" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "arrow-flight", @@ -1462,7 +1462,7 @@ dependencies = [ "session", "snafu", "substrait 0.17.1", - "substrait 0.4.3", + "substrait 0.4.4", "tokio", "tokio-stream", "tonic 0.10.2", @@ -1492,7 +1492,7 @@ dependencies = [ [[package]] name = "cmd" -version = "0.4.3" +version = "0.4.4" dependencies = [ "anymap", "async-trait", @@ -1540,7 +1540,7 @@ dependencies = [ "servers", "session", "snafu", - "substrait 0.4.3", + "substrait 0.4.4", "table", "temp-env", "tikv-jemallocator", @@ -1573,7 +1573,7 @@ checksum = "55b672471b4e9f9e95499ea597ff64941a309b2cdbffcc46f2cc5e2d971fd335" [[package]] name = "common-base" -version = "0.4.3" +version = "0.4.4" dependencies = [ "anymap", "bitvec", @@ -1588,7 +1588,7 @@ dependencies = [ [[package]] name = "common-catalog" -version = "0.4.3" +version = "0.4.4" dependencies = [ "chrono", "common-error", @@ -1599,7 +1599,7 @@ dependencies = [ [[package]] name = "common-config" -version = "0.4.3" +version = "0.4.4" dependencies = [ "common-base", "humantime-serde", @@ -1608,7 +1608,7 @@ dependencies = [ [[package]] name = "common-datasource" -version = "0.4.3" +version = "0.4.4" dependencies = [ "arrow", "arrow-schema", @@ -1639,7 +1639,7 @@ dependencies = [ [[package]] name = "common-decimal" -version = "0.4.3" +version = "0.4.4" dependencies = [ "arrow", "bigdecimal", @@ -1653,7 +1653,7 @@ dependencies = [ [[package]] name = "common-error" -version = "0.4.3" +version = "0.4.4" dependencies = [ "snafu", "strum 0.25.0", @@ -1661,7 +1661,7 @@ dependencies = [ [[package]] name = "common-function" -version = "0.4.3" +version = "0.4.4" dependencies = [ "arc-swap", "chrono-tz 0.6.3", @@ -1684,7 +1684,7 @@ dependencies = [ [[package]] name = "common-greptimedb-telemetry" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-trait", "common-error", @@ -1703,7 +1703,7 @@ dependencies = [ [[package]] name = "common-grpc" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "arrow-flight", @@ -1733,7 +1733,7 @@ dependencies = [ [[package]] name = "common-grpc-expr" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-trait", @@ -1752,7 +1752,7 @@ dependencies = [ [[package]] name = "common-macro" -version = "0.4.3" +version = "0.4.4" dependencies = [ "arc-swap", "common-query", @@ -1767,7 +1767,7 @@ dependencies = [ [[package]] name = "common-mem-prof" -version = "0.4.3" +version = "0.4.4" dependencies = [ "common-error", "common-macro", @@ -1780,7 +1780,7 @@ dependencies = [ [[package]] name = "common-meta" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-recursion", @@ -1819,7 +1819,7 @@ dependencies = [ [[package]] name = "common-procedure" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-stream", "async-trait", @@ -1843,7 +1843,7 @@ dependencies = [ [[package]] name = "common-procedure-test" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-trait", "common-procedure", @@ -1851,7 +1851,7 @@ dependencies = [ [[package]] name = "common-query" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-trait", @@ -1874,7 +1874,7 @@ dependencies = [ [[package]] name = "common-recordbatch" -version = "0.4.3" +version = "0.4.4" dependencies = [ "common-error", "common-macro", @@ -1891,7 +1891,7 @@ dependencies = [ [[package]] name = "common-runtime" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-trait", "common-error", @@ -1909,7 +1909,7 @@ dependencies = [ [[package]] name = "common-telemetry" -version = "0.4.3" +version = "0.4.4" dependencies = [ "backtrace", "common-error", @@ -1934,7 +1934,7 @@ dependencies = [ [[package]] name = "common-test-util" -version = "0.4.3" +version = "0.4.4" dependencies = [ "once_cell", "rand", @@ -1943,7 +1943,7 @@ dependencies = [ [[package]] name = "common-time" -version = "0.4.3" +version = "0.4.4" dependencies = [ "arrow", "chrono", @@ -1958,7 +1958,7 @@ dependencies = [ [[package]] name = "common-version" -version = "0.4.3" +version = "0.4.4" dependencies = [ "build-data", ] @@ -2581,7 +2581,7 @@ dependencies = [ [[package]] name = "datanode" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "arrow-flight", @@ -2640,7 +2640,7 @@ dependencies = [ "snafu", "sql", "store-api", - "substrait 0.4.3", + "substrait 0.4.4", "table", "tokio", "tokio-stream", @@ -2654,7 +2654,7 @@ dependencies = [ [[package]] name = "datatypes" -version = "0.4.3" +version = "0.4.4" dependencies = [ "arrow", "arrow-array", @@ -3092,7 +3092,7 @@ dependencies = [ [[package]] name = "file-engine" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-trait", @@ -3208,7 +3208,7 @@ dependencies = [ [[package]] name = "frontend" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "arc-swap", @@ -3272,7 +3272,7 @@ dependencies = [ "sqlparser 0.38.0 (git+https://github.com/GreptimeTeam/sqlparser-rs.git?rev=0fbae07d0c46dc18e3381c406d8b9b8abef6b1fd)", "store-api", "strfmt", - "substrait 0.4.3", + "substrait 0.4.4", "table", "tokio", "toml 0.7.8", @@ -3910,7 +3910,7 @@ dependencies = [ [[package]] name = "index" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-trait", "common-base", @@ -4365,7 +4365,7 @@ checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" [[package]] name = "log-store" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-stream", "async-trait", @@ -4635,7 +4635,7 @@ dependencies = [ [[package]] name = "meta-client" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-trait", @@ -4665,7 +4665,7 @@ dependencies = [ [[package]] name = "meta-srv" -version = "0.4.3" +version = "0.4.4" dependencies = [ "anymap", "api", @@ -4743,7 +4743,7 @@ dependencies = [ [[package]] name = "metric-engine" -version = "0.4.3" +version = "0.4.4" dependencies = [ "ahash 0.8.6", "api", @@ -4814,7 +4814,7 @@ dependencies = [ [[package]] name = "mito2" -version = "0.4.3" +version = "0.4.4" dependencies = [ "anymap", "api", @@ -5272,7 +5272,7 @@ dependencies = [ [[package]] name = "object-store" -version = "0.4.3" +version = "0.4.4" dependencies = [ "anyhow", "async-trait", @@ -5519,7 +5519,7 @@ dependencies = [ [[package]] name = "operator" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-compat", @@ -5564,7 +5564,7 @@ dependencies = [ "sql", "sqlparser 0.38.0 (git+https://github.com/GreptimeTeam/sqlparser-rs.git?rev=0fbae07d0c46dc18e3381c406d8b9b8abef6b1fd)", "store-api", - "substrait 0.4.3", + "substrait 0.4.4", "table", "tokio", "tonic 0.10.2", @@ -5784,7 +5784,7 @@ dependencies = [ [[package]] name = "partition" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-trait", @@ -6122,7 +6122,7 @@ dependencies = [ [[package]] name = "plugins" -version = "0.4.3" +version = "0.4.4" dependencies = [ "auth", "common-base", @@ -6349,7 +6349,7 @@ dependencies = [ [[package]] name = "promql" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-recursion", "async-trait", @@ -6558,7 +6558,7 @@ dependencies = [ [[package]] name = "puffin" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-trait", "bitflags 2.4.1", @@ -6669,7 +6669,7 @@ dependencies = [ [[package]] name = "query" -version = "0.4.3" +version = "0.4.4" dependencies = [ "ahash 0.8.6", "api", @@ -6727,7 +6727,7 @@ dependencies = [ "stats-cli", "store-api", "streaming-stats", - "substrait 0.4.3", + "substrait 0.4.4", "table", "tokio", "tokio-stream", @@ -7928,7 +7928,7 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "script" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "arc-swap", @@ -8188,7 +8188,7 @@ dependencies = [ [[package]] name = "servers" -version = "0.4.3" +version = "0.4.4" dependencies = [ "aide", "api", @@ -8282,7 +8282,7 @@ dependencies = [ [[package]] name = "session" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "arc-swap", @@ -8543,7 +8543,7 @@ dependencies = [ [[package]] name = "sql" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "common-base", @@ -8595,7 +8595,7 @@ dependencies = [ [[package]] name = "sqlness-runner" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-trait", "clap 4.4.8", @@ -8801,7 +8801,7 @@ dependencies = [ [[package]] name = "store-api" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "aquamarine", @@ -8940,7 +8940,7 @@ dependencies = [ [[package]] name = "substrait" -version = "0.4.3" +version = "0.4.4" dependencies = [ "async-recursion", "async-trait", @@ -9088,7 +9088,7 @@ dependencies = [ [[package]] name = "table" -version = "0.4.3" +version = "0.4.4" dependencies = [ "anymap", "async-trait", @@ -9194,7 +9194,7 @@ dependencies = [ [[package]] name = "tests-integration" -version = "0.4.3" +version = "0.4.4" dependencies = [ "api", "async-trait", @@ -9250,7 +9250,7 @@ dependencies = [ "sql", "sqlx", "store-api", - "substrait 0.4.3", + "substrait 0.4.4", "table", "tempfile", "time", diff --git a/Cargo.toml b/Cargo.toml index 070f7fc186b8..449e17298310 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,7 +58,7 @@ members = [ resolver = "2" [workspace.package] -version = "0.4.3" +version = "0.4.4" edition = "2021" license = "Apache-2.0" From cce5edc88e91f5279ee90a17e72970db9097ad2e Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 29 Nov 2023 18:17:28 +0900 Subject: [PATCH 08/17] feat: add downgrade leader region step (#2792) * feat: add downgrade leader region step * chore: apply suggestions from CR * chore: rename exist to exists * chore: apply suggestions from CR --- src/common/meta/src/instruction.rs | 38 ++ src/datanode/src/heartbeat/handler.rs | 12 + .../src/procedure/region_migration.rs | 26 + .../downgrade_leader_region.rs | 491 +++++++++++++++++- .../region_migration/migration_start.rs | 2 +- .../region_migration/open_candidate_region.rs | 81 +-- .../procedure/region_migration/test_util.rs | 39 +- .../region_migration/update_metadata.rs | 7 +- .../upgrade_candidate_region.rs | 36 ++ 9 files changed, 661 insertions(+), 71 deletions(-) create mode 100644 src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs diff --git a/src/common/meta/src/instruction.rs b/src/common/meta/src/instruction.rs index 5119eedc2e62..7e787b41433e 100644 --- a/src/common/meta/src/instruction.rs +++ b/src/common/meta/src/instruction.rs @@ -48,6 +48,27 @@ impl Display for RegionIdent { } } +/// The result of downgrade leader region. +#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)] +pub struct DowngradeRegionReply { + /// Returns the `last_entry_id` if available. + pub last_entry_id: Option, + /// Indicates whether the region exists. + pub exists: bool, + /// Return error if any during the operation. + pub error: Option, +} + +impl Display for DowngradeRegionReply { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "(last_entry_id={:?}, exists={}, error={:?})", + self.last_entry_id, self.exists, self.error + ) + } +} + #[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)] pub struct SimpleReply { pub result: bool, @@ -87,10 +108,23 @@ impl OpenRegion { } } +/// The instruction of downgrading leader region. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DowngradeRegion { + pub region_id: RegionId, +} + +impl Display for DowngradeRegion { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "DowngradeRegion(region_id={})", self.region_id) + } +} + #[derive(Debug, Clone, Serialize, Deserialize, Display)] pub enum Instruction { OpenRegion(OpenRegion), CloseRegion(RegionIdent), + DowngradeRegion(DowngradeRegion), InvalidateTableIdCache(TableId), InvalidateTableNameCache(TableName), } @@ -101,6 +135,7 @@ pub enum InstructionReply { OpenRegion(SimpleReply), CloseRegion(SimpleReply), InvalidateTableCache(SimpleReply), + DowngradeRegion(DowngradeRegionReply), } impl Display for InstructionReply { @@ -111,6 +146,9 @@ impl Display for InstructionReply { Self::InvalidateTableCache(reply) => { write!(f, "InstructionReply::Invalidate({})", reply) } + Self::DowngradeRegion(reply) => { + write!(f, "InstructionReply::DowngradeRegion({})", reply) + } } } } diff --git a/src/datanode/src/heartbeat/handler.rs b/src/datanode/src/heartbeat/handler.rs index c4a2d57d07ad..9985fd6cf40e 100644 --- a/src/datanode/src/heartbeat/handler.rs +++ b/src/datanode/src/heartbeat/handler.rs @@ -65,6 +65,10 @@ impl RegionHeartbeatResponseHandler { Instruction::InvalidateTableIdCache(_) | Instruction::InvalidateTableNameCache(_) => { InvalidHeartbeatResponseSnafu.fail() } + Instruction::DowngradeRegion(_) => { + // TODO(weny): add it later. + todo!() + } } } @@ -88,6 +92,10 @@ impl RegionHeartbeatResponseHandler { error: None, }) } + Instruction::DowngradeRegion(_) => { + // TODO(weny): add it later. + todo!() + } } } @@ -114,6 +122,10 @@ impl RegionHeartbeatResponseHandler { reply.result = success; reply.error = error; } + InstructionReply::DowngradeRegion(_) => { + // TODO(weny): add it later. + todo!() + } } template diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index cabd1f7805ab..295aab78acbd 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -19,9 +19,11 @@ pub(crate) mod open_candidate_region; #[cfg(test)] pub(crate) mod test_util; pub(crate) mod update_metadata; +pub(crate) mod upgrade_candidate_region; use std::any::Any; use std::fmt::Debug; +use std::time::Duration; use common_meta::key::table_route::TableRouteValue; use common_meta::key::{DeserializedValueWithBytes, TableMetadataManagerRef}; @@ -34,6 +36,7 @@ use common_procedure::{Context as ProcedureContext, LockKey, Procedure, Status}; use serde::{Deserialize, Serialize}; use snafu::{location, Location, OptionExt, ResultExt}; use store_api::storage::RegionId; +use tokio::time::Instant; use self::migration_start::RegionMigrationStart; use crate::error::{self, Error, Result}; @@ -80,6 +83,29 @@ pub struct VolatileContext { opening_region_guard: Option, /// `table_route_info` is stored via previous steps for future use. table_route_info: Option>, + /// The deadline of leader region lease. + leader_region_lease_deadline: Option, + /// The last_entry_id of leader region. + leader_region_last_entry_id: Option, +} + +impl VolatileContext { + /// Sets the `leader_region_lease_deadline` if it does not exist. + pub fn set_leader_region_lease_deadline(&mut self, lease_timeout: Duration) { + if self.leader_region_lease_deadline.is_none() { + self.leader_region_lease_deadline = Some(Instant::now() + lease_timeout); + } + } + + /// Resets the `leader_region_lease_deadline`. + pub fn reset_leader_region_lease_deadline(&mut self) { + self.leader_region_lease_deadline = None; + } + + /// Sets the `leader_region_last_entry_id`. + pub fn set_last_entry_id(&mut self, last_entry_id: u64) { + self.leader_region_last_entry_id = Some(last_entry_id) + } } /// Used to generate new [Context]. diff --git a/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs b/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs index c0ff94330723..2c6a25fa6630 100644 --- a/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs +++ b/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs @@ -13,23 +13,506 @@ // limitations under the License. use std::any::Any; +use std::time::Duration; +use api::v1::meta::MailboxMessage; +use common_meta::distributed_time_constants::REGION_LEASE_SECS; +use common_meta::instruction::{ + DowngradeRegion, DowngradeRegionReply, Instruction, InstructionReply, +}; +use common_telemetry::warn; use serde::{Deserialize, Serialize}; +use snafu::ResultExt; +use tokio::time::sleep; -use crate::error::Result; +use super::upgrade_candidate_region::UpgradeCandidateRegion; +use crate::error::{self, Result}; +use crate::handler::HeartbeatMailbox; use crate::procedure::region_migration::{Context, State}; +use crate::service::mailbox::Channel; + +const DOWNGRADE_LEADER_REGION_TIMEOUT: Duration = Duration::from_secs(1); #[derive(Debug, Serialize, Deserialize)] -pub struct DowngradeLeaderRegion; +pub struct DowngradeLeaderRegion { + // The optimistic retry times. + optimistic_retry: usize, + // The retry initial interval. + retry_initial_interval: Duration, +} + +impl Default for DowngradeLeaderRegion { + fn default() -> Self { + Self { + optimistic_retry: 3, + retry_initial_interval: Duration::from_millis(500), + } + } +} #[async_trait::async_trait] #[typetag::serde] impl State for DowngradeLeaderRegion { - async fn next(&mut self, _ctx: &mut Context) -> Result> { - todo!() + async fn next(&mut self, ctx: &mut Context) -> Result> { + // Ensures the `leader_region_lease_deadline` must exist after recovering. + ctx.volatile_ctx + .set_leader_region_lease_deadline(Duration::from_secs(REGION_LEASE_SECS)); + self.downgrade_region_with_retry(ctx).await; + + // Safety: must exist. + if let Some(deadline) = ctx.volatile_ctx.leader_region_lease_deadline.as_ref() { + tokio::time::sleep_until(*deadline).await; + } + + Ok(Box::new(UpgradeCandidateRegion)) } fn as_any(&self) -> &dyn Any { self } } + +impl DowngradeLeaderRegion { + /// Builds downgrade region instruction. + fn build_downgrade_region_instruction(&self, ctx: &Context) -> Instruction { + let pc = &ctx.persistent_ctx; + let region_id = pc.region_id; + Instruction::DowngradeRegion(DowngradeRegion { region_id }) + } + + /// Tries to downgrade a leader region. + /// + /// Retry: + /// - [MailboxTimeout](error::Error::MailboxTimeout), Timeout. + /// - Failed to downgrade region on the Datanode. + /// + /// Abort: + /// - [PusherNotFound](error::Error::PusherNotFound), The datanode is unreachable. + /// - [PushMessage](error::Error::PushMessage), The receiver is dropped. + /// - [MailboxReceiver](error::Error::MailboxReceiver), The sender is dropped without sending (impossible). + /// - [UnexpectedInstructionReply](error::Error::UnexpectedInstructionReply). + /// - Invalid JSON. + async fn downgrade_region( + &self, + ctx: &mut Context, + downgrade_instruction: &Instruction, + ) -> Result<()> { + let pc = &ctx.persistent_ctx; + let region_id = pc.region_id; + let leader = &pc.from_peer; + + let msg = MailboxMessage::json_message( + &format!("Downgrade leader region: {}", region_id), + &format!("Meta@{}", ctx.server_addr()), + &format!("Datanode-{}@{}", leader.id, leader.addr), + common_time::util::current_time_millis(), + downgrade_instruction, + ) + .with_context(|_| error::SerializeToJsonSnafu { + input: downgrade_instruction.to_string(), + })?; + + let ch = Channel::Datanode(leader.id); + let receiver = ctx + .mailbox + .send(&ch, msg, DOWNGRADE_LEADER_REGION_TIMEOUT) + .await?; + + match receiver.await? { + Ok(msg) => { + let reply = HeartbeatMailbox::json_reply(&msg)?; + let InstructionReply::DowngradeRegion(DowngradeRegionReply { + last_entry_id, + exists, + error, + }) = reply + else { + return error::UnexpectedInstructionReplySnafu { + mailbox_message: msg.to_string(), + reason: "expect downgrade region reply", + } + .fail(); + }; + + if error.is_some() { + return error::RetryLaterSnafu { + reason: format!( + "Failed to downgrade the region {} on Datanode {:?}, error: {:?}", + region_id, leader, error + ), + } + .fail(); + } + + if !exists { + warn!( + "Trying to downgrade the region {} on Datanode {}, but region doesn't exist!", + region_id, leader + ); + } + + if let Some(last_entry_id) = last_entry_id { + ctx.volatile_ctx.set_last_entry_id(last_entry_id); + } + + Ok(()) + } + Err(error::Error::MailboxTimeout { .. }) => { + let reason = format!( + "Mailbox received timeout for downgrade leader region {region_id} on Datanode {:?}", + leader, + ); + error::RetryLaterSnafu { reason }.fail() + } + Err(err) => Err(err), + } + } + + /// Downgrades a leader region. + /// + /// Fast path: + /// - Waits for the reply of downgrade instruction. + /// + /// Slow path: + /// - Waits for the lease of the leader region expired. + async fn downgrade_region_with_retry(&self, ctx: &mut Context) { + let instruction = self.build_downgrade_region_instruction(ctx); + + let mut retry = 0; + + loop { + if let Err(err) = self.downgrade_region(ctx, &instruction).await { + retry += 1; + if err.is_retryable() && retry < self.optimistic_retry { + warn!("Failed to downgrade region, error: {err:?}, retry later"); + sleep(self.retry_initial_interval).await; + } else { + break; + } + } else { + // Resets the deadline. + ctx.volatile_ctx.reset_leader_region_lease_deadline(); + break; + } + } + } +} + +#[cfg(test)] +mod tests { + use std::assert_matches::assert_matches; + + use api::v1::meta::mailbox_message::Payload; + use common_meta::peer::Peer; + use common_time::util::current_time_millis; + use store_api::storage::RegionId; + use tokio::time::Instant; + + use super::*; + use crate::error::Error; + use crate::procedure::region_migration::test_util::{ + new_close_region_reply, send_mock_reply, TestingEnv, + }; + use crate::procedure::region_migration::{ContextFactory, PersistentContext}; + + fn new_persistent_context() -> PersistentContext { + PersistentContext { + from_peer: Peer::empty(1), + to_peer: Peer::empty(2), + region_id: RegionId::new(1024, 1), + cluster_id: 0, + } + } + + fn new_downgrade_region_reply( + id: u64, + last_entry_id: Option, + exist: bool, + error: Option, + ) -> MailboxMessage { + MailboxMessage { + id, + subject: "mock".to_string(), + from: "datanode".to_string(), + to: "meta".to_string(), + timestamp_millis: current_time_millis(), + payload: Some(Payload::Json( + serde_json::to_string(&InstructionReply::DowngradeRegion(DowngradeRegionReply { + last_entry_id, + exists: exist, + error, + })) + .unwrap(), + )), + } + } + + #[tokio::test] + async fn test_datanode_is_unreachable() { + let state = DowngradeLeaderRegion::default(); + let persistent_context = new_persistent_context(); + let env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + + let instruction = &state.build_downgrade_region_instruction(&ctx); + let err = state + .downgrade_region(&mut ctx, instruction) + .await + .unwrap_err(); + + assert_matches!(err, Error::PusherNotFound { .. }); + assert!(!err.is_retryable()); + } + + #[tokio::test] + async fn test_pusher_dropped() { + let state = DowngradeLeaderRegion::default(); + let persistent_context = new_persistent_context(); + let from_peer_id = persistent_context.from_peer.id; + + let mut env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let mailbox_ctx = env.mailbox_context(); + + let (tx, rx) = tokio::sync::mpsc::channel(1); + + mailbox_ctx + .insert_heartbeat_response_receiver(from_peer_id, tx) + .await; + + drop(rx); + + let instruction = &state.build_downgrade_region_instruction(&ctx); + let err = state + .downgrade_region(&mut ctx, instruction) + .await + .unwrap_err(); + + assert_matches!(err, Error::PushMessage { .. }); + assert!(!err.is_retryable()); + } + + #[tokio::test] + async fn test_unexpected_instruction_reply() { + let state = DowngradeLeaderRegion::default(); + let persistent_context = new_persistent_context(); + let from_peer_id = persistent_context.from_peer.id; + + let mut env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let mailbox_ctx = env.mailbox_context(); + let mailbox = mailbox_ctx.mailbox().clone(); + + let (tx, rx) = tokio::sync::mpsc::channel(1); + + mailbox_ctx + .insert_heartbeat_response_receiver(from_peer_id, tx) + .await; + + // Sends an incorrect reply. + send_mock_reply(mailbox, rx, |id| Ok(new_close_region_reply(id))); + + let instruction = &state.build_downgrade_region_instruction(&ctx); + let err = state + .downgrade_region(&mut ctx, instruction) + .await + .unwrap_err(); + + assert_matches!(err, Error::UnexpectedInstructionReply { .. }); + assert!(!err.is_retryable()); + } + + #[tokio::test] + async fn test_instruction_exceeded_deadline() { + let state = DowngradeLeaderRegion::default(); + let persistent_context = new_persistent_context(); + let from_peer_id = persistent_context.from_peer.id; + + let mut env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let mailbox_ctx = env.mailbox_context(); + let mailbox = mailbox_ctx.mailbox().clone(); + + let (tx, rx) = tokio::sync::mpsc::channel(1); + + mailbox_ctx + .insert_heartbeat_response_receiver(from_peer_id, tx) + .await; + + send_mock_reply(mailbox, rx, |id| { + Err(error::MailboxTimeoutSnafu { id }.build()) + }); + + let instruction = &state.build_downgrade_region_instruction(&ctx); + let err = state + .downgrade_region(&mut ctx, instruction) + .await + .unwrap_err(); + + assert_matches!(err, Error::RetryLater { .. }); + assert!(err.is_retryable()); + } + + #[tokio::test] + async fn test_downgrade_region_failed() { + let state = DowngradeLeaderRegion::default(); + let persistent_context = new_persistent_context(); + let from_peer_id = persistent_context.from_peer.id; + + let mut env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let mailbox_ctx = env.mailbox_context(); + let mailbox = mailbox_ctx.mailbox().clone(); + + let (tx, rx) = tokio::sync::mpsc::channel(1); + + mailbox_ctx + .insert_heartbeat_response_receiver(from_peer_id, tx) + .await; + + send_mock_reply(mailbox, rx, |id| { + Ok(new_downgrade_region_reply( + id, + None, + false, + Some("test mocked".to_string()), + )) + }); + + let instruction = &state.build_downgrade_region_instruction(&ctx); + let err = state + .downgrade_region(&mut ctx, instruction) + .await + .unwrap_err(); + + assert_matches!(err, Error::RetryLater { .. }); + assert!(err.is_retryable()); + assert!(err.to_string().contains("test mocked")); + } + + #[tokio::test] + async fn test_downgrade_region_with_retry_fast_path() { + let state = DowngradeLeaderRegion::default(); + let persistent_context = new_persistent_context(); + let from_peer_id = persistent_context.from_peer.id; + + let mut env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let mailbox_ctx = env.mailbox_context(); + let mailbox = mailbox_ctx.mailbox().clone(); + + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + + mailbox_ctx + .insert_heartbeat_response_receiver(from_peer_id, tx) + .await; + + common_runtime::spawn_bg(async move { + // retry: 0. + let resp = rx.recv().await.unwrap().unwrap(); + let reply_id = resp.mailbox_message.unwrap().id; + mailbox + .on_recv( + reply_id, + Err(error::MailboxTimeoutSnafu { id: reply_id }.build()), + ) + .await + .unwrap(); + + // retry: 1. + let resp = rx.recv().await.unwrap().unwrap(); + let reply_id = resp.mailbox_message.unwrap().id; + mailbox + .on_recv( + reply_id, + Ok(new_downgrade_region_reply(reply_id, Some(1), true, None)), + ) + .await + .unwrap(); + }); + + state.downgrade_region_with_retry(&mut ctx).await; + assert_eq!(ctx.volatile_ctx.leader_region_last_entry_id, Some(1)); + assert!(ctx.volatile_ctx.leader_region_lease_deadline.is_none()); + } + + #[tokio::test] + async fn test_downgrade_region_with_retry_slow_path() { + let state = DowngradeLeaderRegion { + optimistic_retry: 3, + retry_initial_interval: Duration::from_millis(100), + }; + let persistent_context = new_persistent_context(); + let from_peer_id = persistent_context.from_peer.id; + + let mut env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let mailbox_ctx = env.mailbox_context(); + let mailbox = mailbox_ctx.mailbox().clone(); + + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + + mailbox_ctx + .insert_heartbeat_response_receiver(from_peer_id, tx) + .await; + + common_runtime::spawn_bg(async move { + for _ in 0..3 { + let resp = rx.recv().await.unwrap().unwrap(); + let reply_id = resp.mailbox_message.unwrap().id; + mailbox + .on_recv( + reply_id, + Err(error::MailboxTimeoutSnafu { id: reply_id }.build()), + ) + .await + .unwrap(); + } + }); + + ctx.volatile_ctx + .set_leader_region_lease_deadline(Duration::from_secs(5)); + let expected_deadline = ctx.volatile_ctx.leader_region_lease_deadline.unwrap(); + state.downgrade_region_with_retry(&mut ctx).await; + assert_eq!(ctx.volatile_ctx.leader_region_last_entry_id, None); + // Should remain no change. + assert_eq!( + ctx.volatile_ctx.leader_region_lease_deadline.unwrap(), + expected_deadline + ) + } + + #[tokio::test] + async fn test_next_upgrade_candidate_state() { + let mut state = Box::::default(); + let persistent_context = new_persistent_context(); + let from_peer_id = persistent_context.from_peer.id; + + let mut env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let mailbox_ctx = env.mailbox_context(); + let mailbox = mailbox_ctx.mailbox().clone(); + + let (tx, rx) = tokio::sync::mpsc::channel(1); + + mailbox_ctx + .insert_heartbeat_response_receiver(from_peer_id, tx) + .await; + + send_mock_reply(mailbox, rx, |id| { + Ok(new_downgrade_region_reply(id, Some(1), true, None)) + }); + + let timer = Instant::now(); + let next = state.next(&mut ctx).await.unwrap(); + let elapsed = timer.elapsed().as_secs(); + assert!(elapsed < REGION_LEASE_SECS / 2); + assert_eq!(ctx.volatile_ctx.leader_region_last_entry_id, Some(1)); + assert!(ctx.volatile_ctx.leader_region_lease_deadline.is_none()); + + let _ = next + .as_any() + .downcast_ref::() + .unwrap(); + } +} diff --git a/src/meta-srv/src/procedure/region_migration/migration_start.rs b/src/meta-srv/src/procedure/region_migration/migration_start.rs index ab61da316c8d..47eb0abb980f 100644 --- a/src/meta-srv/src/procedure/region_migration/migration_start.rs +++ b/src/meta-srv/src/procedure/region_migration/migration_start.rs @@ -47,7 +47,7 @@ impl State for RegionMigrationStart { if self.check_leader_region_on_peer(®ion_route, to_peer)? { Ok(Box::new(RegionMigrationEnd)) } else if self.check_candidate_region_on_peer(®ion_route, to_peer) { - Ok(Box::new(DowngradeLeaderRegion)) + Ok(Box::::default()) } else { Ok(Box::new(OpenCandidateRegion)) } diff --git a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs index 42e057b58f47..fc4d9c7d9a45 100644 --- a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs @@ -41,7 +41,7 @@ impl State for OpenCandidateRegion { let instruction = self.build_open_region_instruction(ctx).await?; self.open_candidate_region(ctx, instruction).await?; - Ok(Box::new(DowngradeLeaderRegion)) + Ok(Box::::default()) } fn as_any(&self) -> &dyn Any { @@ -197,7 +197,9 @@ mod tests { use super::*; use crate::error::Error; use crate::procedure::region_migration::downgrade_leader_region::DowngradeLeaderRegion; - use crate::procedure::region_migration::test_util::TestingEnv; + use crate::procedure::region_migration::test_util::{ + new_close_region_reply, send_mock_reply, TestingEnv, + }; use crate::procedure::region_migration::{ContextFactory, PersistentContext}; fn new_persistent_context() -> PersistentContext { @@ -223,23 +225,6 @@ mod tests { }) } - fn new_close_region_reply(id: u64) -> MailboxMessage { - MailboxMessage { - id, - subject: "mock".to_string(), - from: "datanode".to_string(), - to: "meta".to_string(), - timestamp_millis: current_time_millis(), - payload: Some(Payload::Json( - serde_json::to_string(&InstructionReply::CloseRegion(SimpleReply { - result: false, - error: None, - })) - .unwrap(), - )), - } - } - fn new_open_region_reply(id: u64, result: bool, error: Option) -> MailboxMessage { MailboxMessage { id, @@ -328,21 +313,14 @@ mod tests { let mailbox_ctx = env.mailbox_context(); let mailbox = mailbox_ctx.mailbox().clone(); - let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let (tx, rx) = tokio::sync::mpsc::channel(1); mailbox_ctx .insert_heartbeat_response_receiver(to_peer_id, tx) .await; // Sends an incorrect reply. - common_runtime::spawn_bg(async move { - let resp = rx.recv().await.unwrap().unwrap(); - let reply_id = resp.mailbox_message.unwrap().id; - mailbox - .on_recv(reply_id, Ok(new_close_region_reply(reply_id))) - .await - .unwrap(); - }); + send_mock_reply(mailbox, rx, |id| Ok(new_close_region_reply(id))); let open_instruction = new_mock_open_instruction(to_peer_id, region_id); let err = state @@ -368,23 +346,15 @@ mod tests { let mailbox_ctx = env.mailbox_context(); let mailbox = mailbox_ctx.mailbox().clone(); - let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let (tx, rx) = tokio::sync::mpsc::channel(1); mailbox_ctx .insert_heartbeat_response_receiver(to_peer_id, tx) .await; // Sends an timeout error. - common_runtime::spawn_bg(async move { - let resp = rx.recv().await.unwrap().unwrap(); - let reply_id = resp.mailbox_message.unwrap().id; - mailbox - .on_recv( - reply_id, - Err(error::MailboxTimeoutSnafu { id: reply_id }.build()), - ) - .await - .unwrap(); + send_mock_reply(mailbox, rx, |id| { + Err(error::MailboxTimeoutSnafu { id }.build()) }); let open_instruction = new_mock_open_instruction(to_peer_id, region_id); @@ -411,26 +381,18 @@ mod tests { let mailbox_ctx = env.mailbox_context(); let mailbox = mailbox_ctx.mailbox().clone(); - let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let (tx, rx) = tokio::sync::mpsc::channel(1); mailbox_ctx .insert_heartbeat_response_receiver(to_peer_id, tx) .await; - common_runtime::spawn_bg(async move { - let resp = rx.recv().await.unwrap().unwrap(); - let reply_id = resp.mailbox_message.unwrap().id; - mailbox - .on_recv( - reply_id, - Ok(new_open_region_reply( - reply_id, - false, - Some("test mocked".to_string()), - )), - ) - .await - .unwrap(); + send_mock_reply(mailbox, rx, |id| { + Ok(new_open_region_reply( + id, + false, + Some("test mocked".to_string()), + )) }); let open_instruction = new_mock_open_instruction(to_peer_id, region_id); @@ -471,20 +433,13 @@ mod tests { let mailbox_ctx = env.mailbox_context(); let mailbox = mailbox_ctx.mailbox().clone(); - let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let (tx, rx) = tokio::sync::mpsc::channel(1); mailbox_ctx .insert_heartbeat_response_receiver(to_peer_id, tx) .await; - common_runtime::spawn_bg(async move { - let resp = rx.recv().await.unwrap().unwrap(); - let reply_id = resp.mailbox_message.unwrap().id; - mailbox - .on_recv(reply_id, Ok(new_open_region_reply(reply_id, true, None))) - .await - .unwrap(); - }); + send_mock_reply(mailbox, rx, |id| Ok(new_open_region_reply(id, true, None))); let next = state.next(&mut ctx).await.unwrap(); let vc = ctx.volatile_ctx; diff --git a/src/meta-srv/src/procedure/region_migration/test_util.rs b/src/meta-srv/src/procedure/region_migration/test_util.rs index d880677af03c..277c9f8d90fa 100644 --- a/src/meta-srv/src/procedure/region_migration/test_util.rs +++ b/src/meta-srv/src/procedure/region_migration/test_util.rs @@ -14,20 +14,26 @@ use std::sync::Arc; -use api::v1::meta::{HeartbeatResponse, RequestHeader}; +use api::v1::meta::mailbox_message::Payload; +use api::v1::meta::{HeartbeatResponse, MailboxMessage, RequestHeader}; +use common_meta::instruction::{InstructionReply, SimpleReply}; use common_meta::key::{TableMetadataManager, TableMetadataManagerRef}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::Sequence; use common_meta::DatanodeId; use common_procedure::{Context as ProcedureContext, ProcedureId}; use common_procedure_test::MockContextProvider; -use tokio::sync::mpsc::Sender; +use common_time::util::current_time_millis; +use tokio::sync::mpsc::{Receiver, Sender}; use super::ContextFactoryImpl; +use crate::error::Result; use crate::handler::{HeartbeatMailbox, Pusher, Pushers}; use crate::region::lease_keeper::{OpeningRegionKeeper, OpeningRegionKeeperRef}; use crate::service::mailbox::{Channel, MailboxRef}; +pub type MockHeartbeatReceiver = Receiver>; + /// The context of mailbox. pub struct MailboxContext { mailbox: MailboxRef, @@ -120,3 +126,32 @@ impl TestingEnv { } } } + +pub fn new_close_region_reply(id: u64) -> MailboxMessage { + MailboxMessage { + id, + subject: "mock".to_string(), + from: "datanode".to_string(), + to: "meta".to_string(), + timestamp_millis: current_time_millis(), + payload: Some(Payload::Json( + serde_json::to_string(&InstructionReply::CloseRegion(SimpleReply { + result: false, + error: None, + })) + .unwrap(), + )), + } +} + +pub fn send_mock_reply( + mailbox: MailboxRef, + mut rx: MockHeartbeatReceiver, + msg: impl FnOnce(u64) -> Result + Send + 'static, +) { + common_runtime::spawn_bg(async move { + let resp = rx.recv().await.unwrap().unwrap(); + let reply_id = resp.mailbox_message.unwrap().id; + mailbox.on_recv(reply_id, msg(reply_id)).await.unwrap(); + }); +} diff --git a/src/meta-srv/src/procedure/region_migration/update_metadata.rs b/src/meta-srv/src/procedure/region_migration/update_metadata.rs index f41b66f4c09e..9c8d4b85b7e5 100644 --- a/src/meta-srv/src/procedure/region_migration/update_metadata.rs +++ b/src/meta-srv/src/procedure/region_migration/update_metadata.rs @@ -13,7 +13,9 @@ // limitations under the License. use std::any::Any; +use std::time::Duration; +use common_meta::distributed_time_constants::REGION_LEASE_SECS; use common_meta::rpc::router::RegionStatus; use serde::{Deserialize, Serialize}; use snafu::ResultExt; @@ -36,7 +38,7 @@ impl State for UpdateMetadata { UpdateMetadata::Downgrade => { self.downgrade_leader_region(ctx).await?; - Ok(Box::new(DowngradeLeaderRegion)) + Ok(Box::::default()) } } } @@ -88,6 +90,9 @@ impl UpdateMetadata { debug_assert!(ctx.remove_table_route_value()); + ctx.volatile_ctx + .set_leader_region_lease_deadline(Duration::from_secs(REGION_LEASE_SECS)); + Ok(()) } } diff --git a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs new file mode 100644 index 000000000000..8b15a0730f7a --- /dev/null +++ b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs @@ -0,0 +1,36 @@ +// 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. + +use std::any::Any; + +use serde::{Deserialize, Serialize}; + +use crate::error::Result; +use crate::procedure::region_migration::{Context, State}; +#[derive(Debug, Serialize, Deserialize)] +pub struct UpgradeCandidateRegion; + +#[async_trait::async_trait] +#[typetag::serde] +impl State for UpgradeCandidateRegion { + async fn next(&mut self, _ctx: &mut Context) -> Result> { + todo!(); + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +impl UpgradeCandidateRegion {} From ae8153515bd54a5e347dadc9fbf1aca1f2678bf4 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 29 Nov 2023 20:10:38 +0900 Subject: [PATCH 09/17] feat: add update metadata step for upgrading candidate region (#2811) --- src/common/meta/src/key.rs | 12 +- src/meta-srv/src/error.rs | 9 +- src/meta-srv/src/lib.rs | 2 + .../region_failover/update_metadata.rs | 2 +- .../src/procedure/region_migration.rs | 54 ++- .../region_migration/migration_start.rs | 9 +- .../region_migration/open_candidate_region.rs | 42 +- .../procedure/region_migration/test_util.rs | 15 + .../region_migration/update_metadata.rs | 214 +--------- .../downgrade_leader_region.rs | 210 ++++++++++ .../upgrade_candidate_region.rs | 376 ++++++++++++++++++ 11 files changed, 697 insertions(+), 248 deletions(-) create mode 100644 src/meta-srv/src/procedure/region_migration/update_metadata/downgrade_leader_region.rs create mode 100644 src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index 39a942d5f3bd..850051fa0b93 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -584,7 +584,7 @@ impl TableMetadataManager { &self, table_id: TableId, region_info: RegionInfo, - current_table_route_value: DeserializedValueWithBytes, + current_table_route_value: &DeserializedValueWithBytes, new_region_routes: Vec, new_region_options: &HashMap, ) -> Result<()> { @@ -606,7 +606,7 @@ impl TableMetadataManager { let (update_table_route_txn, on_update_table_route_failure) = self .table_route_manager() - .build_update_txn(table_id, ¤t_table_route_value, &new_table_route_value)?; + .build_update_txn(table_id, current_table_route_value, &new_table_route_value)?; let txn = Txn::merge_all(vec![update_datanode_table_txn, update_table_route_txn]); @@ -1173,7 +1173,7 @@ mod tests { region_storage_path: region_storage_path.to_string(), region_options: HashMap::new(), }, - current_table_route_value.clone(), + ¤t_table_route_value, new_region_routes.clone(), &HashMap::new(), ) @@ -1190,7 +1190,7 @@ mod tests { region_storage_path: region_storage_path.to_string(), region_options: HashMap::new(), }, - current_table_route_value.clone(), + ¤t_table_route_value, new_region_routes.clone(), &HashMap::new(), ) @@ -1212,7 +1212,7 @@ mod tests { region_storage_path: region_storage_path.to_string(), region_options: HashMap::new(), }, - current_table_route_value.clone(), + ¤t_table_route_value, new_region_routes.clone(), &HashMap::new(), ) @@ -1237,7 +1237,7 @@ mod tests { region_storage_path: region_storage_path.to_string(), region_options: HashMap::new(), }, - wrong_table_route_value, + &wrong_table_route_value, new_region_routes, &HashMap::new(), ) diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 3a007b3163c4..abd60a4947cf 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -298,6 +298,12 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to find table route for {region_id}"))] + RegionRouteNotFound { + region_id: RegionId, + location: Location, + }, + #[snafu(display("Table info not found: {}", table_id))] TableInfoNotFound { table_id: TableId, @@ -658,7 +664,8 @@ impl ErrorExt for Error { | Error::Unexpected { .. } | Error::Txn { .. } | Error::TableIdChanged { .. } - | Error::RegionOpeningRace { .. } => StatusCode::Unexpected, + | Error::RegionOpeningRace { .. } + | Error::RegionRouteNotFound { .. } => StatusCode::Unexpected, Error::TableNotFound { .. } => StatusCode::TableNotFound, Error::InvalidateTableCache { source, .. } => source.status_code(), Error::RequestDatanode { source, .. } => source.status_code(), diff --git a/src/meta-srv/src/lib.rs b/src/meta-srv/src/lib.rs index b30c6779b36a..6515ade81ebe 100644 --- a/src/meta-srv/src/lib.rs +++ b/src/meta-srv/src/lib.rs @@ -15,6 +15,8 @@ #![feature(async_closure)] #![feature(result_flattening)] #![feature(assert_matches)] +#![feature(option_take_if)] +#![feature(extract_if)] pub mod bootstrap; mod cache_invalidator; diff --git a/src/meta-srv/src/procedure/region_failover/update_metadata.rs b/src/meta-srv/src/procedure/region_failover/update_metadata.rs index 22c1a89b542e..225588ac0942 100644 --- a/src/meta-srv/src/procedure/region_failover/update_metadata.rs +++ b/src/meta-srv/src/procedure/region_failover/update_metadata.rs @@ -105,7 +105,7 @@ impl UpdateRegionMetadata { region_storage_path: self.region_storage_path.to_string(), region_options: self.region_options.clone(), }, - table_route_value, + &table_route_value, new_region_routes, &self.region_options, ) diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index 295aab78acbd..d5ae7235c79d 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -25,6 +25,7 @@ use std::any::Any; use std::fmt::Debug; use std::time::Duration; +use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_route::TableRouteValue; use common_meta::key::{DeserializedValueWithBytes, TableMetadataManagerRef}; use common_meta::peer::Peer; @@ -81,8 +82,13 @@ pub struct VolatileContext { /// the corresponding [RegionRoute](common_meta::rpc::router::RegionRoute) of the opening region /// was written into [TableRouteValue](common_meta::key::table_route::TableRouteValue). opening_region_guard: Option, - /// `table_route_info` is stored via previous steps for future use. - table_route_info: Option>, + /// `table_route` is stored via previous steps for future use. + table_route: Option>, + /// `table_info` is stored via previous steps for future use. + /// + /// `table_info` should remain unchanged during the procedure; + /// no other DDL procedure executed concurrently for the current table. + table_info: Option>, /// The deadline of leader region lease. leader_region_lease_deadline: Option, /// The last_entry_id of leader region. @@ -153,7 +159,7 @@ impl Context { &self.server_addr } - /// Returns the `table_route_value` of [VolatileContext] if any. + /// Returns the `table_route` of [VolatileContext] if any. /// Otherwise, returns the value retrieved from remote. /// /// Retry: @@ -161,7 +167,7 @@ impl Context { pub async fn get_table_route_value( &mut self, ) -> Result<&DeserializedValueWithBytes> { - let table_route_value = &mut self.volatile_ctx.table_route_info; + let table_route_value = &mut self.volatile_ctx.table_route; if table_route_value.is_none() { let table_id = self.persistent_ctx.region_id.table_id(); @@ -183,9 +189,45 @@ impl Context { Ok(table_route_value.as_ref().unwrap()) } - /// Removes the `table_route_value` of [VolatileContext], returns true if any. + /// Removes the `table_route` of [VolatileContext], returns true if any. pub fn remove_table_route_value(&mut self) -> bool { - let value = self.volatile_ctx.table_route_info.take(); + let value = self.volatile_ctx.table_route.take(); + value.is_some() + } + + /// Returns the `table_info` of [VolatileContext] if any. + /// Otherwise, returns the value retrieved from remote. + /// + /// Retry: + /// - Failed to retrieve the metadata of table. + pub async fn get_table_info_value( + &mut self, + ) -> Result<&DeserializedValueWithBytes> { + let table_info_value = &mut self.volatile_ctx.table_info; + + if table_info_value.is_none() { + let table_id = self.persistent_ctx.region_id.table_id(); + let table_info = self + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .context(error::TableMetadataManagerSnafu) + .map_err(|e| error::Error::RetryLater { + reason: e.to_string(), + location: location!(), + })? + .context(error::TableInfoNotFoundSnafu { table_id })?; + + *table_info_value = Some(table_info); + } + + Ok(table_info_value.as_ref().unwrap()) + } + + /// Removes the `table_info` of [VolatileContext], returns true if any. + pub fn remove_table_info_value(&mut self) -> bool { + let value = self.volatile_ctx.table_info.take(); value.is_some() } diff --git a/src/meta-srv/src/procedure/region_migration/migration_start.rs b/src/meta-srv/src/procedure/region_migration/migration_start.rs index 47eb0abb980f..b10a26886aec 100644 --- a/src/meta-srv/src/procedure/region_migration/migration_start.rs +++ b/src/meta-srv/src/procedure/region_migration/migration_start.rs @@ -137,16 +137,11 @@ mod tests { use super::*; use crate::error::Error; - use crate::procedure::region_migration::test_util::TestingEnv; + use crate::procedure::region_migration::test_util::{self, TestingEnv}; use crate::procedure::region_migration::{ContextFactory, PersistentContext}; fn new_persistent_context() -> PersistentContext { - PersistentContext { - from_peer: Peer::empty(1), - to_peer: Peer::empty(2), - region_id: RegionId::new(1024, 1), - cluster_id: 0, - } + test_util::new_persistent_context(1, 2, RegionId::new(1024, 1)) } #[tokio::test] diff --git a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs index fc4d9c7d9a45..2e126625347f 100644 --- a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs @@ -21,7 +21,7 @@ use common_meta::ddl::utils::region_storage_path; use common_meta::instruction::{Instruction, InstructionReply, OpenRegion, SimpleReply}; use common_meta::RegionIdent; use serde::{Deserialize, Serialize}; -use snafu::{location, Location, OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt}; use crate::error::{self, Result}; use crate::handler::HeartbeatMailbox; @@ -54,38 +54,28 @@ impl OpenCandidateRegion { /// /// Abort(non-retry): /// - Table Info is not found. - async fn build_open_region_instruction(&self, ctx: &Context) -> Result { + async fn build_open_region_instruction(&self, ctx: &mut Context) -> Result { let pc = &ctx.persistent_ctx; let cluster_id = pc.cluster_id; let table_id = pc.region_id.table_id(); let region_number = pc.region_id.region_number(); - let candidate = &pc.to_peer; - let table_info = ctx - .table_metadata_manager - .table_info_manager() - .get(table_id) - .await - .context(error::TableMetadataManagerSnafu) - .map_err(|e| error::Error::RetryLater { - reason: e.to_string(), - location: location!(), - })? - .context(error::TableInfoNotFoundSnafu { table_id })? - .into_inner() - .table_info; + let candidate_id = pc.to_peer.id; + + let table_info_value = ctx.get_table_info_value().await?; + let table_info = &table_info_value.table_info; // The region storage path is immutable after the region is created. // Therefore, it's safe to store it in `VolatileContext` for future use. let region_storage_path = region_storage_path(&table_info.catalog_name, &table_info.schema_name); - let engine = table_info.meta.engine; + let engine = table_info.meta.engine.clone(); let region_options: HashMap = (&table_info.meta.options).into(); let open_instruction = Instruction::OpenRegion(OpenRegion::new( RegionIdent { cluster_id, - datanode_id: candidate.id, + datanode_id: candidate_id, table_id, region_number, engine, @@ -198,17 +188,12 @@ mod tests { use crate::error::Error; use crate::procedure::region_migration::downgrade_leader_region::DowngradeLeaderRegion; use crate::procedure::region_migration::test_util::{ - new_close_region_reply, send_mock_reply, TestingEnv, + self, new_close_region_reply, send_mock_reply, TestingEnv, }; use crate::procedure::region_migration::{ContextFactory, PersistentContext}; fn new_persistent_context() -> PersistentContext { - PersistentContext { - from_peer: Peer::empty(1), - to_peer: Peer::empty(2), - region_id: RegionId::new(1024, 1), - cluster_id: 0, - } + test_util::new_persistent_context(1, 2, RegionId::new(1024, 1)) } fn new_mock_open_instruction(datanode_id: DatanodeId, region_id: RegionId) -> Instruction { @@ -244,9 +229,12 @@ mod tests { let state = OpenCandidateRegion; let persistent_context = new_persistent_context(); let env = TestingEnv::new(); - let ctx = env.context_factory().new_context(persistent_context); + let mut ctx = env.context_factory().new_context(persistent_context); - let err = state.build_open_region_instruction(&ctx).await.unwrap_err(); + let err = state + .build_open_region_instruction(&mut ctx) + .await + .unwrap_err(); assert_matches!(err, Error::TableInfoNotFound { .. }); assert!(!err.is_retryable()); diff --git a/src/meta-srv/src/procedure/region_migration/test_util.rs b/src/meta-srv/src/procedure/region_migration/test_util.rs index 277c9f8d90fa..cc3779b8f54d 100644 --- a/src/meta-srv/src/procedure/region_migration/test_util.rs +++ b/src/meta-srv/src/procedure/region_migration/test_util.rs @@ -19,16 +19,19 @@ use api::v1::meta::{HeartbeatResponse, MailboxMessage, RequestHeader}; use common_meta::instruction::{InstructionReply, SimpleReply}; use common_meta::key::{TableMetadataManager, TableMetadataManagerRef}; use common_meta::kv_backend::memory::MemoryKvBackend; +use common_meta::peer::Peer; use common_meta::sequence::Sequence; use common_meta::DatanodeId; use common_procedure::{Context as ProcedureContext, ProcedureId}; use common_procedure_test::MockContextProvider; use common_time::util::current_time_millis; +use store_api::storage::RegionId; use tokio::sync::mpsc::{Receiver, Sender}; use super::ContextFactoryImpl; use crate::error::Result; use crate::handler::{HeartbeatMailbox, Pusher, Pushers}; +use crate::procedure::region_migration::PersistentContext; use crate::region::lease_keeper::{OpeningRegionKeeper, OpeningRegionKeeperRef}; use crate::service::mailbox::{Channel, MailboxRef}; @@ -127,6 +130,7 @@ impl TestingEnv { } } +/// Generates a [InstructionReply::CloseRegion] reply. pub fn new_close_region_reply(id: u64) -> MailboxMessage { MailboxMessage { id, @@ -144,6 +148,7 @@ pub fn new_close_region_reply(id: u64) -> MailboxMessage { } } +/// Sends a mock reply. pub fn send_mock_reply( mailbox: MailboxRef, mut rx: MockHeartbeatReceiver, @@ -155,3 +160,13 @@ pub fn send_mock_reply( mailbox.on_recv(reply_id, msg(reply_id)).await.unwrap(); }); } + +/// Generates a [PersistentContext]. +pub fn new_persistent_context(from: u64, to: u64, region_id: RegionId) -> PersistentContext { + PersistentContext { + from_peer: Peer::empty(from), + to_peer: Peer::empty(to), + region_id, + cluster_id: 0, + } +} diff --git a/src/meta-srv/src/procedure/region_migration/update_metadata.rs b/src/meta-srv/src/procedure/region_migration/update_metadata.rs index 9c8d4b85b7e5..ba3092548efb 100644 --- a/src/meta-srv/src/procedure/region_migration/update_metadata.rs +++ b/src/meta-srv/src/procedure/region_migration/update_metadata.rs @@ -12,22 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub(crate) mod downgrade_leader_region; +pub(crate) mod upgrade_candidate_region; + use std::any::Any; -use std::time::Duration; -use common_meta::distributed_time_constants::REGION_LEASE_SECS; -use common_meta::rpc::router::RegionStatus; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; -use crate::error::{self, Result}; +use super::migration_end::RegionMigrationEnd; +use crate::error::Result; use crate::procedure::region_migration::downgrade_leader_region::DowngradeLeaderRegion; use crate::procedure::region_migration::{Context, State}; #[derive(Debug, Serialize, Deserialize)] #[serde(tag = "UpdateMetadata")] pub enum UpdateMetadata { + /// Downgrades the leader region. Downgrade, + /// Upgrade the candidate region. + Upgrade, } #[async_trait::async_trait] @@ -40,6 +43,12 @@ impl State for UpdateMetadata { Ok(Box::::default()) } + UpdateMetadata::Upgrade => { + self.upgrade_candidate_region(ctx).await?; + + // TODO(weny): invalidate fe cache. + Ok(Box::new(RegionMigrationEnd)) + } } } @@ -47,198 +56,3 @@ impl State for UpdateMetadata { self } } - -impl UpdateMetadata { - /// Downgrades the leader region. - /// - /// Abort(non-retry): - /// - TableRoute is not found. - /// - /// Retry: - /// - Failed to update [TableRouteValue](common_meta::key::table_region::TableRegionValue). - /// - Failed to retrieve the metadata of table. - /// - /// About the failure of updating the [TableRouteValue](common_meta::key::table_region::TableRegionValue): - /// - /// - There may be another [RegionMigrationProcedure](crate::procedure::region_migration::RegionMigrationProcedure) - /// that is executed concurrently for **other region**. - /// It will only update **other region** info. Therefore, It's safe to retry after failure. - /// - /// - There is no other DDL procedure executed concurrently for the current table. - async fn downgrade_leader_region(&self, ctx: &mut Context) -> Result<()> { - let table_metadata_manager = ctx.table_metadata_manager.clone(); - let region_id = ctx.region_id(); - let table_id = region_id.table_id(); - let current_table_route_value = ctx.get_table_route_value().await?; - - if let Err(err) = table_metadata_manager - .update_leader_region_status(table_id, current_table_route_value, |route| { - if route.region.id == region_id { - Some(Some(RegionStatus::Downgraded)) - } else { - None - } - }) - .await - .context(error::TableMetadataManagerSnafu) - { - debug_assert!(ctx.remove_table_route_value()); - return error::RetryLaterSnafu { - reason: format!("Failed to update the table route during the downgrading leader region, error: {err}") - }.fail(); - } - - debug_assert!(ctx.remove_table_route_value()); - - ctx.volatile_ctx - .set_leader_region_lease_deadline(Duration::from_secs(REGION_LEASE_SECS)); - - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use std::assert_matches::assert_matches; - - use common_meta::key::test_utils::new_test_table_info; - use common_meta::peer::Peer; - use common_meta::rpc::router::{Region, RegionRoute}; - use store_api::storage::RegionId; - - use super::*; - use crate::error::Error; - use crate::procedure::region_migration::test_util::TestingEnv; - use crate::procedure::region_migration::{ContextFactory, PersistentContext}; - - fn new_persistent_context() -> PersistentContext { - PersistentContext { - from_peer: Peer::empty(1), - to_peer: Peer::empty(2), - region_id: RegionId::new(1024, 1), - cluster_id: 0, - } - } - - #[test] - fn test_state_serialization() { - let state = UpdateMetadata::Downgrade; - let expected = r#"{"UpdateMetadata":"Downgrade"}"#; - assert_eq!(expected, serde_json::to_string(&state).unwrap()); - } - - #[tokio::test] - async fn test_table_route_is_not_found_error() { - let state = UpdateMetadata::Downgrade; - let env = TestingEnv::new(); - let persistent_context = new_persistent_context(); - let mut ctx = env.context_factory().new_context(persistent_context); - - let err = state.downgrade_leader_region(&mut ctx).await.unwrap_err(); - - assert_matches!(err, Error::TableRouteNotFound { .. }); - - assert!(!err.is_retryable()); - } - - #[tokio::test] - async fn test_failed_to_update_table_route_error() { - let state = UpdateMetadata::Downgrade; - let persistent_context = new_persistent_context(); - let from_peer = persistent_context.from_peer.clone(); - - let env = TestingEnv::new(); - let mut ctx = env.context_factory().new_context(persistent_context); - let table_id = ctx.region_id().table_id(); - - let table_info = new_test_table_info(1024, vec![1, 2]).into(); - let region_routes = vec![ - RegionRoute { - region: Region::new_test(RegionId::new(1024, 1)), - leader_peer: Some(from_peer.clone()), - ..Default::default() - }, - RegionRoute { - region: Region::new_test(RegionId::new(1024, 2)), - leader_peer: Some(Peer::empty(4)), - ..Default::default() - }, - ]; - - let table_metadata_manager = env.table_metadata_manager(); - table_metadata_manager - .create_table_metadata(table_info, region_routes) - .await - .unwrap(); - - let original_table_route = table_metadata_manager - .table_route_manager() - .get(table_id) - .await - .unwrap() - .unwrap(); - - // modifies the table route. - table_metadata_manager - .update_leader_region_status(table_id, &original_table_route, |route| { - if route.region.id == RegionId::new(1024, 2) { - Some(Some(RegionStatus::Downgraded)) - } else { - None - } - }) - .await - .unwrap(); - - // sets the old table route. - ctx.volatile_ctx.table_route_info = Some(original_table_route); - - let err = state.downgrade_leader_region(&mut ctx).await.unwrap_err(); - - assert_matches!(err, Error::RetryLater { .. }); - - assert!(err.is_retryable()); - assert!(err.to_string().contains("Failed to update the table route")); - } - - #[tokio::test] - async fn test_next_downgrade_leader_region_state() { - let mut state = Box::new(UpdateMetadata::Downgrade); - let persistent_context = new_persistent_context(); - let from_peer = persistent_context.from_peer.clone(); - - let env = TestingEnv::new(); - let mut ctx = env.context_factory().new_context(persistent_context); - let table_id = ctx.region_id().table_id(); - - let table_info = new_test_table_info(1024, vec![1, 2]).into(); - let region_routes = vec![RegionRoute { - region: Region::new_test(RegionId::new(1024, 1)), - leader_peer: Some(from_peer.clone()), - ..Default::default() - }]; - - let table_metadata_manager = env.table_metadata_manager(); - table_metadata_manager - .create_table_metadata(table_info, region_routes) - .await - .unwrap(); - - let next = state.next(&mut ctx).await.unwrap(); - - let _ = next - .as_any() - .downcast_ref::() - .unwrap(); - - let latest_table_route = table_metadata_manager - .table_route_manager() - .get(table_id) - .await - .unwrap() - .unwrap(); - - assert!(latest_table_route.region_routes[0].is_leader_downgraded()); - assert!(ctx.volatile_ctx.table_route_info.is_none()); - } -} diff --git a/src/meta-srv/src/procedure/region_migration/update_metadata/downgrade_leader_region.rs b/src/meta-srv/src/procedure/region_migration/update_metadata/downgrade_leader_region.rs new file mode 100644 index 000000000000..7ff5e59942b7 --- /dev/null +++ b/src/meta-srv/src/procedure/region_migration/update_metadata/downgrade_leader_region.rs @@ -0,0 +1,210 @@ +// 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. + +use common_meta::rpc::router::RegionStatus; +use snafu::ResultExt; + +use crate::error::{self, Result}; +use crate::procedure::region_migration::update_metadata::UpdateMetadata; +use crate::procedure::region_migration::Context; + +impl UpdateMetadata { + /// Downgrades the leader region. + /// + /// Abort(non-retry): + /// - TableRoute is not found. + /// + /// Retry: + /// - Failed to update [TableRouteValue](common_meta::key::table_region::TableRegionValue). + /// - Failed to retrieve the metadata of table. + /// + /// About the failure of updating the [TableRouteValue](common_meta::key::table_region::TableRegionValue): + /// + /// - There may be another [RegionMigrationProcedure](crate::procedure::region_migration::RegionMigrationProcedure) + /// that is executed concurrently for **other region**. + /// It will only update **other region** info. Therefore, It's safe to retry after failure. + /// + /// - There is no other DDL procedure executed concurrently for the current table. + pub async fn downgrade_leader_region(&self, ctx: &mut Context) -> Result<()> { + let table_metadata_manager = ctx.table_metadata_manager.clone(); + let region_id = ctx.region_id(); + let table_id = region_id.table_id(); + let current_table_route_value = ctx.get_table_route_value().await?; + + if let Err(err) = table_metadata_manager + .update_leader_region_status(table_id, current_table_route_value, |route| { + if route.region.id == region_id { + Some(Some(RegionStatus::Downgraded)) + } else { + None + } + }) + .await + .context(error::TableMetadataManagerSnafu) + { + debug_assert!(ctx.remove_table_route_value()); + return error::RetryLaterSnafu { + reason: format!("Failed to update the table route during the downgrading leader region, error: {err}") + }.fail(); + } + + debug_assert!(ctx.remove_table_route_value()); + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::assert_matches::assert_matches; + + use common_meta::key::test_utils::new_test_table_info; + use common_meta::peer::Peer; + use common_meta::rpc::router::{Region, RegionRoute, RegionStatus}; + use store_api::storage::RegionId; + + use crate::error::Error; + use crate::procedure::region_migration::downgrade_leader_region::DowngradeLeaderRegion; + use crate::procedure::region_migration::test_util::{self, TestingEnv}; + use crate::procedure::region_migration::update_metadata::UpdateMetadata; + use crate::procedure::region_migration::{ContextFactory, PersistentContext, State}; + + fn new_persistent_context() -> PersistentContext { + test_util::new_persistent_context(1, 2, RegionId::new(1024, 1)) + } + + #[test] + fn test_state_serialization() { + let state = UpdateMetadata::Downgrade; + let expected = r#"{"UpdateMetadata":"Downgrade"}"#; + assert_eq!(expected, serde_json::to_string(&state).unwrap()); + } + + #[tokio::test] + async fn test_table_route_is_not_found_error() { + let state = UpdateMetadata::Downgrade; + let env = TestingEnv::new(); + let persistent_context = new_persistent_context(); + let mut ctx = env.context_factory().new_context(persistent_context); + + let err = state.downgrade_leader_region(&mut ctx).await.unwrap_err(); + + assert_matches!(err, Error::TableRouteNotFound { .. }); + + assert!(!err.is_retryable()); + } + + #[tokio::test] + async fn test_failed_to_update_table_route_error() { + let state = UpdateMetadata::Downgrade; + let persistent_context = new_persistent_context(); + let from_peer = persistent_context.from_peer.clone(); + + let env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let table_id = ctx.region_id().table_id(); + + let table_info = new_test_table_info(1024, vec![1, 2]).into(); + let region_routes = vec![ + RegionRoute { + region: Region::new_test(RegionId::new(1024, 1)), + leader_peer: Some(from_peer.clone()), + ..Default::default() + }, + RegionRoute { + region: Region::new_test(RegionId::new(1024, 2)), + leader_peer: Some(Peer::empty(4)), + ..Default::default() + }, + ]; + + let table_metadata_manager = env.table_metadata_manager(); + table_metadata_manager + .create_table_metadata(table_info, region_routes) + .await + .unwrap(); + + let original_table_route = table_metadata_manager + .table_route_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + + // modifies the table route. + table_metadata_manager + .update_leader_region_status(table_id, &original_table_route, |route| { + if route.region.id == RegionId::new(1024, 2) { + Some(Some(RegionStatus::Downgraded)) + } else { + None + } + }) + .await + .unwrap(); + + // sets the old table route. + ctx.volatile_ctx.table_route = Some(original_table_route); + + let err = state.downgrade_leader_region(&mut ctx).await.unwrap_err(); + + assert!(ctx.volatile_ctx.table_route.is_none()); + + assert_matches!(err, Error::RetryLater { .. }); + + assert!(err.is_retryable()); + assert!(err.to_string().contains("Failed to update the table route")); + } + + #[tokio::test] + async fn test_next_downgrade_leader_region_state() { + let mut state = Box::new(UpdateMetadata::Downgrade); + let persistent_context = new_persistent_context(); + let from_peer = persistent_context.from_peer.clone(); + + let env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let table_id = ctx.region_id().table_id(); + + let table_info = new_test_table_info(1024, vec![1, 2]).into(); + let region_routes = vec![RegionRoute { + region: Region::new_test(RegionId::new(1024, 1)), + leader_peer: Some(from_peer.clone()), + ..Default::default() + }]; + + let table_metadata_manager = env.table_metadata_manager(); + table_metadata_manager + .create_table_metadata(table_info, region_routes) + .await + .unwrap(); + + let next = state.next(&mut ctx).await.unwrap(); + + let _ = next + .as_any() + .downcast_ref::() + .unwrap(); + + let latest_table_route = table_metadata_manager + .table_route_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + + assert!(latest_table_route.region_routes[0].is_leader_downgraded()); + assert!(ctx.volatile_ctx.table_route.is_none()); + } +} diff --git a/src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs new file mode 100644 index 000000000000..22c732815ea1 --- /dev/null +++ b/src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs @@ -0,0 +1,376 @@ +// 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. + +use std::collections::HashMap; + +use common_meta::ddl::utils::region_storage_path; +use common_meta::key::datanode_table::RegionInfo; +use common_meta::rpc::router::RegionRoute; +use common_telemetry::{info, warn}; +use snafu::{ensure, OptionExt, ResultExt}; + +use crate::error::{self, Result}; +use crate::procedure::region_migration::update_metadata::UpdateMetadata; +use crate::procedure::region_migration::Context; + +impl UpdateMetadata { + /// Returns new [Vec]. + async fn build_upgrade_candidate_region_metadata( + &self, + ctx: &mut Context, + ) -> Result> { + let region_id = ctx.region_id(); + let table_route_value = ctx.get_table_route_value().await?.clone(); + + let mut region_routes = table_route_value.region_routes.clone(); + let region_route = region_routes + .iter_mut() + .find(|route| route.region.id == region_id) + .context(error::RegionRouteNotFoundSnafu { region_id })?; + + // Removes downgraded status. + region_route.set_leader_status(None); + + let candidate = &ctx.persistent_ctx.to_peer; + let expected_old_leader = &ctx.persistent_ctx.from_peer; + + // Upgrades candidate to leader. + ensure!(region_route + .leader_peer + .take_if(|old_leader| old_leader.id == expected_old_leader.id) + .is_some(), + error::UnexpectedSnafu{ + violated: format!("Unexpected region leader: {:?} during the upgrading candidate metadata, expected: {:?}", region_route.leader_peer, expected_old_leader), + } + ); + + region_route.leader_peer = Some(candidate.clone()); + info!( + "Upgrading candidate region to leader region: {:?} for region: {}", + candidate, region_id + ); + + // Removes the candidate region in followers. + let removed = region_route + .follower_peers + .extract_if(|peer| peer.id == candidate.id) + .collect::>(); + + if removed.len() > 1 { + warn!( + "Removes duplicated regions: {removed:?} during the upgrading candidate metadata for region: {region_id}" + ); + } + + Ok(region_routes) + } + + /// Upgrades the candidate region. + /// + /// Abort(non-retry): + /// - TableRoute or RegionRoute is not found. + /// Typically, it's impossible, there is no other DDL procedure executed concurrently for the current table. + /// + /// Retry: + /// - Failed to update [TableRouteValue](common_meta::key::table_region::TableRegionValue). + /// - Failed to retrieve the metadata of table. + pub async fn upgrade_candidate_region(&self, ctx: &mut Context) -> Result<()> { + let region_id = ctx.region_id(); + let table_metadata_manager = ctx.table_metadata_manager.clone(); + + let region_routes = self.build_upgrade_candidate_region_metadata(ctx).await?; + let table_info_value = ctx.get_table_info_value().await?; + + let table_info = &table_info_value.table_info; + let region_storage_path = + region_storage_path(&table_info.catalog_name, &table_info.schema_name); + let engine = table_info.meta.engine.clone(); + let region_options: HashMap = (&table_info.meta.options).into(); + + // No remote fetch. + let table_route_value = ctx.get_table_route_value().await?; + + if let Err(err) = table_metadata_manager + .update_table_route( + region_id.table_id(), + RegionInfo { + engine: engine.to_string(), + region_storage_path: region_storage_path.to_string(), + region_options: region_options.clone(), + }, + table_route_value, + region_routes, + ®ion_options, + ) + .await + .context(error::TableMetadataManagerSnafu) + { + debug_assert!(ctx.remove_table_route_value()); + return error::RetryLaterSnafu { + reason: format!("Failed to update the table route during the upgrading candidate region, error: {err}") + }.fail(); + }; + + debug_assert!(ctx.remove_table_route_value()); + // Consumes the guard. + ctx.volatile_ctx.opening_region_guard.take(); + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::assert_matches::assert_matches; + + use common_meta::key::test_utils::new_test_table_info; + use common_meta::peer::Peer; + use common_meta::rpc::router::{Region, RegionRoute, RegionStatus}; + use store_api::storage::RegionId; + + use crate::error::Error; + use crate::procedure::region_migration::migration_end::RegionMigrationEnd; + use crate::procedure::region_migration::test_util::{self, TestingEnv}; + use crate::procedure::region_migration::update_metadata::UpdateMetadata; + use crate::procedure::region_migration::{ContextFactory, PersistentContext, State}; + use crate::region::lease_keeper::OpeningRegionKeeper; + + fn new_persistent_context() -> PersistentContext { + test_util::new_persistent_context(1, 2, RegionId::new(1024, 1)) + } + + #[tokio::test] + async fn test_table_route_is_not_found_error() { + let state = UpdateMetadata::Upgrade; + + let env = TestingEnv::new(); + let persistent_context = new_persistent_context(); + let mut ctx = env.context_factory().new_context(persistent_context); + + let err = state + .build_upgrade_candidate_region_metadata(&mut ctx) + .await + .unwrap_err(); + + assert_matches!(err, Error::TableRouteNotFound { .. }); + assert!(!err.is_retryable()); + } + + #[tokio::test] + async fn test_region_route_is_not_found() { + let state = UpdateMetadata::Upgrade; + let env = TestingEnv::new(); + let persistent_context = new_persistent_context(); + let mut ctx = env.context_factory().new_context(persistent_context); + + let table_info = new_test_table_info(1024, vec![2]).into(); + let region_routes = vec![RegionRoute { + region: Region::new_test(RegionId::new(1024, 2)), + leader_peer: Some(Peer::empty(4)), + ..Default::default() + }]; + + let table_metadata_manager = env.table_metadata_manager(); + table_metadata_manager + .create_table_metadata(table_info, region_routes) + .await + .unwrap(); + + let err = state + .build_upgrade_candidate_region_metadata(&mut ctx) + .await + .unwrap_err(); + + assert_matches!(err, Error::RegionRouteNotFound { .. }); + assert!(!err.is_retryable()); + } + + #[tokio::test] + async fn test_region_route_expected_leader() { + let state = UpdateMetadata::Upgrade; + let env = TestingEnv::new(); + let persistent_context = new_persistent_context(); + let mut ctx = env.context_factory().new_context(persistent_context); + + let table_info = new_test_table_info(1024, vec![1]).into(); + let region_routes = vec![RegionRoute { + region: Region::new_test(RegionId::new(1024, 1)), + leader_peer: Some(Peer::empty(3)), + ..Default::default() + }]; + + let table_metadata_manager = env.table_metadata_manager(); + table_metadata_manager + .create_table_metadata(table_info, region_routes) + .await + .unwrap(); + + let err = state + .build_upgrade_candidate_region_metadata(&mut ctx) + .await + .unwrap_err(); + + assert_matches!(err, Error::Unexpected { .. }); + assert!(!err.is_retryable()); + assert!(err.to_string().contains("Unexpected region leader")); + } + + #[tokio::test] + async fn test_build_upgrade_candidate_region_metadata() { + let state = UpdateMetadata::Upgrade; + let env = TestingEnv::new(); + let persistent_context = new_persistent_context(); + let mut ctx = env.context_factory().new_context(persistent_context); + + let table_info = new_test_table_info(1024, vec![1]).into(); + let region_routes = vec![RegionRoute { + region: Region::new_test(RegionId::new(1024, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![Peer::empty(2), Peer::empty(3)], + leader_status: Some(RegionStatus::Downgraded), + }]; + + let table_metadata_manager = env.table_metadata_manager(); + table_metadata_manager + .create_table_metadata(table_info, region_routes) + .await + .unwrap(); + + let new_region_routes = state + .build_upgrade_candidate_region_metadata(&mut ctx) + .await + .unwrap(); + + assert!(!new_region_routes[0].is_leader_downgraded()); + assert_eq!(new_region_routes[0].follower_peers, vec![Peer::empty(3)]); + assert_eq!(new_region_routes[0].leader_peer.as_ref().unwrap().id, 2); + } + + #[tokio::test] + async fn test_failed_to_update_table_route_error() { + let state = UpdateMetadata::Upgrade; + let env = TestingEnv::new(); + let persistent_context = new_persistent_context(); + let mut ctx = env.context_factory().new_context(persistent_context); + let opening_keeper = OpeningRegionKeeper::default(); + + let table_id = 1024; + let table_info = new_test_table_info(table_id, vec![1]).into(); + let region_routes = vec![ + RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![Peer::empty(5), Peer::empty(3)], + leader_status: Some(RegionStatus::Downgraded), + }, + RegionRoute { + region: Region::new_test(RegionId::new(table_id, 2)), + leader_peer: Some(Peer::empty(4)), + leader_status: Some(RegionStatus::Downgraded), + ..Default::default() + }, + ]; + + let table_metadata_manager = env.table_metadata_manager(); + table_metadata_manager + .create_table_metadata(table_info, region_routes) + .await + .unwrap(); + + let original_table_route = table_metadata_manager + .table_route_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + + // modifies the table route. + table_metadata_manager + .update_leader_region_status(table_id, &original_table_route, |route| { + if route.region.id == RegionId::new(1024, 2) { + // Removes the status. + Some(None) + } else { + None + } + }) + .await + .unwrap(); + + // sets the old table route. + ctx.volatile_ctx.table_route = Some(original_table_route); + let guard = opening_keeper + .register(2, RegionId::new(table_id, 1)) + .unwrap(); + ctx.volatile_ctx.opening_region_guard = Some(guard); + + let err = state.upgrade_candidate_region(&mut ctx).await.unwrap_err(); + + assert!(ctx.volatile_ctx.table_route.is_none()); + assert!(ctx.volatile_ctx.opening_region_guard.is_some()); + assert_matches!(err, Error::RetryLater { .. }); + + assert!(err.is_retryable()); + assert!(err.to_string().contains("Failed to update the table route")); + } + + #[tokio::test] + async fn test_next_migration_end_state() { + let mut state = Box::new(UpdateMetadata::Upgrade); + let env = TestingEnv::new(); + let persistent_context = new_persistent_context(); + let mut ctx = env.context_factory().new_context(persistent_context); + let opening_keeper = OpeningRegionKeeper::default(); + + let table_id = 1024; + let table_info = new_test_table_info(table_id, vec![1]).into(); + let region_routes = vec![RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::empty(1)), + leader_status: Some(RegionStatus::Downgraded), + ..Default::default() + }]; + + let guard = opening_keeper + .register(2, RegionId::new(table_id, 1)) + .unwrap(); + ctx.volatile_ctx.opening_region_guard = Some(guard); + + let table_metadata_manager = env.table_metadata_manager(); + table_metadata_manager + .create_table_metadata(table_info, region_routes) + .await + .unwrap(); + + let next = state.next(&mut ctx).await.unwrap(); + + let _ = next.as_any().downcast_ref::().unwrap(); + + let region_routes = table_metadata_manager + .table_route_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .region_routes; + + assert!(ctx.volatile_ctx.table_route.is_none()); + assert!(ctx.volatile_ctx.opening_region_guard.is_none()); + assert_eq!(region_routes.len(), 1); + assert!(!region_routes[0].is_leader_downgraded()); + assert!(region_routes[0].follower_peers.is_empty()); + assert_eq!(region_routes[0].leader_peer.as_ref().unwrap().id, 2); + } +} From 9ccd182109ebb96001a528bc19f180ac1e1f2a07 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Thu, 30 Nov 2023 11:17:57 +0800 Subject: [PATCH 10/17] feat: implement PromQL set op AND/UNLESS (#2839) * initial impl Signed-off-by: Ruihang Xia * disable OR for now Signed-off-by: Ruihang Xia --------- Signed-off-by: Ruihang Xia --- src/promql/src/error.rs | 23 +- src/promql/src/planner.rs | 143 +++++++++- .../common/promql/set_operation.result | 270 ++++++++++++++++++ .../common/promql/set_operation.sql | 175 ++++++++++++ 4 files changed, 601 insertions(+), 10 deletions(-) create mode 100644 tests/cases/standalone/common/promql/set_operation.result create mode 100644 tests/cases/standalone/common/promql/set_operation.sql diff --git a/src/promql/src/error.rs b/src/promql/src/error.rs index 31c44e5715e8..4e9b3bb3be6a 100644 --- a/src/promql/src/error.rs +++ b/src/promql/src/error.rs @@ -18,7 +18,7 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; use datafusion::error::DataFusionError; -use promql_parser::parser::{Expr as PromExpr, TokenType}; +use promql_parser::parser::{Expr as PromExpr, TokenType, VectorMatchCardinality}; use snafu::{Location, Snafu}; #[derive(Snafu)] @@ -28,6 +28,12 @@ pub enum Error { #[snafu(display("Unsupported expr type: {}", name))] UnsupportedExpr { name: String, location: Location }, + #[snafu(display("Unsupported vector matches: {:?}", name))] + UnsupportedVectorMatch { + name: VectorMatchCardinality, + location: Location, + }, + #[snafu(display("Unexpected token: {:?}", token))] UnexpectedToken { token: TokenType, @@ -112,6 +118,17 @@ pub enum Error { #[snafu(display("Invalid function argument for {}", fn_name))] FunctionInvalidArgument { fn_name: String, location: Location }, + + #[snafu(display( + "Attempt to combine two tables with different column sets, left: {:?}, right: {:?}", + left, + right + ))] + CombineTableColumnMismatch { + left: Vec, + right: Vec, + location: Location, + }, } impl ErrorExt for Error { @@ -128,7 +145,9 @@ impl ErrorExt for Error { | ZeroRangeSelector { .. } | ColumnNotFound { .. } | Deserialize { .. } - | FunctionInvalidArgument { .. } => StatusCode::InvalidArguments, + | FunctionInvalidArgument { .. } + | UnsupportedVectorMatch { .. } + | CombineTableColumnMismatch { .. } => StatusCode::InvalidArguments, UnknownTable { .. } | DataFusionPlanning { .. } diff --git a/src/promql/src/planner.rs b/src/promql/src/planner.rs index dcadf8c4ebaf..7dfa17a17b3b 100644 --- a/src/promql/src/planner.rs +++ b/src/promql/src/planner.rs @@ -35,19 +35,20 @@ use datafusion::sql::TableReference; use datatypes::arrow::datatypes::DataType as ArrowDataType; use promql_parser::label::{MatchOp, Matcher, Matchers, METRIC_NAME}; use promql_parser::parser::{ - token, AggregateExpr, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr, Function, - LabelModifier, MatrixSelector, NumberLiteral, Offset, ParenExpr, StringLiteral, SubqueryExpr, - TokenType, UnaryExpr, VectorSelector, + token, AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, + Expr as PromExpr, Function, LabelModifier, MatrixSelector, NumberLiteral, Offset, ParenExpr, + StringLiteral, SubqueryExpr, TokenType, UnaryExpr, VectorMatchCardinality, VectorSelector, }; use snafu::{ensure, OptionExt, ResultExt}; use table::table::adapter::DfTableProviderAdapter; use crate::error::{ - CatalogSnafu, ColumnNotFoundSnafu, DataFusionPlanningSnafu, ExpectExprSnafu, - ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, MultipleMetricMatchersSnafu, - MultipleVectorSnafu, NoMetricMatcherSnafu, Result, TableNameNotFoundSnafu, - TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, - UnsupportedExprSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, + CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu, + ExpectExprSnafu, ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, + MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, Result, + TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, + UnknownTableSnafu, UnsupportedExprSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu, + ZeroRangeSelectorSnafu, }; use crate::extension_plan::{ build_special_time_expr, EmptyMetric, HistogramFold, InstantManipulate, Millisecond, @@ -268,14 +269,29 @@ impl PromPlanner { let left_field_columns = self.ctx.field_columns.clone(); let left_table_ref: OwnedTableReference = self.ctx.table_name.clone().unwrap_or_default().into(); + let left_tag_cols = self.ctx.tag_columns.clone(); let right_input = self.prom_expr_to_plan(*rhs.clone()).await?; let right_field_columns = self.ctx.field_columns.clone(); let right_table_ref: OwnedTableReference = self.ctx.table_name.clone().unwrap_or_default().into(); + let right_tag_cols = self.ctx.tag_columns.clone(); // TODO(ruihang): avoid join if left and right are the same table + // set op has "special" join semantics + if Self::is_token_a_set_op(*op) { + return self.set_op_on_non_field_columns( + left_input, + right_input, + left_tag_cols, + right_tag_cols, + *op, + modifier, + ); + } + + // normal join let mut field_columns = left_field_columns.iter().zip(right_field_columns.iter()); let join_plan = self.join_on_non_field_columns( @@ -1310,6 +1326,16 @@ impl PromPlanner { ) } + /// Check if the given op is a set operator (UNION, INTERSECT and EXCEPT in SQL). + fn is_token_a_set_op(token: TokenType) -> bool { + matches!( + token.id(), + token::T_LAND // INTERSECT + | token::T_LOR // UNION + | token::T_LUNLESS // EXCEPT + ) + } + /// Build a inner join on time index column and tag columns to concat two logical plans. fn join_on_non_field_columns( &self, @@ -1351,6 +1377,107 @@ impl PromPlanner { .context(DataFusionPlanningSnafu) } + fn set_op_on_non_field_columns( + &self, + left: LogicalPlan, + right: LogicalPlan, + left_tag_cols: Vec, + right_tag_cols: Vec, + op: TokenType, + modifier: &Option, + ) -> Result { + let mut left_tag_col_set = left_tag_cols.into_iter().collect::>(); + let mut right_tag_col_set = right_tag_cols.into_iter().collect::>(); + + // apply modifier + if let Some(modifier) = modifier { + // one-to-many and many-to-one are not supported + ensure!( + matches!( + modifier.card, + VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany + ), + UnsupportedVectorMatchSnafu { + name: modifier.card.clone(), + }, + ); + // apply label modifier + if let Some(matching) = &modifier.matching { + match matching { + // keeps columns mentioned in `on` + LabelModifier::Include(on) => { + let mask = on.labels.iter().cloned().collect::>(); + left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect(); + right_tag_col_set = + right_tag_col_set.intersection(&mask).cloned().collect(); + } + // removes columns memtioned in `ignoring` + LabelModifier::Exclude(ignoring) => { + // doesn't check existence of label + for label in &ignoring.labels { + let _ = left_tag_col_set.remove(label); + let _ = right_tag_col_set.remove(label); + } + } + } + } + } + // ensure two sides have the same tag columns + if !matches!(op.id(), token::T_LOR) { + ensure!( + left_tag_col_set == right_tag_col_set, + CombineTableColumnMismatchSnafu { + left: left_tag_col_set.into_iter().collect::>(), + right: right_tag_col_set.into_iter().collect::>(), + } + ) + }; + let join_keys = left_tag_col_set + .into_iter() + .chain([self.ctx.time_index_column.clone().unwrap()]) + .collect::>(); + + // Generate join plan. + // All set operations in PromQL are "distinct" + match op.id() { + token::T_LAND => LogicalPlanBuilder::from(left) + .distinct() + .context(DataFusionPlanningSnafu)? + .join_detailed( + right, + JoinType::LeftSemi, + (join_keys.clone(), join_keys), + None, + true, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu), + token::T_LUNLESS => LogicalPlanBuilder::from(left) + .distinct() + .context(DataFusionPlanningSnafu)? + .join_detailed( + right, + JoinType::LeftAnti, + (join_keys.clone(), join_keys), + None, + true, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu), + token::T_LOR => { + // `OR` can not be expressed by `UNION` precisely. + // it will generate unexpceted result when schemas don't match + UnsupportedExprSnafu { + name: "set operation `OR`", + } + .fail() + } + _ => UnexpectedTokenSnafu { token: op }.fail(), + } + } + /// Build a projection that project and perform operation expr for every value columns. /// Non-value columns (tag and timestamp) will be preserved in the projection. /// diff --git a/tests/cases/standalone/common/promql/set_operation.result b/tests/cases/standalone/common/promql/set_operation.result new file mode 100644 index 000000000000..2e0e2b5f5c97 --- /dev/null +++ b/tests/cases/standalone/common/promql/set_operation.result @@ -0,0 +1,270 @@ +-- from promql/testdata/operators.test +-- cases related to AND/OR/UNLESS +-- group_left() and group_right() are not included +create table http_requests ( + ts timestamp time index, + job string, + instance string, + g string, -- for `group` + val double, + primary key (job, instance, g) +); + +Affected Rows: 0 + +insert into http_requests values + (3000000, "api", "0", "production", 100), + (3000000, "api", "1", "production", 200), + (3000000, "api", "0", "canary", 300), + (3000000, "api", "1", "canary", 400), + (3000000, "app", "0", "production", 500), + (3000000, "app", "1", "production", 600), + (3000000, "app", "0", "canary", 700), + (3000000, "app", "1", "canary", 800); + +Affected Rows: 8 + +-- empty metric +create table cpu_count(ts timestamp time index); + +Affected Rows: 0 + +create table vector_matching_a( + ts timestamp time index, + l string primary key, + val double, +); + +Affected Rows: 0 + +insert into vector_matching_a values + (3000000, "x", 10), + (3000000, "y", 20); + +Affected Rows: 2 + +-- eval instant at 50m http_requests{group="canary"} and http_requests{instance="0"} +-- http_requests{group="canary", instance="0", job="api-server"} 300 +-- http_requests{group="canary", instance="0", job="app-server"} 700 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') http_requests{g="canary"} and http_requests{instance="0"}; + ++---------------------+-----+----------+--------+-------+ +| ts | job | instance | g | val | ++---------------------+-----+----------+--------+-------+ +| 1970-01-01T00:50:00 | api | 0 | canary | 300.0 | +| 1970-01-01T00:50:00 | app | 0 | canary | 700.0 | ++---------------------+-----+----------+--------+-------+ + +-- eval instant at 50m (http_requests{group="canary"} + 1) and http_requests{instance="0"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and http_requests{instance="0"}; + ++-----+----------+--------+---------------------+------------------+ +| job | instance | g | ts | val + Float64(1) | ++-----+----------+--------+---------------------+------------------+ +| api | 0 | canary | 1970-01-01T00:50:00 | 301.0 | +| app | 0 | canary | 1970-01-01T00:50:00 | 701.0 | ++-----+----------+--------+---------------------+------------------+ + +-- eval instant at 50m (http_requests{group="canary"} + 1) and on(instance, job) http_requests{instance="0", group="production"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and on(instance, job) http_requests{instance="0", g="production"}; + ++-----+----------+--------+---------------------+------------------+ +| job | instance | g | ts | val + Float64(1) | ++-----+----------+--------+---------------------+------------------+ +| api | 0 | canary | 1970-01-01T00:50:00 | 301.0 | +| app | 0 | canary | 1970-01-01T00:50:00 | 701.0 | ++-----+----------+--------+---------------------+------------------+ + +-- eval instant at 50m (http_requests{group="canary"} + 1) and on(instance) http_requests{instance="0", group="production"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and on(instance) http_requests{instance="0", g="production"}; + ++-----+----------+--------+---------------------+------------------+ +| job | instance | g | ts | val + Float64(1) | ++-----+----------+--------+---------------------+------------------+ +| api | 0 | canary | 1970-01-01T00:50:00 | 301.0 | +| app | 0 | canary | 1970-01-01T00:50:00 | 701.0 | ++-----+----------+--------+---------------------+------------------+ + +-- eval instant at 50m (http_requests{group="canary"} + 1) and ignoring(group) http_requests{instance="0", group="production"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and ignoring(g) http_requests{instance="0", g="production"}; + ++-----+----------+--------+---------------------+------------------+ +| job | instance | g | ts | val + Float64(1) | ++-----+----------+--------+---------------------+------------------+ +| api | 0 | canary | 1970-01-01T00:50:00 | 301.0 | +| app | 0 | canary | 1970-01-01T00:50:00 | 701.0 | ++-----+----------+--------+---------------------+------------------+ + +-- eval instant at 50m (http_requests{group="canary"} + 1) and ignoring(group, job) http_requests{instance="0", group="production"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and ignoring(g, job) http_requests{instance="0", g="production"}; + ++-----+----------+--------+---------------------+------------------+ +| job | instance | g | ts | val + Float64(1) | ++-----+----------+--------+---------------------+------------------+ +| api | 0 | canary | 1970-01-01T00:50:00 | 301.0 | +| app | 0 | canary | 1970-01-01T00:50:00 | 701.0 | ++-----+----------+--------+---------------------+------------------+ + +-- eval instant at 50m http_requests{group="canary"} or http_requests{group="production"} +-- http_requests{group="canary", instance="0", job="api-server"} 300 +-- http_requests{group="canary", instance="0", job="app-server"} 700 +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- http_requests{group="production", instance="0", job="api-server"} 100 +-- http_requests{group="production", instance="0", job="app-server"} 500 +-- http_requests{group="production", instance="1", job="api-server"} 200 +-- http_requests{group="production", instance="1", job="app-server"} 600 +-- NOT SUPPORTED: `or` +tql eval (3000, 3000, '1s') http_requests{g="canary"} or http_requests{g="production"}; + +Error: 1004(InvalidArguments), Unsupported expr type: set operation `OR` + +-- # On overlap the rhs samples must be dropped. +-- eval instant at 50m (http_requests{group="canary"} + 1) or http_requests{instance="1"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- {group="canary", instance="1", job="api-server"} 401 +-- {group="canary", instance="1", job="app-server"} 801 +-- http_requests{group="production", instance="1", job="api-server"} 200 +-- http_requests{group="production", instance="1", job="app-server"} 600 +-- NOT SUPPORTED: `or` +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) or http_requests{instance="1"}; + +Error: 1004(InvalidArguments), Unsupported expr type: set operation `OR` + +-- # Matching only on instance excludes everything that has instance=0/1 but includes +-- # entries without the instance label. +-- eval instant at 50m (http_requests{group="canary"} + 1) or on(instance) (http_requests or cpu_count or vector_matching_a) +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- {group="canary", instance="1", job="api-server"} 401 +-- {group="canary", instance="1", job="app-server"} 801 +-- vector_matching_a{l="x"} 10 +-- vector_matching_a{l="y"} 20 +-- NOT SUPPORTED: union on different schemas +-- NOT SUPPORTED: `or` +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) or on(instance) (http_requests or cpu_count or vector_matching_a); + +Error: 1004(InvalidArguments), Unsupported expr type: set operation `OR` + +-- eval instant at 50m (http_requests{group="canary"} + 1) or ignoring(l, group, job) (http_requests or cpu_count or vector_matching_a) +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- {group="canary", instance="1", job="api-server"} 401 +-- {group="canary", instance="1", job="app-server"} 801 +-- vector_matching_a{l="x"} 10 +-- vector_matching_a{l="y"} 20 +-- NOT SUPPORTED: union on different schemas +-- NOT SUPPORTED: `or` +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) or ignoring(l, g, job) (http_requests or cpu_count or vector_matching_a); + +Error: 1004(InvalidArguments), Unsupported expr type: set operation `OR` + +-- eval instant at 50m http_requests{group="canary"} unless http_requests{instance="0"} +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless http_requests{instance="0"}; + ++---------------------+-----+----------+--------+-------+ +| ts | job | instance | g | val | ++---------------------+-----+----------+--------+-------+ +| 1970-01-01T00:50:00 | api | 1 | canary | 400.0 | +| 1970-01-01T00:50:00 | app | 1 | canary | 800.0 | ++---------------------+-----+----------+--------+-------+ + +-- eval instant at 50m http_requests{group="canary"} unless on(job) http_requests{instance="0"} +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless on(job) http_requests{instance="0"}; + +++ +++ + +-- eval instant at 50m http_requests{group="canary"} unless on(job, instance) http_requests{instance="0"} +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless on(job, instance) http_requests{instance="0"}; + ++---------------------+-----+----------+--------+-------+ +| ts | job | instance | g | val | ++---------------------+-----+----------+--------+-------+ +| 1970-01-01T00:50:00 | api | 1 | canary | 400.0 | +| 1970-01-01T00:50:00 | app | 1 | canary | 800.0 | ++---------------------+-----+----------+--------+-------+ + +-- eval instant at 50m http_requests{group="canary"} unless ignoring(group, instance) http_requests{instance="0"} +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless ignoring(g, instance) http_requests{instance="0"}; + +++ +++ + +-- eval instant at 50m http_requests{group="canary"} unless ignoring(group) http_requests{instance="0"} +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless ignoring(g) http_requests{instance="0"}; + ++---------------------+-----+----------+--------+-------+ +| ts | job | instance | g | val | ++---------------------+-----+----------+--------+-------+ +| 1970-01-01T00:50:00 | api | 1 | canary | 400.0 | +| 1970-01-01T00:50:00 | app | 1 | canary | 800.0 | ++---------------------+-----+----------+--------+-------+ + +-- # https://github.com/prometheus/prometheus/issues/1489 +-- eval instant at 50m http_requests AND ON (dummy) vector(1) +-- http_requests{group="canary", instance="0", job="api-server"} 300 +-- http_requests{group="canary", instance="0", job="app-server"} 700 +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- http_requests{group="production", instance="0", job="api-server"} 100 +-- http_requests{group="production", instance="0", job="app-server"} 500 +-- http_requests{group="production", instance="1", job="api-server"} 200 +-- http_requests{group="production", instance="1", job="app-server"} 600 +-- NOT SUPPORTED: `vector()` +tql eval (3000, 3000, '1s') http_requests AND ON (dummy) vector(1); + +Error: 1004(InvalidArguments), Expect a PromQL expr but not found, input expr: Call(Call { func: Function { name: "vector", arg_types: [Scalar], variadic: false, return_type: Vector }, args: FunctionArgs { args: [NumberLiteral(NumberLiteral { val: 1.0 })] } }) + +-- eval instant at 50m http_requests AND IGNORING (group, instance, job) vector(1) +-- http_requests{group="canary", instance="0", job="api-server"} 300 +-- http_requests{group="canary", instance="0", job="app-server"} 700 +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- http_requests{group="production", instance="0", job="api-server"} 100 +-- http_requests{group="production", instance="0", job="app-server"} 500 +-- http_requests{group="production", instance="1", job="api-server"} 200 +-- http_requests{group="production", instance="1", job="app-server"} 600 +-- NOT SUPPORTED: `vector()` +tql eval (3000, 3000, '1s') http_requests AND IGNORING (g, instance, job) vector(1); + +Error: 1004(InvalidArguments), Expect a PromQL expr but not found, input expr: Call(Call { func: Function { name: "vector", arg_types: [Scalar], variadic: false, return_type: Vector }, args: FunctionArgs { args: [NumberLiteral(NumberLiteral { val: 1.0 })] } }) + +drop table http_requests; + +Affected Rows: 0 + +drop table cpu_count; + +Affected Rows: 0 + +drop table vector_matching_a; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/set_operation.sql b/tests/cases/standalone/common/promql/set_operation.sql new file mode 100644 index 000000000000..e91460df3478 --- /dev/null +++ b/tests/cases/standalone/common/promql/set_operation.sql @@ -0,0 +1,175 @@ +-- from promql/testdata/operators.test +-- cases related to AND/OR/UNLESS +-- group_left() and group_right() are not included + +create table http_requests ( + ts timestamp time index, + job string, + instance string, + g string, -- for `group` + val double, + primary key (job, instance, g) +); + +insert into http_requests values + (3000000, "api", "0", "production", 100), + (3000000, "api", "1", "production", 200), + (3000000, "api", "0", "canary", 300), + (3000000, "api", "1", "canary", 400), + (3000000, "app", "0", "production", 500), + (3000000, "app", "1", "production", 600), + (3000000, "app", "0", "canary", 700), + (3000000, "app", "1", "canary", 800); + +-- empty metric +create table cpu_count(ts timestamp time index); + +create table vector_matching_a( + ts timestamp time index, + l string primary key, + val double, +); + +insert into vector_matching_a values + (3000000, "x", 10), + (3000000, "y", 20); + +-- eval instant at 50m http_requests{group="canary"} and http_requests{instance="0"} +-- http_requests{group="canary", instance="0", job="api-server"} 300 +-- http_requests{group="canary", instance="0", job="app-server"} 700 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') http_requests{g="canary"} and http_requests{instance="0"}; + +-- eval instant at 50m (http_requests{group="canary"} + 1) and http_requests{instance="0"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and http_requests{instance="0"}; + +-- eval instant at 50m (http_requests{group="canary"} + 1) and on(instance, job) http_requests{instance="0", group="production"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and on(instance, job) http_requests{instance="0", g="production"}; + +-- eval instant at 50m (http_requests{group="canary"} + 1) and on(instance) http_requests{instance="0", group="production"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and on(instance) http_requests{instance="0", g="production"}; + +-- eval instant at 50m (http_requests{group="canary"} + 1) and ignoring(group) http_requests{instance="0", group="production"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and ignoring(g) http_requests{instance="0", g="production"}; + +-- eval instant at 50m (http_requests{group="canary"} + 1) and ignoring(group, job) http_requests{instance="0", group="production"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) and ignoring(g, job) http_requests{instance="0", g="production"}; + +-- eval instant at 50m http_requests{group="canary"} or http_requests{group="production"} +-- http_requests{group="canary", instance="0", job="api-server"} 300 +-- http_requests{group="canary", instance="0", job="app-server"} 700 +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- http_requests{group="production", instance="0", job="api-server"} 100 +-- http_requests{group="production", instance="0", job="app-server"} 500 +-- http_requests{group="production", instance="1", job="api-server"} 200 +-- http_requests{group="production", instance="1", job="app-server"} 600 +-- NOT SUPPORTED: `or` +tql eval (3000, 3000, '1s') http_requests{g="canary"} or http_requests{g="production"}; + +-- # On overlap the rhs samples must be dropped. +-- eval instant at 50m (http_requests{group="canary"} + 1) or http_requests{instance="1"} +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- {group="canary", instance="1", job="api-server"} 401 +-- {group="canary", instance="1", job="app-server"} 801 +-- http_requests{group="production", instance="1", job="api-server"} 200 +-- http_requests{group="production", instance="1", job="app-server"} 600 +-- NOT SUPPORTED: `or` +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) or http_requests{instance="1"}; + + +-- # Matching only on instance excludes everything that has instance=0/1 but includes +-- # entries without the instance label. +-- eval instant at 50m (http_requests{group="canary"} + 1) or on(instance) (http_requests or cpu_count or vector_matching_a) +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- {group="canary", instance="1", job="api-server"} 401 +-- {group="canary", instance="1", job="app-server"} 801 +-- vector_matching_a{l="x"} 10 +-- vector_matching_a{l="y"} 20 +-- NOT SUPPORTED: union on different schemas +-- NOT SUPPORTED: `or` +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) or on(instance) (http_requests or cpu_count or vector_matching_a); + +-- eval instant at 50m (http_requests{group="canary"} + 1) or ignoring(l, group, job) (http_requests or cpu_count or vector_matching_a) +-- {group="canary", instance="0", job="api-server"} 301 +-- {group="canary", instance="0", job="app-server"} 701 +-- {group="canary", instance="1", job="api-server"} 401 +-- {group="canary", instance="1", job="app-server"} 801 +-- vector_matching_a{l="x"} 10 +-- vector_matching_a{l="y"} 20 +-- NOT SUPPORTED: union on different schemas +-- NOT SUPPORTED: `or` +tql eval (3000, 3000, '1s') (http_requests{g="canary"} + 1) or ignoring(l, g, job) (http_requests or cpu_count or vector_matching_a); + +-- eval instant at 50m http_requests{group="canary"} unless http_requests{instance="0"} +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless http_requests{instance="0"}; + +-- eval instant at 50m http_requests{group="canary"} unless on(job) http_requests{instance="0"} +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless on(job) http_requests{instance="0"}; + +-- eval instant at 50m http_requests{group="canary"} unless on(job, instance) http_requests{instance="0"} +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless on(job, instance) http_requests{instance="0"}; + +-- eval instant at 50m http_requests{group="canary"} unless ignoring(group, instance) http_requests{instance="0"} +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless ignoring(g, instance) http_requests{instance="0"}; + +-- eval instant at 50m http_requests{group="canary"} unless ignoring(group) http_requests{instance="0"} +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') http_requests{g="canary"} unless ignoring(g) http_requests{instance="0"}; + + +-- # https://github.com/prometheus/prometheus/issues/1489 +-- eval instant at 50m http_requests AND ON (dummy) vector(1) +-- http_requests{group="canary", instance="0", job="api-server"} 300 +-- http_requests{group="canary", instance="0", job="app-server"} 700 +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- http_requests{group="production", instance="0", job="api-server"} 100 +-- http_requests{group="production", instance="0", job="app-server"} 500 +-- http_requests{group="production", instance="1", job="api-server"} 200 +-- http_requests{group="production", instance="1", job="app-server"} 600 +-- NOT SUPPORTED: `vector()` +tql eval (3000, 3000, '1s') http_requests AND ON (dummy) vector(1); + +-- eval instant at 50m http_requests AND IGNORING (group, instance, job) vector(1) +-- http_requests{group="canary", instance="0", job="api-server"} 300 +-- http_requests{group="canary", instance="0", job="app-server"} 700 +-- http_requests{group="canary", instance="1", job="api-server"} 400 +-- http_requests{group="canary", instance="1", job="app-server"} 800 +-- http_requests{group="production", instance="0", job="api-server"} 100 +-- http_requests{group="production", instance="0", job="app-server"} 500 +-- http_requests{group="production", instance="1", job="api-server"} 200 +-- http_requests{group="production", instance="1", job="app-server"} 600 +-- NOT SUPPORTED: `vector()` +tql eval (3000, 3000, '1s') http_requests AND IGNORING (g, instance, job) vector(1); + +drop table http_requests; + +drop table cpu_count; + +drop table vector_matching_a; From 2332305b904a79dcfac70bb74e258c12acb614b9 Mon Sep 17 00:00:00 2001 From: Wei <15172118655@163.com> Date: Thu, 30 Nov 2023 11:27:29 +0800 Subject: [PATCH 11/17] refactor: replace usage of ArrayData by clone (#2827) * refactor: use array clone() * refactor: slice * chore: clippy --- src/datatypes/src/vectors.rs | 7 +- src/datatypes/src/vectors/binary.rs | 12 +- src/datatypes/src/vectors/boolean.rs | 23 +-- src/datatypes/src/vectors/decimal.rs | 39 ++-- src/datatypes/src/vectors/helper.rs | 64 +++---- src/datatypes/src/vectors/list.rs | 23 +-- src/datatypes/src/vectors/primitive.rs | 256 ++++--------------------- src/datatypes/src/vectors/string.rs | 22 +-- 8 files changed, 112 insertions(+), 334 deletions(-) diff --git a/src/datatypes/src/vectors.rs b/src/datatypes/src/vectors.rs index 0fc260970a0a..a1dcafac82ec 100644 --- a/src/datatypes/src/vectors.rs +++ b/src/datatypes/src/vectors.rs @@ -229,17 +229,16 @@ macro_rules! impl_try_from_arrow_array_for_vector { ) -> crate::error::Result<$Vector> { use snafu::OptionExt; - let data = array + let arrow_array = array .as_ref() .as_any() .downcast_ref::<$Array>() .with_context(|| crate::error::ConversionSnafu { from: std::format!("{:?}", array.as_ref().data_type()), })? - .to_data(); + .clone(); - let concrete_array = $Array::from(data); - Ok($Vector::from(concrete_array)) + Ok($Vector::from(arrow_array)) } } }; diff --git a/src/datatypes/src/vectors/binary.rs b/src/datatypes/src/vectors/binary.rs index 57388bf32018..894dbb62cb4e 100644 --- a/src/datatypes/src/vectors/binary.rs +++ b/src/datatypes/src/vectors/binary.rs @@ -15,7 +15,7 @@ use std::any::Any; use std::sync::Arc; -use arrow::array::{Array, ArrayBuilder, ArrayData, ArrayIter, ArrayRef}; +use arrow::array::{Array, ArrayBuilder, ArrayIter, ArrayRef}; use snafu::ResultExt; use crate::arrow_array::{BinaryArray, MutableBinaryArray}; @@ -36,10 +36,6 @@ impl BinaryVector { pub(crate) fn as_arrow(&self) -> &dyn Array { &self.array } - - fn to_array_data(&self) -> ArrayData { - self.array.to_data() - } } impl From for BinaryVector { @@ -74,13 +70,11 @@ impl Vector for BinaryVector { } fn to_arrow_array(&self) -> ArrayRef { - let data = self.to_array_data(); - Arc::new(BinaryArray::from(data)) + Arc::new(self.array.clone()) } fn to_boxed_arrow_array(&self) -> Box { - let data = self.to_array_data(); - Box::new(BinaryArray::from(data)) + Box::new(self.array.clone()) } fn validity(&self) -> Validity { diff --git a/src/datatypes/src/vectors/boolean.rs b/src/datatypes/src/vectors/boolean.rs index 61dd86a434bf..a87bc853e6ec 100644 --- a/src/datatypes/src/vectors/boolean.rs +++ b/src/datatypes/src/vectors/boolean.rs @@ -16,9 +16,7 @@ use std::any::Any; use std::borrow::Borrow; use std::sync::Arc; -use arrow::array::{ - Array, ArrayBuilder, ArrayData, ArrayIter, ArrayRef, BooleanArray, BooleanBuilder, -}; +use arrow::array::{Array, ArrayBuilder, ArrayIter, ArrayRef, BooleanArray, BooleanBuilder}; use snafu::ResultExt; use crate::data_type::ConcreteDataType; @@ -44,16 +42,6 @@ impl BooleanVector { &self.array } - fn to_array_data(&self) -> ArrayData { - self.array.to_data() - } - - fn from_array_data(data: ArrayData) -> BooleanVector { - BooleanVector { - array: BooleanArray::from(data), - } - } - pub(crate) fn false_count(&self) -> usize { self.array.false_count() } @@ -107,13 +95,11 @@ impl Vector for BooleanVector { } fn to_arrow_array(&self) -> ArrayRef { - let data = self.to_array_data(); - Arc::new(BooleanArray::from(data)) + Arc::new(self.array.clone()) } fn to_boxed_arrow_array(&self) -> Box { - let data = self.to_array_data(); - Box::new(BooleanArray::from(data)) + Box::new(self.array.clone()) } fn validity(&self) -> Validity { @@ -133,8 +119,7 @@ impl Vector for BooleanVector { } fn slice(&self, offset: usize, length: usize) -> VectorRef { - let data = self.array.to_data().slice(offset, length); - Arc::new(Self::from_array_data(data)) + Arc::new(Self::from(self.array.slice(offset, length))) } fn get(&self, index: usize) -> Value { diff --git a/src/datatypes/src/vectors/decimal.rs b/src/datatypes/src/vectors/decimal.rs index ebcdb43d6e34..fab5e74cce75 100644 --- a/src/datatypes/src/vectors/decimal.rs +++ b/src/datatypes/src/vectors/decimal.rs @@ -196,8 +196,7 @@ impl Vector for Decimal128Vector { } fn slice(&self, offset: usize, length: usize) -> VectorRef { - let array = self.array.slice(offset, length); - Arc::new(Self { array }) + Arc::new(self.get_slice(offset, length)) } fn get(&self, index: usize) -> Value { @@ -535,23 +534,23 @@ pub mod tests { // because 100 is out of Decimal(3, 1) range, so it will be null assert!(array.is_null(4)); } -} -#[test] -fn test_decimal28_vector_iter_data() { - let vector = Decimal128Vector::from_values(vec![1, 2, 3, 4]) - .with_precision_and_scale(3, 1) - .unwrap(); - let mut iter = vector.iter_data(); - assert_eq!(iter.next(), Some(Some(Decimal128::new(1, 3, 1)))); - assert_eq!(iter.next(), Some(Some(Decimal128::new(2, 3, 1)))); - assert_eq!(iter.next(), Some(Some(Decimal128::new(3, 3, 1)))); - assert_eq!(iter.next(), Some(Some(Decimal128::new(4, 3, 1)))); - assert_eq!(iter.next(), None); - - let values = vector - .iter_data() - .filter_map(|v| v.map(|x| x.val() * 2)) - .collect::>(); - assert_eq!(values, vec![2, 4, 6, 8]); + #[test] + fn test_decimal28_vector_iter_data() { + let vector = Decimal128Vector::from_values(vec![1, 2, 3, 4]) + .with_precision_and_scale(3, 1) + .unwrap(); + let mut iter = vector.iter_data(); + assert_eq!(iter.next(), Some(Some(Decimal128::new(1, 3, 1)))); + assert_eq!(iter.next(), Some(Some(Decimal128::new(2, 3, 1)))); + assert_eq!(iter.next(), Some(Some(Decimal128::new(3, 3, 1)))); + assert_eq!(iter.next(), Some(Some(Decimal128::new(4, 3, 1)))); + assert_eq!(iter.next(), None); + + let values = vector + .iter_data() + .filter_map(|v| v.map(|x| x.val() * 2)) + .collect::>(); + assert_eq!(values, vec![2, 4, 6, 8]); + } } diff --git a/src/datatypes/src/vectors/helper.rs b/src/datatypes/src/vectors/helper.rs index f37048838cc8..1ec9aec45cc5 100644 --- a/src/datatypes/src/vectors/helper.rs +++ b/src/datatypes/src/vectors/helper.rs @@ -284,23 +284,21 @@ impl Helper { ArrowDataType::Date64 => Arc::new(DateTimeVector::try_from_arrow_array(array)?), ArrowDataType::List(_) => Arc::new(ListVector::try_from_arrow_array(array)?), ArrowDataType::Timestamp(unit, _) => match unit { - TimeUnit::Second => Arc::new( - TimestampSecondVector::try_from_arrow_timestamp_array(array)?, - ), - TimeUnit::Millisecond => Arc::new( - TimestampMillisecondVector::try_from_arrow_timestamp_array(array)?, - ), - TimeUnit::Microsecond => Arc::new( - TimestampMicrosecondVector::try_from_arrow_timestamp_array(array)?, - ), - TimeUnit::Nanosecond => Arc::new( - TimestampNanosecondVector::try_from_arrow_timestamp_array(array)?, - ), + TimeUnit::Second => Arc::new(TimestampSecondVector::try_from_arrow_array(array)?), + TimeUnit::Millisecond => { + Arc::new(TimestampMillisecondVector::try_from_arrow_array(array)?) + } + TimeUnit::Microsecond => { + Arc::new(TimestampMicrosecondVector::try_from_arrow_array(array)?) + } + TimeUnit::Nanosecond => { + Arc::new(TimestampNanosecondVector::try_from_arrow_array(array)?) + } }, ArrowDataType::Time32(unit) => match unit { - TimeUnit::Second => Arc::new(TimeSecondVector::try_from_arrow_time_array(array)?), + TimeUnit::Second => Arc::new(TimeSecondVector::try_from_arrow_array(array)?), TimeUnit::Millisecond => { - Arc::new(TimeMillisecondVector::try_from_arrow_time_array(array)?) + Arc::new(TimeMillisecondVector::try_from_arrow_array(array)?) } // Arrow use time32 for second/millisecond. _ => unreachable!( @@ -310,10 +308,10 @@ impl Helper { }, ArrowDataType::Time64(unit) => match unit { TimeUnit::Microsecond => { - Arc::new(TimeMicrosecondVector::try_from_arrow_time_array(array)?) + Arc::new(TimeMicrosecondVector::try_from_arrow_array(array)?) } TimeUnit::Nanosecond => { - Arc::new(TimeNanosecondVector::try_from_arrow_time_array(array)?) + Arc::new(TimeNanosecondVector::try_from_arrow_array(array)?) } // Arrow use time64 for microsecond/nanosecond. _ => unreachable!( @@ -322,29 +320,27 @@ impl Helper { ), }, ArrowDataType::Interval(unit) => match unit { - IntervalUnit::YearMonth => Arc::new( - IntervalYearMonthVector::try_from_arrow_interval_array(array)?, - ), + IntervalUnit::YearMonth => { + Arc::new(IntervalYearMonthVector::try_from_arrow_array(array)?) + } IntervalUnit::DayTime => { - Arc::new(IntervalDayTimeVector::try_from_arrow_interval_array(array)?) + Arc::new(IntervalDayTimeVector::try_from_arrow_array(array)?) + } + IntervalUnit::MonthDayNano => { + Arc::new(IntervalMonthDayNanoVector::try_from_arrow_array(array)?) } - IntervalUnit::MonthDayNano => Arc::new( - IntervalMonthDayNanoVector::try_from_arrow_interval_array(array)?, - ), }, ArrowDataType::Duration(unit) => match unit { - TimeUnit::Second => { - Arc::new(DurationSecondVector::try_from_arrow_duration_array(array)?) + TimeUnit::Second => Arc::new(DurationSecondVector::try_from_arrow_array(array)?), + TimeUnit::Millisecond => { + Arc::new(DurationMillisecondVector::try_from_arrow_array(array)?) + } + TimeUnit::Microsecond => { + Arc::new(DurationMicrosecondVector::try_from_arrow_array(array)?) + } + TimeUnit::Nanosecond => { + Arc::new(DurationNanosecondVector::try_from_arrow_array(array)?) } - TimeUnit::Millisecond => Arc::new( - DurationMillisecondVector::try_from_arrow_duration_array(array)?, - ), - TimeUnit::Microsecond => Arc::new( - DurationMicrosecondVector::try_from_arrow_duration_array(array)?, - ), - TimeUnit::Nanosecond => Arc::new( - DurationNanosecondVector::try_from_arrow_duration_array(array)?, - ), }, ArrowDataType::Decimal128(_, _) => { Arc::new(Decimal128Vector::try_from_arrow_array(array)?) diff --git a/src/datatypes/src/vectors/list.rs b/src/datatypes/src/vectors/list.rs index 21a658f7986a..00ca4398fef9 100644 --- a/src/datatypes/src/vectors/list.rs +++ b/src/datatypes/src/vectors/list.rs @@ -46,17 +46,6 @@ impl ListVector { .map(|value_opt| value_opt.map(Helper::try_into_vector).transpose()) } - fn to_array_data(&self) -> ArrayData { - self.array.to_data() - } - - fn from_array_data_and_type(data: ArrayData, item_type: ConcreteDataType) -> Self { - Self { - array: ListArray::from(data), - item_type, - } - } - pub(crate) fn as_arrow(&self) -> &dyn Array { &self.array } @@ -80,13 +69,11 @@ impl Vector for ListVector { } fn to_arrow_array(&self) -> ArrayRef { - let data = self.to_array_data(); - Arc::new(ListArray::from(data)) + Arc::new(self.array.clone()) } fn to_boxed_arrow_array(&self) -> Box { - let data = self.to_array_data(); - Box::new(ListArray::from(data)) + Box::new(self.array.clone()) } fn validity(&self) -> Validity { @@ -106,8 +93,10 @@ impl Vector for ListVector { } fn slice(&self, offset: usize, length: usize) -> VectorRef { - let data = self.array.to_data().slice(offset, length); - Arc::new(Self::from_array_data_and_type(data, self.item_type.clone())) + Arc::new(Self { + array: self.array.slice(offset, length), + item_type: self.item_type.clone(), + }) } fn get(&self, index: usize) -> Value { diff --git a/src/datatypes/src/vectors/primitive.rs b/src/datatypes/src/vectors/primitive.rs index d570b67c6678..61743b800ccd 100644 --- a/src/datatypes/src/vectors/primitive.rs +++ b/src/datatypes/src/vectors/primitive.rs @@ -16,23 +16,12 @@ use std::any::Any; use std::fmt; use std::sync::Arc; -use arrow::array::{ - Array, ArrayBuilder, ArrayData, ArrayIter, ArrayRef, PrimitiveArray, PrimitiveBuilder, - Time32MillisecondArray as TimeMillisecondArray, Time32SecondArray as TimeSecondArray, - Time64MicrosecondArray as TimeMicrosecondArray, Time64NanosecondArray as TimeNanosecondArray, - TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, - TimestampSecondArray, -}; -use arrow_array::{ - DurationMicrosecondArray, DurationMillisecondArray, DurationNanosecondArray, - DurationSecondArray, IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, -}; -use arrow_schema::DataType; +use arrow::array::{Array, ArrayBuilder, ArrayIter, ArrayRef, PrimitiveArray, PrimitiveBuilder}; use serde_json::Value as JsonValue; use snafu::OptionExt; use crate::data_type::ConcreteDataType; -use crate::error::{self, CastTypeSnafu, Result}; +use crate::error::{self, Result}; use crate::scalars::{Scalar, ScalarRef, ScalarVector, ScalarVectorBuilder}; use crate::serialize::Serializable; use crate::types::{ @@ -66,178 +55,15 @@ impl PrimitiveVector { } pub fn try_from_arrow_array(array: impl AsRef) -> Result { - let data = array + let arrow_array = array .as_ref() .as_any() .downcast_ref::>() .with_context(|| error::ConversionSnafu { from: format!("{:?}", array.as_ref().data_type()), - })? - .to_data(); - let concrete_array = PrimitiveArray::::from(data); - Ok(Self::new(concrete_array)) - } - - /// Converts arrow timestamp array to vectors, ignoring time zone info. - pub fn try_from_arrow_timestamp_array(array: impl AsRef) -> Result { - let array = array.as_ref(); - let array_data = match array.data_type() { - DataType::Timestamp(unit, _) => match unit { - arrow_schema::TimeUnit::Second => array - .as_any() - .downcast_ref::() - .unwrap() - .clone() - .with_timezone_opt(None::) - .to_data(), - arrow_schema::TimeUnit::Millisecond => array - .as_any() - .downcast_ref::() - .unwrap() - .clone() - .with_timezone_opt(None::) - .to_data(), - arrow_schema::TimeUnit::Microsecond => array - .as_any() - .downcast_ref::() - .unwrap() - .clone() - .with_timezone_opt(None::) - .to_data(), - arrow_schema::TimeUnit::Nanosecond => array - .as_any() - .downcast_ref::() - .unwrap() - .clone() - .with_timezone_opt(None::) - .to_data(), - }, - arrow_type => { - return CastTypeSnafu { - msg: format!( - "Failed to cast arrow array {:?} to timestamp vector", - arrow_type, - ), - } - .fail()?; - } - }; - let concrete_array = PrimitiveArray::::from(array_data); - Ok(Self::new(concrete_array)) - } - - /// Converts arrow time array to vectors - pub fn try_from_arrow_time_array(array: impl AsRef) -> Result { - let array = array.as_ref(); - let array_data = match array.data_type() { - DataType::Time32(unit) => match unit { - arrow_schema::TimeUnit::Second => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - arrow_schema::TimeUnit::Millisecond => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - _ => unreachable!(), - }, - DataType::Time64(unit) => match unit { - arrow_schema::TimeUnit::Microsecond => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - arrow_schema::TimeUnit::Nanosecond => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - _ => unreachable!(), - }, - arrow_type => { - return CastTypeSnafu { - msg: format!("Failed to cast arrow array {:?} to time vector", arrow_type,), - } - .fail()?; - } - }; - let concrete_array = PrimitiveArray::::from(array_data); - Ok(Self::new(concrete_array)) - } - - pub fn try_from_arrow_interval_array(array: impl AsRef) -> Result { - let array = array.as_ref(); - let array_data = match array.data_type() { - DataType::Interval(unit) => match unit { - arrow_schema::IntervalUnit::YearMonth => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - arrow_schema::IntervalUnit::DayTime => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - arrow_schema::IntervalUnit::MonthDayNano => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - }, - arrow_type => { - return CastTypeSnafu { - msg: format!( - "Failed to cast arrow array {:?} to interval vector", - arrow_type, - ), - } - .fail()?; - } - }; - let concrete_array = PrimitiveArray::::from(array_data); - Ok(Self::new(concrete_array)) - } - - pub fn try_from_arrow_duration_array(array: impl AsRef) -> Result { - let array = array.as_ref(); - let array_data = match array.data_type() { - DataType::Duration(unit) => match unit { - arrow_schema::TimeUnit::Second => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - arrow_schema::TimeUnit::Millisecond => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - arrow_schema::TimeUnit::Microsecond => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - arrow_schema::TimeUnit::Nanosecond => array - .as_any() - .downcast_ref::() - .unwrap() - .to_data(), - }, - arrow_type => { - return CastTypeSnafu { - msg: format!( - "Failed to cast arrow array {:?} to interval vector", - arrow_type, - ), - } - .fail()?; - } - }; - let concrete_array = PrimitiveArray::::from(array_data); - Ok(Self::new(concrete_array)) + })?; + + Ok(Self::new(arrow_array.clone())) } pub fn from_slice>(slice: P) -> Self { @@ -277,24 +103,15 @@ impl PrimitiveVector { &self.array } - fn to_array_data(&self) -> ArrayData { - self.array.to_data() - } - - fn from_array_data(data: ArrayData) -> Self { - Self { - array: PrimitiveArray::from(data), - } - } - // To distinguish with `Vector::slice()`. /// Slice the vector, returning a new vector. /// /// # Panics /// This function panics if `offset + length > self.len()`. pub fn get_slice(&self, offset: usize, length: usize) -> Self { - let data = self.array.to_data().slice(offset, length); - Self::from_array_data(data) + Self { + array: self.array.slice(offset, length), + } } } @@ -316,13 +133,11 @@ impl Vector for PrimitiveVector { } fn to_arrow_array(&self) -> ArrayRef { - let data = self.to_array_data(); - Arc::new(PrimitiveArray::::from(data)) + Arc::new(self.array.clone()) } fn to_boxed_arrow_array(&self) -> Box { - let data = self.to_array_data(); - Box::new(PrimitiveArray::::from(data)) + Box::new(self.array.clone()) } fn validity(&self) -> Validity { @@ -580,7 +395,12 @@ mod tests { Time64NanosecondArray, }; use arrow::datatypes::DataType as ArrowDataType; - use arrow_array::{DurationSecondArray, IntervalDayTimeArray, IntervalYearMonthArray}; + use arrow_array::{ + DurationMicrosecondArray, DurationMillisecondArray, DurationNanosecondArray, + DurationSecondArray, IntervalDayTimeArray, IntervalYearMonthArray, + TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, + TimestampSecondArray, + }; use serde_json; use super::*; @@ -703,6 +523,14 @@ mod tests { assert_eq!(128, v.memory_size()); } + #[test] + fn test_get_slice() { + let v = Int32Vector::from_slice(vec![1, 2, 3, 4, 5]); + let slice = v.get_slice(1, 3); + assert_eq!(v, Int32Vector::from_slice(vec![1, 2, 3, 4, 5])); + assert_eq!(slice, Int32Vector::from_slice(vec![2, 3, 4])); + } + #[test] fn test_primitive_vector_builder() { let mut builder = Int64Type::default().create_mutable_vector(3); @@ -748,48 +576,48 @@ mod tests { #[test] fn test_try_from_arrow_time_array() { let array: ArrayRef = Arc::new(Time32SecondArray::from(vec![1i32, 2, 3])); - let vector = TimeSecondVector::try_from_arrow_time_array(array).unwrap(); + let vector = TimeSecondVector::try_from_arrow_array(array).unwrap(); assert_eq!(TimeSecondVector::from_values(vec![1, 2, 3]), vector); let array: ArrayRef = Arc::new(Time32MillisecondArray::from(vec![1i32, 2, 3])); - let vector = TimeMillisecondVector::try_from_arrow_time_array(array).unwrap(); + let vector = TimeMillisecondVector::try_from_arrow_array(array).unwrap(); assert_eq!(TimeMillisecondVector::from_values(vec![1, 2, 3]), vector); let array: ArrayRef = Arc::new(Time64MicrosecondArray::from(vec![1i64, 2, 3])); - let vector = TimeMicrosecondVector::try_from_arrow_time_array(array).unwrap(); + let vector = TimeMicrosecondVector::try_from_arrow_array(array).unwrap(); assert_eq!(TimeMicrosecondVector::from_values(vec![1, 2, 3]), vector); let array: ArrayRef = Arc::new(Time64NanosecondArray::from(vec![1i64, 2, 3])); - let vector = TimeNanosecondVector::try_from_arrow_time_array(array).unwrap(); + let vector = TimeNanosecondVector::try_from_arrow_array(array).unwrap(); assert_eq!(TimeNanosecondVector::from_values(vec![1, 2, 3]), vector); // Test convert error let array: ArrayRef = Arc::new(Int32Array::from(vec![1i32, 2, 3])); - assert!(TimeSecondVector::try_from_arrow_time_array(array).is_err()); + assert!(TimeSecondVector::try_from_arrow_array(array).is_err()); } #[test] fn test_try_from_arrow_timestamp_array() { let array: ArrayRef = Arc::new(TimestampSecondArray::from(vec![1i64, 2, 3])); - let vector = TimestampSecondVector::try_from_arrow_timestamp_array(array).unwrap(); + let vector = TimestampSecondVector::try_from_arrow_array(array).unwrap(); assert_eq!(TimestampSecondVector::from_values(vec![1, 2, 3]), vector); let array: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![1i64, 2, 3])); - let vector = TimestampMillisecondVector::try_from_arrow_timestamp_array(array).unwrap(); + let vector = TimestampMillisecondVector::try_from_arrow_array(array).unwrap(); assert_eq!( TimestampMillisecondVector::from_values(vec![1, 2, 3]), vector ); let array: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![1i64, 2, 3])); - let vector = TimestampMicrosecondVector::try_from_arrow_timestamp_array(array).unwrap(); + let vector = TimestampMicrosecondVector::try_from_arrow_array(array).unwrap(); assert_eq!( TimestampMicrosecondVector::from_values(vec![1, 2, 3]), vector ); let array: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![1i64, 2, 3])); - let vector = TimestampNanosecondVector::try_from_arrow_timestamp_array(array).unwrap(); + let vector = TimestampNanosecondVector::try_from_arrow_array(array).unwrap(); assert_eq!( TimestampNanosecondVector::from_values(vec![1, 2, 3]), vector @@ -797,27 +625,27 @@ mod tests { // Test convert error let array: ArrayRef = Arc::new(Int32Array::from(vec![1i32, 2, 3])); - assert!(TimestampSecondVector::try_from_arrow_timestamp_array(array).is_err()); + assert!(TimestampSecondVector::try_from_arrow_array(array).is_err()); } #[test] fn test_try_from_arrow_interval_array() { let array: ArrayRef = Arc::new(IntervalYearMonthArray::from(vec![1000, 2000, 3000])); - let vector = IntervalYearMonthVector::try_from_arrow_interval_array(array).unwrap(); + let vector = IntervalYearMonthVector::try_from_arrow_array(array).unwrap(); assert_eq!( IntervalYearMonthVector::from_values(vec![1000, 2000, 3000]), vector ); let array: ArrayRef = Arc::new(IntervalDayTimeArray::from(vec![1000, 2000, 3000])); - let vector = IntervalDayTimeVector::try_from_arrow_interval_array(array).unwrap(); + let vector = IntervalDayTimeVector::try_from_arrow_array(array).unwrap(); assert_eq!( IntervalDayTimeVector::from_values(vec![1000, 2000, 3000]), vector ); let array: ArrayRef = Arc::new(IntervalYearMonthArray::from(vec![1000, 2000, 3000])); - let vector = IntervalYearMonthVector::try_from_arrow_interval_array(array).unwrap(); + let vector = IntervalYearMonthVector::try_from_arrow_array(array).unwrap(); assert_eq!( IntervalYearMonthVector::from_values(vec![1000, 2000, 3000]), vector @@ -827,28 +655,28 @@ mod tests { #[test] fn test_try_from_arrow_duration_array() { let array: ArrayRef = Arc::new(DurationSecondArray::from(vec![1000, 2000, 3000])); - let vector = DurationSecondVector::try_from_arrow_duration_array(array).unwrap(); + let vector = DurationSecondVector::try_from_arrow_array(array).unwrap(); assert_eq!( DurationSecondVector::from_values(vec![1000, 2000, 3000]), vector ); let array: ArrayRef = Arc::new(DurationMillisecondArray::from(vec![1000, 2000, 3000])); - let vector = DurationMillisecondVector::try_from_arrow_duration_array(array).unwrap(); + let vector = DurationMillisecondVector::try_from_arrow_array(array).unwrap(); assert_eq!( DurationMillisecondVector::from_values(vec![1000, 2000, 3000]), vector ); let array: ArrayRef = Arc::new(DurationMicrosecondArray::from(vec![1000, 2000, 3000])); - let vector = DurationMicrosecondVector::try_from_arrow_duration_array(array).unwrap(); + let vector = DurationMicrosecondVector::try_from_arrow_array(array).unwrap(); assert_eq!( DurationMicrosecondVector::from_values(vec![1000, 2000, 3000]), vector ); let array: ArrayRef = Arc::new(DurationNanosecondArray::from(vec![1000, 2000, 3000])); - let vector = DurationNanosecondVector::try_from_arrow_duration_array(array).unwrap(); + let vector = DurationNanosecondVector::try_from_arrow_array(array).unwrap(); assert_eq!( DurationNanosecondVector::from_values(vec![1000, 2000, 3000]), vector diff --git a/src/datatypes/src/vectors/string.rs b/src/datatypes/src/vectors/string.rs index b0be6ab01da2..bf49491eb1d2 100644 --- a/src/datatypes/src/vectors/string.rs +++ b/src/datatypes/src/vectors/string.rs @@ -15,7 +15,7 @@ use std::any::Any; use std::sync::Arc; -use arrow::array::{Array, ArrayBuilder, ArrayData, ArrayIter, ArrayRef}; +use arrow::array::{Array, ArrayBuilder, ArrayIter, ArrayRef}; use snafu::ResultExt; use crate::arrow_array::{MutableStringArray, StringArray}; @@ -36,16 +36,6 @@ impl StringVector { pub(crate) fn as_arrow(&self) -> &dyn Array { &self.array } - - fn to_array_data(&self) -> ArrayData { - self.array.to_data() - } - - fn from_array_data(data: ArrayData) -> Self { - Self { - array: StringArray::from(data), - } - } } impl From for StringVector { @@ -120,13 +110,11 @@ impl Vector for StringVector { } fn to_arrow_array(&self) -> ArrayRef { - let data = self.to_array_data(); - Arc::new(StringArray::from(data)) + Arc::new(self.array.clone()) } fn to_boxed_arrow_array(&self) -> Box { - let data = self.to_array_data(); - Box::new(StringArray::from(data)) + Box::new(self.array.clone()) } fn validity(&self) -> Validity { @@ -146,8 +134,7 @@ impl Vector for StringVector { } fn slice(&self, offset: usize, length: usize) -> VectorRef { - let data = self.array.to_data().slice(offset, length); - Arc::new(Self::from_array_data(data)) + Arc::new(Self::from(self.array.slice(offset, length))) } fn get(&self, index: usize) -> Value { @@ -256,6 +243,7 @@ vectors::impl_try_from_arrow_array_for_vector!(StringArray, StringVector); #[cfg(test)] mod tests { + use arrow::datatypes::DataType; use super::*; From fe2fc723bc19f37060b15625d363473f1f5bca10 Mon Sep 17 00:00:00 2001 From: Wei <15172118655@163.com> Date: Thu, 30 Nov 2023 11:49:09 +0800 Subject: [PATCH 12/17] refactor: DataType name function (#2836) * refactor: DataType name function * chore: test case --- src/catalog/src/information_schema/columns.rs | 2 +- src/datatypes/src/data_type.rs | 181 +++++++++--------- src/datatypes/src/types/binary_type.rs | 4 +- src/datatypes/src/types/boolean_type.rs | 4 +- src/datatypes/src/types/cast.rs | 2 +- src/datatypes/src/types/date_type.rs | 4 +- src/datatypes/src/types/datetime_type.rs | 4 +- src/datatypes/src/types/decimal_type.rs | 5 +- src/datatypes/src/types/dictionary_type.rs | 8 +- src/datatypes/src/types/duration_type.rs | 4 +- src/datatypes/src/types/interval_type.rs | 4 +- src/datatypes/src/types/list_type.rs | 6 +- src/datatypes/src/types/null_type.rs | 4 +- src/datatypes/src/types/primitive_type.rs | 12 +- src/datatypes/src/types/string_type.rs | 4 +- src/datatypes/src/types/time_type.rs | 4 +- src/datatypes/src/types/timestamp_type.rs | 4 +- src/query/src/sql.rs | 7 +- src/servers/src/http.rs | 2 +- src/store-api/src/metadata.rs | 2 +- 20 files changed, 132 insertions(+), 135 deletions(-) diff --git a/src/catalog/src/information_schema/columns.rs b/src/catalog/src/information_schema/columns.rs index 57880993acfe..34e9c7ef66c5 100644 --- a/src/catalog/src/information_schema/columns.rs +++ b/src/catalog/src/information_schema/columns.rs @@ -202,7 +202,7 @@ impl InformationSchemaColumnsBuilder { &schema_name, &table_name, &column.name, - column.data_type.name(), + &column.data_type.name(), semantic_type, ); } diff --git a/src/datatypes/src/data_type.rs b/src/datatypes/src/data_type.rs index 61470eec6bd7..6c1ffbcc47df 100644 --- a/src/datatypes/src/data_type.rs +++ b/src/datatypes/src/data_type.rs @@ -85,31 +85,48 @@ pub enum ConcreteDataType { impl fmt::Display for ConcreteDataType { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - ConcreteDataType::Null(_) => write!(f, "Null"), - ConcreteDataType::Boolean(_) => write!(f, "Boolean"), - ConcreteDataType::Int8(_) => write!(f, "Int8"), - ConcreteDataType::Int16(_) => write!(f, "Int16"), - ConcreteDataType::Int32(_) => write!(f, "Int32"), - ConcreteDataType::Int64(_) => write!(f, "Int64"), - ConcreteDataType::UInt8(_) => write!(f, "UInt8"), - ConcreteDataType::UInt16(_) => write!(f, "UInt16"), - ConcreteDataType::UInt32(_) => write!(f, "UInt32"), - ConcreteDataType::UInt64(_) => write!(f, "UInt64"), - ConcreteDataType::Float32(_) => write!(f, "Float32"), - ConcreteDataType::Float64(_) => write!(f, "Float64"), - ConcreteDataType::Binary(_) => write!(f, "Binary"), - ConcreteDataType::String(_) => write!(f, "String"), - ConcreteDataType::Date(_) => write!(f, "Date"), - ConcreteDataType::DateTime(_) => write!(f, "DateTime"), - ConcreteDataType::Timestamp(_) => write!(f, "Timestamp"), - ConcreteDataType::Time(_) => write!(f, "Time"), - ConcreteDataType::List(_) => write!(f, "List"), - ConcreteDataType::Dictionary(_) => write!(f, "Dictionary"), - ConcreteDataType::Interval(_) => write!(f, "Interval"), - ConcreteDataType::Duration(_) => write!(f, "Duration"), - ConcreteDataType::Decimal128(d) => { - write!(f, "Decimal128({},{})", d.precision(), d.scale()) - } + ConcreteDataType::Null(v) => write!(f, "{}", v.name()), + ConcreteDataType::Boolean(v) => write!(f, "{}", v.name()), + ConcreteDataType::Int8(v) => write!(f, "{}", v.name()), + ConcreteDataType::Int16(v) => write!(f, "{}", v.name()), + ConcreteDataType::Int32(v) => write!(f, "{}", v.name()), + ConcreteDataType::Int64(v) => write!(f, "{}", v.name()), + ConcreteDataType::UInt8(v) => write!(f, "{}", v.name()), + ConcreteDataType::UInt16(v) => write!(f, "{}", v.name()), + ConcreteDataType::UInt32(v) => write!(f, "{}", v.name()), + ConcreteDataType::UInt64(v) => write!(f, "{}", v.name()), + ConcreteDataType::Float32(v) => write!(f, "{}", v.name()), + ConcreteDataType::Float64(v) => write!(f, "{}", v.name()), + ConcreteDataType::Binary(v) => write!(f, "{}", v.name()), + ConcreteDataType::String(v) => write!(f, "{}", v.name()), + ConcreteDataType::Date(v) => write!(f, "{}", v.name()), + ConcreteDataType::DateTime(v) => write!(f, "{}", v.name()), + ConcreteDataType::Timestamp(t) => match t { + TimestampType::Second(v) => write!(f, "{}", v.name()), + TimestampType::Millisecond(v) => write!(f, "{}", v.name()), + TimestampType::Microsecond(v) => write!(f, "{}", v.name()), + TimestampType::Nanosecond(v) => write!(f, "{}", v.name()), + }, + ConcreteDataType::Time(t) => match t { + TimeType::Second(v) => write!(f, "{}", v.name()), + TimeType::Millisecond(v) => write!(f, "{}", v.name()), + TimeType::Microsecond(v) => write!(f, "{}", v.name()), + TimeType::Nanosecond(v) => write!(f, "{}", v.name()), + }, + ConcreteDataType::Interval(i) => match i { + IntervalType::YearMonth(v) => write!(f, "{}", v.name()), + IntervalType::DayTime(v) => write!(f, "{}", v.name()), + IntervalType::MonthDayNano(v) => write!(f, "{}", v.name()), + }, + ConcreteDataType::Duration(d) => match d { + DurationType::Second(v) => write!(f, "{}", v.name()), + DurationType::Millisecond(v) => write!(f, "{}", v.name()), + DurationType::Microsecond(v) => write!(f, "{}", v.name()), + DurationType::Nanosecond(v) => write!(f, "{}", v.name()), + }, + ConcreteDataType::Decimal128(v) => write!(f, "{}", v.name()), + ConcreteDataType::List(v) => write!(f, "{}", v.name()), + ConcreteDataType::Dictionary(v) => write!(f, "{}", v.name()), } } } @@ -492,7 +509,7 @@ impl ConcreteDataType { #[enum_dispatch::enum_dispatch] pub trait DataType: std::fmt::Debug + Send + Sync { /// Name of this data type. - fn name(&self) -> &str; + fn name(&self) -> String; /// Returns id of the Logical data type. fn logical_type_id(&self) -> LogicalTypeId; @@ -523,7 +540,7 @@ mod tests { fn test_concrete_type_as_datatype_trait() { let concrete_type = ConcreteDataType::boolean_datatype(); - assert_eq!("Boolean", concrete_type.name()); + assert_eq!("Boolean", concrete_type.to_string()); assert_eq!(Value::Boolean(false), concrete_type.default_value()); assert_eq!(LogicalTypeId::Boolean, concrete_type.logical_type_id()); assert_eq!(ArrowDataType::Boolean, concrete_type.as_arrow_type()); @@ -767,94 +784,68 @@ mod tests { #[test] fn test_display_concrete_data_type() { + assert_eq!(ConcreteDataType::null_datatype().to_string(), "Null"); + assert_eq!(ConcreteDataType::boolean_datatype().to_string(), "Boolean"); + assert_eq!(ConcreteDataType::binary_datatype().to_string(), "Binary"); + assert_eq!(ConcreteDataType::int8_datatype().to_string(), "Int8"); + assert_eq!(ConcreteDataType::int16_datatype().to_string(), "Int16"); + assert_eq!(ConcreteDataType::int32_datatype().to_string(), "Int32"); + assert_eq!(ConcreteDataType::int64_datatype().to_string(), "Int64"); + assert_eq!(ConcreteDataType::uint8_datatype().to_string(), "UInt8"); + assert_eq!(ConcreteDataType::uint16_datatype().to_string(), "UInt16"); + assert_eq!(ConcreteDataType::uint32_datatype().to_string(), "UInt32"); + assert_eq!(ConcreteDataType::uint64_datatype().to_string(), "UInt64"); + assert_eq!(ConcreteDataType::float32_datatype().to_string(), "Float32"); + assert_eq!(ConcreteDataType::float64_datatype().to_string(), "Float64"); + assert_eq!(ConcreteDataType::string_datatype().to_string(), "String"); + assert_eq!(ConcreteDataType::date_datatype().to_string(), "Date"); assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Null).to_string(), - "Null" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Boolean).to_string(), - "Boolean" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Binary).to_string(), - "Binary" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::LargeBinary).to_string(), - "Binary" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Int8).to_string(), - "Int8" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Int16).to_string(), - "Int16" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Int32).to_string(), - "Int32" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Int64).to_string(), - "Int64" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::UInt8).to_string(), - "UInt8" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::UInt16).to_string(), - "UInt16" + ConcreteDataType::timestamp_millisecond_datatype().to_string(), + "TimestampMillisecond" ); assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::UInt32).to_string(), - "UInt32" + ConcreteDataType::time_millisecond_datatype().to_string(), + "TimeMillisecond" ); assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::UInt64).to_string(), - "UInt64" + ConcreteDataType::interval_month_day_nano_datatype().to_string(), + "IntervalMonthDayNano" ); assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Float32).to_string(), - "Float32" + ConcreteDataType::duration_second_datatype().to_string(), + "DurationSecond" ); assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Float64).to_string(), - "Float64" + ConcreteDataType::decimal128_datatype(10, 2).to_string(), + "Decimal(10, 2)" ); + // Nested types assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Utf8).to_string(), - "String" + ConcreteDataType::list_datatype(ConcreteDataType::int32_datatype()).to_string(), + "List" ); assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::List(Arc::new(Field::new( - "item", - ArrowDataType::Int32, - true, - )))) + ConcreteDataType::list_datatype(ConcreteDataType::Dictionary(DictionaryType::new( + ConcreteDataType::int32_datatype(), + ConcreteDataType::string_datatype() + ))) .to_string(), - "List" - ); - assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Date32).to_string(), - "Date" + "List>" ); - assert_eq!(ConcreteDataType::time_second_datatype().to_string(), "Time"); assert_eq!( - ConcreteDataType::from_arrow_type(&ArrowDataType::Interval( - arrow_schema::IntervalUnit::MonthDayNano, + ConcreteDataType::list_datatype(ConcreteDataType::list_datatype( + ConcreteDataType::list_datatype(ConcreteDataType::int32_datatype()) )) .to_string(), - "Interval" - ); - assert_eq!( - ConcreteDataType::duration_second_datatype().to_string(), - "Duration" + "List>>" ); assert_eq!( - ConcreteDataType::decimal128_datatype(10, 2).to_string(), - "Decimal128(10,2)" + ConcreteDataType::dictionary_datatype( + ConcreteDataType::int32_datatype(), + ConcreteDataType::string_datatype() + ) + .to_string(), + "Dictionary" ); } } diff --git a/src/datatypes/src/types/binary_type.rs b/src/datatypes/src/types/binary_type.rs index 6f4c1f6bc0b5..7213489da102 100644 --- a/src/datatypes/src/types/binary_type.rs +++ b/src/datatypes/src/types/binary_type.rs @@ -34,8 +34,8 @@ impl BinaryType { } impl DataType for BinaryType { - fn name(&self) -> &str { - "Binary" + fn name(&self) -> String { + "Binary".to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/boolean_type.rs b/src/datatypes/src/types/boolean_type.rs index 1d4b9e80a2b9..d3e465a57c70 100644 --- a/src/datatypes/src/types/boolean_type.rs +++ b/src/datatypes/src/types/boolean_type.rs @@ -34,8 +34,8 @@ impl BooleanType { } impl DataType for BooleanType { - fn name(&self) -> &str { - "Boolean" + fn name(&self) -> String { + "Boolean".to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/cast.rs b/src/datatypes/src/types/cast.rs index 80ca310ac893..299d0d625066 100644 --- a/src/datatypes/src/types/cast.rs +++ b/src/datatypes/src/types/cast.rs @@ -248,7 +248,7 @@ mod tests { assert!(res.is_err()); assert_eq!( res.unwrap_err().to_string(), - "Type Timestamp with value 1970-01-01 08:00:10+0800 can't be cast to the destination type Int8" + "Type TimestampSecond with value 1970-01-01 08:00:10+0800 can't be cast to the destination type Int8" ); } diff --git a/src/datatypes/src/types/date_type.rs b/src/datatypes/src/types/date_type.rs index 8bbde3a7c7f0..fe32f00b26c8 100644 --- a/src/datatypes/src/types/date_type.rs +++ b/src/datatypes/src/types/date_type.rs @@ -32,8 +32,8 @@ use crate::vectors::{DateVector, DateVectorBuilder, MutableVector, Vector}; pub struct DateType; impl DataType for DateType { - fn name(&self) -> &str { - "Date" + fn name(&self) -> String { + "Date".to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/datetime_type.rs b/src/datatypes/src/types/datetime_type.rs index 17326ef9ae4a..abed366264fd 100644 --- a/src/datatypes/src/types/datetime_type.rs +++ b/src/datatypes/src/types/datetime_type.rs @@ -30,8 +30,8 @@ use crate::vectors::{DateTimeVector, DateTimeVectorBuilder, PrimitiveVector}; pub struct DateTimeType; impl DataType for DateTimeType { - fn name(&self) -> &str { - "DateTime" + fn name(&self) -> String { + "DateTime".to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/decimal_type.rs b/src/datatypes/src/types/decimal_type.rs index 48ede0c44136..e397ba197cf8 100644 --- a/src/datatypes/src/types/decimal_type.rs +++ b/src/datatypes/src/types/decimal_type.rs @@ -56,9 +56,8 @@ impl Decimal128Type { } impl DataType for Decimal128Type { - fn name(&self) -> &str { - // TODO(QuenKar): support precision and scale information in name - "decimal" + fn name(&self) -> String { + format!("Decimal({}, {})", self.precision, self.scale) } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/dictionary_type.rs b/src/datatypes/src/types/dictionary_type.rs index cc29c41403df..dcedb5d6d052 100644 --- a/src/datatypes/src/types/dictionary_type.rs +++ b/src/datatypes/src/types/dictionary_type.rs @@ -62,8 +62,12 @@ impl DictionaryType { } impl DataType for DictionaryType { - fn name(&self) -> &str { - "Dictionary" + fn name(&self) -> String { + format!( + "Dictionary<{}, {}>", + self.key_type.name(), + self.value_type.name() + ) } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/duration_type.rs b/src/datatypes/src/types/duration_type.rs index ffc8fe92467b..db668451d32c 100644 --- a/src/datatypes/src/types/duration_type.rs +++ b/src/datatypes/src/types/duration_type.rs @@ -78,8 +78,8 @@ macro_rules! impl_data_type_for_duration { pub struct []; impl DataType for [] { - fn name(&self) -> &str { - stringify!([]) + fn name(&self) -> String { + stringify!([]).to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/interval_type.rs b/src/datatypes/src/types/interval_type.rs index 1acc506cfce0..7ee796498211 100644 --- a/src/datatypes/src/types/interval_type.rs +++ b/src/datatypes/src/types/interval_type.rs @@ -66,8 +66,8 @@ macro_rules! impl_data_type_for_interval { pub struct []; impl DataType for [] { - fn name(&self) -> &str { - stringify!([]) + fn name(&self) -> String { + stringify!([]).to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/list_type.rs b/src/datatypes/src/types/list_type.rs index 37d620620297..18115837f570 100644 --- a/src/datatypes/src/types/list_type.rs +++ b/src/datatypes/src/types/list_type.rs @@ -52,8 +52,8 @@ impl ListType { } impl DataType for ListType { - fn name(&self) -> &str { - "List" + fn name(&self) -> String { + format!("List<{}>", self.item_type.name()) } fn logical_type_id(&self) -> LogicalTypeId { @@ -92,7 +92,7 @@ mod tests { #[test] fn test_list_type() { let t = ListType::new(ConcreteDataType::boolean_datatype()); - assert_eq!("List", t.name()); + assert_eq!("List", t.name()); assert_eq!(LogicalTypeId::List, t.logical_type_id()); assert_eq!( Value::List(ListValue::new(None, ConcreteDataType::boolean_datatype())), diff --git a/src/datatypes/src/types/null_type.rs b/src/datatypes/src/types/null_type.rs index 04c44c38c573..d07f1b38a8b2 100644 --- a/src/datatypes/src/types/null_type.rs +++ b/src/datatypes/src/types/null_type.rs @@ -32,8 +32,8 @@ impl NullType { } impl DataType for NullType { - fn name(&self) -> &str { - "Null" + fn name(&self) -> String { + "Null".to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/primitive_type.rs b/src/datatypes/src/types/primitive_type.rs index 52e1bd30a7c6..1893b8ac27eb 100644 --- a/src/datatypes/src/types/primitive_type.rs +++ b/src/datatypes/src/types/primitive_type.rs @@ -251,8 +251,8 @@ macro_rules! define_non_timestamp_primitive { define_logical_primitive_type!($Native, $TypeId, $DataType, $Largest); impl DataType for $DataType { - fn name(&self) -> &str { - stringify!($TypeId) + fn name(&self) -> String { + stringify!($TypeId).to_string() } fn logical_type_id(&self) -> LogicalTypeId { @@ -350,8 +350,8 @@ define_logical_primitive_type!(i64, Int64, Int64Type, Int64Type); define_logical_primitive_type!(i32, Int32, Int32Type, Int64Type); impl DataType for Int64Type { - fn name(&self) -> &str { - "Int64" + fn name(&self) -> String { + "Int64".to_string() } fn logical_type_id(&self) -> LogicalTypeId { @@ -397,8 +397,8 @@ impl DataType for Int64Type { } impl DataType for Int32Type { - fn name(&self) -> &str { - "Int32" + fn name(&self) -> String { + "Int32".to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/string_type.rs b/src/datatypes/src/types/string_type.rs index 5f01207a6f74..38045a600adb 100644 --- a/src/datatypes/src/types/string_type.rs +++ b/src/datatypes/src/types/string_type.rs @@ -34,8 +34,8 @@ impl StringType { } impl DataType for StringType { - fn name(&self) -> &str { - "String" + fn name(&self) -> String { + "String".to_string() } fn logical_type_id(&self) -> LogicalTypeId { diff --git a/src/datatypes/src/types/time_type.rs b/src/datatypes/src/types/time_type.rs index a8d48a7f586d..bb360f5a08df 100644 --- a/src/datatypes/src/types/time_type.rs +++ b/src/datatypes/src/types/time_type.rs @@ -92,8 +92,8 @@ macro_rules! impl_data_type_for_time { pub struct [