mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-04 06:18:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Merge pull request #50 from nickolay/window-functions

Browse source 
Support OVER clause for window/analytic functions, add support for qualified function names
This commit is contained in:
Nickolay Ponomarev 2019-04-27 21:12:08 +03:00 • committed by GitHub
commit 6bb2acc9f8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 295 additions and 89 deletions
Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -30,7 +30,7 @@ println!("AST: {:?}", ast);
This outputs
```rust
AST: [SQLSelect(SQLSelect { projection: [SQLIdentifier("a"), SQLIdentifier("b"), SQLValue(Long(123)), SQLFunction { id: "myfunc", args: [SQLIdentifier("b")] }], relation: Some(Table { name: SQLObjectName(["table_1"]), alias: None }), joins: [], selection: Some(SQLBinaryExpr { left: SQLBinaryExpr { left: SQLIdentifier("a"), op: Gt, right: SQLIdentifier("b") }, op: And, right: SQLBinaryExpr { left: SQLIdentifier("b"), op: Lt, right: SQLValue(Long(100)) } }), order_by: Some([SQLOrderByExpr { expr: SQLIdentifier("a"), asc: Some(false) }, SQLOrderByExpr { expr: SQLIdentifier("b"), asc: None }]), group_by: None, having: None, limit: None })]
AST: [SQLSelect(SQLQuery { ctes: [], body: Select(SQLSelect { distinct: false, projection: [UnnamedExpression(SQLIdentifier("a")), UnnamedExpression(SQLIdentifier("b")), UnnamedExpression(SQLValue(Long(123))), UnnamedExpression(SQLFunction { name: SQLObjectName(["myfunc"]), args: [SQLIdentifier("b")], over: None })], relation: Some(Table { name: SQLObjectName(["table_1"]), alias: None }), joins: [], selection: Some(SQLBinaryExpr { left: SQLBinaryExpr { left: SQLIdentifier("a"), op: Gt, right: SQLIdentifier("b") }, op: And, right: SQLBinaryExpr { left: SQLIdentifier("b"), op: Lt, right: SQLValue(Long(100)) } }), group_by: None, having: None }), order_by: Some([SQLOrderByExpr { expr: SQLIdentifier("a"), asc: Some(false) }, SQLOrderByExpr { expr: SQLIdentifier("b"), asc: None }]), limit: None })]
```
## Design

6
src/dialect/keywords.rs
View file

