Support distribution styles for redshift.

src/ast/helpers/stmt_create_table.rs

@@ -1,7 +1,7 @@
 #[cfg(not(feature = "std"))]
 use alloc::{boxed::Box, format, string::String, vec, vec::Vec};
 
-use crate::ast::WithSpan;
+use crate::ast::{SortKey, WithSpan};
 
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
@@ -10,8 +10,8 @@ use serde::{Deserialize, Serialize};
 use sqlparser_derive::{Visit, VisitMut};
 
 use crate::ast::{
-    ColumnDef, EngineSpec, Expr, FileFormat, HiveDistributionStyle, HiveFormat, Ident, ObjectName,
-    OnCommit, Query, SqlOption, Statement, TableConstraint,
+    ColumnDef, DistributionStyle, EngineSpec, Expr, FileFormat, HiveDistributionStyle, HiveFormat,
+    Ident, ObjectName, OnCommit, Query, SqlOption, Statement, TableConstraint,
 };
 use crate::parser::ParserError;
 
@@ -58,6 +58,9 @@ pub struct CreateTableBuilder {
     pub constraints: Vec<TableConstraint>,
     pub hive_distribution: HiveDistributionStyle,
     pub hive_formats: Option<HiveFormat>,
+    pub dist_style: Option<DistributionStyle>,
+    pub dist_key: Option<WithSpan<Ident>>,
+    pub sort_key: Option<SortKey>,
     pub table_options: Vec<SqlOption>,
     pub table_properties: Vec<SqlOption>,
     pub with_options: Vec<SqlOption>,
@@ -98,6 +101,9 @@ impl CreateTableBuilder {
             constraints: vec![],
             hive_distribution: HiveDistributionStyle::NONE,
             hive_formats: None,
+            dist_style: None,
+            dist_key: None,
+            sort_key: None,
             table_options: vec![],
             table_properties: vec![],
             with_options: vec![],
@@ -174,6 +180,21 @@ impl CreateTableBuilder {
         self
     }
 
+    pub fn dist_style(mut self, dist_style: Option<DistributionStyle>) -> Self {
+        self.dist_style = dist_style;
+        self
+    }
+
+    pub fn dist_key(mut self, dist_key: Option<WithSpan<Ident>>) -> Self {
+        self.dist_key = dist_key;
+        self
+    }
+
+    pub fn sort_key(mut self, sort_key: Option<SortKey>) -> Self {
+        self.sort_key = sort_key;
+        self
+    }
+
     pub fn table_options(mut self, table_options: Vec<SqlOption>) -> Self {
         self.table_options = table_options;
         self
@@ -305,6 +326,9 @@ impl CreateTableBuilder {
             constraints: self.constraints,
             hive_distribution: self.hive_distribution,
             hive_formats: self.hive_formats,
+            dist_style: self.dist_style,
+            dist_key: self.dist_key,
+            sort_key: self.sort_key,
             table_options: self.table_options,
             table_properties: self.table_properties,
             with_options: self.with_options,
@@ -352,6 +376,9 @@ impl TryFrom<Statement> for CreateTableBuilder {
                 constraints,
                 hive_distribution,
                 hive_formats,
+                dist_style,
+                dist_key,
+                sort_key,
                 table_options,
                 table_properties,
                 with_options,
@@ -388,6 +415,9 @@ impl TryFrom<Statement> for CreateTableBuilder {
                 constraints,
                 hive_distribution,
                 hive_formats,
+                dist_style,
+                dist_key,
+                sort_key,
                 table_options,
                 table_properties,
                 with_options,

src/ast/mod.rs

@@ -1648,6 +1648,9 @@ pub enum Statement {
         constraints: Vec<TableConstraint>,
         hive_distribution: HiveDistributionStyle,
         hive_formats: Option<HiveFormat>,
+        dist_style: Option<DistributionStyle>,
+        dist_key: Option<WithSpan<Ident>>,
+        sort_key: Option<SortKey>,
         table_properties: Vec<SqlOption>,
         table_options: Vec<SqlOption>,
         with_options: Vec<SqlOption>,
@@ -2762,6 +2765,9 @@ impl fmt::Display for Statement {
                 transient,
                 hive_distribution,
                 hive_formats,
+                dist_style,
+                dist_key,
+                sort_key,
                 external,
                 global,
                 temporary,
@@ -2831,6 +2837,22 @@ impl fmt::Display for Statement {
                     write!(f, " ()")?;
                 }
 
+                if let Some(dist_style) = dist_style {
+                    write!(f, " DISTSTYLE {style}", style = dist_style)?;
+                }
+
+                if let Some(dist_key) = dist_key {
+                    write!(f, " DISTKEY({dist_key})")?;
+                }
+
+                if let Some(SortKey { compound, columns }) = sort_key {
+                    if *compound {
+                        write!(f, " COMPOUND SORTKEY({})", display_comma_separated(columns))?;
+                    } else {
+                        write!(f, " SORTKEY({})", display_comma_separated(columns))?;
+                    }
+                }
+
                 if let Some(using) = using {
                     write!(f, " USING {using}")?;
                 }
@@ -4366,6 +4388,35 @@ impl fmt::Display for KillType {
     }
 }
 
+#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum DistributionStyle {
+    Auto,
+    Even,
+    Key,
+    All,
+}
+
+impl fmt::Display for DistributionStyle {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            DistributionStyle::Auto => "AUTO",
+            DistributionStyle::Even => "EVEN",
+            DistributionStyle::Key => "KEY",
+            DistributionStyle::All => "ALL",
+        })
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct SortKey {
+    pub compound: bool,
+    pub columns: Vec<WithSpan<Ident>>,
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/keywords.rs

@@ -157,6 +157,7 @@ define_keywords!(
     COMMENT,
     COMMIT,
     COMMITTED,
+    COMPOUND,
     COMPRESSION,
     COMPUTE,
     CONCURRENTLY,
@@ -223,7 +224,9 @@ define_keywords!(
     DISCARD,
     DISCONNECT,
     DISTINCT,
+    DISTKEY,
     DISTRIBUTE,
+    DISTSTYLE,
     DIV,
     DO,
     DOUBLE,
@@ -250,6 +253,7 @@ define_keywords!(
     EQUALS,
     ERROR,
     ESCAPE,
+    EVEN,
     EVENT,
     EVERY,
     EXCEPT,
@@ -581,6 +585,7 @@ define_keywords!(
     SNAPSHOT,
     SOME,
     SORT,
+    SORTKEY,
     SPATIAL,
     SPECIFIC,
     SPECIFICTYPE,

src/parser/mod.rs

@@ -4168,6 +4168,48 @@ impl<'a> Parser<'a> {
             None
         };
 
+        // Redshift allows specifying DISTSTYLE after column definitions
+        let dist_style = if self.parse_keywords(&[Keyword::DISTSTYLE]) {
+            match self.parse_one_of_keywords(&[
+                Keyword::EVEN,
+                Keyword::ALL,
+                Keyword::AUTO,
+                Keyword::KEY,
+            ]) {
+                Some(Keyword::EVEN) => Some(DistributionStyle::Even),
+                Some(Keyword::ALL) => Some(DistributionStyle::All),
+                Some(Keyword::AUTO) => Some(DistributionStyle::Auto),
+                Some(Keyword::KEY) => Some(DistributionStyle::Key),
+                _ => self.expected("KEY, EVEN, ALL or AUTO", self.peek_token())?,
+            }
+        } else {
+            None
+        };
+
+        // Redshift allows specifying DISTKEY after column definitions
+        let dist_key = if self.parse_keywords(&[Keyword::DISTKEY]) {
+            self.expect_token(&Token::LParen)?;
+            let key = self.parse_identifier(false)?;
+            self.expect_token(&Token::RParen)?;
+            Some(key)
+        } else {
+            None
+        };
+
+        // Redshift allows specifying SORTKEY after column definitions
+        let compound_sort_key = self.parse_keywords(&[Keyword::COMPOUND]);
+        let sort_key = if self.parse_keywords(&[Keyword::SORTKEY]) {
+            self.expect_token(&Token::LParen)?;
+            let columns = self.parse_comma_separated(|p| p.parse_identifier(false))?;
+            self.expect_token(&Token::RParen)?;
+            Some(SortKey {
+                compound: compound_sort_key,
+                columns,
+            })
+        } else {
+            None
+        };
+
         // SQLite supports `WITHOUT ROWID` at the end of `CREATE TABLE`
         let without_rowid = self.parse_keywords(&[Keyword::WITHOUT, Keyword::ROWID]);
 
@@ -4355,6 +4397,9 @@ impl<'a> Parser<'a> {
             .transient(transient)
             .hive_distribution(hive_distribution)
             .hive_formats(Some(hive_formats))
+            .dist_style(dist_style)
+            .dist_key(dist_key)
+            .sort_key(sort_key)
             .global(global)
             .query(query)
             .without_rowid(without_rowid)

tests/sqlparser_redshift.rs

@@ -377,6 +377,12 @@ fn test_escape_string() {
     redshift_unescaped().verified_stmt(r#"SELECT '[\'\\[\\]]'"#);
 }
 
+#[test]
+fn test_distribution_styles() {
+    let sql = "CREATE TABLE foo (id VARCHAR(32)) DISTSTYLE KEY DISTKEY(id) COMPOUND SORTKEY(id)";
+    redshift().verified_stmt(sql);
+}
+
 #[test]
 fn test_utf8_column_names() {
     redshift().verified_stmt("SELECT financing_cost_€k FROM tbl");