DuckDB: Support CREATE TABLE PARTITIONED BY expressions

src/ast/ddl.rs

@@ -2991,6 +2991,8 @@ pub struct CreateTable {
     /// Hive: Table clustering column list.
     /// <https://cwiki.apache.org/confluence/display/Hive/LanguageManual+DDL#LanguageManualDDL-CreateTable>
     pub clustered_by: Option<ClusteredBy>,
+    /// DuckDB partition expressions, distinct from Hive partition columns.
+    pub partitioned_by: Option<Vec<Expr>>,
     /// Postgres `INHERITs` clause, which contains the list of tables from which
     /// the new table inherits.
     /// <https://www.postgresql.org/docs/current/ddl-inherit.html>
@@ -3215,6 +3217,14 @@ impl fmt::Display for CreateTable {
             _ => (),
         }
 
+        if let Some(partitioned_by) = &self.partitioned_by {
+            write!(
+                f,
+                " PARTITIONED BY ({})",
+                display_comma_separated(partitioned_by)
+            )?;
+        }
+
         if let Some(clustered_by) = &self.clustered_by {
             write!(f, " {clustered_by}")?;
         }

src/ast/helpers/stmt_create_table.rs

@@ -127,6 +127,8 @@ pub struct CreateTableBuilder {
     pub cluster_by: Option<WrappedCollection<Vec<Expr>>>,
     /// Optional `CLUSTERED BY` clause.
     pub clustered_by: Option<ClusteredBy>,
+    /// Optional expression-based `PARTITIONED BY` clause.
+    pub partitioned_by: Option<Vec<Expr>>,
     /// Optional parent tables (`INHERITS`).
     pub inherits: Option<Vec<ObjectName>>,
     /// Optional partitioned table (`PARTITION OF`)
@@ -231,6 +233,7 @@ impl CreateTableBuilder {
             partition_by: None,
             cluster_by: None,
             clustered_by: None,
+            partitioned_by: None,
             inherits: None,
             partition_of: None,
             for_values: None,
@@ -417,6 +420,11 @@ impl CreateTableBuilder {
         self.clustered_by = clustered_by;
         self
     }
+    /// Set expression-based partitioning.
+    pub fn partitioned_by(mut self, partitioned_by: Option<Vec<Expr>>) -> Self {
+        self.partitioned_by = partitioned_by;
+        self
+    }
     /// Set parent tables via `INHERITS`.
     pub fn inherits(mut self, inherits: Option<Vec<ObjectName>>) -> Self {
         self.inherits = inherits;
@@ -632,6 +640,7 @@ impl CreateTableBuilder {
             partition_by: self.partition_by,
             cluster_by: self.cluster_by,
             clustered_by: self.clustered_by,
+            partitioned_by: self.partitioned_by,
             inherits: self.inherits,
             partition_of: self.partition_of,
             for_values: self.for_values,
@@ -718,6 +727,7 @@ impl From<CreateTable> for CreateTableBuilder {
             partition_by: table.partition_by,
             cluster_by: table.cluster_by,
             clustered_by: table.clustered_by,
+            partitioned_by: table.partitioned_by,
             inherits: table.inherits,
             partition_of: table.partition_of,
             for_values: table.for_values,

src/ast/spans.rs

@@ -580,7 +580,8 @@ impl Spanned for CreateTable {
             partition_by: _, // todo, BigQuery specific
             cluster_by: _,   // todo, BigQuery specific
             clustered_by: _, // todo, Hive specific
-            inherits: _,     // todo, PostgreSQL specific
+            partitioned_by,
+            inherits: _, // todo, PostgreSQL specific
             partition_of,
             for_values,
             strict: _,                          // bool
@@ -624,7 +625,8 @@ impl Spanned for CreateTable {
                 .chain(query.iter().map(|i| i.span()))
                 .chain(clone.iter().map(|i| i.span()))
                 .chain(partition_of.iter().map(|i| i.span()))
-                .chain(for_values.iter().map(|i| i.span())),
+                .chain(for_values.iter().map(|i| i.span()))
+                .chain(partitioned_by.iter().flatten().map(|i| i.span())),
         )
     }
 }

src/dialect/duckdb.rs

@@ -24,6 +24,10 @@ pub struct DuckDbDialect;
 
 // In most cases the redshift dialect is identical to [`PostgresSqlDialect`].
 impl Dialect for DuckDbDialect {
+    fn supports_create_table_partitioned_by_expressions(&self) -> bool {
+        true
+    }
+
     fn supports_trailing_commas(&self) -> bool {
         true
     }

src/dialect/mod.rs

@@ -763,6 +763,11 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Uses expressions rather than column declarations in `PARTITIONED BY`.
+    fn supports_create_table_partitioned_by_expressions(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports MySQL-specific SELECT modifiers
     /// like `HIGH_PRIORITY`, `STRAIGHT_JOIN`, `SQL_SMALL_RESULT`, etc.
     ///

src/parser/mod.rs

@@ -8756,6 +8756,7 @@ impl<'a> Parser<'a> {
         };
 
         // parse optional column list (schema)
+        let has_columns = self.peek_token_ref().token == Token::LParen;
         let (columns, constraints) = self.parse_columns()?;
         let comment_after_column_def =
             if dialect_of!(self is HiveDialect) && self.parse_keyword(Keyword::COMMENT) {
@@ -8785,7 +8786,22 @@ impl<'a> Parser<'a> {
         // SQLite supports `WITHOUT ROWID` at the end of `CREATE TABLE`
         let without_rowid = self.parse_keywords(&[Keyword::WITHOUT, Keyword::ROWID]);
 
-        let hive_distribution = self.parse_hive_distribution()?;
+        let (partitioned_by, hive_distribution) = if self
+            .dialect
+            .supports_create_table_partitioned_by_expressions()
+        {
+            let expressions = if self.parse_keywords(&[Keyword::PARTITIONED, Keyword::BY]) {
+                self.expect_token(&Token::LParen)?;
+                let expressions = self.parse_comma_separated(Parser::parse_expr)?;
+                self.expect_token(&Token::RParen)?;
+                Some(expressions)
+            } else {
+                None
+            };
+            (expressions, HiveDistributionStyle::NONE)
+        } else {
+            (None, self.parse_hive_distribution()?)
+        };
         let clustered_by = self.parse_optional_clustered_by()?;
         let hive_formats = self.parse_hive_formats()?;
 
@@ -8879,6 +8895,10 @@ impl<'a> Parser<'a> {
             None
         };
 
+        if query.is_none() && !has_columns && partitioned_by.is_some() {
+            return self.expected_ref("AS query or a table schema", self.peek_token_ref());
+        }
+
         // `WITH DATA` clause only applies if there is a query body.
         let with_data = if query.is_some() {
             self.maybe_parse_with_data()?
@@ -8909,6 +8929,7 @@ impl<'a> Parser<'a> {
             .on_commit(on_commit)
             .on_cluster(on_cluster)
             .clustered_by(clustered_by)
+            .partitioned_by(partitioned_by)
             .partition_by(partition_by)
             .cluster_by(create_table_config.cluster_by)
             .inherits(create_table_config.inherits)

tests/sqlparser_common.rs

@@ -20015,3 +20015,25 @@ fn parse_function_arg_call_chain_no_exponential_blowup() {
     rx.recv_timeout(Duration::from_secs(5))
         .expect("parser should reject this quickly, not loop exponentially");
 }
+
+#[test]
+fn create_table_partitioned_by_dialect_isolation() {
+    let sql = "CREATE TABLE t (id INTEGER) PARTITIONED BY (id + 1)";
+    all_dialects_where(|dialect| dialect.supports_create_table_partitioned_by_expressions())
+        .verified_stmt(sql);
+    all_dialects_where(|dialect| !dialect.supports_create_table_partitioned_by_expressions())
+        .one_of_identical_results(|dialect| assert!(Parser::parse_sql(dialect, sql).is_err()));
+    let hive = TestedDialects::new(vec![Box::new(HiveDialect {}), Box::new(GenericDialect {})]);
+    let Statement::CreateTable(table) = hive.verified_stmt(
+        "CREATE TABLE t (id INT) PARTITIONED BY (category STRING) CLUSTERED BY (id) SORTED BY (id DESC) INTO 4 BUCKETS",
+    ) else {
+        unreachable!()
+    };
+    let HiveDistributionStyle::PARTITIONED { columns } = table.hive_distribution else {
+        unreachable!()
+    };
+    assert_eq!(columns[0].name, Ident::new("category"));
+    assert_eq!(columns[0].data_type, DataType::String(None));
+    assert!(table.partitioned_by.is_none());
+    assert!(table.clustered_by.unwrap().sorted_by.is_some());
+}

tests/sqlparser_duckdb.rs

@@ -764,6 +764,7 @@ fn test_duckdb_union_datatype() {
             partition_by: Default::default(),
             cluster_by: Default::default(),
             clustered_by: Default::default(),
+            partitioned_by: None,
             inherits: Default::default(),
             partition_of: Default::default(),
             for_values: Default::default(),
@@ -910,3 +911,173 @@ fn test_duckdb_lambda_function() {
     let sql_transform = "SELECT list_transform([1, 2, 3], lambda x : x * 2)";
     duckdb().verified_stmt(sql_transform);
 }
+
+#[test]
+fn create_table_partitioned_by_expressions() {
+    for sql in [
+        "CREATE TABLE events (id INTEGER) PARTITIONED BY (id + 1)",
+        "CREATE TABLE events (id INTEGER) PARTITIONED BY (abs(id), id % 2)",
+        "CREATE TABLE events PARTITIONED BY (id + 1) AS SELECT 1 AS id",
+        "CREATE TABLE events (id INTEGER) PARTITIONED BY (id + 1) WITH (flag = true)",
+    ] {
+        let statement = duckdb().verified_stmt(sql);
+        assert_eq!(statement, duckdb().verified_stmt(&statement.to_string()));
+    }
+    duckdb().one_statement_parses_to(
+        "CREATE TABLE events (id INTEGER) PARTITIONED /* before BY */ BY (id + 1,);",
+        "CREATE TABLE events (id INTEGER) PARTITIONED BY (id + 1)",
+    );
+}
+
+#[test]
+fn create_table_partitioned_by_errors() {
+    for (sql, expected) in [
+        (
+            "CREATE TABLE events (id INTEGER) PARTITIONED BY ()",
+            "Expected: an expression, found: )",
+        ),
+        (
+            "CREATE TABLE events (id INTEGER) PARTITIONED BY id",
+            "Expected: (, found: id",
+        ),
+        (
+            "CREATE TABLE events (id INTEGER) PARTITIONED BY (id",
+            "Expected: ), found: EOF",
+        ),
+        (
+            "CREATE TABLE events (id INTEGER) PARTITIONED BY (id INTEGER)",
+            "Expected: ), found: INTEGER",
+        ),
+        (
+            "CREATE TABLE events (id INTEGER) PARTITIONED BY (id) PARTITIONED BY (id)",
+            "Expected: end of statement, found: PARTITIONED",
+        ),
+        (
+            "CREATE TABLE events PARTITIONED BY (id)",
+            "Expected: AS query or a table schema, found: EOF",
+        ),
+        (
+            "CREATE TABLE events (id INTEGER) PARTITIONED BY (id ASC)",
+            "Expected: ), found: ASC",
+        ),
+        (
+            "CREATE TABLE events (id INTEGER) PARTITIONED BY (id DESC NULLS LAST)",
+            "Expected: ), found: DESC",
+        ),
+        (
+            "CREATE TABLE events (id INTEGER) PARTITIONED BY (id AS alias)",
+            "Expected: ), found: AS",
+        ),
+        (
+            "CREATE TABLE events (id INTEGER) WITH (flag = true) PARTITIONED BY (id)",
+            "Expected: end of statement, found: PARTITIONED",
+        ),
+    ] {
+        assert_eq!(
+            duckdb().parse_sql_statements(sql).unwrap_err(),
+            ParserError::ParserError(expected.to_owned()),
+            "{sql}"
+        );
+    }
+}
+
+#[test]
+fn create_table_partitioned_by_ast_and_builder() {
+    let sql = "CREATE TABLE events (id INTEGER) PARTITIONED BY (id + 1)";
+    let Statement::CreateTable(table) = duckdb().verified_stmt(sql) else {
+        unreachable!()
+    };
+    let expressions = vec![Expr::BinaryOp {
+        left: Box::new(Expr::Identifier(Ident::new("id"))),
+        op: BinaryOperator::Plus,
+        right: Box::new(Expr::Value(
+            Value::Number("1".parse().unwrap(), false).into(),
+        )),
+    }];
+    assert_eq!(table.partitioned_by, Some(expressions.clone()));
+    assert_eq!(table.hive_distribution, HiveDistributionStyle::NONE);
+    assert!(table.partition_by.is_none());
+    let rebuilt = helpers::stmt_create_table::CreateTableBuilder::from(table.clone()).build();
+    assert_eq!(rebuilt, table);
+    let built = helpers::stmt_create_table::CreateTableBuilder::new(table.name.clone())
+        .columns(table.columns.clone())
+        .partitioned_by(Some(expressions.clone()))
+        .build();
+    assert_eq!(built.partitioned_by, Some(expressions));
+    assert_eq!(built.to_string(), sql);
+    let Statement::CreateTable(unpartitioned) =
+        duckdb().verified_stmt("CREATE TABLE events (id INTEGER)")
+    else {
+        unreachable!()
+    };
+    assert!(unpartitioned.partitioned_by.is_none());
+}
+
+#[test]
+fn create_table_partitioned_by_spans() {
+    use sqlparser::parser::Parser;
+    use sqlparser::tokenizer::Location;
+    let sql = "CREATE TABLE events (id INTEGER) PARTITIONED BY (id + 2)";
+    let Statement::CreateTable(table) =
+        Parser::parse_sql(&DuckDbDialect {}, sql).unwrap().remove(0)
+    else {
+        unreachable!()
+    };
+    let expression_start = sql.find("id + 2").unwrap() as u64 + 1;
+    let expression_end = sql.len() as u64;
+    assert_eq!(
+        table.partitioned_by.as_ref().unwrap()[0].span(),
+        Span::new(
+            Location::new(1, expression_start),
+            Location::new(1, expression_end)
+        )
+    );
+    assert_eq!(
+        table.span(),
+        Span::new(Location::new(1, 14), Location::new(1, expression_end))
+    );
+}
+
+#[test]
+#[cfg(feature = "json_example")]
+fn create_table_partitioned_by_serialization() {
+    let statement =
+        duckdb().verified_stmt("CREATE TABLE events (id INTEGER) PARTITIONED BY (id + 1)");
+    let json = serde_json::to_string(&statement).unwrap();
+    assert_eq!(statement, serde_json::from_str::<Statement>(&json).unwrap());
+    let mut table = serde_json::to_value(
+        helpers::stmt_create_table::CreateTableBuilder::new(Ident::new("events").into()).build(),
+    )
+    .unwrap();
+    table.as_object_mut().unwrap().remove("partitioned_by");
+    assert!(serde_json::from_value::<CreateTable>(table)
+        .unwrap()
+        .partitioned_by
+        .is_none());
+}
+
+#[test]
+#[cfg(feature = "visitor")]
+fn create_table_partitioned_by_visitors() {
+    use core::ops::ControlFlow;
+    let mut statement =
+        duckdb().verified_stmt("CREATE TABLE events (id INTEGER) PARTITIONED BY (id + 2)");
+    let mut expressions = vec![];
+    let _ = visit_expressions(&statement, |expression| {
+        expressions.push(expression.to_string());
+        ControlFlow::<()>::Continue(())
+    });
+    assert_eq!(expressions, ["id + 2", "id", "2"]);
+    let _ = visit_expressions_mut(&mut statement, |expression| {
+        if let Expr::Identifier(ident) = expression {
+            if ident.value == "id" {
+                ident.value = "value".to_owned();
+            }
+        }
+        ControlFlow::<()>::Continue(())
+    });
+    assert_eq!(
+        statement.to_string(),
+        "CREATE TABLE events (id INTEGER) PARTITIONED BY (value + 2)"
+    );
+}

tests/sqlparser_mssql.rs

@@ -1982,6 +1982,7 @@ fn parse_create_table_with_valid_options() {
                 partition_by: None,
                 cluster_by: None,
                 clustered_by: None,
+                partitioned_by: None,
                 inherits: None,
                 partition_of: None,
                 for_values: None,
@@ -2163,6 +2164,7 @@ fn parse_create_table_with_identity_column() {
                 partition_by: None,
                 cluster_by: None,
                 clustered_by: None,
+                partitioned_by: None,
                 inherits: None,
                 partition_of: None,
                 for_values: None,

tests/sqlparser_postgres.rs

@@ -7031,6 +7031,7 @@ fn parse_trigger_related_functions() {
             partition_by: None,
             cluster_by: None,
             clustered_by: None,
+            partitioned_by: None,
             inherits: None,
             partition_of: None,
             for_values: None,