Add support for cluster by expressions (apache#1883)

Co-authored-by: osipovartem <[email protected]>

Author: osipovartem

src/ast/dml.rs

@@ -175,7 +175,9 @@ pub struct CreateTable {
     pub partition_by: Option<Box<Expr>>,
     /// BigQuery: Table clustering column list.
     /// <https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#table_option_list>
-    pub cluster_by: Option<WrappedCollection<Vec<Ident>>>,
+    /// Snowflake: Table clustering list which contains base column, expressions on base columns.
+    /// <https://docs.snowflake.com/en/user-guide/tables-clustering-keys#defining-a-clustering-key-for-a-table>
+    pub cluster_by: Option<WrappedCollection<Vec<Expr>>>,
     /// Hive: Table clustering column list.
     /// <https://cwiki.apache.org/confluence/display/Hive/LanguageManual+DDL#LanguageManualDDL-CreateTable>
     pub clustered_by: Option<ClusteredBy>,

src/ast/helpers/stmt_create_table.rs

@@ -90,7 +90,7 @@ pub struct CreateTableBuilder {
     pub primary_key: Option<Box<Expr>>,
     pub order_by: Option<OneOrManyWithParens<Expr>>,
     pub partition_by: Option<Box<Expr>>,
-    pub cluster_by: Option<WrappedCollection<Vec<Ident>>>,
+    pub cluster_by: Option<WrappedCollection<Vec<Expr>>>,
     pub clustered_by: Option<ClusteredBy>,
     pub inherits: Option<Vec<ObjectName>>,
     pub strict: bool,
@@ -279,7 +279,7 @@ impl CreateTableBuilder {
         self
     }
 
-    pub fn cluster_by(mut self, cluster_by: Option<WrappedCollection<Vec<Ident>>>) -> Self {
+    pub fn cluster_by(mut self, cluster_by: Option<WrappedCollection<Vec<Expr>>>) -> Self {
         self.cluster_by = cluster_by;
         self
     }
@@ -542,7 +542,7 @@ impl TryFrom<Statement> for CreateTableBuilder {
 #[derive(Default)]
 pub(crate) struct CreateTableConfiguration {
     pub partition_by: Option<Box<Expr>>,
-    pub cluster_by: Option<WrappedCollection<Vec<Ident>>>,
+    pub cluster_by: Option<WrappedCollection<Vec<Expr>>>,
     pub inherits: Option<Vec<ObjectName>>,
     pub table_options: CreateTableOptions,
 }

src/dialect/snowflake.rs

@@ -453,7 +453,7 @@ pub fn parse_create_table(
                     parser.expect_keyword_is(Keyword::BY)?;
                     parser.expect_token(&Token::LParen)?;
                     let cluster_by = Some(WrappedCollection::Parentheses(
-                        parser.parse_comma_separated(|p| p.parse_identifier())?,
+                        parser.parse_comma_separated(|p| p.parse_expr())?,
                     ));
                     parser.expect_token(&Token::RParen)?;
 

src/parser/mod.rs

@@ -7316,7 +7316,7 @@ impl<'a> Parser<'a> {
         if dialect_of!(self is BigQueryDialect | GenericDialect) {
             if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) {
                 cluster_by = Some(WrappedCollection::NoWrapping(
-                    self.parse_comma_separated(|p| p.parse_identifier())?,
+                    self.parse_comma_separated(|p| p.parse_expr())?,
                 ));
             };
 

tests/sqlparser_bigquery.rs

@@ -536,8 +536,8 @@ fn parse_create_table_with_options() {
                 (
                     Some(Box::new(Expr::Identifier(Ident::new("_PARTITIONDATE")))),
                     Some(WrappedCollection::NoWrapping(vec![
-                        Ident::new("userid"),
-                        Ident::new("age"),
+                        Expr::Identifier(Ident::new("userid")),
+                        Expr::Identifier(Ident::new("age")),
                     ])),
                     CreateTableOptions::Options(vec![
                         SqlOption::KeyValue {

tests/sqlparser_snowflake.rs

@@ -471,15 +471,31 @@ fn test_snowflake_create_table_if_not_exists() {
 
 #[test]
 fn test_snowflake_create_table_cluster_by() {
-    match snowflake().verified_stmt("CREATE TABLE my_table (a INT) CLUSTER BY (a, b)") {
+    match snowflake().verified_stmt("CREATE TABLE my_table (a INT) CLUSTER BY (a, b, my_func(c))") {
         Statement::CreateTable(CreateTable {
             name, cluster_by, ..
         }) => {
             assert_eq!("my_table", name.to_string());
             assert_eq!(
                 Some(WrappedCollection::Parentheses(vec![
-                    Ident::new("a"),
-                    Ident::new("b"),
+                    Expr::Identifier(Ident::new("a")),
+                    Expr::Identifier(Ident::new("b")),
+                    Expr::Function(Function {
+                        name: ObjectName::from(vec![Ident::new("my_func")]),
+                        uses_odbc_syntax: false,
+                        parameters: FunctionArguments::None,
+                        args: FunctionArguments::List(FunctionArgumentList {
+                            args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                                Expr::Identifier(Ident::new("c"))
+                            ))],
+                            duplicate_treatment: None,
+                            clauses: vec![],
+                        }),
+                        filter: None,
+                        null_treatment: None,
+                        over: None,
+                        within_group: vec![],
+                    }),
                 ])),
                 cluster_by
             )
@@ -903,8 +919,8 @@ fn test_snowflake_create_iceberg_table_all_options() {
             assert_eq!("my_table", name.to_string());
             assert_eq!(
                 Some(WrappedCollection::Parentheses(vec![
-                    Ident::new("a"),
-                    Ident::new("b"),
+                    Expr::Identifier(Ident::new("a")),
+                    Expr::Identifier(Ident::new("b")),
                 ])),
                 cluster_by
             );