feat: add repartition partition count hint (#9080)

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-09-10 05:04:21 +00:00
committed by GitHub
parent 2d78220a4f
commit adda50e03f
8 changed files with 963 additions and 123 deletions
+133 -30
View File
@@ -49,33 +49,14 @@ const LOCATION_TYPE_AFTER: i32 = LocationType::After as i32;
/// mixed batch — annotation alters skip region dispatch, so they cannot share
/// a statement with options that regions must see.
fn annotation_family_of_keys<'a>(
mut keys: impl Iterator<Item = &'a str>,
keys: impl Iterator<Item = &'a str>,
) -> Result<Option<AnnotationFamily>> {
let Some(first) = keys.next() else {
return Ok(None);
};
let family = AnnotationFamily::of_key(first);
let mut count = 1usize;
for key in keys {
count += 1;
let this = AnnotationFamily::of_key(key);
if this != family {
return error::InvalidTableOptionRequestSnafu {
err_msg: family.or(this).unwrap().mixed_batch_error(),
}
.fail();
table::requests::validate_annotation_keys(keys).map_err(|err| {
error::InvalidTableOptionRequestSnafu {
err_msg: err.to_string(),
}
}
if let Some(family) = family
&& family.requires_single_key()
&& count > 1
{
return error::InvalidTableOptionRequestSnafu {
err_msg: family.mixed_batch_error(),
}
.fail();
}
Ok(family)
.build()
})
}
/// Returns the annotation family when `kind` is a SET/UNSET whose keys all
@@ -583,6 +564,131 @@ mod tests {
assert_eq!("mem_usage".to_string(), drop_names.pop().unwrap());
}
#[test]
fn test_repartition_hints_together() {
let keys = [
REPARTITION_COLUMN_HINT_KEY,
table::requests::REPARTITION_PARTITION_NUM_HINT_KEY,
];
let options = vec![
(keys[0].to_string(), "host".to_string()),
(keys[1].to_string(), "10".to_string()),
];
let set = Kind::SetTableOptions(SetTableOptions {
table_options: options
.iter()
.map(|(key, value)| PbOption {
key: key.clone(),
value: value.clone(),
})
.collect(),
});
let unset = Kind::UnsetTableOptions(UnsetTableOptions {
keys: keys.map(str::to_string).to_vec(),
});
for kind in [set, unset] {
assert_eq!(
annotation_alter_family(&kind).unwrap(),
Some(AnnotationFamily::RepartitionHint)
);
let request = alter_expr_to_request(
1,
AlterTableExpr {
kind: Some(kind),
..Default::default()
},
None,
)
.unwrap();
match request.alter_kind {
AlterKind::SetAnnotations {
family,
options: actual,
} => {
assert_eq!(family, AnnotationFamily::RepartitionHint);
assert_eq!(actual, options);
}
AlterKind::UnsetAnnotations {
family,
keys: actual,
} => {
assert_eq!(family, AnnotationFamily::RepartitionHint);
assert_eq!(actual, keys);
}
other => panic!("unexpected alter kind: {other:?}"),
}
}
}
#[test]
fn test_repartition_partition_num_hint_expr() {
let key = table::requests::REPARTITION_PARTITION_NUM_HINT_KEY;
let set = Kind::SetTableOptions(SetTableOptions {
table_options: vec![PbOption {
key: key.to_string(),
value: "8".to_string(),
}],
});
let unset = Kind::UnsetTableOptions(UnsetTableOptions {
keys: vec![key.to_string()],
});
for kind in [set, unset] {
assert_eq!(
annotation_alter_family(&kind).unwrap(),
Some(AnnotationFamily::RepartitionHint)
);
let request = alter_expr_to_request(
1,
AlterTableExpr {
kind: Some(kind),
..Default::default()
},
None,
)
.unwrap();
match request.alter_kind {
AlterKind::SetAnnotations { family, options } => {
assert_eq!(family, AnnotationFamily::RepartitionHint);
assert_eq!(options, vec![(key.to_string(), "8".to_string())]);
}
AlterKind::UnsetAnnotations { family, keys } => {
assert_eq!(family, AnnotationFamily::RepartitionHint);
assert_eq!(keys, vec![key.to_string()]);
}
other => panic!("unexpected alter kind: {other:?}"),
}
}
for other in [key, table::requests::TTL_KEY] {
for kind in [
Kind::SetTableOptions(SetTableOptions {
table_options: [key, other]
.into_iter()
.map(|key| PbOption {
key: key.to_string(),
value: "8".to_string(),
})
.collect(),
}),
Kind::UnsetTableOptions(UnsetTableOptions {
keys: vec![key.to_string(), other.to_string()],
}),
] {
assert!(annotation_alter_family(&kind).is_err());
assert!(
alter_expr_to_request(
1,
AlterTableExpr {
kind: Some(kind),
..Default::default()
},
None
)
.is_err()
);
}
}
}
#[test]
fn test_set_repartition_column_hint_expr() {
let expr = AlterTableExpr {
@@ -633,7 +739,7 @@ mod tests {
let err = alter_expr_to_request(1, expr, None).unwrap_err();
assert!(
err.to_string()
.contains("repartition.column.hint must be altered separately")
.contains("repartition hints must be altered separately")
);
// Duplicate hint entries are not a meaningful batch either.
@@ -655,10 +761,7 @@ mod tests {
})),
};
let err = alter_expr_to_request(1, dup, None).unwrap_err();
assert!(
err.to_string()
.contains("repartition.column.hint must be altered separately")
);
assert!(err.to_string().contains("duplicate repartition hint keys"));
}
#[test]
+78 -51
View File
@@ -1123,58 +1123,85 @@ async fn test_on_submit_alter_request_with_exist_poison() {
}
#[tokio::test]
async fn test_semantic_annotation_alter_is_metadata_only() {
let (tx, mut rx) = mpsc::channel(8);
let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx)));
let ddl_context = new_ddl_context(node_manager);
let table_id = 1024;
let table_name = "foo";
let task = test_create_table_task(table_name, table_id);
ddl_context
.table_metadata_manager
.create_table_metadata(
task.table_info.clone(),
prepare_table_route(table_id),
HashMap::new(),
)
.await
.unwrap();
async fn test_annotation_alter_is_metadata_only() {
for (key, value) in [
("greptime.semantic.signal_type", "metric"),
(table::requests::REPARTITION_PARTITION_NUM_HINT_KEY, "8"),
] {
let (tx, mut rx) = mpsc::channel(8);
let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx)));
let ddl_context = new_ddl_context(node_manager);
let table_id = 1024;
let table_name = "foo";
let task = test_create_table_task(table_name, table_id);
ddl_context
.table_metadata_manager
.create_table_metadata(
task.table_info.clone(),
prepare_table_route(table_id),
HashMap::new(),
)
.await
.unwrap();
let alter_table_task = AlterTableTask {
alter_table: AlterTableExpr {
catalog_name: DEFAULT_CATALOG_NAME.to_string(),
schema_name: DEFAULT_SCHEMA_NAME.to_string(),
table_name: table_name.to_string(),
kind: Some(Kind::SetTableOptions(SetTableOptions {
table_options: vec![api::v1::Option {
key: "greptime.semantic.signal_type".to_string(),
value: "metric".to_string(),
}],
})),
},
};
let mut procedure =
AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap();
execute_procedure_until_done(&mut procedure).await;
let alter_table_task = AlterTableTask {
alter_table: AlterTableExpr {
catalog_name: DEFAULT_CATALOG_NAME.to_string(),
schema_name: DEFAULT_SCHEMA_NAME.to_string(),
table_name: table_name.to_string(),
kind: Some(Kind::SetTableOptions(SetTableOptions {
table_options: vec![api::v1::Option {
key: key.to_string(),
value: value.to_string(),
}],
})),
},
};
let mut procedure =
AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap();
execute_procedure_until_done(&mut procedure).await;
// Metadata-only: no region request reaches any datanode.
rx.try_recv().unwrap_err();
// Metadata-only: no region request reaches any datanode.
rx.try_recv().unwrap_err();
let table_info = ddl_context
.table_metadata_manager
.table_info_manager()
.get(table_id)
.await
.unwrap()
.unwrap()
.into_inner()
.table_info;
assert_eq!(
table_info
.meta
.options
.extra_options
.get("greptime.semantic.signal_type"),
Some(&"metric".to_string())
);
let table_info = ddl_context
.table_metadata_manager
.table_info_manager()
.get(table_id)
.await
.unwrap()
.unwrap()
.into_inner()
.table_info;
assert_eq!(
table_info.meta.options.extra_options.get(key),
Some(&value.to_string())
);
for _ in 0..2 {
let task = AlterTableTask {
alter_table: AlterTableExpr {
catalog_name: DEFAULT_CATALOG_NAME.to_string(),
schema_name: DEFAULT_SCHEMA_NAME.to_string(),
table_name: table_name.to_string(),
kind: Some(Kind::UnsetTableOptions(api::v1::UnsetTableOptions {
keys: vec![key.to_string()],
})),
},
};
let mut procedure =
AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap();
execute_procedure_until_done(&mut procedure).await;
rx.try_recv().unwrap_err();
let info = ddl_context
.table_metadata_manager
.table_info_manager()
.get(table_id)
.await
.unwrap()
.unwrap()
.into_inner()
.table_info;
assert!(!info.meta.options.extra_options.contains_key(key));
}
}
}
+20 -15
View File
@@ -2036,20 +2036,25 @@ mod tests {
// The repartition hint drives physical repartitioning; on a logical
// route it stays rejected by the guard.
let err = ddl_manager
.submit_ddl_task(
ExecutorContext {
query_context: Some(QueryContext::default()),
..Default::default()
},
SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr(
"logical",
&[("repartition.column.hint", "host")],
))),
)
.await
.unwrap_err();
let msg = common_error::ext::ErrorExt::output_msg(&err);
assert!(msg.contains("non-physical TableRouteValue"), "{msg}");
for (key, value) in [
(table::requests::REPARTITION_COLUMN_HINT_KEY, "host"),
(table::requests::REPARTITION_PARTITION_NUM_HINT_KEY, "8"),
] {
let err = ddl_manager
.submit_ddl_task(
ExecutorContext {
query_context: Some(QueryContext::default()),
..Default::default()
},
SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr(
"logical",
&[(key, value)],
))),
)
.await
.unwrap_err();
let msg = common_error::ext::ErrorExt::output_msg(&err);
assert!(msg.contains("non-physical TableRouteValue"), "{msg}");
}
}
}
+31
View File
@@ -2587,6 +2587,7 @@ fn validate_and_normalize_annotations(
CheckError::ColumnNotFound { column } => ColumnNotFoundSnafu { msg: column }.build(),
e @ (CheckError::UnknownKey { .. }
| CheckError::InvalidValue { .. }
| CheckError::InvalidPartitionNumHint { .. }
| CheckError::ColumnNotStringForm { .. }) => InvalidSqlSnafu {
err_msg: e.to_string(),
}
@@ -3525,6 +3526,36 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;";
}
}
#[test]
fn test_create_table_with_repartition_partition_num_hint() {
let key = table::requests::REPARTITION_PARTITION_NUM_HINT_KEY;
for (value, expected) in [
(" 8 ", Some("8")),
("0", None),
("-1", None),
("4294967296", None),
] {
let expr = create_expr_from_sql(&format!(
"CREATE TABLE metrics (host STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host)) WITH ('{key}' = '{value}')"
));
let result = create_table_info(&expr, vec![]);
if let Some(expected) = expected {
let info = result.unwrap();
assert_eq!(
info.meta.options.extra_options.get(key).map(String::as_str),
Some(expected)
);
} else {
assert!(
result
.unwrap_err()
.to_string()
.contains("expects a positive integer")
);
}
}
}
#[test]
fn test_create_table_with_repartition_column_hint() {
let expr = create_expr_from_sql(
+180 -15
View File
@@ -40,9 +40,10 @@ use store_api::storage::{ColumnDescriptor, ColumnDescriptorBuilder, ColumnId};
use crate::error::{self, Result};
use crate::requests::{
AddColumnRequest, AlterKind, AnnotationContext, AnnotationFamily, AnnotationValidationError,
ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, SetDefaultRequest, SetIndexOption,
TableOptions, UnsetIndexOption, has_stable_string_form, parse_entity_columns,
parse_entity_option_key, validate_and_normalize_annotation,
ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, REPARTITION_PARTITION_NUM_HINT_KEY,
SetDefaultRequest, SetIndexOption, TableOptions, UnsetIndexOption, has_stable_string_form,
parse_entity_columns, parse_entity_option_key, validate_and_normalize_annotation,
validate_annotation_keys,
};
use crate::table_reference::TableReference;
@@ -461,13 +462,13 @@ impl TableMeta {
family: AnnotationFamily,
options: &[(String, String)],
) -> Result<TableMetaBuilder> {
ensure!(
!family.requires_single_key() || options.len() == 1,
validate_annotation_keys(options.iter().map(|(key, _)| key.as_str())).map_err(|err| {
error::InvalidAlterRequestSnafu {
table: table_name,
err: family.mixed_batch_error(),
err: err.to_string(),
}
);
.build()
})?;
let cx = AnnotationContext {
schema: &self.schema,
partition_key_indices: &self.partition_key_indices,
@@ -516,13 +517,13 @@ impl TableMeta {
family: AnnotationFamily,
keys: &[String],
) -> Result<TableMetaBuilder> {
ensure!(
!family.requires_single_key() || keys.len() == 1,
validate_annotation_keys(keys.iter().map(String::as_str)).map_err(|err| {
error::InvalidAlterRequestSnafu {
table: table_name,
err: family.mixed_batch_error(),
err: err.to_string(),
}
);
.build()
})?;
let mut new_options = self.options.clone();
for key in keys {
ensure!(
@@ -1491,6 +1492,7 @@ impl TableInfo {
pub fn to_region_options(&self) -> HashMap<String, String> {
let mut options = HashMap::from(&self.meta.options);
options.remove(REPARTITION_COLUMN_HINT_KEY);
options.remove(REPARTITION_PARTITION_NUM_HINT_KEY);
options
}
@@ -1943,6 +1945,160 @@ mod tests {
);
}
#[test]
fn test_set_unset_repartition_hints_together() {
let meta = TableMetaBuilder::empty()
.schema(Arc::new(new_test_schema()))
.primary_key_indices(vec![0])
.engine("engine")
.next_column_id(3)
.build()
.unwrap();
let options = vec![
(REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string()),
(
REPARTITION_PARTITION_NUM_HINT_KEY.to_string(),
"10".to_string(),
),
];
for reverse in [false, true] {
let mut options = options.clone();
if reverse {
options.reverse();
}
let updated = meta
.builder_with_alter_kind(
"t",
&AlterKind::SetAnnotations {
family: AnnotationFamily::RepartitionHint,
options: options.clone(),
},
)
.unwrap()
.build()
.unwrap();
assert_eq!(
updated.options.extra_options,
options.iter().cloned().collect()
);
for (key, value) in [
(REPARTITION_COLUMN_HINT_KEY, "missing"),
(REPARTITION_PARTITION_NUM_HINT_KEY, "0"),
] {
let invalid = options
.iter()
.map(|(k, v)| {
(
k.clone(),
if k == key {
value.to_string()
} else {
v.clone()
},
)
})
.collect();
assert!(
updated
.builder_with_alter_kind(
"t",
&AlterKind::SetAnnotations {
family: AnnotationFamily::RepartitionHint,
options: invalid,
}
)
.is_err()
);
assert_eq!(
updated.options.extra_options,
options.iter().cloned().collect()
);
}
let cleared = updated
.builder_with_alter_kind(
"t",
&AlterKind::UnsetAnnotations {
family: AnnotationFamily::RepartitionHint,
keys: options.into_iter().map(|(key, _)| key).collect(),
},
)
.unwrap()
.build()
.unwrap();
assert!(cleared.options.extra_options.is_empty());
}
}
#[test]
fn test_repartition_partition_num_hint() {
for partition_key_indices in [vec![], vec![0]] {
let mut meta = TableMetaBuilder::empty()
.schema(Arc::new(new_test_schema()))
.primary_key_indices(vec![0])
.partition_key_indices(partition_key_indices)
.engine("engine")
.next_column_id(3)
.build()
.unwrap();
for value in ["", " ", "0", "-1", "1.5", "abc", "4294967296"] {
let err = meta
.builder_with_alter_kind(
"t",
&AlterKind::SetAnnotations {
family: AnnotationFamily::RepartitionHint,
options: vec![(
REPARTITION_PARTITION_NUM_HINT_KEY.to_string(),
value.to_string(),
)],
},
)
.err()
.unwrap();
assert!(
err.to_string().contains("expects a positive integer"),
"{err}"
);
}
for (value, expected) in [(" 8 ", "8"), ("1", "1"), ("4294967295", "4294967295")] {
meta = meta
.builder_with_alter_kind(
"t",
&AlterKind::SetAnnotations {
family: AnnotationFamily::RepartitionHint,
options: vec![(
REPARTITION_PARTITION_NUM_HINT_KEY.to_string(),
value.to_string(),
)],
},
)
.unwrap()
.build()
.unwrap();
assert_eq!(
meta.options.extra_options,
HashMap::from([(
REPARTITION_PARTITION_NUM_HINT_KEY.to_string(),
expected.to_string()
)])
);
}
for _ in 0..2 {
meta = meta
.builder_with_alter_kind(
"t",
&AlterKind::UnsetAnnotations {
family: AnnotationFamily::RepartitionHint,
keys: vec![REPARTITION_PARTITION_NUM_HINT_KEY.to_string()],
},
)
.unwrap()
.build()
.unwrap();
assert!(meta.options.extra_options.is_empty());
}
}
}
#[test]
fn test_set_repartition_column_hint() {
let meta = TableMetaBuilder::empty()
@@ -2141,11 +2297,15 @@ mod tests {
}
#[test]
fn test_repartition_column_hint_is_not_region_option() {
fn test_repartition_hints_are_not_region_options() {
let mut table_options = TableOptions::default();
table_options
.extra_options
.insert(REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string());
table_options.extra_options.insert(
REPARTITION_PARTITION_NUM_HINT_KEY.to_string(),
"8".to_string(),
);
let table_info = TableInfoBuilder::default()
.table_id(1)
.table_version(0)
@@ -2165,6 +2325,11 @@ mod tests {
.build()
.unwrap();
assert!(
!table_info
.to_region_options()
.contains_key(REPARTITION_PARTITION_NUM_HINT_KEY)
);
assert!(
!table_info
.to_region_options()
@@ -2972,7 +3137,7 @@ mod tests {
}
#[test]
fn test_repartition_hint_batch_must_be_single_key() {
fn test_repartition_hint_batch_rejects_duplicate_keys() {
let meta = TableMetaBuilder::empty()
.schema(Arc::new(new_test_schema()))
.primary_key_indices(vec![0])
@@ -2995,7 +3160,7 @@ mod tests {
.err()
.unwrap();
assert!(
err.to_string().contains("must be altered separately"),
err.to_string().contains("duplicate repartition hint keys"),
"{err}"
);
@@ -3011,7 +3176,7 @@ mod tests {
.err()
.unwrap();
assert!(
err.to_string().contains("must be altered separately"),
err.to_string().contains("duplicate repartition hint keys"),
"{err}"
);
}
+125 -12
View File
@@ -77,7 +77,7 @@ pub fn is_trace_v1_table(table_info: &crate::metadata::TableInfo) -> bool {
pub const OTLP_METRIC_COMPAT_KEY: &str = "otlp_metric_compat";
pub const OTLP_METRIC_COMPAT_PROM: &str = "prom";
pub const VALID_TABLE_OPTION_KEYS: [&str; 14] = [
pub const VALID_TABLE_OPTION_KEYS: [&str; 15] = [
// common keys:
WRITE_BUFFER_SIZE_KEY,
TTL_KEY,
@@ -96,6 +96,7 @@ pub const VALID_TABLE_OPTION_KEYS: [&str; 14] = [
TABLE_DATA_MODEL,
OTLP_METRIC_COMPAT_KEY,
REPARTITION_COLUMN_HINT_KEY,
REPARTITION_PARTITION_NUM_HINT_KEY,
];
pub const DDL_TIMEOUT: &str = "timeout";
@@ -209,6 +210,9 @@ pub const SKIP_WAL_KEY: &str = store_api::mito_engine_options::SKIP_WAL_KEY;
pub const TRACE_TABLE_PARTITIONS_HINT_KEY: &str = "trace_table_partitions";
pub const REPARTITION_COLUMN_HINT_KEY: &str = "repartition.column.hint";
/// Table-level partition count hint consumed by the auto-repartition planner.
pub const REPARTITION_PARTITION_NUM_HINT_KEY: &str = "repartition.partition.num.hint";
impl TableOptions {
pub fn try_from_iter<T: ToString, U: IntoIterator<Item = (T, T)>>(
iter: U,
@@ -350,24 +354,26 @@ pub struct ModifyColumnTypeRequest {
pub enum AnnotationFamily {
/// `greptime.semantic.*` options (see the [`semantic`] module).
Semantic,
/// `repartition.column.hint`, consumed by the auto-repartition planner.
/// Column and partition count hints consumed by the auto-repartition planner.
RepartitionHint,
}
impl AnnotationFamily {
/// The key namespace: a prefix for [`Self::Semantic`], the exact key for
/// [`Self::RepartitionHint`].
/// The key namespace used in diagnostics; accepted keys are classified by [`Self::of_key`].
pub fn namespace(self) -> &'static str {
match self {
Self::Semantic => SEMANTIC_PREFIX,
Self::RepartitionHint => REPARTITION_COLUMN_HINT_KEY,
Self::RepartitionHint => "repartition.",
}
}
pub fn of_key(key: &str) -> Option<Self> {
if key.starts_with(SEMANTIC_PREFIX) {
Some(Self::Semantic)
} else if key == REPARTITION_COLUMN_HINT_KEY {
} else if matches!(
key,
REPARTITION_COLUMN_HINT_KEY | REPARTITION_PARTITION_NUM_HINT_KEY
) {
Some(Self::RepartitionHint)
} else {
None
@@ -384,10 +390,8 @@ impl AnnotationFamily {
}
}
/// Whether this family's SET/UNSET batch must contain exactly one key.
/// The repartition hint is a single marker; a batch with several hint
/// entries (duplicates included) has no meaningful order.
pub fn requires_single_key(self) -> bool {
/// Whether duplicate keys are rejected in this family's SET/UNSET batch.
pub fn requires_unique_keys(self) -> bool {
matches!(self, Self::RepartitionHint)
}
@@ -398,12 +402,60 @@ impl AnnotationFamily {
"`{SEMANTIC_PREFIX}*` options must be altered separately from other table options"
),
Self::RepartitionHint => {
format!("{REPARTITION_COLUMN_HINT_KEY} must be altered separately")
"repartition hints must be altered separately from other table options".to_string()
}
}
}
}
/// Why an annotation SET/UNSET key batch was rejected.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AnnotationKeyError {
MixedFamilies { family: AnnotationFamily },
DuplicateKey,
}
impl fmt::Display for AnnotationKeyError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::MixedFamilies { family } => f.write_str(&family.mixed_batch_error()),
Self::DuplicateKey => f.write_str("duplicate repartition hint keys"),
}
}
}
/// Validates a SET/UNSET batch and returns its annotation family.
///
/// Empty batches and batches containing only non-annotation keys return `None`.
/// Annotation keys cannot share a batch with another family or region options.
/// Duplicate keys are rejected only for families that require unique keys.
pub fn validate_annotation_keys<'a>(
keys: impl IntoIterator<Item = &'a str>,
) -> std::result::Result<Option<AnnotationFamily>, AnnotationKeyError> {
let mut keys = keys.into_iter();
let Some(first) = keys.next() else {
return Ok(None);
};
let family = AnnotationFamily::of_key(first);
let reject_duplicates = family.is_some_and(|family| family.requires_unique_keys());
let mut seen = HashSet::new();
if reject_duplicates {
seen.insert(first);
}
for key in keys {
let this = AnnotationFamily::of_key(key);
if this != family
&& let Some(family) = family.or(this)
{
return Err(AnnotationKeyError::MixedFamilies { family });
}
if reject_duplicates && !seen.insert(key) {
return Err(AnnotationKeyError::DuplicateKey);
}
}
Ok(family)
}
/// Table shape an annotation option is validated against.
pub struct AnnotationContext<'a> {
pub schema: &'a Schema,
@@ -431,6 +483,9 @@ pub enum AnnotationValidationError {
column: String,
ty: ConcreteDataType,
},
InvalidPartitionNumHint {
value: String,
},
NotSingleColumn,
PartitionMetadataConflict,
TimeIndexConflict,
@@ -449,6 +504,10 @@ impl fmt::Display for AnnotationValidationError {
"entity column `{column}` (option `{key}`) has type `{ty}`, \
which cannot render as a string"
),
Self::InvalidPartitionNumHint { value } => write!(
f,
"{REPARTITION_PARTITION_NUM_HINT_KEY} expects a positive integer within u32 range, got `{value}`"
),
Self::NotSingleColumn => write!(
f,
"{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name"
@@ -466,7 +525,7 @@ impl fmt::Display for AnnotationValidationError {
}
/// Validates one annotation option and returns the value to store — the
/// repartition hint is trimmed to the bare column name, semantic values pass
/// repartition hints are trimmed, semantic values pass
/// through unchanged.
pub(crate) fn validate_and_normalize_annotation(
family: AnnotationFamily,
@@ -505,6 +564,15 @@ pub(crate) fn validate_and_normalize_annotation(
}
Ok(value.to_string())
}
AnnotationFamily::RepartitionHint if key == REPARTITION_PARTITION_NUM_HINT_KEY => {
let value = value.trim();
if !matches!(value.parse::<u32>(), Ok(1..)) {
return Err(AnnotationValidationError::InvalidPartitionNumHint {
value: value.to_string(),
});
}
Ok(value.to_string())
}
AnnotationFamily::RepartitionHint => {
let column_name = value.trim();
if column_name.is_empty() || column_name.contains(',') {
@@ -763,6 +831,49 @@ mod tests {
use super::*;
#[test]
fn test_validate_annotation_keys() {
let column = REPARTITION_COLUMN_HINT_KEY;
let count = REPARTITION_PARTITION_NUM_HINT_KEY;
for (keys, expected) in [
(vec![], None),
(vec![TTL_KEY, TTL_KEY], None),
(vec!["repartition.unknown.hint"], None),
(vec![column], Some(AnnotationFamily::RepartitionHint)),
(vec![count], Some(AnnotationFamily::RepartitionHint)),
(vec![column, count], Some(AnnotationFamily::RepartitionHint)),
(vec![count, column], Some(AnnotationFamily::RepartitionHint)),
(
vec!["greptime.semantic.source", "greptime.semantic.source"],
Some(AnnotationFamily::Semantic),
),
] {
assert_eq!(validate_annotation_keys(keys), Ok(expected));
}
for keys in [
vec![column, column],
vec![count, count],
vec![column, count, column],
] {
assert_eq!(
validate_annotation_keys(keys),
Err(AnnotationKeyError::DuplicateKey)
);
}
for keys in [
vec![column, TTL_KEY],
vec![TTL_KEY, count],
vec![column, "greptime.semantic.source"],
] {
assert_eq!(
validate_annotation_keys(keys),
Err(AnnotationKeyError::MixedFamilies {
family: AnnotationFamily::RepartitionHint,
})
);
}
}
#[test]
fn test_validate_table_option() {
assert!(validate_table_option(FILE_TABLE_LOCATION_KEY));
@@ -773,6 +884,8 @@ mod tests {
assert!(validate_table_option(STORAGE_KEY));
assert!(validate_table_option(MEMTABLE_BULK_MERGE_THRESHOLD));
assert!(validate_table_option(REPARTITION_COLUMN_HINT_KEY));
assert!(validate_table_option(REPARTITION_PARTITION_NUM_HINT_KEY));
assert_eq!(AnnotationFamily::of_key("repartition.unknown.hint"), None);
assert!(!validate_table_option("foo"));
// Only whitelisted semantic keys are accepted.
@@ -180,6 +180,310 @@ DROP TABLE alter_hint_table;
Affected Rows: 0
CREATE TABLE partition_num_hint_table (
host STRING PRIMARY KEY,
ts TIMESTAMP TIME INDEX
)
PARTITION ON COLUMNS (host) (
host < 'm',
host >= 'm'
)
WITH ('repartition.partition.num.hint' = ' 8 ');
Affected Rows: 0
SHOW CREATE TABLE partition_num_hint_table;
+--------------------------+---------------------------------------------------------+
| Table | Create Table |
+--------------------------+---------------------------------------------------------+
| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( |
| | "host" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | PARTITION ON COLUMNS ("host") ( |
| | host < 'm', |
| | host >= 'm' |
| | ) |
| | ENGINE=mito |
| | WITH( |
| | 'repartition.partition.num.hint' = '8' |
| | ) |
+--------------------------+---------------------------------------------------------+
ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = 16;
Affected Rows: 0
SHOW CREATE TABLE partition_num_hint_table;
+--------------------------+---------------------------------------------------------+
| Table | Create Table |
+--------------------------+---------------------------------------------------------+
| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( |
| | "host" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | PARTITION ON COLUMNS ("host") ( |
| | host < 'm', |
| | host >= 'm' |
| | ) |
| | ENGINE=mito |
| | WITH( |
| | 'repartition.partition.num.hint' = '16' |
| | ) |
+--------------------------+---------------------------------------------------------+
ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = '0';
Error: 1004(InvalidArguments), Invalid alter table(partition_num_hint_table) request: repartition.partition.num.hint expects a positive integer within u32 range, got `0`
ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint';
Affected Rows: 0
ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint';
Affected Rows: 0
SHOW CREATE TABLE partition_num_hint_table;
+--------------------------+---------------------------------------------------------+
| Table | Create Table |
+--------------------------+---------------------------------------------------------+
| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( |
| | "host" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | PARTITION ON COLUMNS ("host") ( |
| | host < 'm', |
| | host >= 'm' |
| | ) |
| | ENGINE=mito |
| | |
+--------------------------+---------------------------------------------------------+
DROP TABLE partition_num_hint_table;
Affected Rows: 0
CREATE TABLE partition_num_hint_table (
host STRING PRIMARY KEY,
ts TIMESTAMP TIME INDEX
)
WITH ('repartition.partition.num.hint' = '10');
Affected Rows: 0
SHOW CREATE TABLE partition_num_hint_table;
+--------------------------+---------------------------------------------------------+
| Table | Create Table |
+--------------------------+---------------------------------------------------------+
| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( |
| | "host" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | |
| | ENGINE=mito |
| | WITH( |
| | 'repartition.partition.num.hint' = '10' |
| | ) |
+--------------------------+---------------------------------------------------------+
DROP TABLE partition_num_hint_table;
Affected Rows: 0
CREATE TABLE partition_num_hint_table (
host STRING PRIMARY KEY,
ts TIMESTAMP TIME INDEX
)
WITH ('repartition.partition.num.hint' = 10);
Affected Rows: 0
SHOW CREATE TABLE partition_num_hint_table;
+--------------------------+---------------------------------------------------------+
| Table | Create Table |
+--------------------------+---------------------------------------------------------+
| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( |
| | "host" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | |
| | ENGINE=mito |
| | WITH( |
| | 'repartition.partition.num.hint' = '10' |
| | ) |
+--------------------------+---------------------------------------------------------+
DROP TABLE partition_num_hint_table;
Affected Rows: 0
CREATE TABLE combined_hint_table (
host STRING PRIMARY KEY,
`service` STRING,
ts TIMESTAMP TIME INDEX
);
Affected Rows: 0
ALTER TABLE combined_hint_table SET
'repartition.column.hint' = 'host',
'repartition.partition.num.hint' = 10;
Affected Rows: 0
SHOW CREATE TABLE combined_hint_table;
+---------------------+----------------------------------------------------+
| Table | Create Table |
+---------------------+----------------------------------------------------+
| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( |
| | "host" STRING NULL, |
| | "service" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | |
| | ENGINE=mito |
| | WITH( |
| | 'repartition.column.hint' = 'host', |
| | 'repartition.partition.num.hint' = '10' |
| | ) |
+---------------------+----------------------------------------------------+
ALTER TABLE combined_hint_table SET
'repartition.column.hint' = 'service',
'repartition.partition.num.hint' = 0;
Error: 1004(InvalidArguments), Invalid alter table(combined_hint_table) request: repartition.partition.num.hint expects a positive integer within u32 range, got `0`
SHOW CREATE TABLE combined_hint_table;
+---------------------+----------------------------------------------------+
| Table | Create Table |
+---------------------+----------------------------------------------------+
| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( |
| | "host" STRING NULL, |
| | "service" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | |
| | ENGINE=mito |
| | WITH( |
| | 'repartition.column.hint' = 'host', |
| | 'repartition.partition.num.hint' = '10' |
| | ) |
+---------------------+----------------------------------------------------+
ALTER TABLE combined_hint_table SET
'repartition.partition.num.hint' = 20,
'repartition.column.hint' = 'missing';
Error: 4002(TableColumnNotFound), Column missing not exists in table combined_hint_table
SHOW CREATE TABLE combined_hint_table;
+---------------------+----------------------------------------------------+
| Table | Create Table |
+---------------------+----------------------------------------------------+
| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( |
| | "host" STRING NULL, |
| | "service" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | |
| | ENGINE=mito |
| | WITH( |
| | 'repartition.column.hint' = 'host', |
| | 'repartition.partition.num.hint' = '10' |
| | ) |
+---------------------+----------------------------------------------------+
ALTER TABLE combined_hint_table SET
'repartition.partition.num.hint' = 20,
'repartition.column.hint' = 'service';
Affected Rows: 0
SHOW CREATE TABLE combined_hint_table;
+---------------------+----------------------------------------------------+
| Table | Create Table |
+---------------------+----------------------------------------------------+
| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( |
| | "host" STRING NULL, |
| | "service" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | |
| | ENGINE=mito |
| | WITH( |
| | 'repartition.column.hint' = 'service', |
| | 'repartition.partition.num.hint' = '20' |
| | ) |
+---------------------+----------------------------------------------------+
ALTER TABLE combined_hint_table SET
'repartition.partition.num.hint' = 30,
'repartition.partition.num.hint' = 40;
Error: 1004(InvalidArguments), Invalid table option request: duplicate repartition hint keys
ALTER TABLE combined_hint_table SET
'repartition.partition.num.hint' = 30,
'ttl' = '7d';
Error: 1004(InvalidArguments), Invalid table option request: repartition hints must be altered separately from other table options
ALTER TABLE combined_hint_table UNSET
'repartition.column.hint',
'repartition.partition.num.hint';
Affected Rows: 0
SHOW CREATE TABLE combined_hint_table;
+---------------------+----------------------------------------------------+
| Table | Create Table |
+---------------------+----------------------------------------------------+
| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( |
| | "host" STRING NULL, |
| | "service" STRING NULL, |
| | "ts" TIMESTAMP(3) NOT NULL, |
| | TIME INDEX ("ts"), |
| | PRIMARY KEY ("host") |
| | ) |
| | |
| | ENGINE=mito |
| | |
+---------------------+----------------------------------------------------+
DROP TABLE combined_hint_table;
Affected Rows: 0
CREATE TABLE not_supported_table_storage_option (
`id` INT UNSIGNED,
host STRING,
@@ -61,6 +61,98 @@ SHOW CREATE TABLE alter_hint_table;
DROP TABLE alter_hint_table;
CREATE TABLE partition_num_hint_table (
host STRING PRIMARY KEY,
ts TIMESTAMP TIME INDEX
)
PARTITION ON COLUMNS (host) (
host < 'm',
host >= 'm'
)
WITH ('repartition.partition.num.hint' = ' 8 ');
SHOW CREATE TABLE partition_num_hint_table;
ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = 16;
SHOW CREATE TABLE partition_num_hint_table;
ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = '0';
ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint';
ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint';
SHOW CREATE TABLE partition_num_hint_table;
DROP TABLE partition_num_hint_table;
CREATE TABLE partition_num_hint_table (
host STRING PRIMARY KEY,
ts TIMESTAMP TIME INDEX
)
WITH ('repartition.partition.num.hint' = '10');
SHOW CREATE TABLE partition_num_hint_table;
DROP TABLE partition_num_hint_table;
CREATE TABLE partition_num_hint_table (
host STRING PRIMARY KEY,
ts TIMESTAMP TIME INDEX
)
WITH ('repartition.partition.num.hint' = 10);
SHOW CREATE TABLE partition_num_hint_table;
DROP TABLE partition_num_hint_table;
CREATE TABLE combined_hint_table (
host STRING PRIMARY KEY,
`service` STRING,
ts TIMESTAMP TIME INDEX
);
ALTER TABLE combined_hint_table SET
'repartition.column.hint' = 'host',
'repartition.partition.num.hint' = 10;
SHOW CREATE TABLE combined_hint_table;
ALTER TABLE combined_hint_table SET
'repartition.column.hint' = 'service',
'repartition.partition.num.hint' = 0;
SHOW CREATE TABLE combined_hint_table;
ALTER TABLE combined_hint_table SET
'repartition.partition.num.hint' = 20,
'repartition.column.hint' = 'missing';
SHOW CREATE TABLE combined_hint_table;
ALTER TABLE combined_hint_table SET
'repartition.partition.num.hint' = 20,
'repartition.column.hint' = 'service';
SHOW CREATE TABLE combined_hint_table;
ALTER TABLE combined_hint_table SET
'repartition.partition.num.hint' = 30,
'repartition.partition.num.hint' = 40;
ALTER TABLE combined_hint_table SET
'repartition.partition.num.hint' = 30,
'ttl' = '7d';
ALTER TABLE combined_hint_table UNSET
'repartition.column.hint',
'repartition.partition.num.hint';
SHOW CREATE TABLE combined_hint_table;
DROP TABLE combined_hint_table;
CREATE TABLE not_supported_table_storage_option (
`id` INT UNSIGNED,
host STRING,