From 68b59e0e5e28b1d6e2e0a89f570545319c1edf86 Mon Sep 17 00:00:00 2001
From: Ruihang Xia <waynestxia@gmail.com>
Date: Sat, 31 Aug 2024 11:16:01 +0800
Subject: [PATCH] feat: remove the requirement that partition column must be PK
 (#4647)

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
---
 src/operator/src/error.rs            | 13 -------
 src/operator/src/statement/ddl.rs    | 52 +++-------------------------
 src/sql/src/parsers/create_parser.rs |  2 +-
 3 files changed, 6 insertions(+), 61 deletions(-)

diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs
index e7526c4092f..2473128137d 100644
--- a/src/operator/src/error.rs
+++ b/src/operator/src/error.rs
@@ -676,18 +676,6 @@ pub enum Error {
         location: Location,
     },
 
-    #[snafu(display(
-        "Invalid partition columns when creating table '{}', reason: {}",
-        table,
-        reason
-    ))]
-    InvalidPartitionColumns {
-        table: String,
-        reason: String,
-        #[snafu(implicit)]
-        location: Location,
-    },
-
     #[snafu(display("Failed to prepare file table"))]
     PrepareFileTable {
         #[snafu(implicit)]
@@ -812,7 +800,6 @@ impl ErrorExt for Error {
             | Error::ProjectSchema { .. }
             | Error::UnsupportedFormat { .. }
             | Error::ColumnNoneDefaultValue { .. }
-            | Error::InvalidPartitionColumns { .. }
             | Error::PrepareFileTable { .. }
             | Error::InferFileTableSchema { .. }
             | Error::SchemaIncompatible { .. }
diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs
index 6e646bc5fe5..afef6d590d5 100644
--- a/src/operator/src/statement/ddl.rs
+++ b/src/operator/src/statement/ddl.rs
@@ -70,11 +70,11 @@ use super::StatementExecutor;
 use crate::error::{
     self, AlterExprToRequestSnafu, CatalogSnafu, ColumnDataTypeSnafu, ColumnNotFoundSnafu,
     CreateLogicalTablesSnafu, CreateTableInfoSnafu, DeserializePartitionSnafu, EmptyDdlExprSnafu,
-    ExtractTableNamesSnafu, FlowNotFoundSnafu, InvalidPartitionColumnsSnafu,
-    InvalidPartitionRuleSnafu, InvalidPartitionSnafu, InvalidTableNameSnafu, InvalidViewNameSnafu,
-    InvalidViewStmtSnafu, ParseSqlValueSnafu, Result, SchemaInUseSnafu, SchemaNotFoundSnafu,
-    SchemaReadOnlySnafu, SubstraitCodecSnafu, TableAlreadyExistsSnafu, TableMetadataManagerSnafu,
-    TableNotFoundSnafu, UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu,
+    ExtractTableNamesSnafu, FlowNotFoundSnafu, InvalidPartitionRuleSnafu, InvalidPartitionSnafu,
+    InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu, ParseSqlValueSnafu, Result,
+    SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu, SubstraitCodecSnafu,
+    TableAlreadyExistsSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu,
+    UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu,
 };
 use crate::expr_factory;
 use crate::statement::show::create_partitions_stmt;
@@ -239,7 +239,6 @@ impl StatementExecutor {
         );
 
         let (partitions, partition_cols) = parse_partitions(create_table, partitions, &query_ctx)?;
-        validate_partition_columns(create_table, &partition_cols)?;
         let mut table_info = create_table_info(create_table, partition_cols, schema_opts)?;
 
         let resp = self
@@ -1209,22 +1208,6 @@ impl StatementExecutor {
     }
 }
 
-fn validate_partition_columns(
-    create_table: &CreateTableExpr,
-    partition_cols: &[String],
-) -> Result<()> {
-    ensure!(
-        partition_cols
-            .iter()
-            .all(|col| &create_table.time_index == col || create_table.primary_keys.contains(col)),
-        InvalidPartitionColumnsSnafu {
-            table: &create_table.table_name,
-            reason: "partition column must belongs to primary keys or equals to time index"
-        }
-    );
-    Ok(())
-}
-
 /// Parse partition statement [Partitions] into [MetaPartition] and partition columns.
 fn parse_partitions(
     create_table: &CreateTableExpr,
@@ -1519,31 +1502,6 @@ mod test {
         assert!(NAME_PATTERN_REG.is_match("hello"));
     }
 
-    #[test]
-    fn test_validate_partition_columns() {
-        let create_table = CreateTableExpr {
-            table_name: "my_table".to_string(),
-            time_index: "ts".to_string(),
-            primary_keys: vec!["a".to_string(), "b".to_string()],
-            ..Default::default()
-        };
-
-        assert!(validate_partition_columns(&create_table, &[]).is_ok());
-        assert!(validate_partition_columns(&create_table, &["ts".to_string()]).is_ok());
-        assert!(validate_partition_columns(&create_table, &["a".to_string()]).is_ok());
-        assert!(
-            validate_partition_columns(&create_table, &["b".to_string(), "a".to_string()]).is_ok()
-        );
-
-        assert_eq!(
-            validate_partition_columns(&create_table, &["a".to_string(), "c".to_string()])
-                .unwrap_err()
-                .to_string(),
-            "Invalid partition columns when creating table 'my_table', \
-             reason: partition column must belongs to primary keys or equals to time index",
-        );
-    }
-
     #[tokio::test]
     #[ignore = "TODO(ruihang): WIP new partition rule"]
     async fn test_parse_partitions() {
diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs
index 98638df3f21..de81bd1a214 100644
--- a/src/sql/src/parsers/create_parser.rs
+++ b/src/sql/src/parsers/create_parser.rs
@@ -380,7 +380,7 @@ impl<'a> ParserContext<'a> {
         Ok(options.into())
     }
 
-    /// "PARTITION BY ..." clause
+    /// "PARTITION ON COLUMNS (...)" clause
     fn parse_partitions(&mut self) -> Result<Option<Partitions>> {
         if !self.parser.parse_keyword(Keyword::PARTITION) {
             return Ok(None);