Skip to content

Commit 4b3a22b

Browse files
TCeasonTCeason
committed
support array_sort([...] (, ASC|DESC) (, NULLS FIRST|LAST))
1 parent 79aba03 commit 4b3a22b

File tree

18 files changed

+479
-166
lines changed

18 files changed

+479
-166
lines changed

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

+27
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,12 @@ pub enum Expr {
179179
},
180180
/// The `Array` expr
181181
Array { span: Span, exprs: Vec<Expr> },
182+
ArraySort {
183+
span: Span,
184+
expr: Box<Expr>,
185+
asc: bool,
186+
null_first: bool,
187+
},
182188
/// The `Interval 1 DAY` expr
183189
Interval {
184190
span: Span,
@@ -383,6 +389,7 @@ impl Expr {
383389
| Expr::Subquery { span, .. }
384390
| Expr::MapAccess { span, .. }
385391
| Expr::Array { span, .. }
392+
| Expr::ArraySort { span, .. }
386393
| Expr::Interval { span, .. }
387394
| Expr::DateAdd { span, .. }
388395
| Expr::DateSub { span, .. }
@@ -866,6 +873,26 @@ impl Display for Expr {
866873
write_comma_separated_list(f, exprs)?;
867874
write!(f, "]")?;
868875
}
876+
Expr::ArraySort {
877+
expr,
878+
asc,
879+
null_first,
880+
..
881+
} => {
882+
write!(f, "ARRAY_SORT(")?;
883+
write!(f, "{expr})")?;
884+
if *asc {
885+
write!(f, " , ASC")?;
886+
} else {
887+
write!(f, " , DESC")?;
888+
}
889+
if *null_first {
890+
write!(f, " , NULLS FIRST")?;
891+
} else {
892+
write!(f, " , NULLS LAST")?;
893+
}
894+
write!(f, ")")?;
895+
}
869896
Expr::Interval { expr, unit, .. } => {
870897
write!(f, "INTERVAL {expr} {unit}")?;
871898
}

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

+32
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,38 @@ pub(crate) fn pretty_expr(expr: Expr) -> RcDoc<'static> {
318318
Expr::Array { exprs, .. } => RcDoc::text("[")
319319
.append(inline_comma(exprs.into_iter().map(pretty_expr)))
320320
.append(RcDoc::text("]")),
321+
Expr::ArraySort {
322+
expr,
323+
asc,
324+
null_first,
325+
..
326+
} => {
327+
let res = pretty_expr(*expr);
328+
if asc {
329+
res.clone()
330+
.append(RcDoc::text(","))
331+
.append(RcDoc::space())
332+
.append(RcDoc::text("ASC"));
333+
} else {
334+
res.clone()
335+
.append(RcDoc::text(","))
336+
.append(RcDoc::space())
337+
.append(RcDoc::text("DESC"));
338+
}
339+
if null_first {
340+
res.clone()
341+
.append(RcDoc::text(","))
342+
.append(RcDoc::space())
343+
.append(RcDoc::text("NULL FIRST"))
344+
.append(RcDoc::space());
345+
} else {
346+
res.clone()
347+
.append(RcDoc::text(","))
348+
.append(RcDoc::space())
349+
.append(RcDoc::text("NULL LAST"));
350+
}
351+
res.clone().append(RcDoc::text(")"))
352+
}
321353
Expr::Interval { expr, unit, .. } => RcDoc::text("INTERVAL")
322354
.append(RcDoc::space())
323355
.append(pretty_expr(*expr))

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

+43
Original file line numberDiff line numberDiff line change
@@ -278,6 +278,14 @@ pub enum ExprElement {
278278
Array {
279279
exprs: Vec<Expr>,
280280
},
281+
/// ARRAY_SORT([1,2,3], ASC|DESC, NULLS FIRST|LAST)
282+
ArraySort {
283+
expr: Box<Expr>,
284+
// Optional `ASC` or `DESC`
285+
asc: Option<bool>,
286+
// Optional `NULLS FIRST` or `NULLS LAST`
287+
nulls_first: Option<bool>,
288+
},
281289
Interval {
282290
expr: Expr,
283291
unit: IntervalKind,
@@ -461,6 +469,24 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
461469
span: transform_span(elem.span.0),
462470
exprs,
463471
},
472+
ExprElement::ArraySort {
473+
expr,
474+
asc,
475+
nulls_first,
476+
} => {
477+
let asc = if let Some(asc) = asc { asc } else { true };
478+
let null_first = if let Some(nulls_first) = nulls_first {
479+
nulls_first
480+
} else {
481+
true
482+
};
483+
Expr::ArraySort {
484+
span: transform_span(elem.span.0),
485+
expr,
486+
asc,
487+
null_first,
488+
}
489+
}
464490
ExprElement::Interval { expr, unit } => Expr::Interval {
465491
span: transform_span(elem.span.0),
466492
expr: Box::new(expr),
@@ -830,6 +856,22 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
830856
ExprElement::Array { exprs }
831857
},
832858
);
859+
// ARRAY_SORT([...], ASC | DESC, NULLS FIRST | LAST)
860+
let array_sort = map(
861+
rule! {
862+
( ARRAY_SORT )
863+
~ "("
864+
~ #subexpr(0)
865+
~ ( "," ~ ( ASC | DESC ) )?
866+
~ ( "," ~ NULLS ~ ( FIRST | LAST ) )?
867+
~ ")"
868+
},
869+
|(_, _, expr, opt_asc, opt_null_first, _)| ExprElement::ArraySort {
870+
expr: Box::new(expr),
871+
asc: opt_asc.map(|(_, asc)| asc.kind == ASC),
872+
nulls_first: opt_null_first.map(|(_, _, first_last)| first_last.kind == FIRST),
873+
},
874+
);
833875
let date_add = map(
834876
rule! {
835877
DATE_ADD ~ "(" ~ #interval_kind ~ "," ~ #subexpr(0) ~ "," ~ #subexpr(0) ~ ")"
@@ -890,6 +932,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
890932
| #extract : "`EXTRACT((YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND) FROM ...)`"
891933
| #position : "`POSITION(... IN ...)`"
892934
| #substring : "`SUBSTRING(... [FROM ...] [FOR ...])`"
935+
| #array_sort : "`ARRAY_SORT([...], ASC | DESC, NULLS FIRST | LAST)`"
893936
| #trim : "`TRIM(...)`"
894937
| #trim_from : "`TRIM([(BOTH | LEADEING | TRAILING) ... FROM ...)`"
895938
),

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

+3
Original file line numberDiff line numberDiff line change
@@ -751,6 +751,8 @@ pub enum TokenKind {
751751
TRANSIENT,
752752
#[token("TRIM", ignore(ascii_case))]
753753
TRIM,
754+
#[token("ARRAY_SORT", ignore(ascii_case))]
755+
ARRAY_SORT,
754756
#[token("TRUE", ignore(ascii_case))]
755757
TRUE,
756758
#[token("TRUNCATE", ignore(ascii_case))]
@@ -981,6 +983,7 @@ impl TokenKind {
981983
| TokenKind::TRAILING
982984
// | TokenKind::TREAT
983985
| TokenKind::TRIM
986+
| TokenKind::ARRAY_SORT
984987
| TokenKind::TRUE
985988
| TokenKind::TRY_CAST
986989
// | TokenKind::UNIQUE

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

+4
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,10 @@ pub trait Visitor<'ast>: Sized {
257257
}
258258
}
259259

260+
fn visit_array_sort(&mut self, _span: Span, expr: &'ast Expr, _asc: bool, _null_first: bool) {
261+
walk_expr(self, expr);
262+
}
263+
260264
fn visit_interval(&mut self, _span: Span, expr: &'ast Expr, _unit: &'ast IntervalKind) {
261265
walk_expr(self, expr);
262266
}

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

+4
Original file line numberDiff line numberDiff line change
@@ -261,6 +261,10 @@ pub trait VisitorMut: Sized {
261261
}
262262
}
263263

264+
fn visit_array_sort(&mut self, _span: Span, expr: &mut Expr, _asc: bool, _null_first: bool) {
265+
walk_expr_mut(self, expr);
266+
}
267+
264268
fn visit_interval(&mut self, _span: Span, expr: &mut Expr, _unit: &mut IntervalKind) {
265269
walk_expr_mut(self, expr);
266270
}

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

+6
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,12 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expr: &'a Expr) {
117117
accessor,
118118
} => visitor.visit_map_access(*span, expr, accessor),
119119
Expr::Array { span, exprs } => visitor.visit_array(*span, exprs),
120+
Expr::ArraySort {
121+
span,
122+
expr,
123+
asc,
124+
null_first,
125+
} => visitor.visit_array_sort(*span, expr, *asc, *null_first),
120126
Expr::Interval { span, expr, unit } => visitor.visit_interval(*span, expr, unit),
121127
Expr::DateAdd {
122128
span,

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

+6
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,12 @@ pub fn walk_expr_mut<V: VisitorMut>(visitor: &mut V, expr: &mut Expr) {
117117
accessor,
118118
} => visitor.visit_map_access(*span, expr, accessor),
119119
Expr::Array { span, exprs } => visitor.visit_array(*span, exprs),
120+
Expr::ArraySort {
121+
span,
122+
expr,
123+
asc,
124+
null_first,
125+
} => visitor.visit_array_sort(*span, expr, *asc, *null_first),
120126
Expr::Interval { span, expr, unit } => visitor.visit_interval(*span, expr, unit),
121127
Expr::DateAdd {
122128
span,

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

+3
Original file line numberDiff line numberDiff line change
@@ -540,6 +540,9 @@ fn test_expr() {
540540
r#"position('a' in str)"#,
541541
r#"substring(a from b for c)"#,
542542
r#"substring(a, b, c)"#,
543+
r#"array_sort([2])"#,
544+
r#"array_sort([2,0.1], ASC)"#,
545+
r#"array_sort([3,2], DESC, NULLS FIRST)"#,
543546
r#"col1::UInt8"#,
544547
r#"(arr[0]:a).b"#,
545548
r#"arr[4]["k"]"#,

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ error:
5555
--> SQL:1:10
5656
|
5757
1 | CAST(col1)
58-
| ---- ^ expected `AS`, `,`, `(`, `.`, `IS`, `NOT`, or 55 more ...
58+
| ---- ^ expected `AS`, `,`, `(`, `.`, `IS`, `NOT`, or 56 more ...
5959
| |
6060
| while parsing `CAST(... AS ...)`
6161
| while parsing expression

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

+103
Original file line numberDiff line numberDiff line change
@@ -1646,6 +1646,109 @@ Substring {
16461646
}
16471647

16481648

1649+
---------- Input ----------
1650+
array_sort([2])
1651+
---------- Output ---------
1652+
ARRAY_SORT([2]) , ASC , NULLS FIRST)
1653+
---------- AST ------------
1654+
ArraySort {
1655+
span: Some(
1656+
0..15,
1657+
),
1658+
expr: Array {
1659+
span: Some(
1660+
11..14,
1661+
),
1662+
exprs: [
1663+
Literal {
1664+
span: Some(
1665+
12..13,
1666+
),
1667+
lit: Integer(
1668+
2,
1669+
),
1670+
},
1671+
],
1672+
},
1673+
asc: true,
1674+
null_first: true,
1675+
}
1676+
1677+
1678+
---------- Input ----------
1679+
array_sort([2,0.1], ASC)
1680+
---------- Output ---------
1681+
ARRAY_SORT([2, 0.1]) , ASC , NULLS FIRST)
1682+
---------- AST ------------
1683+
ArraySort {
1684+
span: Some(
1685+
0..24,
1686+
),
1687+
expr: Array {
1688+
span: Some(
1689+
11..18,
1690+
),
1691+
exprs: [
1692+
Literal {
1693+
span: Some(
1694+
12..13,
1695+
),
1696+
lit: Integer(
1697+
2,
1698+
),
1699+
},
1700+
Literal {
1701+
span: Some(
1702+
14..17,
1703+
),
1704+
lit: Float(
1705+
0.1,
1706+
),
1707+
},
1708+
],
1709+
},
1710+
asc: true,
1711+
null_first: true,
1712+
}
1713+
1714+
1715+
---------- Input ----------
1716+
array_sort([3,2], DESC, NULLS FIRST)
1717+
---------- Output ---------
1718+
ARRAY_SORT([3, 2]) , DESC , NULLS FIRST)
1719+
---------- AST ------------
1720+
ArraySort {
1721+
span: Some(
1722+
0..36,
1723+
),
1724+
expr: Array {
1725+
span: Some(
1726+
11..16,
1727+
),
1728+
exprs: [
1729+
Literal {
1730+
span: Some(
1731+
12..13,
1732+
),
1733+
lit: Integer(
1734+
3,
1735+
),
1736+
},
1737+
Literal {
1738+
span: Some(
1739+
14..15,
1740+
),
1741+
lit: Integer(
1742+
2,
1743+
),
1744+
},
1745+
],
1746+
},
1747+
asc: false,
1748+
null_first: true,
1749+
}
1750+
1751+
16491752
---------- Input ----------
16501753
col1::UInt8
16511754
---------- Output ---------

0 commit comments

Comments
 (0)