Add support for PostgreSQL's ORDER BY ... USING <operator> clause (apache#2246)

Author: LucaCappelletti94

src/ast/mod.rs

@@ -99,7 +99,7 @@ pub use self::query::{
     JsonTableNestedColumn, LateralView, LimitClause, LockClause, LockType, MatchRecognizePattern,
     MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
     OffsetRows, OpenJsonTableColumn, OrderBy, OrderByExpr, OrderByKind, OrderByOptions,
-    PipeOperator, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
+    OrderBySort, PipeOperator, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
     RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
     SelectFlavor, SelectInto, SelectItem, SelectItemQualifiedWildcardKind, SelectModifiers,
     SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table, TableAlias,

src/ast/query.rs

@@ -2927,7 +2927,7 @@ impl fmt::Display for OrderBy {
 pub struct OrderByExpr {
     /// The expression to order by.
     pub expr: Expr,
-    /// Ordering options such as `ASC`/`DESC` and `NULLS` behavior.
+    /// Ordering options such as `ASC`/`DESC`/`USING <operator>` and `NULLS` behavior.
     pub options: OrderByOptions,
     /// Optional `WITH FILL` clause (ClickHouse extension) which specifies how to fill gaps.
     pub with_fill: Option<WithFill>,
@@ -2945,7 +2945,8 @@ impl From<Ident> for OrderByExpr {
 
 impl fmt::Display for OrderByExpr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{}{}", self.expr, self.options)?;
+        write!(f, "{}", self.expr)?;
+        write!(f, "{}", self.options)?;
         if let Some(ref with_fill) = self.with_fill {
             write!(f, " {with_fill}")?
         }
@@ -3020,22 +3021,47 @@ impl fmt::Display for InterpolateExpr {
     }
 }
 
-#[derive(Default, Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+/// The sort order for an `ORDER BY` expression.
+///
+/// See PostgreSQL `USING` operator:
+/// <https://www.postgresql.org/docs/current/sql-select.html>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-/// Options for an `ORDER BY` expression (ASC/DESC and NULLS FIRST/LAST).
+pub enum OrderBySort {
+    /// `ASC`
+    Asc,
+    /// `DESC`
+    Desc,
+    /// PostgreSQL `USING <operator>` ordering.
+    ///
+    /// See <https://www.postgresql.org/docs/current/sql-select.html>
+    Using(ObjectName),
+}
+
+#[derive(Default, Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+/// Options for an `ORDER BY` expression.
 pub struct OrderByOptions {
-    /// Optional `ASC` (`Some(true)`) or `DESC` (`Some(false)`).
-    pub asc: Option<bool>,
+    /// Optional sort order: `ASC`, `DESC`, or `USING <operator>`.
+    pub sort: Option<OrderBySort>,
     /// Optional `NULLS FIRST` (`Some(true)`) or `NULLS LAST` (`Some(false)`).
     pub nulls_first: Option<bool>,
 }
 
 impl fmt::Display for OrderByOptions {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self.asc {
-            Some(true) => write!(f, " ASC")?,
-            Some(false) => write!(f, " DESC")?,
+        match &self.sort {
+            Some(OrderBySort::Asc) => write!(f, " ASC")?,
+            Some(OrderBySort::Desc) => write!(f, " DESC")?,
+            Some(OrderBySort::Using(op)) => {
+                if op.0.len() > 1 {
+                    write!(f, " USING OPERATOR({op})")?;
+                } else {
+                    write!(f, " USING {op}")?;
+                }
+            }
             None => (),
         }
         match self.nulls_first {

src/dialect/mod.rs

@@ -1414,6 +1414,14 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports PostgreSQL-style ordering operators:
+    /// `ORDER BY expr USING <operator>`.
+    ///
+    /// For example: `SELECT * FROM t ORDER BY a USING <`.
+    fn supports_order_by_using_operator(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports `SET NAMES <charset_name> [COLLATE <collation_name>]`.
     ///
     /// - [MySQL](https://dev.mysql.com/doc/refman/8.4/en/set-names.html)

src/dialect/postgresql.rs

@@ -278,6 +278,10 @@ impl Dialect for PostgreSqlDialect {
         true
     }
 
+    fn supports_order_by_using_operator(&self) -> bool {
+        true
+    }
+
     fn supports_set_names(&self) -> bool {
         true
     }

src/parser/mod.rs

@@ -1378,7 +1378,7 @@ impl<'a> Parser<'a> {
         }
         let alias = self.parse_optional_alias_inner(None, validator)?;
         let order_by = OrderByOptions {
-            asc: self.parse_asc_desc(),
+            sort: self.parse_optional_order_by_sort(),
             nulls_first: None,
         };
         Ok(ExprWithAliasAndOrderBy {
@@ -18852,6 +18852,15 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// Parse ASC or DESC and map to [OrderBySort].
+    fn parse_optional_order_by_sort(&mut self) -> Option<OrderBySort> {
+        match self.parse_asc_desc() {
+            Some(true) => Some(OrderBySort::Asc),
+            Some(false) => Some(OrderBySort::Desc),
+            None => None,
+        }
+    }
+
     /// Parse an [OrderByExpr] expression.
     pub fn parse_order_by_expr(&mut self) -> Result<OrderByExpr, ParserError> {
         self.parse_order_by_expr_inner(false)
@@ -18888,7 +18897,18 @@ impl<'a> Parser<'a> {
             None
         };
 
-        let options = self.parse_order_by_options()?;
+        let options = if !with_operator_class
+            && self.dialect.supports_order_by_using_operator()
+            && self.parse_keyword(Keyword::USING)
+        {
+            let op = self.parse_order_by_using_operator()?;
+            OrderByOptions {
+                sort: Some(OrderBySort::Using(op)),
+                nulls_first: self.parse_null_ordering_modifier(),
+            }
+        } else {
+            self.parse_order_by_options()?
+        };
 
         let with_fill = if self.dialect.supports_with_fill()
             && self.parse_keywords(&[Keyword::WITH, Keyword::FILL])
@@ -18908,18 +18928,33 @@ impl<'a> Parser<'a> {
         ))
     }
 
-    fn parse_order_by_options(&mut self) -> Result<OrderByOptions, ParserError> {
-        let asc = self.parse_asc_desc();
+    fn parse_order_by_using_operator(&mut self) -> Result<ObjectName, ParserError> {
+        if self.parse_keyword(Keyword::OPERATOR) {
+            self.expect_token(&Token::LParen)?;
+            let operator_name = self.parse_operator_name()?;
+            self.expect_token(&Token::RParen)?;
+            return Ok(operator_name);
+        }
+
+        let token = self.next_token();
+        Ok(ObjectName::from(vec![Ident::new(token.token.to_string())]))
+    }
 
-        let nulls_first = if self.parse_keywords(&[Keyword::NULLS, Keyword::FIRST]) {
+    fn parse_null_ordering_modifier(&mut self) -> Option<bool> {
+        if self.parse_keywords(&[Keyword::NULLS, Keyword::FIRST]) {
             Some(true)
         } else if self.parse_keywords(&[Keyword::NULLS, Keyword::LAST]) {
             Some(false)
         } else {
             None
-        };
+        }
+    }
+
+    fn parse_order_by_options(&mut self) -> Result<OrderByOptions, ParserError> {
+        let sort = self.parse_optional_order_by_sort();
+        let nulls_first = self.parse_null_ordering_modifier();
 
-        Ok(OrderByOptions { asc, nulls_first })
+        Ok(OrderByOptions { sort, nulls_first })
     }
 
     // Parse a WITH FILL clause (ClickHouse dialect)
@@ -21170,7 +21205,7 @@ mod tests {
                 column: OrderByExpr {
                     expr: Expr::Identifier(name.into()),
                     options: OrderByOptions {
-                        asc: None,
+                        sort: None,
                         nulls_first: None,
                     },
                     with_fill: None,

tests/sqlparser_bigquery.rs

@@ -2731,7 +2731,7 @@ fn test_export_data() {
                     kind: OrderByKind::Expressions(vec![OrderByExpr {
                         expr: Expr::Identifier(Ident::new("field1")),
                         options: OrderByOptions {
-                            asc: None,
+                            sort: None,
                             nulls_first: None,
                         },
                         with_fill: None,
@@ -2837,7 +2837,7 @@ fn test_export_data() {
                     kind: OrderByKind::Expressions(vec![OrderByExpr {
                         expr: Expr::Identifier(Ident::new("field1")),
                         options: OrderByOptions {
-                            asc: None,
+                            sort: None,
                             nulls_first: None,
                         },
                         with_fill: None,

tests/sqlparser_clickhouse.rs

@@ -390,7 +390,7 @@ fn parse_alter_table_add_projection() {
                             kind: OrderByKind::Expressions(vec![OrderByExpr {
                                 expr: Identifier(Ident::new("b")),
                                 options: OrderByOptions {
-                                    asc: None,
+                                    sort: None,
                                     nulls_first: None,
                                 },
                                 with_fill: None,
@@ -1218,7 +1218,7 @@ fn parse_select_order_by_with_fill_interpolate() {
                 OrderByExpr {
                     expr: Expr::Identifier(Ident::new("fname")),
                     options: OrderByOptions {
-                        asc: Some(true),
+                        sort: Some(OrderBySort::Asc),
                         nulls_first: Some(true),
                     },
                     with_fill: Some(WithFill {
@@ -1230,7 +1230,7 @@ fn parse_select_order_by_with_fill_interpolate() {
                 OrderByExpr {
                     expr: Expr::Identifier(Ident::new("lname")),
                     options: OrderByOptions {
-                        asc: Some(false),
+                        sort: Some(OrderBySort::Desc),
                         nulls_first: Some(false),
                     },
                     with_fill: Some(WithFill {