feat: repartition grammar candy (#7518)

* feat: repartition grammar candy

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* align keyword

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2026-01-06 12:44:13 +08:00
committed by GitHub
parent 522ca99cd6
commit 5162c1de4d
3 changed files with 167 additions and 3 deletions

View File

@@ -20,7 +20,7 @@ use std::collections::HashMap;
use common_query::AddColumnLocation;
use datatypes::schema::COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE;
use snafu::{ResultExt, ensure};
use sqlparser::ast::{Expr, Ident};
use sqlparser::ast::{BinaryOperator, Expr, Ident};
use sqlparser::keywords::Keyword;
use sqlparser::parser::{Parser, ParserError};
use sqlparser::tokenizer::{Token, TokenWithSpan};
@@ -127,6 +127,10 @@ impl ParserContext<'_> {
self.parse_alter_table_unset()?
} else if w.value.eq_ignore_ascii_case("REPARTITION") {
self.parse_alter_table_repartition()?
} else if w.value.eq_ignore_ascii_case("SPLIT") {
self.parse_alter_table_split_partition()?
} else if w.value.eq_ignore_ascii_case("MERGE") {
self.parse_alter_table_merge_partition()?
} else {
match w.keyword {
Keyword::ADD => self.parse_alter_table_add()?,
@@ -169,7 +173,7 @@ impl ParserContext<'_> {
AlterTableOperation::SetTableOptions { options }
}
_ => self.expected(
"ADD or DROP or MODIFY or RENAME or SET or REPARTITION after ALTER TABLE",
"ADD or DROP or MODIFY or RENAME or SET or REPARTITION or SPLIT or MERGE after ALTER TABLE",
self.parser.peek_token(),
)?,
}
@@ -210,6 +214,63 @@ impl ParserContext<'_> {
})
}
fn parse_alter_table_split_partition(&mut self) -> Result<AlterTableOperation> {
let _ = self.parser.next_token();
self.parser
.expect_keyword(Keyword::PARTITION)
.context(error::SyntaxSnafu)?;
let from_exprs = self.parse_repartition_expr_list()?;
if from_exprs.len() != 1 {
return self.expected(
"single partition expression inside SPLIT PARTITION clause",
self.parser.peek_token(),
);
}
self.parser
.expect_keyword(Keyword::INTO)
.context(error::SyntaxSnafu)?;
let into_exprs = self.parse_repartition_expr_list()?;
if matches!(self.parser.peek_token().token, Token::Comma) {
return self.expected("end of SPLIT PARTITION clause", self.parser.peek_token());
}
Ok(AlterTableOperation::Repartition {
operation: RepartitionOperation::new(from_exprs, into_exprs),
})
}
fn parse_alter_table_merge_partition(&mut self) -> Result<AlterTableOperation> {
let _ = self.parser.next_token();
self.parser
.expect_keyword(Keyword::PARTITION)
.context(error::SyntaxSnafu)?;
let from_exprs = self.parse_repartition_expr_list()?;
let mut expr_iter = from_exprs.iter().cloned();
let Some(first) = expr_iter.next() else {
return self.expected(
"expression inside MERGE PARTITION clause",
self.parser.peek_token(),
);
};
let merged_expr = expr_iter.fold(first, |left, right| Expr::BinaryOp {
left: Box::new(left),
op: BinaryOperator::Or,
right: Box::new(right),
});
if matches!(self.parser.peek_token().token, Token::Comma) {
return self.expected("end of MERGE PARTITION clause", self.parser.peek_token());
}
Ok(AlterTableOperation::Repartition {
operation: RepartitionOperation::new(from_exprs, vec![merged_expr]),
})
}
fn parse_repartition_expr_list(&mut self) -> Result<Vec<Expr>> {
self.parser
.expect_token(&Token::LParen)
@@ -911,6 +972,77 @@ ALTER TABLE t REPARTITION (
}
}
#[test]
fn test_parse_alter_table_split_partition() {
let sql = r#"
ALTER TABLE metrics SPLIT PARTITION (
device_id < 100
) 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.into_exprs[0].to_string(),
"device_id < 100 AND area < 'South'"
);
assert_eq!(
operation.into_exprs[1].to_string(),
"device_id < 100 AND area >= 'South'"
);
}
}
}
#[test]
fn test_parse_alter_table_merge_partition() {
let sql = r#"
ALTER TABLE metrics MERGE PARTITION (
device_id < 100,
device_id >= 100
);"#;
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(), 2);
assert_eq!(operation.from_exprs[0].to_string(), "device_id < 100");
assert_eq!(operation.from_exprs[1].to_string(), "device_id >= 100");
assert_eq!(operation.into_exprs.len(), 1);
assert_eq!(
operation.into_exprs[0].to_string(),
"device_id < 100 OR device_id >= 100"
);
}
}
}
#[test]
fn test_parse_alter_table_repartition_multiple() {
let sql = r#"
@@ -1094,7 +1226,7 @@ ALTER TABLE metrics REPARTITION
let err = result.output_msg();
assert_eq!(
err,
"Invalid SQL syntax: sql parser error: Expected ADD or DROP or MODIFY or RENAME or SET or REPARTITION after ALTER TABLE, found: table_t"
"Invalid SQL syntax: sql parser error: Expected ADD or DROP or MODIFY or RENAME or SET or REPARTITION or SPLIT or MERGE after ALTER TABLE, found: table_t"
);
let sql = "ALTER TABLE test_table RENAME table_t";