From 61f802cdf25e653ec5569264060a5349fd66f7d8 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 12 Jun 2026 15:36:53 +0800 Subject: [PATCH] feat: support repartition with target partition columns (#8278) * feat(sql): support repartition on columns syntax Signed-off-by: WenyXu * feat(meta): support repartition target partition columns Signed-off-by: WenyXu * feat(operator): wire repartition target columns Signed-off-by: WenyXu * fix(operator): validate repartition target columns Signed-off-by: WenyXu * test: cover repartition on columns integration Signed-off-by: WenyXu * fix(operator): validate duplicate target partition columns Signed-off-by: WenyXu * chore: upgrade proto Signed-off-by: WenyXu * fix: fix unit tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/common/meta/src/ddl_manager.rs | 26 +- src/meta-srv/src/procedure/repartition.rs | 46 ++- .../repartition/repartition_start.rs | 158 +++++++- .../repartition/update_partition_metadata.rs | 211 ++++++++-- src/operator/src/expr_helper.rs | 56 ++- src/operator/src/statement/ddl.rs | 199 ++++++++- src/sql/src/parsers/alter_parser.rs | 244 ++++++++++- src/sql/src/statements/alter.rs | 29 +- tests-integration/tests/repartition.rs | 383 ++++++++++++++++++ 11 files changed, 1267 insertions(+), 89 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ee4d0b82ba2..8188a8aa494 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5934,7 +5934,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=912f6c21017d5ab6c529568d36e7b8c30a94d84b#912f6c21017d5ab6c529568d36e7b8c30a94d84b" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=00e1d312ecc4b11ffb3c48161d61b03ed0e848a5#00e1d312ecc4b11ffb3c48161d61b03ed0e848a5" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index e45f1da0814..4e95472c515 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -158,7 +158,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "912f6c21017d5ab6c529568d36e7b8c30a94d84b" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "00e1d312ecc4b11ffb3c48161d61b03ed0e848a5" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 8dceeb2e5a6..a6192f06bab 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -15,9 +15,9 @@ use std::sync::Arc; use std::time::Duration; +use api::v1::Repartition; use api::v1::alter_table_expr::Kind; use api::v1::repartition::Source as PbRepartitionSource; -use api::v1::{PartitionExprs, Repartition}; use common_error::ext::BoxedError; use common_procedure::{ BoxedProcedure, BoxedProcedureLoader, Output, ProcedureId, ProcedureManagerRef, @@ -153,8 +153,18 @@ macro_rules! procedure_loader { pub type RepartitionProcedureFactoryRef = Arc; pub enum RepartitionSource { - Partitioned { exprs: Vec }, - Unpartitioned { partition_columns: Vec }, + Partitioned { + exprs: Vec, + /// Full target partition columns to overwrite table metadata. + /// + /// `None` means the repartition keeps using the current table + /// partition columns, so the procedure won't update + /// `partition_key_indices`. + target_partition_columns: Option>, + }, + Unpartitioned { + partition_columns: Vec, + }, } pub trait RepartitionProcedureFactory: Send + Sync { @@ -296,9 +306,12 @@ impl DdlManager { let source = repartition.source; let source = match source { - Some(PbRepartitionSource::PartitionExprs(PartitionExprs { exprs })) => { - RepartitionSource::Partitioned { exprs } - } + Some(PbRepartitionSource::PartitionExprs(source)) => RepartitionSource::Partitioned { + exprs: source.exprs, + target_partition_columns: source + .target_partition_columns + .map(|columns| columns.columns), + }, Some(PbRepartitionSource::Unpartitioned(source)) => RepartitionSource::Unpartitioned { partition_columns: source.partition_columns, }, @@ -307,6 +320,7 @@ impl DdlManager { #[allow(deprecated)] RepartitionSource::Partitioned { exprs: repartition.from_partition_exprs, + target_partition_columns: None, } } }; diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index 8c7c1dfeff0..1d91bd8ea88 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -96,7 +96,7 @@ pub struct PersistentContext { #[serde(with = "humantime_serde", default = "default_timeout")] pub timeout: Duration, #[serde(default)] - /// Records table-level partition metadata added by this repartition. + /// Records table-level partition metadata updated by this repartition. pub partition_metadata_update: Option, } @@ -634,20 +634,15 @@ impl RepartitionProcedure { else { return Ok(()); }; - if update.partition_key_indices.is_empty() { - return Ok(()); - } - let table_info_value = self.context.get_raw_table_info_value().await?; - let mut new_partition_key_indices = table_info_value - .table_info - .meta - .partition_key_indices - .clone(); - new_partition_key_indices.retain(|idx| !update.partition_key_indices.contains(idx)); - if new_partition_key_indices == table_info_value.table_info.meta.partition_key_indices { + let current_partition_key_indices = &table_info_value.table_info.meta.partition_key_indices; + let Some(new_partition_key_indices) = update.rollback_partition_key_indices( + self.context.persistent_ctx.table_id, + current_partition_key_indices, + )? + else { return Ok(()); - } + }; let mut new_table_info = table_info_value.table_info.clone(); new_table_info.meta.partition_key_indices = new_partition_key_indices; @@ -824,7 +819,10 @@ impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { ) -> std::result::Result { let persistent_ctx = PersistentContext::new(table_name, table_id, timeout); let from = match source { - RepartitionSource::Partitioned { exprs } => { + RepartitionSource::Partitioned { + exprs, + target_partition_columns, + } => { let exprs = exprs .iter() .map(|e| { @@ -834,7 +832,10 @@ impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { }) .collect::>>() .map_err(BoxedError::new)?; - RepartitionFrom::Partitioned { exprs } + RepartitionFrom::Partitioned { + exprs, + target_partition_columns, + } } RepartitionSource::Unpartitioned { partition_columns } => { RepartitionFrom::Unpartitioned { partition_columns } @@ -1071,7 +1072,10 @@ mod tests { let procedure = test_procedure( Box::new(RepartitionStart::new( - RepartitionFrom::Partitioned { exprs: vec![] }, + RepartitionFrom::Partitioned { + exprs: vec![], + target_partition_columns: None, + }, vec![], )), test_context(&env, table_id), @@ -1193,9 +1197,9 @@ mod tests { .update_table_info(¤t, current.update(table_info)) .await .unwrap(); - context.persistent_ctx.partition_metadata_update = Some(PartitionMetadataUpdate { - partition_key_indices: vec![0], - }); + context.persistent_ctx.partition_metadata_update = Some( + PartitionMetadataUpdate::from_partitioned(vec![1], vec![0, 1]), + ); let mut procedure = RepartitionProcedure { state: Box::new(UpdatePartitionMetadata::new(vec![])), context, @@ -1881,6 +1885,7 @@ mod tests { let mut procedure = RepartitionProcedure::new( RepartitionFrom::Partitioned { exprs: vec![range_expr("x", 0, 100)], + target_partition_columns: None, }, vec![range_expr("x", 0, 50), range_expr("x", 50, 100)], context, @@ -2019,7 +2024,7 @@ mod tests { .partition_metadata_update .as_ref() .unwrap() - .partition_key_indices, + .target_partition_key_indices, vec![0] ); @@ -2247,6 +2252,7 @@ mod tests { let mut procedure = RepartitionProcedure::new( RepartitionFrom::Partitioned { exprs: vec![range_expr("x", 0, 100)], + target_partition_columns: None, }, vec![range_expr("x", 0, 50), range_expr("x", 50, 100)], context, diff --git a/src/meta-srv/src/procedure/repartition/repartition_start.rs b/src/meta-srv/src/procedure/repartition/repartition_start.rs index b6f0ec9c0a0..11ed2d1840d 100644 --- a/src/meta-srv/src/procedure/repartition/repartition_start.rs +++ b/src/meta-srv/src/procedure/repartition/repartition_start.rs @@ -36,8 +36,18 @@ use crate::procedure::repartition::{Context, State}; #[derive(Debug, Clone, Serialize)] pub enum RepartitionFrom { - Partitioned { exprs: Vec }, - Unpartitioned { partition_columns: Vec }, + Partitioned { + exprs: Vec, + /// Full target partition columns to overwrite table metadata. + /// + /// `None` means the repartition keeps using the current table + /// partition columns, so the procedure won't update + /// `partition_key_indices`. + target_partition_columns: Option>, + }, + Unpartitioned { + partition_columns: Vec, + }, } impl<'de> Deserialize<'de> for RepartitionFrom { @@ -47,8 +57,14 @@ impl<'de> Deserialize<'de> for RepartitionFrom { { #[derive(Deserialize)] enum CurrentRepartitionFrom { - Partitioned { exprs: Vec }, - Unpartitioned { partition_columns: Vec }, + Partitioned { + exprs: Vec, + #[serde(default)] + target_partition_columns: Option>, + }, + Unpartitioned { + partition_columns: Vec, + }, } #[derive(Deserialize)] @@ -59,13 +75,20 @@ impl<'de> Deserialize<'de> for RepartitionFrom { } match RepartitionFromRepr::deserialize(deserializer)? { - RepartitionFromRepr::Current(CurrentRepartitionFrom::Partitioned { exprs }) => { - Ok(Self::Partitioned { exprs }) - } + RepartitionFromRepr::Current(CurrentRepartitionFrom::Partitioned { + exprs, + target_partition_columns, + }) => Ok(Self::Partitioned { + exprs, + target_partition_columns, + }), RepartitionFromRepr::Current(CurrentRepartitionFrom::Unpartitioned { partition_columns, }) => Ok(Self::Unpartitioned { partition_columns }), - RepartitionFromRepr::Legacy(exprs) => Ok(Self::Partitioned { exprs }), + RepartitionFromRepr::Legacy(exprs) => Ok(Self::Partitioned { + exprs, + target_partition_columns: None, + }), } } } @@ -157,7 +180,13 @@ impl State for RepartitionStart { impl RepartitionStart { async fn prepare_from<'a>(&'a self, ctx: &mut Context) -> Result<&'a [PartitionExpr]> { match &self.from { - RepartitionFrom::Partitioned { exprs } => Ok(exprs), + RepartitionFrom::Partitioned { + exprs, + target_partition_columns, + } => { + Self::prepare_partitioned(ctx, target_partition_columns.as_deref()).await?; + Ok(exprs) + } RepartitionFrom::Unpartitioned { partition_columns } => { Self::prepare_unpartitioned(ctx, partition_columns).await?; Ok(&[]) @@ -208,8 +237,52 @@ impl RepartitionStart { }) }) .collect::>>()?; + ctx.persistent_ctx.partition_metadata_update = Some( + PartitionMetadataUpdate::from_unpartitioned(partition_key_indices), + ); + + Ok(()) + } + + async fn prepare_partitioned( + ctx: &mut Context, + target_partition_columns: Option<&[String]>, + ) -> Result<()> { + let Some(target_partition_columns) = target_partition_columns else { + return Ok(()); + }; + if ctx.persistent_ctx.partition_metadata_update.is_some() { + return Ok(()); + } + + ensure!( + !target_partition_columns.is_empty(), + error::InvalidArgumentsSnafu { + err_msg: "Partitioned source expects non-empty target partition columns" + .to_string(), + } + ); + + let table_info_value = ctx.get_table_info_value().await?; + let schema = &table_info_value.table_info.meta.schema; + let target_partition_key_indices = target_partition_columns + .iter() + .map(|column_name| { + schema.column_index_by_name(column_name).with_context(|| { + error::InvalidArgumentsSnafu { + err_msg: format!( + "Target partition column {} not found in table {}", + column_name, ctx.persistent_ctx.table_id + ), + } + }) + }) + .collect::>>()?; ctx.persistent_ctx.partition_metadata_update = - Some(PartitionMetadataUpdate::new(partition_key_indices)); + Some(PartitionMetadataUpdate::from_partitioned( + table_info_value.table_info.meta.partition_key_indices, + target_partition_key_indices, + )); Ok(()) } @@ -507,10 +580,15 @@ mod tests { let state: RepartitionStart = serde_json::from_str(&json).unwrap(); - let RepartitionFrom::Partitioned { exprs } = state.from else { + let RepartitionFrom::Partitioned { + exprs, + target_partition_columns, + } = state.from + else { panic!("expected partition source"); }; assert_eq!(exprs, vec![range_expr("x", 0, 100)]); + assert!(target_partition_columns.is_none()); } #[test] @@ -549,6 +627,7 @@ mod tests { let mut state = RepartitionStart::new( RepartitionFrom::Partitioned { exprs: vec![range_expr("x", 0, 100)], + target_partition_columns: None, }, vec![range_expr("x", 0, 50), range_expr("x", 50, 100)], ); @@ -563,6 +642,52 @@ mod tests { assert!(ctx.persistent_ctx.partition_metadata_update.is_none()); } + #[tokio::test] + async fn test_partitioned_source_initializes_target_partition_metadata_update() { + let env = TestingEnv::new(); + let table_id = 1024; + env.create_physical_table_metadata_for_repartition( + table_id, + vec![test_region_route( + RegionId::new(table_id, 1), + &range_expr("x", 0, 100).as_json_str().unwrap(), + )], + test_region_wal_options(&[1]), + ) + .await; + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ctx = new_parent_context(&env, node_manager, table_id); + let current = ctx.get_raw_table_info_value().await.unwrap(); + let mut table_info = current.table_info.clone(); + table_info.meta.partition_key_indices = vec![0]; + ctx.update_table_info(¤t, current.update(table_info)) + .await + .unwrap(); + let mut state = RepartitionStart::new( + RepartitionFrom::Partitioned { + exprs: vec![range_expr("x", 0, 100)], + target_partition_columns: Some(vec!["col2".to_string(), "col1".to_string()]), + }, + vec![range_expr("x", 0, 50), range_expr("x", 50, 100)], + ); + + let (next, status) = state + .next(&mut ctx, &TestingEnv::procedure_context()) + .await + .unwrap(); + + assert!(status.need_persist()); + assert!(next.as_any().is::()); + let update = ctx + .persistent_ctx + .partition_metadata_update + .as_ref() + .unwrap(); + assert_eq!(update.original_partition_key_indices, vec![0]); + assert_eq!(update.target_partition_key_indices, vec![2, 0]); + assert!(!update.expect_empty_partition_key_indices); + } + #[tokio::test] async fn test_unpartitioned_source_initializes_partition_metadata_update() { let env = TestingEnv::new(); @@ -587,7 +712,7 @@ mod tests { .partition_metadata_update .as_ref() .unwrap() - .partition_key_indices, + .target_partition_key_indices, vec![2, 0] ); } @@ -624,8 +749,13 @@ mod tests { let env = TestingEnv::new(); let table_id = 1024; let mut ctx = new_test_context(&env, table_id).await; - let mut state = - RepartitionStart::new(RepartitionFrom::Partitioned { exprs: vec![] }, vec![]); + let mut state = RepartitionStart::new( + RepartitionFrom::Partitioned { + exprs: vec![], + target_partition_columns: None, + }, + vec![], + ); let err = state .next(&mut ctx, &TestingEnv::procedure_context()) diff --git a/src/meta-srv/src/procedure/repartition/update_partition_metadata.rs b/src/meta-srv/src/procedure/repartition/update_partition_metadata.rs index a213b8b225f..c59e4ca5a83 100644 --- a/src/meta-srv/src/procedure/repartition/update_partition_metadata.rs +++ b/src/meta-srv/src/procedure/repartition/update_partition_metadata.rs @@ -18,6 +18,7 @@ use common_meta::lock_key::TableLock; use common_procedure::{Context as ProcedureContext, Status}; use serde::{Deserialize, Serialize}; use snafu::ensure; +use store_api::storage::TableId; use crate::error::{self, Result}; use crate::procedure::repartition::allocate_region::AllocateRegion; @@ -26,15 +27,107 @@ use crate::procedure::repartition::{Context, State}; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct PartitionMetadataUpdate { - pub partition_key_indices: Vec, + #[serde(default)] + pub original_partition_key_indices: Vec, + #[serde(alias = "partition_key_indices")] + pub target_partition_key_indices: Vec, + #[serde(default)] + pub expect_empty_partition_key_indices: bool, } impl PartitionMetadataUpdate { - pub fn new(partition_key_indices: Vec) -> Self { + pub fn from_unpartitioned(target_partition_key_indices: Vec) -> Self { Self { - partition_key_indices, + original_partition_key_indices: vec![], + target_partition_key_indices, + expect_empty_partition_key_indices: true, } } + + pub fn from_partitioned( + original_partition_key_indices: Vec, + target_partition_key_indices: Vec, + ) -> Self { + Self { + original_partition_key_indices, + target_partition_key_indices, + expect_empty_partition_key_indices: false, + } + } + + fn validate_target_partition_key_indices(&self) -> Result<()> { + ensure!( + !self.target_partition_key_indices.is_empty(), + error::InvalidArgumentsSnafu { + err_msg: + "Repartition partition metadata update expects non-empty target partition key indices" + .to_string(), + } + ); + + Ok(()) + } + + pub fn target_partition_key_indices( + &self, + table_id: TableId, + current_partition_key_indices: &[usize], + ) -> Result>> { + self.validate_target_partition_key_indices()?; + if current_partition_key_indices == self.target_partition_key_indices.as_slice() { + return Ok(None); + } + if current_partition_key_indices == self.original_partition_key_indices.as_slice() { + return Ok(Some(self.target_partition_key_indices.clone())); + } + if self.expect_empty_partition_key_indices { + ensure!( + current_partition_key_indices.is_empty(), + error::InvalidArgumentsSnafu { + err_msg: format!( + "Repartition partition metadata update expects an unpartitioned table, but table {} has partition key indices: {:?}", + table_id, current_partition_key_indices + ), + } + ); + } + + error::InvalidArgumentsSnafu { + err_msg: format!( + "Repartition partition metadata update expects partition key indices {:?} or {:?}, but table {} has partition key indices: {:?}", + self.original_partition_key_indices, + self.target_partition_key_indices, + table_id, + current_partition_key_indices + ), + } + .fail() + } + + pub fn rollback_partition_key_indices( + &self, + table_id: TableId, + current_partition_key_indices: &[usize], + ) -> Result>> { + self.validate_target_partition_key_indices()?; + if current_partition_key_indices == self.original_partition_key_indices.as_slice() { + return Ok(None); + } + if current_partition_key_indices == self.target_partition_key_indices.as_slice() { + return Ok(Some(self.original_partition_key_indices.clone())); + } + + error::InvalidArgumentsSnafu { + err_msg: format!( + "Repartition partition metadata rollback expects partition key indices {:?} or {:?}, but table {} has partition key indices: {:?}", + self.original_partition_key_indices, + self.target_partition_key_indices, + table_id, + current_partition_key_indices + ), + } + .fail() + } } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -62,39 +155,22 @@ impl State for UpdatePartitionMetadata { Status::executing(false), )); }; - let partition_key_indices = update.partition_key_indices.clone(); - ensure!( - !partition_key_indices.is_empty(), - error::InvalidArgumentsSnafu { - err_msg: - "Repartition partition metadata update expects non-empty partition key indices" - .to_string(), - } - ); - let table_id = ctx.persistent_ctx.table_id; let table_lock = TableLock::Write(table_id).into(); let _guard = procedure_ctx.provider.acquire_lock(&table_lock).await; let table_info_value = ctx.get_raw_table_info_value().await?; let current_partition_key_indices = &table_info_value.table_info.meta.partition_key_indices; - if current_partition_key_indices == &partition_key_indices { + let Some(target_partition_key_indices) = + update.target_partition_key_indices(table_id, current_partition_key_indices)? + else { return Ok(( Box::new(AllocateRegion::new(self.plan_entries.clone())), Status::executing(true), )); - } - ensure!( - current_partition_key_indices.is_empty(), - error::InvalidArgumentsSnafu { - err_msg: format!( - "Repartition partition metadata update expects an unpartitioned table, but table {} has partition key indices: {:?}", - table_id, current_partition_key_indices - ), - } - ); + }; let mut new_table_info = table_info_value.table_info.clone(); - new_table_info.meta.partition_key_indices = partition_key_indices; + new_table_info.meta.partition_key_indices = target_partition_key_indices; common_telemetry::info!( "Update table partition metadata, table_id: {}, partition_key_indices: {:?}, partition_columns: {:?}", table_id, @@ -141,7 +217,8 @@ mod tests { .await; let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); let mut ctx = new_parent_context(env, node_manager, table_id); - ctx.persistent_ctx.partition_metadata_update = Some(PartitionMetadataUpdate::new(vec![0])); + ctx.persistent_ctx.partition_metadata_update = + Some(PartitionMetadataUpdate::from_unpartitioned(vec![0])); ctx } @@ -203,7 +280,8 @@ mod tests { let env = TestingEnv::new(); let table_id = 1024; let mut ctx = new_test_context(&env, table_id).await; - ctx.persistent_ctx.partition_metadata_update = Some(PartitionMetadataUpdate::new(vec![])); + ctx.persistent_ctx.partition_metadata_update = + Some(PartitionMetadataUpdate::from_unpartitioned(vec![])); let mut state = UpdatePartitionMetadata::new(vec![]); let err = state @@ -211,7 +289,10 @@ mod tests { .await .unwrap_err(); - assert!(err.to_string().contains("non-empty partition key indices")); + assert!( + err.to_string() + .contains("non-empty target partition key indices") + ); assert!(partition_key_indices(&ctx).await.is_empty()); } @@ -256,4 +337,78 @@ mod tests { assert!(next.as_any().is::()); } + + #[tokio::test] + async fn test_update_partition_metadata_overwrites_partitioned_table() { + let env = TestingEnv::new(); + let table_id = 1024; + let mut ctx = new_test_context(&env, table_id).await; + set_partition_key_indices(&ctx, vec![0]).await; + ctx.persistent_ctx.partition_metadata_update = Some( + PartitionMetadataUpdate::from_partitioned(vec![0], vec![2, 1]), + ); + let mut state = UpdatePartitionMetadata::new(vec![]); + + let (next, status) = state + .next(&mut ctx, &TestingEnv::procedure_context()) + .await + .unwrap(); + + assert!(status.need_persist()); + assert!(next.as_any().is::()); + assert_eq!(partition_key_indices(&ctx).await, vec![2, 1]); + } + + #[tokio::test] + async fn test_update_partition_metadata_rejects_unexpected_partition_keys() { + let env = TestingEnv::new(); + let table_id = 1024; + let mut ctx = new_test_context(&env, table_id).await; + set_partition_key_indices(&ctx, vec![1]).await; + ctx.persistent_ctx.partition_metadata_update = + Some(PartitionMetadataUpdate::from_partitioned(vec![0], vec![2])); + let mut state = UpdatePartitionMetadata::new(vec![]); + + let err = state + .next(&mut ctx, &TestingEnv::procedure_context()) + .await + .unwrap_err(); + + assert!(err.to_string().contains("expects partition key indices")); + assert_eq!(partition_key_indices(&ctx).await, vec![1]); + } + + #[test] + fn test_partition_metadata_update_rollback_target_to_original() { + let update = PartitionMetadataUpdate::from_partitioned(vec![0], vec![2, 1]); + + let rollback_indices = update + .rollback_partition_key_indices(1024, &[2, 1]) + .unwrap(); + + assert_eq!(rollback_indices, Some(vec![0])); + } + + #[test] + fn test_partition_metadata_update_rollback_original_is_noop() { + let update = PartitionMetadataUpdate::from_partitioned(vec![0], vec![2, 1]); + + let rollback_indices = update.rollback_partition_key_indices(1024, &[0]).unwrap(); + + assert_eq!(rollback_indices, None); + } + + #[test] + fn test_partition_metadata_update_rollback_rejects_unexpected_state() { + let update = PartitionMetadataUpdate::from_partitioned(vec![0], vec![2, 1]); + + let err = update + .rollback_partition_key_indices(1024, &[1]) + .unwrap_err(); + + assert!( + err.to_string() + .contains("rollback expects partition key indices") + ); + } } diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index af6e7d10320..b7e9ba39e6e 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -696,8 +696,13 @@ pub struct RepartitionRequest { #[derive(Debug, Clone, PartialEq, Eq)] pub enum RepartitionSource { - Partitions { from_exprs: Vec }, - Unpartitioned { partition_columns: Vec }, + Partitions { + from_exprs: Vec, + target_partition_columns: Option>, + }, + Unpartitioned { + partition_columns: Vec, + }, } pub(crate) fn to_repartition_request( @@ -718,6 +723,12 @@ pub(crate) fn to_repartition_request( AlterTableOperation::Repartition { operation } => ( RepartitionSource::Partitions { from_exprs: operation.from_exprs, + target_partition_columns: operation.partition_columns.map(|columns| { + columns + .into_iter() + .map(|ident| ident.value) + .collect::>() + }), }, operation.into_exprs, ), @@ -1717,9 +1728,14 @@ ALTER TABLE metrics REPARTITION ( assert_eq!("greptime", request.catalog_name); assert_eq!("public", request.schema_name); assert_eq!("metrics", request.table_name); - let RepartitionSource::Partitions { from_exprs } = request.source else { + let RepartitionSource::Partitions { + from_exprs, + target_partition_columns, + } = request.source + else { unreachable!() }; + assert!(target_partition_columns.is_none()); assert_eq!( from_exprs .into_iter() @@ -1740,6 +1756,40 @@ ALTER TABLE metrics REPARTITION ( ); } + #[test] + fn test_to_repartition_request_with_target_partition_columns() { + let sql = r#" +ALTER TABLE metrics REPARTITION ( + device_id < 100 +) ON COLUMNS (device_id, area) INTO ( + device_id < 100 AND area < 'South', + device_id < 100 AND area >= 'South' +);"#; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + + let request = to_repartition_request(alter_table, &QueryContext::arc()).unwrap(); + let RepartitionSource::Partitions { + target_partition_columns, + .. + } = request.source + else { + unreachable!() + }; + + assert_eq!( + target_partition_columns, + Some(vec!["device_id".to_string(), "area".to_string()]) + ); + } + #[test] fn test_to_repartition_request_with_unpartitioned_source() { let sql = r#" diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index cadb1cde666..e7ddd7bfbb9 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -22,7 +22,7 @@ use api::v1::meta::CreateFlowTask as PbCreateFlowTask; use api::v1::repartition::Source; use api::v1::{ AlterDatabaseExpr, AlterTableExpr, CreateFlowExpr, CreateTableExpr, CreateViewExpr, - PartitionExprs, Repartition, UnpartitionedSource, column_def, + PartitionedSource, Repartition, TargetPartitionColumns, UnpartitionedSource, column_def, }; #[cfg(feature = "enterprise")] use api::v1::{ @@ -1517,8 +1517,17 @@ impl StatementExecutor { let table_info = table.table_info(); let existing_partition_columns = table_info.meta.partition_columns().collect::>(); - let partition_columns = match &request.source { - RepartitionSource::Partitions { .. } => { + let column_schemas = table_info.meta.schema.column_schemas(); + // `REPARTITION ... ON COLUMNS` uses overwrite semantics: the provided + // columns are the full target partition columns, not an extension of the + // current ones. Therefore source expressions are converted with the + // existing partition columns, while target expressions and the final + // partition rule are validated against this effective target column set. + let target_partition_columns = match &request.source { + RepartitionSource::Partitions { + target_partition_columns, + .. + } => { ensure!( !existing_partition_columns.is_empty(), InvalidPartitionRuleSnafu { @@ -1528,7 +1537,21 @@ impl StatementExecutor { ) } ); - existing_partition_columns + + if let Some(target_partition_columns) = target_partition_columns { + ensure!( + !target_partition_columns.is_empty(), + InvalidPartitionRuleSnafu { + reason: "ON COLUMNS requires at least one partition column" + } + ); + validate_and_collect_partition_columns( + target_partition_columns, + column_schemas, + )? + } else { + existing_partition_columns.clone() + } } RepartitionSource::Unpartitioned { partition_columns } => { ensure!( @@ -1543,7 +1566,6 @@ impl StatementExecutor { reason: format!("table {} already has partition columns", table_ref) } ); - let column_schemas = table_info.meta.schema.column_schemas(); partition_columns .iter() .map(|column_name| { @@ -1556,16 +1578,18 @@ impl StatementExecutor { } }; - let column_name_and_type = partition_columns + let from_column_name_and_type = column_name_and_type(&existing_partition_columns); + let target_column_name_and_type = column_name_and_type(&target_partition_columns); + let target_partition_column_names = target_partition_columns .iter() - .map(|column| (&column.name, column.data_type.clone())) - .collect(); + .map(|column| column.name.clone()) + .collect::>(); let timezone = query_context.timezone(); // Convert SQL Exprs to PartitionExprs. let from_partition_exprs = match &request.source { - RepartitionSource::Partitions { from_exprs } => from_exprs + RepartitionSource::Partitions { from_exprs, .. } => from_exprs .iter() - .map(|expr| convert_one_expr(expr, &column_name_and_type, &timezone)) + .map(|expr| convert_one_expr(expr, &from_column_name_and_type, &timezone)) .collect::>>()?, RepartitionSource::Unpartitioned { .. } => vec![], }; @@ -1573,7 +1597,7 @@ impl StatementExecutor { let mut into_partition_exprs = request .into_exprs .iter() - .map(|expr| convert_one_expr(expr, &column_name_and_type, &timezone)) + .map(|expr| convert_one_expr(expr, &target_column_name_and_type, &timezone)) .collect::>>()?; // `MERGE PARTITION` (and some `REPARTITION`) generates a single `OR` expression from @@ -1630,12 +1654,16 @@ impl StatementExecutor { .collect(), RepartitionSource::Unpartitioned { .. } => into_partition_exprs.clone(), }; + ensure_partition_expr_columns_in_target( + &new_partition_exprs, + &target_partition_column_names.iter().collect(), + )?; let new_partition_exprs_len = new_partition_exprs.len(); let from_partition_exprs_len = from_partition_exprs.len(); // Validate the new partition expressions using MultiDimPartitionRule and PartitionChecker. let _ = MultiDimPartitionRule::try_new( - partition_columns.iter().map(|c| c.name.clone()).collect(), + target_partition_column_names, vec![], new_partition_exprs, true, @@ -1653,8 +1681,14 @@ impl StatementExecutor { let from_partition_exprs_json = serialize_exprs(from_partition_exprs)?; let into_partition_exprs_json = serialize_exprs(into_partition_exprs)?; let source = match &request.source { - RepartitionSource::Partitions { .. } => Source::PartitionExprs(PartitionExprs { + RepartitionSource::Partitions { + target_partition_columns, + .. + } => Source::PartitionExprs(PartitionedSource { exprs: from_partition_exprs_json, + target_partition_columns: target_partition_columns + .clone() + .map(|columns| TargetPartitionColumns { columns }), }), RepartitionSource::Unpartitioned { partition_columns } => { Source::Unpartitioned(UnpartitionedSource { @@ -2360,6 +2394,73 @@ fn find_partition_entries( Ok(partition_exprs) } +fn column_name_and_type<'a>( + partition_columns: &'a [&'a ColumnSchema], +) -> HashMap<&'a String, ConcreteDataType> { + partition_columns + .iter() + .map(|column| (&column.name, column.data_type.clone())) + .collect() +} + +fn validate_and_collect_partition_columns<'a>( + column_names: &[String], + column_schemas: &'a [ColumnSchema], +) -> Result> { + let mut seen = HashSet::with_capacity(column_names.len()); + column_names + .iter() + .map(|column_name| { + ensure!( + seen.insert(column_name), + InvalidPartitionRuleSnafu { + reason: format!("duplicate partition column '{}'", column_name) + } + ); + column_schemas + .iter() + .find(|column| &column.name == column_name) + .with_context(|| ColumnNotFoundSnafu { msg: column_name }) + }) + .collect() +} + +fn ensure_partition_expr_columns_in_target( + partition_exprs: &[PartitionExpr], + target_partition_columns: &HashSet<&String>, +) -> Result<()> { + for expr in partition_exprs { + ensure_partition_operand_columns_in_target(&expr.lhs, target_partition_columns)?; + ensure_partition_operand_columns_in_target(&expr.rhs, target_partition_columns)?; + } + + Ok(()) +} + +fn ensure_partition_operand_columns_in_target( + operand: &Operand, + target_partition_columns: &HashSet<&String>, +) -> Result<()> { + match operand { + Operand::Column(column) => ensure!( + target_partition_columns.contains(column), + InvalidPartitionRuleSnafu { + reason: format!( + "partition expression references column '{}' that is not in target partition columns", + column + ) + } + ), + Operand::Expr(expr) => { + ensure_partition_operand_columns_in_target(&expr.lhs, target_partition_columns)?; + ensure_partition_operand_columns_in_target(&expr.rhs, target_partition_columns)?; + } + Operand::Value(_) => {} + } + + Ok(()) +} + fn convert_one_expr( expr: &Expr, column_name_and_type: &HashMap<&String, ConcreteDataType>, @@ -2648,6 +2749,78 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; assert_eq!(physical_partition_exprs, logical_partition_exprs); } + #[test] + fn test_repartition_target_partition_columns_are_overwrite_context() { + let device_id = ColumnSchema::new("device_id", ConcreteDataType::int32_datatype(), true); + let area = ColumnSchema::new("area", ConcreteDataType::string_datatype(), true); + let existing_partition_columns = vec![&device_id]; + let target_partition_columns = vec![&device_id, &area]; + let existing_column_name_and_type = column_name_and_type(&existing_partition_columns); + let target_column_name_and_type = column_name_and_type(&target_partition_columns); + let timezone = Timezone::from_tz_string("UTC").unwrap(); + let dialect = GreptimeDbDialect {}; + + let mut parser = Parser::new(&dialect) + .try_with_sql("device_id < 100 AND area < 'South'") + .unwrap(); + let expr = parser.parse_expr().unwrap(); + + let err = convert_one_expr(&expr, &existing_column_name_and_type, &timezone).unwrap_err(); + assert!(err.to_string().contains("area")); + + let partition_expr = convert_one_expr(&expr, &target_column_name_and_type, &timezone) + .expect("target columns should overwrite the conversion context"); + let partition_expr = partition_expr.to_string(); + assert!(partition_expr.contains("device_id")); + assert!(partition_expr.contains("area")); + assert!(partition_expr.contains("South")); + } + + #[test] + fn test_repartition_rejects_remaining_expr_outside_target_columns() { + let device_id = "device_id".to_string(); + let area = "area".to_string(); + let timezone = Timezone::from_tz_string("UTC").unwrap(); + let column_name_and_type = HashMap::from([ + (&device_id, ConcreteDataType::int32_datatype()), + (&area, ConcreteDataType::string_datatype()), + ]); + let dialect = GreptimeDbDialect {}; + let mut parser = Parser::new(&dialect) + .try_with_sql("device_id >= 100") + .unwrap(); + let remaining_old_expr = convert_one_expr( + &parser.parse_expr().unwrap(), + &column_name_and_type, + &timezone, + ) + .unwrap(); + let target_partition_columns = HashSet::from([&area]); + + let err = ensure_partition_expr_columns_in_target( + &[remaining_old_expr], + &target_partition_columns, + ) + .unwrap_err(); + + assert!(err.to_string().contains("device_id")); + assert!(err.to_string().contains("target partition columns")); + } + + #[test] + fn test_repartition_rejects_duplicate_target_partition_columns() { + let device_id = ColumnSchema::new("device_id", ConcreteDataType::int32_datatype(), true); + let column_schemas = vec![device_id]; + let target_partition_columns = vec!["device_id".to_string(), "device_id".to_string()]; + + let err = + validate_and_collect_partition_columns(&target_partition_columns, &column_schemas) + .unwrap_err(); + + assert!(err.to_string().contains("duplicate partition column")); + assert!(err.to_string().contains("device_id")); + } + #[tokio::test] #[ignore = "TODO(ruihang): WIP new partition rule"] async fn test_parse_partitions() { diff --git a/src/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index 4aa571005a7..6c71702f5bf 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -22,6 +22,7 @@ use datatypes::schema::COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE; use snafu::{ResultExt, ensure}; use sqlparser::ast::{BinaryOperator, Expr, Ident}; use sqlparser::keywords::Keyword; +use sqlparser::parser::IsOptional::Mandatory; use sqlparser::parser::{Parser, ParserError}; use sqlparser::tokenizer::{Token, TokenWithSpan}; @@ -205,6 +206,8 @@ impl ParserContext<'_> { let _ = self.parser.next_token(); let from_exprs = self.parse_repartition_expr_list()?; + let partition_columns = self.parse_optional_repartition_columns()?; + self.parser .expect_keyword(Keyword::INTO) .context(error::SyntaxSnafu)?; @@ -215,7 +218,14 @@ impl ParserContext<'_> { } Ok(AlterTableOperation::Repartition { - operation: RepartitionOperation::new(from_exprs, into_exprs), + operation: match partition_columns { + Some(partition_columns) => RepartitionOperation::with_partition_columns( + from_exprs, + into_exprs, + partition_columns, + ), + None => RepartitionOperation::new(from_exprs, into_exprs), + }, }) } @@ -246,6 +256,8 @@ impl ParserContext<'_> { ); } + let partition_columns = self.parse_optional_repartition_columns()?; + self.parser .expect_keyword(Keyword::INTO) .context(error::SyntaxSnafu)?; @@ -254,12 +266,39 @@ impl ParserContext<'_> { if matches!(self.parser.peek_token().token, Token::Comma) { return self.expected("end of SPLIT PARTITION clause", self.parser.peek_token()); } + if matches!(&self.parser.peek_token().token, Token::Word(w) if w.keyword == Keyword::ON) { + return self.expected("end of SPLIT PARTITION clause", self.parser.peek_token()); + } Ok(AlterTableOperation::Repartition { - operation: RepartitionOperation::new(from_exprs, into_exprs), + operation: match partition_columns { + Some(partition_columns) => RepartitionOperation::with_partition_columns( + from_exprs, + into_exprs, + partition_columns, + ), + None => RepartitionOperation::new(from_exprs, into_exprs), + }, }) } + fn parse_optional_repartition_columns(&mut self) -> Result>> { + if !self.parser.parse_keywords(&[Keyword::ON, Keyword::COLUMNS]) { + return Ok(None); + } + + 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(); + + Ok(Some(column_list)) + } + fn parse_alter_table_merge_partition(&mut self) -> Result { let _ = self.parser.next_token(); self.parser @@ -978,6 +1017,7 @@ ALTER TABLE t REPARTITION ( assert_eq!(operation.from_exprs.len(), 1); assert_eq!(operation.from_exprs[0].to_string(), "device_id < 100"); assert_eq!(operation.into_exprs.len(), 2); + assert!(operation.partition_columns.is_none()); assert_eq!( operation.into_exprs[0].to_string(), "device_id < 100 AND area < 'South'" @@ -990,6 +1030,100 @@ ALTER TABLE t REPARTITION ( } } + #[test] + fn test_parse_alter_table_repartition_on_columns() { + let sql = r#" +ALTER TABLE t REPARTITION ( + device_id < 100 +) +ON COLUMNS (device_id, area) +INTO ( + device_id < 100 AND area < 'South', + device_id < 100 AND area >= 'South' +);"#; + let mut result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + assert_eq!(1, result.len()); + + let statement = result.remove(0); + assert_matches!(statement, Statement::AlterTable { .. }); + if let Statement::AlterTable(alter_table) = statement { + assert_matches!( + alter_table.alter_operation(), + AlterTableOperation::Repartition { .. } + ); + + if let AlterTableOperation::Repartition { operation } = alter_table.alter_operation() { + assert_eq!(operation.from_exprs.len(), 1); + assert_eq!(operation.from_exprs[0].to_string(), "device_id < 100"); + assert_eq!(operation.into_exprs.len(), 2); + assert_eq!( + operation + .partition_columns + .as_ref() + .unwrap() + .iter() + .map(|ident| ident.value.as_str()) + .collect::>(), + vec!["device_id", "area"] + ); + assert_eq!( + operation.to_string(), + "(device_id < 100) ON COLUMNS (device_id, area) INTO (device_id < 100 AND area < 'South', device_id < 100 AND area >= 'South')" + ); + } + } + } + + #[test] + fn test_parse_alter_table_repartition_on_columns_with_options() { + let sql = r#" +ALTER TABLE t REPARTITION ( + device_id < 100 +) +ON COLUMNS (device_id, area) +INTO ( + device_id < 100 AND area < 'South', + device_id < 100 AND area >= 'South' +) +WITH ( + TIMEOUT = '5m', + WAIT = false +);"#; + let mut result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + assert_eq!(1, result.len()); + + let statement = result.remove(0); + assert_matches!(statement, Statement::AlterTable { .. }); + if let Statement::AlterTable(alter_table) = statement { + assert_matches!( + alter_table.alter_operation(), + AlterTableOperation::Repartition { .. } + ); + + if let AlterTableOperation::Repartition { operation } = alter_table.alter_operation() { + assert_eq!( + operation + .partition_columns + .as_ref() + .unwrap() + .iter() + .map(|ident| ident.value.as_str()) + .collect::>(), + vec!["device_id", "area"] + ); + } + + let options = alter_table.options().to_str_map(); + assert_eq!(options.get("timeout").unwrap(), &"5m"); + assert_eq!(options.get("wait").unwrap(), &"false"); + assert_eq!(options.len(), 2); + } + } + #[test] fn test_parse_alter_table_partition_on_columns() { let sql = r#" @@ -1110,6 +1244,7 @@ ALTER TABLE metrics SPLIT PARTITION ( assert_eq!(operation.from_exprs.len(), 1); assert_eq!(operation.from_exprs[0].to_string(), "device_id < 100"); assert_eq!(operation.into_exprs.len(), 2); + assert!(operation.partition_columns.is_none()); assert_eq!( operation.into_exprs[0].to_string(), "device_id < 100 AND area < 'South'" @@ -1122,6 +1257,90 @@ ALTER TABLE metrics SPLIT PARTITION ( } } + #[test] + fn test_parse_alter_table_split_partition_on_columns() { + let sql = r#" +ALTER TABLE metrics SPLIT PARTITION ( + device_id < 100 +) +ON COLUMNS (device_id, area) +INTO ( + device_id < 100 AND area < 'South', + device_id < 100 AND area >= 'South' +);"#; + let mut result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + assert_eq!(1, result.len()); + + let statement = result.remove(0); + assert_matches!(statement, Statement::AlterTable { .. }); + if let Statement::AlterTable(alter_table) = statement { + assert_matches!( + alter_table.alter_operation(), + AlterTableOperation::Repartition { .. } + ); + + if let AlterTableOperation::Repartition { operation } = alter_table.alter_operation() { + assert_eq!(operation.from_exprs.len(), 1); + assert_eq!(operation.from_exprs[0].to_string(), "device_id < 100"); + assert_eq!(operation.into_exprs.len(), 2); + assert_eq!( + operation + .partition_columns + .as_ref() + .unwrap() + .iter() + .map(|ident| ident.value.as_str()) + .collect::>(), + vec!["device_id", "area"] + ); + assert_eq!( + operation.to_string(), + "(device_id < 100) ON COLUMNS (device_id, area) INTO (device_id < 100 AND area < 'South', device_id < 100 AND area >= 'South')" + ); + } + } + } + + #[test] + fn test_parse_alter_table_split_partition_on_columns_empty_columns() { + let sql = r#" +ALTER TABLE metrics SPLIT PARTITION ( + device_id < 100 +) +ON COLUMNS () +INTO ( + device_id < 100 AND area < 'South', + device_id < 100 AND area >= 'South' +);"#; + let result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); + + assert!(result.is_err()); + } + + #[test] + fn test_parse_alter_table_split_partition_on_columns_wrong_order() { + let sql = r#" +ALTER TABLE metrics SPLIT PARTITION ( + device_id < 100 +) +INTO ( + device_id < 100 AND area < 'South', + device_id < 100 AND area >= 'South' +) +ON COLUMNS (device_id, area);"#; + let result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap_err(); + + assert_eq!( + result.output_msg(), + "Invalid SQL syntax: sql parser error: Expected end of SPLIT PARTITION clause, found: ON" + ); + } + #[test] fn test_parse_alter_table_merge_partition() { let sql = r#" @@ -1155,6 +1374,27 @@ ALTER TABLE metrics MERGE PARTITION ( } } + #[test] + fn test_parse_alter_table_merge_partition_on_columns_rejected() { + let sql = r#" +ALTER TABLE metrics MERGE PARTITION ( + device_id < 100, + device_id >= 100 +) +ON COLUMNS (device_id, area) +INTO ( + device_id >= 0 +);"#; + let result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap_err(); + + assert_eq!( + result.output_msg(), + "SQL statement is not supported, keyword: ON" + ); + } + #[test] fn test_parse_alter_table_merge_partition_with_options() { let sql = r#" diff --git a/src/sql/src/statements/alter.rs b/src/sql/src/statements/alter.rs index 72182a4b60a..dd1f8e19d26 100644 --- a/src/sql/src/statements/alter.rs +++ b/src/sql/src/statements/alter.rs @@ -140,6 +140,12 @@ pub struct SetDefaultsOperation { pub struct RepartitionOperation { pub from_exprs: Vec, pub into_exprs: Vec, + /// Optional new partition columns for `REPARTITION ... ON COLUMNS (...) INTO (...)` and + /// `SPLIT PARTITION ... ON COLUMNS (...) INTO (...)`. + /// + /// This is `Some` only when the statement explicitly carries `ON COLUMNS`. + /// Legacy `REPARTITION`, `SPLIT PARTITION`, and `MERGE PARTITION` keep this as `None`. + pub partition_columns: Option>, } impl RepartitionOperation { @@ -147,6 +153,19 @@ impl RepartitionOperation { Self { from_exprs, into_exprs, + partition_columns: None, + } + } + + pub fn with_partition_columns( + from_exprs: Vec, + into_exprs: Vec, + partition_columns: Vec, + ) -> Self { + Self { + from_exprs, + into_exprs, + partition_columns: Some(partition_columns), } } } @@ -164,7 +183,15 @@ impl Display for RepartitionOperation { .map(|expr| expr.to_string()) .join(", "); - write!(f, "({from}) INTO ({into})") + if let Some(partition_columns) = &self.partition_columns { + let partition_columns = partition_columns + .iter() + .map(|ident| ident.to_string()) + .join(", "); + write!(f, "({from}) ON COLUMNS ({partition_columns}) INTO ({into})") + } else { + write!(f, "({from}) INTO ({into})") + } } } diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index ef59d1b9109..2adec87edcb 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use std::time::Duration; use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_error::root_source; use common_meta::key::table_name::TableNameKey; use common_procedure::{ProcedureWithId, watcher}; use common_query::Output; @@ -64,6 +65,24 @@ macro_rules! repartition_tests { } } + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_on_columns_metadata_mito >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + common_telemetry::init_default_ut_logging(); + $crate::repartition::test_repartition_on_columns_metadata_mito(store_type).await; + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_on_columns_data_correctness_mito >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + common_telemetry::init_default_ut_logging(); + $crate::repartition::test_repartition_on_columns_data_correctness_mito(store_type).await; + } + } + #[tokio::test(flavor = "multi_thread")] async fn [< test_partition_unpartitioned_metric >]() { let store_type = tests_integration::test_util::StorageType::$service; @@ -223,6 +242,146 @@ ORDER BY partition_ordinal_position;", .unwrap(); } +pub async fn test_repartition_on_columns_metadata_mito(store_type: StorageType) { + info!( + "test_repartition_on_columns_metadata_mito: store_type: {:?}", + store_type + ); + let cluster_name = "test_repartition_on_columns_metadata_mito"; + let (store_config, _guard) = get_test_store_config(&store_type); + let datanodes = 3u64; + let mut builder = GreptimeDbClusterBuilder::new(cluster_name).await; + if matches!(store_type, StorageType::File) { + let home_dir = create_temp_dir("test_repartition_on_columns_metadata_mito_data_home"); + builder = builder.with_shared_home_dir(Arc::new(home_dir)); + } + + let cluster = builder + .with_datanodes(datanodes as u32) + .with_store_config(store_config) + .with_datanode_wal_config(DatanodeWalConfig::Noop) + .build(true) + .await; + + let query_ctx = QueryContext::arc(); + let instance = cluster.fe_instance(); + + assert_on_columns_metadata_overwrite( + instance, + "repartition_on_columns_metadata_table", + &repartition_on_columns_sql("repartition_on_columns_metadata_table"), + query_ctx.clone(), + ) + .await; + assert_on_columns_metadata_overwrite( + instance, + "split_on_columns_metadata_table", + &split_on_columns_sql("split_on_columns_metadata_table"), + query_ctx.clone(), + ) + .await; + + assert_on_columns_rejects_removed_remaining_column( + instance, + "repartition_on_columns_error_table", + &repartition_on_columns_removed_column_sql("repartition_on_columns_error_table"), + query_ctx.clone(), + ) + .await; + assert_on_columns_rejects_removed_remaining_column( + instance, + "split_on_columns_error_table", + &split_on_columns_removed_column_sql("split_on_columns_error_table"), + query_ctx.clone(), + ) + .await; + + run_sql( + instance, + "DROP TABLE `repartition_on_columns_metadata_table`", + query_ctx.clone(), + ) + .await + .unwrap(); + run_sql( + instance, + "DROP TABLE `split_on_columns_metadata_table`", + query_ctx.clone(), + ) + .await + .unwrap(); + run_sql( + instance, + "DROP TABLE `repartition_on_columns_error_table`", + query_ctx.clone(), + ) + .await + .unwrap(); + run_sql( + instance, + "DROP TABLE `split_on_columns_error_table`", + query_ctx.clone(), + ) + .await + .unwrap(); +} + +pub async fn test_repartition_on_columns_data_correctness_mito(store_type: StorageType) { + info!( + "test_repartition_on_columns_data_correctness_mito: store_type: {:?}", + store_type + ); + let cluster_name = "test_repartition_on_columns_data_correctness_mito"; + let (store_config, _guard) = get_test_store_config(&store_type); + let datanodes = 3u64; + let mut builder = GreptimeDbClusterBuilder::new(cluster_name).await; + if matches!(store_type, StorageType::File) { + let home_dir = + create_temp_dir("test_repartition_on_columns_data_correctness_mito_data_home"); + builder = builder.with_shared_home_dir(Arc::new(home_dir)); + } + + let cluster = builder + .with_datanodes(datanodes as u32) + .with_store_config(store_config) + .with_datanode_wal_config(DatanodeWalConfig::Noop) + .build(true) + .await; + + let query_ctx = QueryContext::arc(); + let instance = cluster.fe_instance(); + + assert_on_columns_data_correctness( + instance, + "repartition_on_columns_data_table", + &repartition_on_columns_sql("repartition_on_columns_data_table"), + query_ctx.clone(), + ) + .await; + assert_on_columns_data_correctness( + instance, + "split_on_columns_data_table", + &split_on_columns_sql("split_on_columns_data_table"), + query_ctx.clone(), + ) + .await; + + run_sql( + instance, + "DROP TABLE `repartition_on_columns_data_table`", + query_ctx.clone(), + ) + .await + .unwrap(); + run_sql( + instance, + "DROP TABLE `split_on_columns_data_table`", + query_ctx.clone(), + ) + .await + .unwrap(); +} + pub async fn test_partition_unpartitioned_metric(store_type: StorageType) { info!( "test_partition_unpartitioned_metric: store_type: {:?}", @@ -1165,6 +1324,230 @@ pub async fn test_repartition_metric( .unwrap(); } +async fn create_on_columns_test_table( + instance: &Arc, + table_name: &str, + query_ctx: QueryContextRef, +) { + let sql = format!( + r#" + CREATE TABLE `{table_name}`( + `id` INT, + `city` STRING, + `ts` TIMESTAMP TIME INDEX, + PRIMARY KEY(`id`, `city`) + ) PARTITION ON COLUMNS (`id`) ( + `id` < 10, + `id` >= 10 AND `id` < 20, + `id` >= 20 + ) ENGINE = mito; + "# + ); + run_sql(instance, &sql, query_ctx).await.unwrap(); +} + +async fn insert_on_columns_old_rows( + instance: &Arc, + table_name: &str, + query_ctx: QueryContextRef, +) { + let sql = format!( + r#" + INSERT INTO `{table_name}` VALUES + (1, 'London', '2022-01-01 00:00:00'), + (5, 'New York', '2022-01-01 00:00:00'), + (10, 'Paris', '2022-01-01 00:00:00'), + (15, 'Tokyo', '2022-01-01 00:00:00'), + (20, 'Beijing', '2022-01-01 00:00:00'); + "# + ); + run_sql(instance, &sql, query_ctx).await.unwrap(); +} + +async fn assert_on_columns_metadata_overwrite( + instance: &Arc, + table_name: &str, + alter_sql: &str, + query_ctx: QueryContextRef, +) { + create_on_columns_test_table(instance, table_name, query_ctx.clone()).await; + run_sql(instance, alter_sql, query_ctx.clone()) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(500)).await; + assert_partition_expression(instance, table_name, "id, city", query_ctx).await; +} + +async fn assert_on_columns_rejects_removed_remaining_column( + instance: &Arc, + table_name: &str, + alter_sql: &str, + query_ctx: QueryContextRef, +) { + create_on_columns_test_table(instance, table_name, query_ctx.clone()).await; + assert_sql_err_eq( + instance, + alter_sql, + "Invalid partition rule: partition expression references column 'id' that is not in target partition columns", + query_ctx, + ) + .await; +} + +async fn assert_on_columns_data_correctness( + instance: &Arc, + table_name: &str, + alter_sql: &str, + query_ctx: QueryContextRef, +) { + create_on_columns_test_table(instance, table_name, query_ctx.clone()).await; + insert_on_columns_old_rows(instance, table_name, query_ctx.clone()).await; + run_sql(instance, alter_sql, query_ctx.clone()) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(500)).await; + + let result = run_sql( + instance, + &format!("SELECT * FROM `{table_name}` ORDER BY `id`, `city`"), + query_ctx.clone(), + ) + .await + .unwrap(); + let expected = "\ ++----+----------+---------------------+ +| id | city | ts | ++----+----------+---------------------+ +| 1 | London | 2022-01-01T00:00:00 | +| 5 | New York | 2022-01-01T00:00:00 | +| 10 | Paris | 2022-01-01T00:00:00 | +| 15 | Tokyo | 2022-01-01T00:00:00 | +| 20 | Beijing | 2022-01-01T00:00:00 | ++----+----------+---------------------+"; + check_output_stream(result.data, expected).await; + + let sql = format!( + r#" + INSERT INTO `{table_name}` VALUES + (2, 'Amsterdam', '2022-01-02 00:00:00'), + (7, 'Zurich', '2022-01-02 00:00:00'), + (12, 'Berlin', '2022-01-02 00:00:00'); + "# + ); + run_sql(instance, &sql, query_ctx.clone()).await.unwrap(); + + let result = run_sql( + instance, + &format!("SELECT * FROM `{table_name}` ORDER BY `id`, `city`"), + query_ctx, + ) + .await + .unwrap(); + let expected = "\ ++----+-----------+---------------------+ +| id | city | ts | ++----+-----------+---------------------+ +| 1 | London | 2022-01-01T00:00:00 | +| 2 | Amsterdam | 2022-01-02T00:00:00 | +| 5 | New York | 2022-01-01T00:00:00 | +| 7 | Zurich | 2022-01-02T00:00:00 | +| 10 | Paris | 2022-01-01T00:00:00 | +| 12 | Berlin | 2022-01-02T00:00:00 | +| 15 | Tokyo | 2022-01-01T00:00:00 | +| 20 | Beijing | 2022-01-01T00:00:00 | ++----+-----------+---------------------+"; + check_output_stream(result.data, expected).await; +} + +fn repartition_on_columns_sql(table_name: &str) -> String { + format!( + r#" + ALTER TABLE `{table_name}` REPARTITION ( + `id` < 10 + ) ON COLUMNS (`id`, `city`) INTO ( + `id` < 10 AND `city` < 'M', + `id` < 10 AND `city` >= 'M' + ); + "# + ) +} + +fn split_on_columns_sql(table_name: &str) -> String { + format!( + r#" + ALTER TABLE `{table_name}` SPLIT PARTITION ( + `id` < 10 + ) ON COLUMNS (`id`, `city`) INTO ( + `id` < 10 AND `city` < 'M', + `id` < 10 AND `city` >= 'M' + ); + "# + ) +} + +fn repartition_on_columns_removed_column_sql(table_name: &str) -> String { + format!( + r#" + ALTER TABLE `{table_name}` REPARTITION ( + `id` < 10 + ) ON COLUMNS (`city`) INTO ( + `city` < 'M', + `city` >= 'M' + ); + "# + ) +} + +fn split_on_columns_removed_column_sql(table_name: &str) -> String { + format!( + r#" + ALTER TABLE `{table_name}` SPLIT PARTITION ( + `id` < 10 + ) ON COLUMNS (`city`) INTO ( + `city` < 'M', + `city` >= 'M' + ); + "# + ) +} + +async fn assert_partition_expression( + instance: &Arc, + table_name: &str, + expected_partition_expression: &str, + query_ctx: QueryContextRef, +) { + let result = run_sql( + instance, + &format!( + "SELECT DISTINCT partition_expression FROM information_schema.partitions WHERE table_name = '{table_name}' ORDER BY partition_expression" + ), + query_ctx, + ) + .await + .unwrap(); + let expected = format!( + "\ ++----------------------+ +| partition_expression | ++----------------------+ +| {expected_partition_expression:<20} | ++----------------------+" + ); + check_output_stream(result.data, &expected).await; +} + +async fn assert_sql_err_eq( + instance: &Arc, + sql: &str, + expected: &str, + query_ctx: QueryContextRef, +) { + let err = run_sql(instance, sql, query_ctx).await.unwrap_err(); + let root = root_source(&err).unwrap_or(&err); + assert_eq!(root.to_string(), expected); +} + async fn run_sql( instance: &Arc, sql: &str,