Skip to content

Commit bab168f

Browse files
zdenalzdenal
authored
Support distribution styles for redshift. (#39)
* Support distribution styles for redshift. * Working.
1 parent ecaed29 commit bab168f

File tree

5 files changed

+140
-3
lines changed

5 files changed

+140
-3
lines changed

‎src/ast/helpers/stmt_create_table.rs

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
#[cfg(not(feature = "std"))]
22
use alloc::{boxed::Box, format, string::String, vec, vec::Vec};
33

4-
use crate::ast::WithSpan;
4+
use crate::ast::{SortKey, WithSpan};
55

66
#[cfg(feature = "serde")]
77
use serde::{Deserialize, Serialize};
@@ -10,8 +10,8 @@ use serde::{Deserialize, Serialize};
1010
use sqlparser_derive::{Visit, VisitMut};
1111

1212
use crate::ast::{
13-
ColumnDef, EngineSpec, Expr, FileFormat, HiveDistributionStyle, HiveFormat, Ident, ObjectName,
14-
OnCommit, Query, SqlOption, Statement, TableConstraint,
13+
ColumnDef, DistributionStyle, EngineSpec, Expr, FileFormat, HiveDistributionStyle, HiveFormat,
14+
Ident, ObjectName, OnCommit, Query, SqlOption, Statement, TableConstraint,
1515
};
1616
use crate::parser::ParserError;
1717

@@ -58,6 +58,9 @@ pub struct CreateTableBuilder {
5858
pub constraints: Vec<TableConstraint>,
5959
pub hive_distribution: HiveDistributionStyle,
6060
pub hive_formats: Option<HiveFormat>,
61+
pub dist_style: Option<DistributionStyle>,
62+
pub dist_key: Option<WithSpan<Ident>>,
63+
pub sort_key: Option<SortKey>,
6164
pub table_options: Vec<SqlOption>,
6265
pub table_properties: Vec<SqlOption>,
6366
pub with_options: Vec<SqlOption>,
@@ -98,6 +101,9 @@ impl CreateTableBuilder {
98101
constraints: vec![],
99102
hive_distribution: HiveDistributionStyle::NONE,
100103
hive_formats: None,
104+
dist_style: None,
105+
dist_key: None,
106+
sort_key: None,
101107
table_options: vec![],
102108
table_properties: vec![],
103109
with_options: vec![],
@@ -174,6 +180,21 @@ impl CreateTableBuilder {
174180
self
175181
}
176182

183+
pub fn dist_style(mut self, dist_style: Option<DistributionStyle>) -> Self {
184+
self.dist_style = dist_style;
185+
self
186+
}
187+
188+
pub fn dist_key(mut self, dist_key: Option<WithSpan<Ident>>) -> Self {
189+
self.dist_key = dist_key;
190+
self
191+
}
192+
193+
pub fn sort_key(mut self, sort_key: Option<SortKey>) -> Self {
194+
self.sort_key = sort_key;
195+
self
196+
}
197+
177198
pub fn table_options(mut self, table_options: Vec<SqlOption>) -> Self {
178199
self.table_options = table_options;
179200
self
@@ -305,6 +326,9 @@ impl CreateTableBuilder {
305326
constraints: self.constraints,
306327
hive_distribution: self.hive_distribution,
307328
hive_formats: self.hive_formats,
329+
dist_style: self.dist_style,
330+
dist_key: self.dist_key,
331+
sort_key: self.sort_key,
308332
table_options: self.table_options,
309333
table_properties: self.table_properties,
310334
with_options: self.with_options,
@@ -352,6 +376,9 @@ impl TryFrom<Statement> for CreateTableBuilder {
352376
constraints,
353377
hive_distribution,
354378
hive_formats,
379+
dist_style,
380+
dist_key,
381+
sort_key,
355382
table_options,
356383
table_properties,
357384
with_options,
@@ -388,6 +415,9 @@ impl TryFrom<Statement> for CreateTableBuilder {
388415
constraints,
389416
hive_distribution,
390417
hive_formats,
418+
dist_style,
419+
dist_key,
420+
sort_key,
391421
table_options,
392422
table_properties,
393423
with_options,

‎src/ast/mod.rs

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1648,6 +1648,9 @@ pub enum Statement {
16481648
constraints: Vec<TableConstraint>,
16491649
hive_distribution: HiveDistributionStyle,
16501650
hive_formats: Option<HiveFormat>,
1651+
dist_style: Option<DistributionStyle>,
1652+
dist_key: Option<WithSpan<Ident>>,
1653+
sort_key: Option<SortKey>,
16511654
table_properties: Vec<SqlOption>,
16521655
table_options: Vec<SqlOption>,
16531656
with_options: Vec<SqlOption>,
@@ -2762,6 +2765,9 @@ impl fmt::Display for Statement {
27622765
transient,
27632766
hive_distribution,
27642767
hive_formats,
2768+
dist_style,
2769+
dist_key,
2770+
sort_key,
27652771
external,
27662772
global,
27672773
temporary,
@@ -2831,6 +2837,22 @@ impl fmt::Display for Statement {
28312837
write!(f, " ()")?;
28322838
}
28332839

2840+
if let Some(dist_style) = dist_style {
2841+
write!(f, " DISTSTYLE {style}", style = dist_style)?;
2842+
}
2843+
2844+
if let Some(dist_key) = dist_key {
2845+
write!(f, " DISTKEY({dist_key})")?;
2846+
}
2847+
2848+
if let Some(SortKey { compound, columns }) = sort_key {
2849+
if *compound {
2850+
write!(f, " COMPOUND SORTKEY({})", display_comma_separated(columns))?;
2851+
} else {
2852+
write!(f, " SORTKEY({})", display_comma_separated(columns))?;
2853+
}
2854+
}
2855+
28342856
if let Some(using) = using {
28352857
write!(f, " USING {using}")?;
28362858
}
@@ -4366,6 +4388,35 @@ impl fmt::Display for KillType {
43664388
}
43674389
}
43684390

4391+
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
4392+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
4393+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
4394+
pub enum DistributionStyle {
4395+
Auto,
4396+
Even,
4397+
Key,
4398+
All,
4399+
}
4400+
4401+
impl fmt::Display for DistributionStyle {
4402+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
4403+
f.write_str(match self {
4404+
DistributionStyle::Auto => "AUTO",
4405+
DistributionStyle::Even => "EVEN",
4406+
DistributionStyle::Key => "KEY",
4407+
DistributionStyle::All => "ALL",
4408+
})
4409+
}
4410+
}
4411+
4412+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
4413+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
4414+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
4415+
pub struct SortKey {
4416+
pub compound: bool,
4417+
pub columns: Vec<WithSpan<Ident>>,
4418+
}
4419+
43694420
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
43704421
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
43714422
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

‎src/keywords.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,7 @@ define_keywords!(
157157
COMMENT,
158158
COMMIT,
159159
COMMITTED,
160+
COMPOUND,
160161
COMPRESSION,
161162
COMPUTE,
162163
CONCURRENTLY,
@@ -223,7 +224,9 @@ define_keywords!(
223224
DISCARD,
224225
DISCONNECT,
225226
DISTINCT,
227+
DISTKEY,
226228
DISTRIBUTE,
229+
DISTSTYLE,
227230
DIV,
228231
DO,
229232
DOUBLE,
@@ -250,6 +253,7 @@ define_keywords!(
250253
EQUALS,
251254
ERROR,
252255
ESCAPE,
256+
EVEN,
253257
EVENT,
254258
EVERY,
255259
EXCEPT,
@@ -581,6 +585,7 @@ define_keywords!(
581585
SNAPSHOT,
582586
SOME,
583587
SORT,
588+
SORTKEY,
584589
SPATIAL,
585590
SPECIFIC,
586591
SPECIFICTYPE,

‎src/parser/mod.rs

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4168,6 +4168,48 @@ impl<'a> Parser<'a> {
41684168
None
41694169
};
41704170

4171+
// Redshift allows specifying DISTSTYLE after column definitions
4172+
let dist_style = if self.parse_keywords(&[Keyword::DISTSTYLE]) {
4173+
match self.parse_one_of_keywords(&[
4174+
Keyword::EVEN,
4175+
Keyword::ALL,
4176+
Keyword::AUTO,
4177+
Keyword::KEY,
4178+
]) {
4179+
Some(Keyword::EVEN) => Some(DistributionStyle::Even),
4180+
Some(Keyword::ALL) => Some(DistributionStyle::All),
4181+
Some(Keyword::AUTO) => Some(DistributionStyle::Auto),
4182+
Some(Keyword::KEY) => Some(DistributionStyle::Key),
4183+
_ => self.expected("KEY, EVEN, ALL or AUTO", self.peek_token())?,
4184+
}
4185+
} else {
4186+
None
4187+
};
4188+
4189+
// Redshift allows specifying DISTKEY after column definitions
4190+
let dist_key = if self.parse_keywords(&[Keyword::DISTKEY]) {
4191+
self.expect_token(&Token::LParen)?;
4192+
let key = self.parse_identifier(false)?;
4193+
self.expect_token(&Token::RParen)?;
4194+
Some(key)
4195+
} else {
4196+
None
4197+
};
4198+
4199+
// Redshift allows specifying SORTKEY after column definitions
4200+
let compound_sort_key = self.parse_keywords(&[Keyword::COMPOUND]);
4201+
let sort_key = if self.parse_keywords(&[Keyword::SORTKEY]) {
4202+
self.expect_token(&Token::LParen)?;
4203+
let columns = self.parse_comma_separated(|p| p.parse_identifier(false))?;
4204+
self.expect_token(&Token::RParen)?;
4205+
Some(SortKey {
4206+
compound: compound_sort_key,
4207+
columns,
4208+
})
4209+
} else {
4210+
None
4211+
};
4212+
41714213
// SQLite supports `WITHOUT ROWID` at the end of `CREATE TABLE`
41724214
let without_rowid = self.parse_keywords(&[Keyword::WITHOUT, Keyword::ROWID]);
41734215

@@ -4355,6 +4397,9 @@ impl<'a> Parser<'a> {
43554397
.transient(transient)
43564398
.hive_distribution(hive_distribution)
43574399
.hive_formats(Some(hive_formats))
4400+
.dist_style(dist_style)
4401+
.dist_key(dist_key)
4402+
.sort_key(sort_key)
43584403
.global(global)
43594404
.query(query)
43604405
.without_rowid(without_rowid)

‎tests/sqlparser_redshift.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,6 +377,12 @@ fn test_escape_string() {
377377
redshift_unescaped().verified_stmt(r#"SELECT '[\'\\[\\]]'"#);
378378
}
379379

380+
#[test]
381+
fn test_distribution_styles() {
382+
let sql = "CREATE TABLE foo (id VARCHAR(32)) DISTSTYLE KEY DISTKEY(id) COMPOUND SORTKEY(id)";
383+
redshift().verified_stmt(sql);
384+
}
385+
380386
#[test]
381387
fn test_utf8_column_names() {
382388
redshift().verified_stmt("SELECT financing_cost_€k FROM tbl");

0 commit comments

Comments
 (0)