From d929c5e9e9535da585fd163ade8a707f12170d94 Mon Sep 17 00:00:00 2001
From: osipovartem <artem@PC.localdomain>
Date: Thu, 12 Jun 2025 13:09:43 +0300
Subject: [PATCH] Add support for cluster by expressions

---
 src/ast/dml.rs                       |  4 +++-
 src/ast/helpers/stmt_create_table.rs |  6 +++---
 src/dialect/snowflake.rs             |  2 +-
 src/parser/mod.rs                    |  2 +-
 tests/sqlparser_bigquery.rs          |  4 ++--
 tests/sqlparser_snowflake.rs         | 26 +++++++++++++++++++++-----
 6 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/src/ast/dml.rs b/src/ast/dml.rs
index da82a4ede..292650c88 100644
--- a/src/ast/dml.rs
+++ b/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>,
diff --git a/src/ast/helpers/stmt_create_table.rs b/src/ast/helpers/stmt_create_table.rs
index 542d30ea9..d66a869bf 100644
--- a/src/ast/helpers/stmt_create_table.rs
+++ b/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,
 }
diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs
index 990e2ea29..ea0b94a63 100644
--- a/src/dialect/snowflake.rs
+++ b/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)?;
 
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 2c208e2e5..73cc3e0ed 100644
--- a/src/parser/mod.rs
+++ b/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())?,
                 ));
             };
 
diff --git a/tests/sqlparser_bigquery.rs b/tests/sqlparser_bigquery.rs
index b64f190f6..6a303577f 100644
--- a/tests/sqlparser_bigquery.rs
+++ b/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 {
diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs
index 7a164c248..7d734ca26 100644
--- a/tests/sqlparser_snowflake.rs
+++ b/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
             );