feat(query): support array_sort

src/query/ast/src/ast/expr.rs

@@ -179,6 +179,12 @@ pub enum Expr {
     },
     /// The `Array` expr
     Array { span: Span, exprs: Vec<Expr> },
+    ArraySort {
+        span: Span,
+        expr: Box<Expr>,
+        asc: bool,
+        null_first: bool,
+    },
     /// The `Interval 1 DAY` expr
     Interval {
         span: Span,
@@ -383,6 +389,7 @@ impl Expr {
             | Expr::Subquery { span, .. }
             | Expr::MapAccess { span, .. }
             | Expr::Array { span, .. }
+            | Expr::ArraySort { span, .. }
             | Expr::Interval { span, .. }
             | Expr::DateAdd { span, .. }
             | Expr::DateSub { span, .. }
@@ -866,6 +873,26 @@ impl Display for Expr {
                 write_comma_separated_list(f, exprs)?;
                 write!(f, "]")?;
             }
+            Expr::ArraySort {
+                expr,
+                asc,
+                null_first,
+                ..
+            } => {
+                write!(f, "ARRAY_SORT(")?;
+                write!(f, "{expr})")?;
+                if *asc {
+                    write!(f, " , 'ASC'")?;
+                } else {
+                    write!(f, " , 'DESC'")?;
+                }
+                if *null_first {
+                    write!(f, " , 'NULLS FIRST'")?;
+                } else {
+                    write!(f, " , 'NULLS LAST'")?;
+                }
+                write!(f, ")")?;
+            }
             Expr::Interval { expr, unit, .. } => {
                 write!(f, "INTERVAL {expr} {unit}")?;
             }

src/query/ast/src/ast/format/syntax/expr.rs

@@ -318,6 +318,45 @@ pub(crate) fn pretty_expr(expr: Expr) -> RcDoc<'static> {
         Expr::Array { exprs, .. } => RcDoc::text("[")
             .append(inline_comma(exprs.into_iter().map(pretty_expr)))
             .append(RcDoc::text("]")),
+        Expr::ArraySort {
+            expr,
+            asc,
+            null_first,
+            ..
+        } => {
+            let res = pretty_expr(*expr);
+            if asc {
+                res.clone()
+                    .append(RcDoc::text(","))
+                    .append(RcDoc::space())
+                    .append(RcDoc::text("'"))
+                    .append(RcDoc::text("ASC"))
+                    .append(RcDoc::text("'"));
+            } else {
+                res.clone()
+                    .append(RcDoc::text(","))
+                    .append(RcDoc::space())
+                    .append(RcDoc::text("'"))
+                    .append(RcDoc::text("DESC"))
+                    .append(RcDoc::text("'"));
+            }
+            if null_first {
+                res.clone()
+                    .append(RcDoc::text(","))
+                    .append(RcDoc::space())
+                    .append(RcDoc::text("'"))
+                    .append(RcDoc::text("NULL FIRST"))
+                    .append(RcDoc::text("'"));
+            } else {
+                res.clone()
+                    .append(RcDoc::text(","))
+                    .append(RcDoc::space())
+                    .append(RcDoc::text("'"))
+                    .append(RcDoc::text("NULL LAST"))
+                    .append(RcDoc::text("'"));
+            }
+            res.clone().append(RcDoc::text(")"))
+        }
         Expr::Interval { expr, unit, .. } => RcDoc::text("INTERVAL")
             .append(RcDoc::space())
             .append(pretty_expr(*expr))

src/query/ast/src/parser/expr.rs

@@ -278,6 +278,14 @@ pub enum ExprElement {
     Array {
         exprs: Vec<Expr>,
     },
+    /// ARRAY_SORT([1,2,3], ASC|DESC, NULLS FIRST|LAST)
+    ArraySort {
+        expr: Box<Expr>,
+        // Optional `ASC` or `DESC`
+        asc: Option<String>,
+        // Optional `NULLS FIRST` or `NULLS LAST`
+        nulls_first: Option<String>,
+    },
     Interval {
         expr: Expr,
         unit: IntervalKind,
@@ -461,6 +469,41 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
                 span: transform_span(elem.span.0),
                 exprs,
             },
+            ExprElement::ArraySort {
+                expr,
+                asc,
+                nulls_first,
+            } => {
+                let asc = if let Some(asc) = asc {
+                    if asc.to_lowercase() == "asc" {
+                        true
+                    } else if asc.to_lowercase() == "desc" {
+                        false
+                    } else {
+                        return Err("Sorting order must be either ASC or DESC");
+                    }
+                } else {
+                    true
+                };
+                let null_first = if let Some(nulls_first) = nulls_first {
+                    let null_first = nulls_first.trim().to_lowercase();
+                    if null_first == "nulls first" {
+                        true
+                    } else if null_first == "nulls last" {
+                        false
+                    } else {
+                        return Err("Null sorting order must be either NULLS FIRST or NULLS LAST");
+                    }
+                } else {
+                    true
+                };
+                Expr::ArraySort {
+                    span: transform_span(elem.span.0),
+                    expr,
+                    asc,
+                    null_first,
+                }
+            }
             ExprElement::Interval { expr, unit } => Expr::Interval {
                 span: transform_span(elem.span.0),
                 expr: Box::new(expr),
@@ -830,6 +873,22 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
             ExprElement::Array { exprs }
         },
     );
+    // ARRAY_SORT([...], ASC | DESC, NULLS FIRST | LAST)
+    let array_sort = map(
+        rule! {
+            ( ARRAY_SORT )
+            ~ "("
+            ~ #subexpr(0)
+            ~ ( "," ~ #literal_string )?
+            ~ ( "," ~ #literal_string )?
+            ~ ")"
+        },
+        |(_, _, expr, opt_asc, opt_null_first, _)| ExprElement::ArraySort {
+            expr: Box::new(expr),
+            asc: opt_asc.map(|(_, asc)| asc),
+            nulls_first: opt_null_first.map(|(_, first_last)| first_last),
+        },
+    );
     let date_add = map(
         rule! {
             DATE_ADD ~ "(" ~ #interval_kind ~ "," ~ #subexpr(0) ~ "," ~ #subexpr(0) ~ ")"
@@ -890,6 +949,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
             | #extract : "`EXTRACT((YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND) FROM ...)`"
             | #position : "`POSITION(... IN ...)`"
             | #substring : "`SUBSTRING(... [FROM ...] [FOR ...])`"
+            | #array_sort : "`ARRAY_SORT([...], 'ASC' | 'DESC', 'NULLS FIRST' | 'NULLS LAST')`"
             | #trim : "`TRIM(...)`"
             | #trim_from : "`TRIM([(BOTH | LEADEING | TRAILING) ... FROM ...)`"
         ),

src/query/ast/src/parser/token.rs

@@ -751,6 +751,8 @@ pub enum TokenKind {
     TRANSIENT,
     #[token("TRIM", ignore(ascii_case))]
     TRIM,
+    #[token("ARRAY_SORT", ignore(ascii_case))]
+    ARRAY_SORT,
     #[token("TRUE", ignore(ascii_case))]
     TRUE,
     #[token("TRUNCATE", ignore(ascii_case))]
@@ -981,6 +983,7 @@ impl TokenKind {
             | TokenKind::TRAILING
             // | TokenKind::TREAT
             | TokenKind::TRIM
+            | TokenKind::ARRAY_SORT
             | TokenKind::TRUE
             | TokenKind::TRY_CAST
             // | TokenKind::UNIQUE

src/query/ast/src/visitors/visitor.rs

@@ -257,6 +257,10 @@ pub trait Visitor<'ast>: Sized {
         }
     }
 
+    fn visit_array_sort(&mut self, _span: Span, expr: &'ast Expr, _asc: bool, _null_first: bool) {
+        walk_expr(self, expr);
+    }
+
     fn visit_interval(&mut self, _span: Span, expr: &'ast Expr, _unit: &'ast IntervalKind) {
         walk_expr(self, expr);
     }

src/query/ast/src/visitors/visitor_mut.rs

@@ -261,6 +261,10 @@ pub trait VisitorMut: Sized {
         }
     }
 
+    fn visit_array_sort(&mut self, _span: Span, expr: &mut Expr, _asc: bool, _null_first: bool) {
+        walk_expr_mut(self, expr);
+    }
+
     fn visit_interval(&mut self, _span: Span, expr: &mut Expr, _unit: &mut IntervalKind) {
         walk_expr_mut(self, expr);
     }

src/query/ast/src/visitors/walk.rs

@@ -117,6 +117,12 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expr: &'a Expr) {
             accessor,
         } => visitor.visit_map_access(*span, expr, accessor),
         Expr::Array { span, exprs } => visitor.visit_array(*span, exprs),
+        Expr::ArraySort {
+            span,
+            expr,
+            asc,
+            null_first,
+        } => visitor.visit_array_sort(*span, expr, *asc, *null_first),
         Expr::Interval { span, expr, unit } => visitor.visit_interval(*span, expr, unit),
         Expr::DateAdd {
             span,

src/query/ast/src/visitors/walk_mut.rs

@@ -117,6 +117,12 @@ pub fn walk_expr_mut<V: VisitorMut>(visitor: &mut V, expr: &mut Expr) {
             accessor,
         } => visitor.visit_map_access(*span, expr, accessor),
         Expr::Array { span, exprs } => visitor.visit_array(*span, exprs),
+        Expr::ArraySort {
+            span,
+            expr,
+            asc,
+            null_first,
+        } => visitor.visit_array_sort(*span, expr, *asc, *null_first),
         Expr::Interval { span, expr, unit } => visitor.visit_interval(*span, expr, unit),
         Expr::DateAdd {
             span,

src/query/ast/tests/it/parser.rs

@@ -540,6 +540,9 @@ fn test_expr() {
         r#"position('a' in str)"#,
         r#"substring(a from b for c)"#,
         r#"substring(a, b, c)"#,
+        r#"array_sort([2])"#,
+        r#"array_sort([2,0.1], 'ASC')"#,
+        r#"array_sort([3,2], 'DESC', 'NULLS FIRST')"#,
         r#"col1::UInt8"#,
         r#"(arr[0]:a).b"#,
         r#"arr[4]["k"]"#,

src/query/ast/tests/it/testdata/expr-error.txt

@@ -55,7 +55,7 @@ error:
   --> SQL:1:10
   |
 1 | CAST(col1)
-  | ----     ^ expected `AS`, `,`, `(`, `.`, `IS`, `NOT`, or 55 more ...
+  | ----     ^ expected `AS`, `,`, `(`, `.`, `IS`, `NOT`, or 56 more ...
   | |         
   | while parsing `CAST(... AS ...)`
   | while parsing expression

src/query/ast/tests/it/testdata/expr.txt

@@ -1646,6 +1646,109 @@ Substring {
 }
 
 
+---------- Input ----------
+array_sort([2])
+---------- Output ---------
+ARRAY_SORT([2]) , 'ASC' , 'NULLS FIRST')
+---------- AST ------------
+ArraySort {
+    span: Some(
+        0..15,
+    ),
+    expr: Array {
+        span: Some(
+            11..14,
+        ),
+        exprs: [
+            Literal {
+                span: Some(
+                    12..13,
+                ),
+                lit: Integer(
+                    2,
+                ),
+            },
+        ],
+    },
+    asc: true,
+    null_first: true,
+}
+
+
+---------- Input ----------
+array_sort([2,0.1], 'ASC')
+---------- Output ---------
+ARRAY_SORT([2, 0.1]) , 'ASC' , 'NULLS FIRST')
+---------- AST ------------
+ArraySort {
+    span: Some(
+        0..26,
+    ),
+    expr: Array {
+        span: Some(
+            11..18,
+        ),
+        exprs: [
+            Literal {
+                span: Some(
+                    12..13,
+                ),
+                lit: Integer(
+                    2,
+                ),
+            },
+            Literal {
+                span: Some(
+                    14..17,
+                ),
+                lit: Float(
+                    0.1,
+                ),
+            },
+        ],
+    },
+    asc: true,
+    null_first: true,
+}
+
+
+---------- Input ----------
+array_sort([3,2], 'DESC', 'NULLS FIRST')
+---------- Output ---------
+ARRAY_SORT([3, 2]) , 'DESC' , 'NULLS FIRST')
+---------- AST ------------
+ArraySort {
+    span: Some(
+        0..40,
+    ),
+    expr: Array {
+        span: Some(
+            11..16,
+        ),
+        exprs: [
+            Literal {
+                span: Some(
+                    12..13,
+                ),
+                lit: Integer(
+                    3,
+                ),
+            },
+            Literal {
+                span: Some(
+                    14..15,
+                ),
+                lit: Integer(
+                    2,
+                ),
+            },
+        ],
+    },
+    asc: false,
+    null_first: true,
+}
+
+
 ---------- Input ----------
 col1::UInt8
 ---------- Output ---------