From 815ee5bcb36910a1ad28cdde1a53a6aa8c5e1cdc Mon Sep 17 00:00:00 2001 From: Nick Presta Date: Wed, 3 Jul 2024 11:07:47 -0400 Subject: [PATCH] Add support for WITH FILL to OrderByExpr ClickHouse supports the ORDER BY ... WITH FILL modifier: https://clickhouse.com/docs/en/sql-reference/statements/select/order-by#order-by-expr-with-fill-modifier WITH FILL itself supports a simple "from", "to", and "step" parameters, and a more sophisticated INTERPOLATE option. --- src/ast/mod.rs | 16 +++--- src/ast/query.rs | 59 ++++++++++++++++++++ src/keywords.rs | 3 ++ src/parser/mod.rs | 62 +++++++++++++++++++++ tests/sqlparser_common.rs | 111 ++++++++++++++++++++++++++++++++++++++ tests/sqlparser_mysql.rs | 1 + 6 files changed, 244 insertions(+), 8 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index c7f461418..e9b0e61a4 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -43,14 +43,14 @@ pub use self::operator::{BinaryOperator, UnaryOperator}; pub use self::query::{ AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode, ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml, - 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, SymbolDefinition, Table, - TableAlias, TableFactor, TableVersion, TableWithJoins, Top, TopQuantity, ValueTableMode, - Values, WildcardAdditionalOptions, With, + GroupByExpr, GroupByWithModifier, IdentWithAlias, IlikeSelectItem, Interpolation, 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, 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, diff --git a/src/ast/query.rs b/src/ast/query.rs index d00a0dfcc..75957a789 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -1627,6 +1627,9 @@ pub struct OrderByExpr { pub asc: Option, /// Optional `NULLS FIRST` or `NULLS LAST` pub nulls_first: Option, + /// Optional: `WITH FILL` + /// Supported by [ClickHouse syntax]: + pub with_fill: Option, } impl fmt::Display for OrderByExpr { @@ -1642,6 +1645,62 @@ 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. +#[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, + pub to: Option, + pub step: Option, + pub interpolate: Vec, +} + +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)?; + } + if !self.interpolate.is_empty() { + write!( + f, + " INTERPOLATE ({})", + display_comma_separated(&self.interpolate) + )?; + } + Ok(()) + } +} + +/// ClickHouse `INTERPOLATE` clause for use in `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 Interpolation { + pub column: Expr, + pub formula: Option, +} + +impl fmt::Display for Interpolation { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", self.column)?; + if let Some(ref formula) = self.formula { + write!(f, " AS {}", formula)?; + } Ok(()) } } diff --git a/src/keywords.rs b/src/keywords.rs index 5db55e9da..b4ae2a2cf 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -297,6 +297,7 @@ define_keywords!( FILE, FILES, FILE_FORMAT, + FILL, FILTER, FIRST, FIRST_VALUE, @@ -382,6 +383,7 @@ define_keywords!( INT64, INT8, INTEGER, + INTERPOLATE, INTERSECT, INTERSECTION, INTERVAL, @@ -678,6 +680,7 @@ define_keywords!( STDDEV_SAMP, STDIN, STDOUT, + STEP, STORAGE_INTEGRATION, STORED, STRICT, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 4e9c3836b..ffc70779c 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -10408,13 +10408,75 @@ impl<'a> Parser<'a> { None }; + let with_fill = if 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 { + 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 + }; + + let interpolate = + if self.parse_keyword(Keyword::INTERPOLATE) && self.consume_token(&Token::LParen) { + let interpolations = self.parse_interpolations()?; + self.expect_token(&Token::RParen)?; + interpolations + } else { + vec![] + }; + + Ok(WithFill { + from, + to, + step, + interpolate, }) } + // Parse a set of comma seperated INTERPOLATE expressions (ClickHouse dialect) + // that follow the INTERPOLATE keyword in a WITH FILL clause + pub fn parse_interpolations(&mut self) -> Result, ParserError> { + self.parse_comma_separated(|p| p.parse_interpolation()) + } + + // Parse a INTERPOLATE expression (ClickHouse dialect) + pub fn parse_interpolation(&mut self) -> Result { + let column = self.parse_expr()?; + let formula = if self.parse_keyword(Keyword::AS) { + Some(self.parse_expr()?) + } else { + None + }; + Ok(Interpolation { column, formula }) + } + /// Parse a TOP clause, MSSQL equivalent of LIMIT, /// that follows after `SELECT [DISTINCT]`. pub fn parse_top(&mut self) -> Result { diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index ac2133946..b4eca6491 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -2048,16 +2048,19 @@ fn parse_select_order_by() { expr: Expr::Identifier(Ident::new("lname")), asc: Some(true), nulls_first: None, + with_fill: None, }, OrderByExpr { expr: Expr::Identifier(Ident::new("fname")), asc: Some(false), nulls_first: None, + with_fill: None, }, OrderByExpr { expr: Expr::Identifier(Ident::new("id")), asc: None, nulls_first: None, + with_fill: None, }, ], select.order_by @@ -2080,11 +2083,13 @@ fn parse_select_order_by_limit() { expr: Expr::Identifier(Ident::new("lname")), asc: Some(true), nulls_first: None, + with_fill: None, }, OrderByExpr { expr: Expr::Identifier(Ident::new("fname")), asc: Some(false), nulls_first: None, + with_fill: None, }, ], select.order_by @@ -2103,11 +2108,13 @@ fn parse_select_order_by_nulls_order() { expr: Expr::Identifier(Ident::new("lname")), asc: Some(true), nulls_first: Some(true), + with_fill: None, }, OrderByExpr { expr: Expr::Identifier(Ident::new("fname")), asc: Some(false), nulls_first: Some(false), + with_fill: None, }, ], select.order_by @@ -2115,6 +2122,100 @@ fn parse_select_order_by_nulls_order() { assert_eq!(Some(Expr::Value(number("2"))), select.limit); } +#[test] +fn parse_select_order_by_with_fill() { + let sql = "SELECT id, fname, lname FROM customer WHERE id < 5 \ + ORDER BY lname ASC WITH FILL, \ + fname DESC NULLS LAST WITH FILL FROM 10 TO 20 STEP 2 INTERPOLATE (col1 AS col1 + 1) LIMIT 2"; + let select = verified_query(sql); + assert_eq!( + vec![ + OrderByExpr { + expr: Expr::Identifier(Ident::new("lname")), + asc: Some(true), + nulls_first: None, + with_fill: Some(WithFill { + from: None, + to: None, + step: None, + interpolate: vec![], + }), + }, + OrderByExpr { + expr: Expr::Identifier(Ident::new("fname")), + asc: Some(false), + nulls_first: Some(false), + with_fill: Some(WithFill { + from: Some(Expr::Value(number("10"))), + to: Some(Expr::Value(number("20"))), + step: Some(Expr::Value(number("2"))), + interpolate: vec![Interpolation { + column: Expr::Identifier(Ident::new("col1")), + formula: Some(Expr::BinaryOp { + left: Box::new(Expr::Identifier(Ident::new("col1"))), + op: BinaryOperator::Plus, + right: Box::new(Expr::Value(number("1"))), + }), + }] + }), + }, + ], + select.order_by + ); + assert_eq!(Some(Expr::Value(number("2"))), select.limit); +} + +#[test] +fn parse_with_fill() { + let sql = "SELECT fname FROM customer ORDER BY fname WITH FILL FROM 10 TO 20 STEP 2 INTERPOLATE (col1 AS col1 + 1, col2 AS col3)"; + let select = verified_query(sql); + assert_eq!( + Some(WithFill { + from: Some(Expr::Value(number("10"))), + to: Some(Expr::Value(number("20"))), + step: Some(Expr::Value(number("2"))), + interpolate: vec![ + Interpolation { + column: Expr::Identifier(Ident::new("col1")), + formula: Some(Expr::BinaryOp { + left: Box::new(Expr::Identifier(Ident::new("col1"))), + op: BinaryOperator::Plus, + right: Box::new(Expr::Value(number("1"))), + }), + }, + Interpolation { + column: Expr::Identifier(Ident::new("col2")), + formula: Some(Expr::Identifier(Ident::new("col3"))), + }, + ] + }), + select.order_by[0].with_fill + ); +} + +#[test] +fn parse_interpolation() { + let sql = "SELECT fname FROM customer ORDER BY fname WITH FILL INTERPOLATE (col1 AS col1 + 1, col2 AS col3)"; + let select = verified_query(sql); + assert_eq!( + vec![ + Interpolation { + column: Expr::Identifier(Ident::new("col1")), + formula: Some(Expr::BinaryOp { + left: Box::new(Expr::Identifier(Ident::new("col1"))), + op: BinaryOperator::Plus, + right: Box::new(Expr::Value(number("1"))), + }), + }, + Interpolation { + column: Expr::Identifier(Ident::new("col2")), + formula: Some(Expr::Identifier(Ident::new("col3"))), + }, + ], + select.order_by[0].with_fill.as_ref().unwrap().interpolate + ); +} + #[test] fn parse_select_group_by() { let sql = "SELECT id, fname, lname FROM customer GROUP BY lname, fname"; @@ -2202,6 +2303,7 @@ fn parse_select_qualify() { expr: Expr::Identifier(Ident::new("o")), asc: None, nulls_first: None, + with_fill: None, }], window_frame: None, })), @@ -2562,6 +2664,7 @@ fn parse_listagg() { }), asc: None, nulls_first: None, + with_fill: None, }, OrderByExpr { expr: Expr::Identifier(Ident { @@ -2570,6 +2673,7 @@ fn parse_listagg() { }), asc: None, nulls_first: None, + with_fill: None, }, ] }), @@ -4363,6 +4467,7 @@ fn parse_window_functions() { expr: Expr::Identifier(Ident::new("dt")), asc: Some(false), nulls_first: None, + with_fill: None, }], window_frame: None, })), @@ -4570,6 +4675,7 @@ fn test_parse_named_window() { }), asc: None, nulls_first: None, + with_fill: None, }], window_frame: None, }), @@ -7254,11 +7360,13 @@ fn parse_create_index() { expr: Expr::Identifier(Ident::new("name")), asc: None, nulls_first: None, + with_fill: None, }, OrderByExpr { expr: Expr::Identifier(Ident::new("age")), asc: Some(false), nulls_first: None, + with_fill: None, }, ]; match verified_stmt(sql) { @@ -7288,11 +7396,13 @@ fn test_create_index_with_using_function() { expr: Expr::Identifier(Ident::new("name")), asc: None, nulls_first: None, + with_fill: None, }, OrderByExpr { expr: Expr::Identifier(Ident::new("age")), asc: Some(false), nulls_first: None, + with_fill: None, }, ]; match verified_stmt(sql) { @@ -9547,6 +9657,7 @@ fn test_match_recognize() { expr: Expr::Identifier(Ident::new("price_date")), asc: None, nulls_first: None, + with_fill: None, }], measures: vec![ Measure { diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs index 4c18d4a75..3475f5013 100644 --- a/tests/sqlparser_mysql.rs +++ b/tests/sqlparser_mysql.rs @@ -1883,6 +1883,7 @@ fn parse_delete_with_order_by() { }), asc: Some(false), nulls_first: None, + with_fill: None, }], order_by );