@ -139,6 +139,7 @@ keyword!(
FIRST_VALUE,
FLOAT,
FLOOR,
FOLLOWING,
FOR,
FOREIGN,
FRAME_ROW,
@ -246,6 +247,7 @@ keyword!(
POSITION_REGEX,
POWER,
PRECEDES,
PRECEDING,
PRECISION,
PREPARE,
PRIMARY,
@ -333,6 +335,7 @@ keyword!(
TRIM_ARRAY,
TRUE,
UESCAPE,
UNBOUNDED,
UNION,
UNIQUE,
UNKNOWN,
@ -488,6 +491,7 @@ pub const ALL_KEYWORDS: &[&str] = &[
FIRST_VALUE,
FLOAT,
FLOOR,
FOLLOWING,
FOR,
FOREIGN,
FRAME_ROW,
@ -595,6 +599,7 @@ pub const ALL_KEYWORDS: &[&str] = &[
POSITION_REGEX,
POWER,
PRECEDES,
PRECEDING,
PRECISION,
PREPARE,
PRIMARY,
@ -682,6 +687,7 @@ pub const ALL_KEYWORDS: &[&str] = &[
TRIM_ARRAY,
TRUE,
UESCAPE,
UNBOUNDED,
UNION,
UNIQUE,
UNKNOWN,

189
src/sqlast/mod.rs
View file

@ -30,6 +30,14 @@ pub use self::value::Value;
pub use self::sql_operator::SQLOperator;
/// Like `vec.join(", ")`, but for any types implementing ToString.
fn comma_separated_string<T: ToString>(vec: &[T]) -> String {
vec.iter()
.map(T::to_string)
.collect::<Vec<String>>()
.join(", ")
}
/// Identifier name, in the originally quoted form (e.g. `"id"`)
pub type SQLIdent = String;
@ -46,7 +54,7 @@ pub enum ASTNode {
/// Qualified wildcard, e.g. `alias.*` or `schema.table.*`.
/// (Same caveats apply to SQLQualifiedWildcard as to SQLWildcard.)
SQLQualifiedWildcard(Vec<SQLIdent>),
/// Multi part identifier e.g. `myschema.dbo.mytable`
/// Multi-part identifier, e.g. `table_alias.column` or `schema.table.col`
SQLCompoundIdentifier(Vec<SQLIdent>),
/// `IS NULL` expression
SQLIsNull(Box<ASTNode>),
@ -92,8 +100,11 @@ pub enum ASTNode {
/// SQLValue
SQLValue(Value),
/// Scalar function call e.g. `LEFT(foo, 5)`
/// TODO: this can be a compound SQLObjectName as well (for UDFs)
SQLFunction { id: SQLIdent, args: Vec<ASTNode> },
SQLFunction {
name: SQLObjectName,
args: Vec<ASTNode>,
over: Option<SQLWindowSpec>,
},
/// CASE [<operand>] WHEN <condition> THEN <result> ... [ELSE <result>] END
SQLCase {
// TODO: support optional operand for "simple case"
@ -123,10 +134,7 @@ impl ToString for ASTNode {
"{} {}IN ({})",
expr.as_ref().to_string(),
if *negated { "NOT " } else { "" },
list.iter()
.map(|a| a.to_string())
.collect::<Vec<String>>()
.join(", ")
comma_separated_string(list)
),
ASTNode::SQLInSubquery {
expr,
@ -166,14 +174,13 @@ impl ToString for ASTNode {
format!("{} {}", operator.to_string(), expr.as_ref().to_string())
}
ASTNode::SQLValue(v) => v.to_string(),
ASTNode::SQLFunction { id, args } => format!(
"{}({})",
id,
args.iter()
.map(|a| a.to_string())
.collect::<Vec<String>>()
.join(", ")
),
ASTNode::SQLFunction { name, args, over } => {
let mut s = format!("{}({})", name.to_string(), comma_separated_string(args));
if let Some(o) = over {
s += &format!(" OVER ({})", o.to_string())
}
s
}
ASTNode::SQLCase {
conditions,
results,
@ -198,6 +205,116 @@ impl ToString for ASTNode {
}
}
/// A window specification (i.e. `OVER (PARTITION BY .. ORDER BY .. etc.)`)
#[derive(Debug, Clone, PartialEq)]
pub struct SQLWindowSpec {
pub partition_by: Vec<ASTNode>,
pub order_by: Vec<SQLOrderByExpr>,
pub window_frame: Option<SQLWindowFrame>,
}
impl ToString for SQLWindowSpec {
fn to_string(&self) -> String {
let mut clauses = vec![];
if !self.partition_by.is_empty() {
clauses.push(format!(
"PARTITION BY {}",
comma_separated_string(&self.partition_by)
))
};
if !self.order_by.is_empty() {
clauses.push(format!(
"ORDER BY {}",
comma_separated_string(&self.order_by)
))
};
if let Some(window_frame) = &self.window_frame {
if let Some(end_bound) = &window_frame.end_bound {
clauses.push(format!(
"{} BETWEEN {} AND {}",
window_frame.units.to_string(),
window_frame.start_bound.to_string(),
end_bound.to_string()
));
} else {
clauses.push(format!(
"{} {}",
window_frame.units.to_string(),
window_frame.start_bound.to_string()
));
}
}
clauses.join(" ")
}
}
/// Specifies the data processed by a window function, e.g.
/// `RANGE UNBOUNDED PRECEDING` or `ROWS BETWEEN 5 PRECEDING AND CURRENT ROW`.
#[derive(Debug, Clone, PartialEq)]
pub struct SQLWindowFrame {
pub units: SQLWindowFrameUnits,
pub start_bound: SQLWindowFrameBound,
/// The right bound of the `BETWEEN .. AND` clause.
pub end_bound: Option<SQLWindowFrameBound>,
// TBD: EXCLUDE
}
#[derive(Debug, Clone, PartialEq)]
pub enum SQLWindowFrameUnits {
Rows,
Range,
Groups,
}
impl ToString for SQLWindowFrameUnits {
fn to_string(&self) -> String {
match self {
SQLWindowFrameUnits::Rows => "ROWS".to_string(),
SQLWindowFrameUnits::Range => "RANGE".to_string(),
SQLWindowFrameUnits::Groups => "GROUPS".to_string(),
}
}
}
impl FromStr for SQLWindowFrameUnits {
type Err = ParserError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"ROWS" => Ok(SQLWindowFrameUnits::Rows),
"RANGE" => Ok(SQLWindowFrameUnits::Range),
"GROUPS" => Ok(SQLWindowFrameUnits::Groups),
_ => Err(ParserError::ParserError(format!(
"Expected ROWS, RANGE, or GROUPS, found: {}",
s
))),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum SQLWindowFrameBound {
/// "CURRENT ROW"
CurrentRow,
/// "<N> PRECEDING" or "UNBOUNDED PRECEDING"
Preceding(Option<u64>),
/// "<N> FOLLOWING" or "UNBOUNDED FOLLOWING". This can only appear in
/// SQLWindowFrame::end_bound.
Following(Option<u64>),
}
impl ToString for SQLWindowFrameBound {
fn to_string(&self) -> String {
match self {
SQLWindowFrameBound::CurrentRow => "CURRENT ROW".to_string(),
SQLWindowFrameBound::Preceding(None) => "UNBOUNDED PRECEDING".to_string(),
SQLWindowFrameBound::Following(None) => "UNBOUNDED FOLLOWING".to_string(),
SQLWindowFrameBound::Preceding(Some(n)) => format!("{} PRECEDING", n),
SQLWindowFrameBound::Following(Some(n)) => format!("{} FOLLOWING", n),
}
}
}
/// A top-level statement (SELECT, INSERT, CREATE, etc.)
#[derive(Debug, Clone, PartialEq)]
pub enum SQLStatement {
@ -279,11 +396,7 @@ impl ToString for SQLStatement {
" VALUES({})",
values
.iter()
.map(|row| row
.iter()
.map(|c| c.to_string())
.collect::<Vec<String>>()
.join(", "))
.map(|row| comma_separated_string(row))
.collect::<Vec<String>>()
.join(", ")
);
@ -296,15 +409,8 @@ impl ToString for SQLStatement {
values,
} => {
let mut s = format!("COPY {}", table_name.to_string());
if columns.len() > 0 {
s += &format!(
" ({})",
columns
.iter()
.map(|c| c.to_string())
.collect::<Vec<String>>()
.join(", ")
);
if !columns.is_empty() {
s += &format!(" ({})", comma_separated_string(columns));
}
s += " FROM stdin; ";
if !values.is_empty() {
@ -326,15 +432,8 @@ impl ToString for SQLStatement {
selection,
} => {
let mut s = format!("UPDATE {}", table_name.to_string());
if assignments.len() > 0 {
s += &format!(
"{}",
assignments
.iter()
.map(|ass| ass.to_string())
.collect::<Vec<String>>()
.join(", ")
);
if !assignments.is_empty() {
s += &comma_separated_string(assignments);
}
if let Some(selection) = selection {
s += &format!(" WHERE {}", selection.to_string());
@ -373,22 +472,14 @@ impl ToString for SQLStatement {
} if *external => format!(
"CREATE EXTERNAL TABLE {} ({}) STORED AS {} LOCATION '{}'",
name.to_string(),
columns
.iter()
.map(|c| c.to_string())
.collect::<Vec<String>>()
.join(", "),
file_format.as_ref().map(|f| f.to_string()).unwrap(),
comma_separated_string(columns),
file_format.as_ref().unwrap().to_string(),
location.as_ref().unwrap()
),
SQLStatement::SQLCreateTable { name, columns, .. } => format!(
"CREATE TABLE {} ({})",
name.to_string(),
columns
.iter()
.map(|c| c.to_string())
.collect::<Vec<String>>()
.join(", ")
comma_separated_string(columns)
),
SQLStatement::SQLAlterTable { name, operation } => {
format!("ALTER TABLE {} {}", name.to_string(), operation.to_string())

28
src/sqlast/query.rs
View file

@ -29,14 +29,7 @@ impl ToString for SQLQuery {
}
s += &self.body.to_string();
if let Some(ref order_by) = self.order_by {
s += &format!(
" ORDER BY {}",
order_by
.iter()
.map(|o| o.to_string())
.collect::<Vec<String>>()
.join(", ")
);
s += &format!(" ORDER BY {}", comma_separated_string(order_by));
}
if let Some(ref limit) = self.limit {
s += &format!(" LIMIT {}", limit.to_string());
@ -130,11 +123,7 @@ impl ToString for SQLSelect {
let mut s = format!(
"SELECT{} {}",
if self.distinct { " DISTINCT" } else { "" },
self.projection
.iter()
.map(|p| p.to_string())
.collect::<Vec<String>>()
.join(", ")
comma_separated_string(&self.projection)
);
if let Some(ref relation) = self.relation {
s += &format!(" FROM {}", relation.to_string());
@ -146,14 +135,7 @@ impl ToString for SQLSelect {
s += &format!(" WHERE {}", selection.to_string());
}
if let Some(ref group_by) = self.group_by {
s += &format!(
" GROUP BY {}",
group_by
.iter()
.map(|g| g.to_string())
.collect::<Vec<String>>()
.join(", ")
);
s += &format!(" GROUP BY {}", comma_separated_string(group_by));
}
if let Some(ref having) = self.having {
s += &format!(" HAVING {}", having.to_string());
@ -175,7 +157,7 @@ pub enum SQLSelectItem {
/// Any expression, not followed by `[ AS ] alias`
UnnamedExpression(ASTNode),
/// An expression, followed by `[ AS ] alias`
ExpressionWithAlias(ASTNode, SQLIdent),
ExpressionWithAlias { expr: ASTNode, alias: SQLIdent },
/// `alias.*` or even `schema.table.*`
QualifiedWildcard(SQLObjectName),
/// An unqualified `*`
@ -186,7 +168,7 @@ impl ToString for SQLSelectItem {
fn to_string(&self) -> String {
match &self {
SQLSelectItem::UnnamedExpression(expr) => expr.to_string(),
SQLSelectItem::ExpressionWithAlias(expr, alias) => {
SQLSelectItem::ExpressionWithAlias { expr, alias } => {
format!("{} AS {}", expr.to_string(), alias)
}
SQLSelectItem::QualifiedWildcard(prefix) => format!("{}.*", prefix.to_string()),

111
src/sqlparser.rs
View file

@ -174,10 +174,10 @@ impl Parser {
expr: Box::new(self.parse_subexpr(p)?),
})
}
// another SQLWord:
// Here `w` is a word, check if it's a part of a multi-part
// identifier, a function call, or a simple identifier:
_ => match self.peek_token() {
Some(Token::LParen) => self.parse_function(w.as_sql_ident()),
Some(Token::Period) => {
Some(Token::LParen) | Some(Token::Period) => {
let mut id_parts: Vec<SQLIdent> = vec![w.as_sql_ident()];
let mut ends_with_wildcard = false;
while self.consume_token(&Token::Period) {
@ -187,15 +187,19 @@ impl Parser {
ends_with_wildcard = true;
break;
}
_ => {
unexpected => {
return parser_err!(format!(
"Error parsing compound identifier"
"Expected an identifier or a '*' after '.', got: {:?}",
unexpected
));
}
}
}
if ends_with_wildcard {
Ok(ASTNode::SQLQualifiedWildcard(id_parts))
} else if self.consume_token(&Token::LParen) {
self.prev_token();
self.parse_function(SQLObjectName(id_parts))
} else {
Ok(ASTNode::SQLCompoundIdentifier(id_parts))
}
@ -236,7 +240,7 @@ impl Parser {
}
}
pub fn parse_function(&mut self, id: SQLIdent) -> Result<ASTNode, ParserError> {
pub fn parse_function(&mut self, name: SQLObjectName) -> Result<ASTNode, ParserError> {
self.expect_token(&Token::LParen)?;
let args = if self.consume_token(&Token::RParen) {
vec![]
@ -245,7 +249,98 @@ impl Parser {
self.expect_token(&Token::RParen)?;
args
};
Ok(ASTNode::SQLFunction { id, args })
let over = if self.parse_keyword("OVER") {
// TBD: support window names (`OVER mywin`) in place of inline specification
self.expect_token(&Token::LParen)?;
let partition_by = if self.parse_keywords(vec!["PARTITION", "BY"]) {
// a list of possibly-qualified column names
self.parse_expr_list()?
} else {
vec![]
};
let order_by = if self.parse_keywords(vec!["ORDER", "BY"]) {
self.parse_order_by_expr_list()?
} else {
vec![]
};
let window_frame = self.parse_window_frame()?;
Some(SQLWindowSpec {
partition_by,
order_by,
window_frame,
})
} else {
None
};
Ok(ASTNode::SQLFunction { name, args, over })
}
pub fn parse_window_frame(&mut self) -> Result<Option<SQLWindowFrame>, ParserError> {
let window_frame = match self.peek_token() {
Some(Token::SQLWord(w)) => {
let units = w.keyword.parse::<SQLWindowFrameUnits>()?;
self.next_token();
if self.parse_keyword("BETWEEN") {
let start_bound = self.parse_window_frame_bound()?;
self.expect_keyword("AND")?;
let end_bound = Some(self.parse_window_frame_bound()?);
Some(SQLWindowFrame {
units,
start_bound,
end_bound,
})
} else {
let start_bound = self.parse_window_frame_bound()?;
let end_bound = None;
Some(SQLWindowFrame {
units,
start_bound,
end_bound,
})
}
}
Some(Token::RParen) => None,
unexpected => {
return parser_err!(format!(
"Expected 'ROWS', 'RANGE', 'GROUPS', or ')', got {:?}",
unexpected
));
}
};
self.expect_token(&Token::RParen)?;
Ok(window_frame)
}
/// "CURRENT ROW" | ( (<positive number> | "UNBOUNDED") ("PRECEDING" | FOLLOWING) )
pub fn parse_window_frame_bound(&mut self) -> Result<SQLWindowFrameBound, ParserError> {
if self.parse_keywords(vec!["CURRENT", "ROW"]) {
Ok(SQLWindowFrameBound::CurrentRow)
} else {
let rows = if self.parse_keyword("UNBOUNDED") {
None
} else {
let rows = self.parse_literal_int()?;
if rows < 0 {
parser_err!(format!(
"The number of rows must be non-negative, got {}",
rows
))?;
}
Some(rows as u64)
};
if self.parse_keyword("PRECEDING") {
Ok(SQLWindowFrameBound::Preceding(rows))
} else if self.parse_keyword("FOLLOWING") {
Ok(SQLWindowFrameBound::Following(rows))
} else {
parser_err!(format!(
"Expected PRECEDING or FOLLOWING, found {:?}",
self.peek_token()
))
}
}
}
pub fn parse_case_expression(&mut self) -> Result<ASTNode, ParserError> {
@ -1527,7 +1622,7 @@ impl Parser {
if let Some(alias) =
self.parse_optional_alias(keywords::RESERVED_FOR_COLUMN_ALIAS)?
{
projections.push(SQLSelectItem::ExpressionWithAlias(expr, alias));
projections.push(SQLSelectItem::ExpressionWithAlias { expr, alias });
} else {
projections.push(SQLSelectItem::UnnamedExpression(expr));
}

48
tests/sqlparser_generic.rs
View file

@ -101,12 +101,12 @@ fn parse_count_wildcard() {
fn parse_column_aliases() {
let sql = "SELECT a.col + 1 AS newname FROM foo AS a";
let select = verified_only_select(sql);
if let SQLSelectItem::ExpressionWithAlias(
ASTNode::SQLBinaryExpr {
if let SQLSelectItem::ExpressionWithAlias {
expr: ASTNode::SQLBinaryExpr {
ref op, ref right, ..
},
ref alias,
) = only(&select.projection)
} = only(&select.projection)
{
assert_eq!(&SQLOperator::Plus, op);
assert_eq!(&ASTNode::SQLValue(Value::Long(1)), right.as_ref());
@ -125,8 +125,9 @@ fn parse_select_count_wildcard() {
let select = verified_only_select(sql);
assert_eq!(
&ASTNode::SQLFunction {
id: "COUNT".to_string(),
name: SQLObjectName(vec!["COUNT".to_string()]),
args: vec![ASTNode::SQLWildcard],
over: None,
},
expr_from_projection(only(&select.projection))
);
@ -551,13 +552,43 @@ fn parse_scalar_function_in_projection() {
let select = verified_only_select(sql);
assert_eq!(
&ASTNode::SQLFunction {
id: String::from("sqrt"),
name: SQLObjectName(vec![String::from("sqrt")]),
args: vec![ASTNode::SQLIdentifier(String::from("id"))],
over: None,
},
expr_from_projection(only(&select.projection))
);
}
#[test]
fn parse_window_functions() {
let sql = "SELECT row_number() OVER (ORDER BY dt DESC), \
sum(foo) OVER (PARTITION BY a, b ORDER BY c, d \
ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), \
avg(bar) OVER (ORDER BY a \
RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING), \
max(baz) OVER (ORDER BY a \
ROWS UNBOUNDED PRECEDING) \
FROM foo";
let select = verified_only_select(sql);
assert_eq!(4, select.projection.len());
assert_eq!(
&ASTNode::SQLFunction {
name: SQLObjectName(vec!["row_number".to_string()]),
args: vec![],
over: Some(SQLWindowSpec {
partition_by: vec![],
order_by: vec![SQLOrderByExpr {
expr: ASTNode::SQLIdentifier("dt".to_string()),
asc: Some(false)
}],
window_frame: None,
})
},
expr_from_projection(&select.projection[0])
);
}
#[test]
fn parse_aggregate_with_group_by() {
let sql = "SELECT a, COUNT(1), MIN(b), MAX(b) FROM foo GROUP BY a";
@ -624,13 +655,14 @@ fn parse_delimited_identifiers() {
);
assert_eq!(
&ASTNode::SQLFunction {
id: r#""myfun""#.to_string(),
args: vec![]
name: SQLObjectName(vec![r#""myfun""#.to_string()]),
args: vec![],
over: None,
},
expr_from_projection(&select.projection[1]),
);
match &select.projection[2] {
&SQLSelectItem::ExpressionWithAlias(ref expr, ref alias) => {
SQLSelectItem::ExpressionWithAlias { expr, alias } => {
assert_eq!(&ASTNode::SQLIdentifier(r#""simple id""#.to_string()), expr);
assert_eq!(r#""column alias""#, alias);
}