Skip to content

Commit

Permalink
[ClickHouse] Add support for WITH FILL to OrderByExpr (apache#1330)
Browse files Browse the repository at this point in the history
Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
  • Loading branch information
2 people authored and ayman-sigma committed Nov 19, 2024
1 parent e710836 commit 7ccedd4
Show file tree
Hide file tree
Showing 9 changed files with 397 additions and 47 deletions.
17 changes: 9 additions & 8 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,14 +43,15 @@ pub use self::operator::{BinaryOperator, UnaryOperator};
pub use self::query::{
AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml,
FormatClause, GroupByExpr, GroupByWithModifier, IdentWithAlias, IlikeSelectItem, Join,
JoinConstraint, JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling, LateralView,
LockClause, LockType, MatchRecognizePattern, MatchRecognizeSymbol, Measure,
NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset, OffsetRows, OrderByExpr,
PivotValueSource, Query, RenameSelectItem, RepetitionQuantifier, ReplaceSelectElement,
ReplaceSelectItem, RowsPerMatch, Select, SelectInto, SelectItem, SetExpr, SetOperator,
SetQuantifier, Setting, SymbolDefinition, Table, TableAlias, TableFactor, TableVersion,
TableWithJoins, Top, TopQuantity, ValueTableMode, Values, WildcardAdditionalOptions, With,
FormatClause, GroupByExpr, GroupByWithModifier, IdentWithAlias, IlikeSelectItem, Interpolate,
InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonTableColumn,
JsonTableColumnErrorHandling, LateralView, LockClause, LockType, MatchRecognizePattern,
MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
OffsetRows, OrderBy, OrderByExpr, PivotValueSource, Query, RenameSelectItem,
RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table,
TableAlias, TableFactor, TableVersion, TableWithJoins, Top, TopQuantity, ValueTableMode,
Values, WildcardAdditionalOptions, With, WithFill,
};
pub use self::value::{
escape_double_quote_string, escape_quoted_string, DateTimeField, DollarQuotedString,
Expand Down
91 changes: 88 additions & 3 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ pub struct Query {
/// SELECT or UNION / EXCEPT / INTERSECT
pub body: Box<SetExpr>,
/// ORDER BY
pub order_by: Vec<OrderByExpr>,
pub order_by: Option<OrderBy>,
/// `LIMIT { <N> | ALL }`
pub limit: Option<Expr>,

Expand Down Expand Up @@ -67,8 +67,17 @@ impl fmt::Display for Query {
write!(f, "{with} ")?;
}
write!(f, "{}", self.body)?;
if !self.order_by.is_empty() {
write!(f, " ORDER BY {}", display_comma_separated(&self.order_by))?;
if let Some(ref order_by) = self.order_by {
write!(f, " ORDER BY")?;
if !order_by.exprs.is_empty() {
write!(f, " {}", display_comma_separated(&order_by.exprs))?;
}
if let Some(ref interpolate) = order_by.interpolate {
match &interpolate.exprs {
Some(exprs) => write!(f, " INTERPOLATE ({})", display_comma_separated(exprs))?,
None => write!(f, " INTERPOLATE")?,
}
}
}
if let Some(ref limit) = self.limit {
write!(f, " LIMIT {limit}")?;
Expand Down Expand Up @@ -1680,6 +1689,18 @@ pub enum JoinConstraint {
None,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct OrderBy {
pub exprs: Vec<OrderByExpr>,
/// Optional: `INTERPOLATE`
/// Supported by [ClickHouse syntax]
///
/// [ClickHouse syntax]: <https://clickhouse.com/docs/en/sql-reference/statements/select/order-by#order-by-expr-with-fill-modifier>
pub interpolate: Option<Interpolate>,
}

/// An `ORDER BY` expression
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
Expand All @@ -1690,6 +1711,9 @@ pub struct OrderByExpr {
pub asc: Option<bool>,
/// Optional `NULLS FIRST` or `NULLS LAST`
pub nulls_first: Option<bool>,
/// Optional: `WITH FILL`
/// Supported by [ClickHouse syntax]: <https://clickhouse.com/docs/en/sql-reference/statements/select/order-by#order-by-expr-with-fill-modifier>
pub with_fill: Option<WithFill>,
}

impl fmt::Display for OrderByExpr {
Expand All @@ -1705,6 +1729,67 @@ impl fmt::Display for OrderByExpr {
Some(false) => write!(f, " NULLS LAST")?,
None => (),
}
if let Some(ref with_fill) = self.with_fill {
write!(f, " {}", with_fill)?
}
Ok(())
}
}

/// ClickHouse `WITH FILL` modifier for `ORDER BY` clause.
/// Supported by [ClickHouse syntax]
///
/// [ClickHouse syntax]: <https://clickhouse.com/docs/en/sql-reference/statements/select/order-by#order-by-expr-with-fill-modifier>
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct WithFill {
pub from: Option<Expr>,
pub to: Option<Expr>,
pub step: Option<Expr>,
}

impl fmt::Display for WithFill {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "WITH FILL")?;
if let Some(ref from) = self.from {
write!(f, " FROM {}", from)?;
}
if let Some(ref to) = self.to {
write!(f, " TO {}", to)?;
}
if let Some(ref step) = self.step {
write!(f, " STEP {}", step)?;
}
Ok(())
}
}

/// ClickHouse `INTERPOLATE` clause for use in `ORDER BY` clause when using `WITH FILL` modifier.
/// Supported by [ClickHouse syntax]
///
/// [ClickHouse syntax]: <https://clickhouse.com/docs/en/sql-reference/statements/select/order-by#order-by-expr-with-fill-modifier>
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct InterpolateExpr {
pub column: Ident,
pub expr: Option<Expr>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct Interpolate {
pub exprs: Option<Vec<InterpolateExpr>>,
}

impl fmt::Display for InterpolateExpr {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}", self.column)?;
if let Some(ref expr) = self.expr {
write!(f, " AS {}", expr)?;
}
Ok(())
}
}
Expand Down
3 changes: 3 additions & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -297,6 +297,7 @@ define_keywords!(
FILE,
FILES,
FILE_FORMAT,
FILL,
FILTER,
FIRST,
FIRST_VALUE,
Expand Down Expand Up @@ -382,6 +383,7 @@ define_keywords!(
INT64,
INT8,
INTEGER,
INTERPOLATE,
INTERSECT,
INTERSECTION,
INTERVAL,
Expand Down Expand Up @@ -683,6 +685,7 @@ define_keywords!(
STDDEV_SAMP,
STDIN,
STDOUT,
STEP,
STORAGE_INTEGRATION,
STORED,
STRICT,
Expand Down
84 changes: 79 additions & 5 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7948,7 +7948,7 @@ impl<'a> Parser<'a> {
body: self.parse_insert_setexpr_boxed()?,
limit: None,
limit_by: vec![],
order_by: vec![],
order_by: None,
offset: None,
fetch: None,
locks: vec![],
Expand All @@ -7962,7 +7962,7 @@ impl<'a> Parser<'a> {
body: self.parse_update_setexpr_boxed()?,
limit: None,
limit_by: vec![],
order_by: vec![],
order_by: None,
offset: None,
fetch: None,
locks: vec![],
Expand All @@ -7974,9 +7974,19 @@ impl<'a> Parser<'a> {
let body = self.parse_boxed_query_body(0)?;

let order_by = if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
self.parse_comma_separated(Parser::parse_order_by_expr)?
let order_by_exprs = self.parse_comma_separated(Parser::parse_order_by_expr)?;
let interpolate = if dialect_of!(self is ClickHouseDialect | GenericDialect) {
self.parse_interpolations()?
} else {
None
};

Some(OrderBy {
exprs: order_by_exprs,
interpolate,
})
} else {
vec![]
None
};

let mut limit = None;
Expand Down Expand Up @@ -9208,7 +9218,7 @@ impl<'a> Parser<'a> {
subquery: Box::new(Query {
with: None,
body: Box::new(values),
order_by: vec![],
order_by: None,
limit: None,
limit_by: vec![],
offset: None,
Expand Down Expand Up @@ -10548,13 +10558,77 @@ impl<'a> Parser<'a> {
None
};

let with_fill = if dialect_of!(self is ClickHouseDialect | GenericDialect)
&& self.parse_keywords(&[Keyword::WITH, Keyword::FILL])
{
Some(self.parse_with_fill()?)
} else {
None
};

Ok(OrderByExpr {
expr,
asc,
nulls_first,
with_fill,
})
}

// Parse a WITH FILL clause (ClickHouse dialect)
// that follow the WITH FILL keywords in a ORDER BY clause
pub fn parse_with_fill(&mut self) -> Result<WithFill, ParserError> {
let from = if self.parse_keyword(Keyword::FROM) {
Some(self.parse_expr()?)
} else {
None
};

let to = if self.parse_keyword(Keyword::TO) {
Some(self.parse_expr()?)
} else {
None
};

let step = if self.parse_keyword(Keyword::STEP) {
Some(self.parse_expr()?)
} else {
None
};

Ok(WithFill { from, to, step })
}

// Parse a set of comma seperated INTERPOLATE expressions (ClickHouse dialect)
// that follow the INTERPOLATE keyword in an ORDER BY clause with the WITH FILL modifier
pub fn parse_interpolations(&mut self) -> Result<Option<Interpolate>, ParserError> {
if !self.parse_keyword(Keyword::INTERPOLATE) {
return Ok(None);
}

if self.consume_token(&Token::LParen) {
let interpolations = self.parse_comma_separated0(|p| p.parse_interpolation())?;
self.expect_token(&Token::RParen)?;
// INTERPOLATE () and INTERPOLATE ( ... ) variants
return Ok(Some(Interpolate {
exprs: Some(interpolations),
}));
}

// INTERPOLATE
Ok(Some(Interpolate { exprs: None }))
}

// Parse a INTERPOLATE expression (ClickHouse dialect)
pub fn parse_interpolation(&mut self) -> Result<InterpolateExpr, ParserError> {
let column = self.parse_identifier(false)?;
let expr = if self.parse_keyword(Keyword::AS) {
Some(self.parse_expr()?)
} else {
None
};
Ok(InterpolateExpr { column, expr })
}

/// Parse a TOP clause, MSSQL equivalent of LIMIT,
/// that follows after `SELECT [DISTINCT]`.
pub fn parse_top(&mut self) -> Result<Top, ParserError> {
Expand Down
Loading

0 comments on commit 7ccedd4

Please sign in to comment.