[ClickHouse] Add support for WITH FILL to OrderByExpr (apache#1330)

Co-authored-by: Andrew Lamb <[email protected]>

Author: 2 people

src/ast/mod.rs

@@ -41,12 +41,12 @@ pub use self::ddl::{
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     AggregateItem, Cte, Distinct, ExceptSelectItem, ExcludeSelectItem, Fetch, GroupByExpr,
-    IdentWithAlias, Join, JoinConstraint, JoinOperator, LateralView, LockClause, LockType,
-    NamedWindowDefinition, NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem,
-    ReplaceSelectElement, ReplaceSelectItem, SamplingMethod, Select, SelectInto, SelectItem,
-    SelectionCount, SetExpr, SetOperator, SetQuantifier, Table, TableAlias, TableFactor,
-    TableSampleSeed, TableVersion, TableWithJoins, Top, ValueTableMode, Values,
-    WildcardAdditionalOptions, With,
+    IdentWithAlias, Interpolate, InterpolateExpr, Join, JoinConstraint, JoinOperator, LateralView,
+    LockClause, LockType, NamedWindowDefinition, NonBlock, Offset, OffsetRows, OrderBy,
+    OrderByExpr, Query, RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, SamplingMethod,
+    Select, SelectInto, SelectItem, SelectionCount, SetExpr, SetOperator, SetQuantifier, Table,
+    TableAlias, TableFactor, TableSampleSeed, TableVersion, TableWithJoins, Top, ValueTableMode,
+    Values, WildcardAdditionalOptions, With, WithFill,
 };
 pub use self::value::{
     escape_quoted_string, DateTimeField, DollarQuotedString, ObjectConstantKeyValue,

src/ast/query.rs

@@ -32,7 +32,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>,
 
@@ -53,8 +53,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}")?;
@@ -100,6 +109,17 @@ pub enum SetExpr {
     Table(Box<Table>),
 }
 
+impl SetExpr {
+    /// If this `SetExpr` is a `SELECT`, returns the [`Select`].
+    pub fn as_select(&self) -> Option<&Select> {
+        if let Self::Select(select) = self {
+            Some(&**select)
+        } else {
+            None
+        }
+    }
+}
+
 impl fmt::Display for SetExpr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
@@ -1223,6 +1243,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))]
@@ -1233,6 +1265,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 {
@@ -1248,6 +1283,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: WithSpan<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(())
     }
 }

src/keywords.rs

@@ -284,6 +284,7 @@ define_keywords!(
     FILE,
     FILES,
     FILE_FORMAT,
+    FILL,
     FILTER,
     FINAL,
     FIRST,
@@ -358,6 +359,7 @@ define_keywords!(
     INT64,
     INT8,
     INTEGER,
+    INTERPOLATE,
     INTERSECT,
     INTERSECTION,
     INTERVAL,
@@ -620,6 +622,7 @@ define_keywords!(
     STDDEV_SAMP,
     STDIN,
     STDOUT,
+    STEP,
     STORAGE_INTEGRATION,
     STORED,
     STRICT,