rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-30 15:36:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support repartition with target partition columns (#8278)

Browse Source 
* feat(sql): support repartition on columns syntax

Signed-off-by: WenyXu <wenymedia@gmail.com>

* feat(meta): support repartition target partition columns

Signed-off-by: WenyXu <wenymedia@gmail.com>

* feat(operator): wire repartition target columns

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(operator): validate repartition target columns

Signed-off-by: WenyXu <wenymedia@gmail.com>

* test: cover repartition on columns integration

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(operator): validate duplicate target partition columns

Signed-off-by: WenyXu <wenymedia@gmail.com>

* chore: upgrade proto

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix: fix unit tests

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-06-12 07:36:53 +00:00
committed by GitHub
parent c567a6562a
commit 61f802cdf2
11 changed files with 1267 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1 -1
View File
@@ -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",
Cargo.toml
+1 -1
View File
@@ -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"
src/common/meta/src/ddl_manager.rs
+20 -6
View File
@@ -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<dyn RepartitionProcedureFactory>;
pub enum RepartitionSource {
Partitioned { exprs: Vec<String> },
Unpartitioned { partition_columns: Vec<String> },
Partitioned {
exprs: Vec<String>,
/// 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<Vec<String>>,
},
Unpartitioned {
partition_columns: Vec<String>,
},
}
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,
}
}
};
src/meta-srv/src/procedure/repartition.rs
+26 -20
View File
@@ -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<PartitionMetadataUpdate>,
}
@@ -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<BoxedProcedure, BoxedError> {
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::<Result<Vec<_>>>()
.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(&current, 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,
src/meta-srv/src/procedure/repartition/repartition_start.rs
+144 -14
View File
@@ -36,8 +36,18 @@ use crate::procedure::repartition::{Context, State};
#[derive(Debug, Clone, Serialize)]
pub enum RepartitionFrom {
Partitioned { exprs: Vec<PartitionExpr> },
Unpartitioned { partition_columns: Vec<String> },
Partitioned {
exprs: Vec<PartitionExpr>,
/// 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<Vec<String>>,
},
Unpartitioned {
partition_columns: Vec<String>,
},
}
impl<'de> Deserialize<'de> for RepartitionFrom {
@@ -47,8 +57,14 @@ impl<'de> Deserialize<'de> for RepartitionFrom {
{
#[derive(Deserialize)]
enum CurrentRepartitionFrom {
Partitioned { exprs: Vec<PartitionExpr> },
Unpartitioned { partition_columns: Vec<String> },
Partitioned {
exprs: Vec<PartitionExpr>,
#[serde(default)]
target_partition_columns: Option<Vec<String>>,
},
Unpartitioned {
partition_columns: Vec<String>,
},
}
#[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::<Result<Vec<_>>>()?;
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::<Result<Vec<_>>>()?;
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(&current, 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::<UpdatePartitionMetadata>());
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())
src/meta-srv/src/procedure/repartition/update_partition_metadata.rs
+183 -28
View File
@@ -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<usize>,
#[serde(default)]
pub original_partition_key_indices: Vec<usize>,
#[serde(alias = "partition_key_indices")]
pub target_partition_key_indices: Vec<usize>,
#[serde(default)]
pub expect_empty_partition_key_indices: bool,
}
impl PartitionMetadataUpdate {
pub fn new(partition_key_indices: Vec<usize>) -> Self {
pub fn from_unpartitioned(target_partition_key_indices: Vec<usize>) -> 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<usize>,
target_partition_key_indices: Vec<usize>,
) -> 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<Option<Vec<usize>>> {
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<Option<Vec<usize>>> {
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::<AllocateRegion>());
}
#[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::<AllocateRegion>());
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")
);
}
}
src/operator/src/expr_helper.rs
+53 -3
View File
@@ -696,8 +696,13 @@ pub struct RepartitionRequest {
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RepartitionSource {
Partitions { from_exprs: Vec<Expr> },
Unpartitioned { partition_columns: Vec<String> },
Partitions {
from_exprs: Vec<Expr>,
target_partition_columns: Option<Vec<String>>,
},
Unpartitioned {
partition_columns: Vec<String>,
},
}
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::<Vec<_>>()
}),
},
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#"
src/operator/src/statement/ddl.rs
+186 -13
View File
@@ -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::<Vec<_>>();
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::<Vec<_>>();
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::<Result<Vec<_>>>()?,
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::<Result<Vec<_>>>()?;
// `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<Vec<&'a ColumnSchema>> {
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() {
src/sql/src/parsers/alter_parser.rs
+242 -2
View File
@@ -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<Option<Vec<Ident>>> {
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<AlterTableOperation> {
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<_>>(),
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<_>>(),
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<_>>(),
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#"
src/sql/src/statements/alter.rs
+28 -1
View File
@@ -140,6 +140,12 @@ pub struct SetDefaultsOperation {
pub struct RepartitionOperation {
pub from_exprs: Vec<Expr>,
pub into_exprs: Vec<Expr>,
/// 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<Vec<Ident>>,
}
impl RepartitionOperation {
@@ -147,6 +153,19 @@ impl RepartitionOperation {
Self {
from_exprs,
into_exprs,
partition_columns: None,
}
}
pub fn with_partition_columns(
from_exprs: Vec<Expr>,
into_exprs: Vec<Expr>,
partition_columns: Vec<Ident>,
) -> 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})")
}
}
}
tests-integration/tests/repartition.rs
+383
View File
@@ -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<Instance>,
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<Instance>,
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<Instance>,
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<Instance>,
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<Instance>,
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<Instance>,
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<Instance>,
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<Instance>,
sql: &str,