Merge pull request #5718 from zhyass/feature_cluster_table

feat(query): alter table cluster key

common/ast/src/ast/statement.rs

@@ -294,6 +294,7 @@ pub enum AlterDatabaseAction<'a> {
 #[derive(Debug, Clone, PartialEq)]
 pub enum AlterTableAction<'a> {
     RenameTable { new_table: Identifier<'a> },
+    AlterClusterKey { cluster_by: Vec<Expr<'a>> },
     // TODO(wuzhiguo): AddColumn etc
 }
 
@@ -579,6 +580,10 @@ impl<'a> Display for Statement<'a> {
                     AlterTableAction::RenameTable { new_table } => {
                         write!(f, " RENAME TO {new_table}")?;
                     }
+                    AlterTableAction::AlterClusterKey { cluster_by } => {
+                        write!(f, " CLUSTER BY ")?;
+                        write_comma_separated_list(f, cluster_by)?;
+                    }
                 }
             }
             Statement::RenameTable {

common/ast/src/parser/statement.rs

@@ -567,15 +567,23 @@ pub fn alter_database_action(i: Input) -> IResult<AlterDatabaseAction> {
 }
 
 pub fn alter_table_action(i: Input) -> IResult<AlterTableAction> {
-    let mut rename_table = map(
+    let rename_table = map(
         rule! {
             RENAME ~ TO ~ #ident
         },
         |(_, _, new_table)| AlterTableAction::RenameTable { new_table },
     );
 
+    let cluster_by = map(
+        rule! {
+            CLUSTER ~ ^BY ~ ^"(" ~ ^#comma_separated_list1(expr) ~ ^")"
+        },
+        |(_, _, _, cluster_by, _)| AlterTableAction::AlterClusterKey { cluster_by },
+    );
+
     rule!(
         #rename_table
+        | #cluster_by
     )(i)
 }
 

common/meta/app/src/schema/table.rs

@@ -174,7 +174,11 @@ pub struct TableMeta {
     pub engine: String,
     pub engine_options: BTreeMap<String, String>,
     pub options: BTreeMap<String, String>,
-    pub cluster_keys: Option<String>,
+    pub cluster_key: Option<String>,
+    // The vector of cluster keys.
+    pub cluster_keys: Vec<String>,
+    // The default cluster keys id.
+    pub default_cluster_key_id: Option<u32>,
     pub created_on: DateTime<Utc>,
     pub updated_on: DateTime<Utc>,
     pub comment: String,
@@ -237,7 +241,9 @@ impl Default for TableMeta {
             engine: "".to_string(),
             engine_options: BTreeMap::new(),
             options: BTreeMap::new(),
-            cluster_keys: None,
+            cluster_key: None,
+            cluster_keys: vec![],
+            default_cluster_key_id: None,
             created_on: Default::default(),
             updated_on: Default::default(),
             comment: "".to_string(),
@@ -247,6 +253,19 @@ impl Default for TableMeta {
     }
 }
 
+impl TableMeta {
+    pub fn push_cluster_key(mut self, cluster_key: String) -> Self {
+        self.cluster_keys.push(cluster_key.clone());
+        self.cluster_key = Some(cluster_key);
+        self.default_cluster_key_id = Some(self.cluster_keys.len() as u32 - 1);
+        self
+    }
+
+    pub fn cluster_key(&self) -> Option<(u32, String)> {
+        self.default_cluster_key_id.zip(self.cluster_key.clone())
+    }
+}
+
 impl Display for TableMeta {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         write!(

common/planners/src/lib.rs

@@ -14,6 +14,7 @@
 
 mod plan_aggregator_final;
 mod plan_aggregator_partial;
+mod plan_alter_cluster_key;
 mod plan_broadcast;
 mod plan_call;
 mod plan_copy;
@@ -105,6 +106,7 @@ mod plan_view_drop;
 
 pub use plan_aggregator_final::AggregatorFinalPlan;
 pub use plan_aggregator_partial::AggregatorPartialPlan;
+pub use plan_alter_cluster_key::AlterClusterKeyPlan;
 pub use plan_broadcast::BroadcastPlan;
 pub use plan_call::CallPlan;
 pub use plan_copy::CopyMode;

common/planners/src/plan_alter_cluster_key.rs

@@ -0,0 +1,35 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use common_datavalues::DataSchema;
+use common_datavalues::DataSchemaRef;
+
+use crate::Expression;
+
+#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq)]
+pub struct AlterClusterKeyPlan {
+    pub tenant: String,
+    pub catalog_name: String,
+    pub database_name: String,
+    pub table_name: String,
+    pub cluster_keys: Vec<Expression>,
+}
+
+impl AlterClusterKeyPlan {
+    pub fn schema(&self) -> DataSchemaRef {
+        Arc::new(DataSchema::empty())
+    }
+}

common/planners/src/plan_node.rs

@@ -19,6 +19,7 @@ use common_datavalues::DataSchemaRef;
 use crate::plan_table_undrop::UnDropTablePlan;
 use crate::AggregatorFinalPlan;
 use crate::AggregatorPartialPlan;
+use crate::AlterClusterKeyPlan;
 use crate::AlterUserPlan;
 use crate::AlterUserUDFPlan;
 use crate::AlterViewPlan;
@@ -113,6 +114,9 @@ pub enum PlanNode {
     // List
     List(ListPlan),
 
+    // Alter.
+    AlterClusterKey(AlterClusterKeyPlan),
+
     // Show.
     Show(ShowPlan),
 
@@ -275,6 +279,9 @@ impl PlanNode {
 
             // Kill.
             PlanNode::Kill(v) => v.schema(),
+
+            // Alter
+            PlanNode::AlterClusterKey(v) => v.schema(),
         }
     }
 
@@ -376,6 +383,9 @@ impl PlanNode {
 
             // Kill.
             PlanNode::Kill(_) => "KillQuery",
+
+            // Alter.
+            PlanNode::AlterClusterKey(_) => "AlterClusterKeyPlan",
         }
     }
 

common/planners/src/plan_node_display_indent.rs

@@ -19,6 +19,7 @@ use common_datavalues::DataType;
 
 use crate::AggregatorFinalPlan;
 use crate::AggregatorPartialPlan;
+use crate::AlterClusterKeyPlan;
 use crate::BroadcastPlan;
 use crate::CallPlan;
 use crate::CopyPlan;
@@ -85,6 +86,7 @@ impl<'a> fmt::Display for PlanNodeIndentFormatDisplay<'a> {
             PlanNode::DropRole(plan) => Self::format_drop_role(f, plan),
             PlanNode::Copy(plan) => Self::format_copy(f, plan),
             PlanNode::Call(plan) => Self::format_call(f, plan),
+            PlanNode::AlterClusterKey(plan) => Self::format_alter_cluster_key(f, plan),
             _ => {
                 let mut printed = true;
 
@@ -356,4 +358,13 @@ impl<'a> PlanNodeIndentFormatDisplay<'a> {
         write!(f, "Call {:}", plan.name)?;
         write!(f, " args: {:?}", plan.args)
     }
+
+    fn format_alter_cluster_key(f: &mut Formatter, plan: &AlterClusterKeyPlan) -> fmt::Result {
+        write!(
+            f,
+            "Alter table {:}.{:}",
+            plan.database_name, plan.table_name
+        )?;
+        write!(f, " cluster by {:?}", plan.cluster_keys)
+    }
 }

common/planners/src/plan_node_rewriter.rs

@@ -26,6 +26,7 @@ use crate::plan_subqueries_set::SubQueriesSetPlan;
 use crate::plan_table_undrop::UnDropTablePlan;
 use crate::AggregatorFinalPlan;
 use crate::AggregatorPartialPlan;
+use crate::AlterClusterKeyPlan;
 use crate::AlterUserPlan;
 use crate::AlterUserUDFPlan;
 use crate::AlterViewPlan;
@@ -199,6 +200,9 @@ pub trait PlanRewriter: Sized {
 
             // Kill.
             PlanNode::Kill(plan) => self.rewrite_kill(plan),
+
+            // Alter.
+            PlanNode::AlterClusterKey(plan) => self.rewrite_alter_cluster_key(plan),
         }
     }
 
@@ -511,6 +515,10 @@ pub trait PlanRewriter: Sized {
     fn rewrite_alter_user_udf(&mut self, plan: &AlterUserUDFPlan) -> Result<PlanNode> {
         Ok(PlanNode::AlterUserUDF(plan.clone()))
     }
+
+    fn rewrite_alter_cluster_key(&mut self, plan: &AlterClusterKeyPlan) -> Result<PlanNode> {
+        Ok(PlanNode::AlterClusterKey(plan.clone()))
+    }
 }
 
 pub struct RewriteHelper {}

common/planners/src/plan_node_visitor.rs

@@ -19,6 +19,7 @@ use crate::plan_subqueries_set::SubQueriesSetPlan;
 use crate::plan_table_undrop::UnDropTablePlan;
 use crate::AggregatorFinalPlan;
 use crate::AggregatorPartialPlan;
+use crate::AlterClusterKeyPlan;
 use crate::AlterUserPlan;
 use crate::AlterUserUDFPlan;
 use crate::AlterViewPlan;
@@ -211,6 +212,9 @@ pub trait PlanVisitor {
 
             // Kill.
             PlanNode::Kill(plan) => self.visit_kill_query(plan),
+
+            // Alter.
+            PlanNode::AlterClusterKey(plan) => self.visit_alter_cluster_key(plan),
         }
     }
 
@@ -476,4 +480,8 @@ pub trait PlanVisitor {
     fn visit_alter_user_udf(&mut self, _: &AlterUserUDFPlan) -> Result<()> {
         Ok(())
     }
+
+    fn visit_alter_cluster_key(&mut self, _: &AlterClusterKeyPlan) -> Result<()> {
+        Ok(())
+    }
 }

common/proto-conv/src/table_from_to_protobuf_impl.rs

@@ -130,7 +130,9 @@ impl FromToProto<pb::TableMeta> for mt::TableMeta {
             engine: p.engine,
             engine_options: p.engine_options,
             options: p.options,
+            cluster_key: p.cluster_key,
             cluster_keys: p.cluster_keys,
+            default_cluster_key_id: p.default_cluster_key_id,
             created_on: DateTime::<Utc>::from_pb(p.created_on)?,
             updated_on: DateTime::<Utc>::from_pb(p.updated_on)?,
             drop_on: match p.drop_on {
@@ -155,7 +157,9 @@ impl FromToProto<pb::TableMeta> for mt::TableMeta {
             engine: self.engine.clone(),
             engine_options: self.engine_options.clone(),
             options: self.options.clone(),
+            cluster_key: self.cluster_key.clone(),
             cluster_keys: self.cluster_keys.clone(),
+            default_cluster_key_id: self.default_cluster_key_id,
             created_on: self.created_on.to_pb()?,
             updated_on: self.updated_on.to_pb()?,
             drop_on: match self.drop_on {

common/proto-conv/tests/it/proto_conv.rs

@@ -108,7 +108,9 @@ fn new_table_info() -> mt::TableInfo {
             engine: "44".to_string(),
             engine_options: btreemap! {s("abc") => s("def")},
             options: btreemap! {s("xyz") => s("foo")},
-            cluster_keys: Some("(a + 2, b)".to_string()),
+            cluster_key: Some("(a + 2, b)".to_string()),
+            cluster_keys: vec!["(a + 2, b)".to_string()],
+            default_cluster_key_id: Some(0),
             created_on: Utc.ymd(2014, 11, 28).and_hms(12, 0, 9),
             updated_on: Utc.ymd(2014, 11, 29).and_hms(12, 0, 10),
             comment: s("table_comment"),
@@ -329,7 +331,9 @@ fn test_load_old() -> anyhow::Result<()> {
                 engine: "44".to_string(),
                 engine_options: btreemap! {s("abc") => s("def")},
                 options: btreemap! {s("xyz") => s("foo")},
-                cluster_keys: Some("(a + 2, b)".to_string()),
+                cluster_key: Some("(a + 2, b)".to_string()),
+                cluster_keys: vec![],
+                default_cluster_key_id: None,
                 created_on: Utc.ymd(2014, 11, 28).and_hms(12, 0, 9),
                 updated_on: Utc.ymd(2014, 11, 29).and_hms(12, 0, 10),
                 comment: s("table_comment"),

common/protos/proto/metadata.proto

@@ -164,7 +164,13 @@ message TableMeta {
   map<string, string> options = 5;
 
   // Keys to sort rows in table.
-  optional string cluster_keys = 9;
+   optional string cluster_key = 9;
+
+  // The vector of cluster keys.
+  repeated string cluster_keys = 4;
+
+  // The default cluster keys id.
+  optional uint32 default_cluster_key_id = 8;
 
   // The time table created.
   string created_on = 20;