Skip to content

Commit 7bdb6b7

Browse files
TCeasonTCeason
committed
support array_sort([...] (, 'ASC'|'DESC') (, 'NULLS FIRST' | 'NULLS LAST'))
1 parent 578f85c commit 7bdb6b7

File tree

10 files changed

+117
-110
lines changed

10 files changed

+117
-110
lines changed

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

+4-4
Original file line numberDiff line numberDiff line change
@@ -882,14 +882,14 @@ impl Display for Expr {
882882
write!(f, "ARRAY_SORT(")?;
883883
write!(f, "{expr})")?;
884884
if *asc {
885-
write!(f, " , ASC")?;
885+
write!(f, " , 'ASC'")?;
886886
} else {
887-
write!(f, " , DESC")?;
887+
write!(f, " , 'DESC'")?;
888888
}
889889
if *null_first {
890-
write!(f, " , NULLS FIRST")?;
890+
write!(f, " , 'NULLS FIRST'")?;
891891
} else {
892-
write!(f, " , NULLS LAST")?;
892+
write!(f, " , 'NULLS LAST'")?;
893893
}
894894
write!(f, ")")?;
895895
}

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

+26-9
Original file line numberDiff line numberDiff line change
@@ -282,9 +282,9 @@ pub enum ExprElement {
282282
ArraySort {
283283
expr: Box<Expr>,
284284
// Optional `ASC` or `DESC`
285-
asc: Option<bool>,
285+
asc: Option<String>,
286286
// Optional `NULLS FIRST` or `NULLS LAST`
287-
nulls_first: Option<bool>,
287+
nulls_first: Option<String>,
288288
},
289289
Interval {
290290
expr: Expr,
@@ -474,9 +474,26 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
474474
asc,
475475
nulls_first,
476476
} => {
477-
let asc = if let Some(asc) = asc { asc } else { true };
477+
let asc = if let Some(asc) = asc {
478+
if asc.to_lowercase() == "asc" {
479+
true
480+
} else if asc.to_lowercase() == "desc" {
481+
false
482+
} else {
483+
return Err("Sorting order must be either ASC or DESC");
484+
}
485+
} else {
486+
true
487+
};
478488
let null_first = if let Some(nulls_first) = nulls_first {
479-
nulls_first
489+
let null_first = nulls_first.trim().to_lowercase();
490+
if null_first == "nulls first" {
491+
true
492+
} else if null_first == "nulls last" {
493+
false
494+
} else {
495+
return Err("Null sorting order must be either NULLS FIRST or NULLS LAST");
496+
}
480497
} else {
481498
true
482499
};
@@ -862,14 +879,14 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
862879
( ARRAY_SORT )
863880
~ "("
864881
~ #subexpr(0)
865-
~ ( "," ~ ( ASC | DESC ) )?
866-
~ ( "," ~ NULLS ~ ( FIRST | LAST ) )?
882+
~ ( "," ~ #literal_string )?
883+
~ ( "," ~ #literal_string )?
867884
~ ")"
868885
},
869886
|(_, _, expr, opt_asc, opt_null_first, _)| ExprElement::ArraySort {
870887
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),
888+
asc: opt_asc.map(|(_, asc)| asc),
889+
nulls_first: opt_null_first.map(|(_, first_last)| first_last),
873890
},
874891
);
875892
let date_add = map(
@@ -932,7 +949,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
932949
| #extract : "`EXTRACT((YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND) FROM ...)`"
933950
| #position : "`POSITION(... IN ...)`"
934951
| #substring : "`SUBSTRING(... [FROM ...] [FOR ...])`"
935-
| #array_sort : "`ARRAY_SORT([...], ASC | DESC, NULLS FIRST | LAST)`"
952+
| #array_sort : "`ARRAY_SORT([...], 'ASC' | 'DESC', 'NULLS FIRST' | 'NULLS LAST')`"
936953
| #trim : "`TRIM(...)`"
937954
| #trim_from : "`TRIM([(BOTH | LEADEING | TRAILING) ... FROM ...)`"
938955
),

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -541,8 +541,8 @@ fn test_expr() {
541541
r#"substring(a from b for c)"#,
542542
r#"substring(a, b, c)"#,
543543
r#"array_sort([2])"#,
544-
r#"array_sort([2,0.1], ASC)"#,
545-
r#"array_sort([3,2], DESC, NULLS FIRST)"#,
544+
r#"array_sort([2,0.1], 'ASC')"#,
545+
r#"array_sort([3,2], 'DESC', 'NULLS FIRST')"#,
546546
r#"col1::UInt8"#,
547547
r#"(arr[0]:a).b"#,
548548
r#"arr[4]["k"]"#,

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

+7-7
Original file line numberDiff line numberDiff line change
@@ -1649,7 +1649,7 @@ Substring {
16491649
---------- Input ----------
16501650
array_sort([2])
16511651
---------- Output ---------
1652-
ARRAY_SORT([2]) , ASC , NULLS FIRST)
1652+
ARRAY_SORT([2]) , 'ASC' , 'NULLS FIRST')
16531653
---------- AST ------------
16541654
ArraySort {
16551655
span: Some(
@@ -1676,13 +1676,13 @@ ArraySort {
16761676

16771677

16781678
---------- Input ----------
1679-
array_sort([2,0.1], ASC)
1679+
array_sort([2,0.1], 'ASC')
16801680
---------- Output ---------
1681-
ARRAY_SORT([2, 0.1]) , ASC , NULLS FIRST)
1681+
ARRAY_SORT([2, 0.1]) , 'ASC' , 'NULLS FIRST')
16821682
---------- AST ------------
16831683
ArraySort {
16841684
span: Some(
1685-
0..24,
1685+
0..26,
16861686
),
16871687
expr: Array {
16881688
span: Some(
@@ -1713,13 +1713,13 @@ ArraySort {
17131713

17141714

17151715
---------- Input ----------
1716-
array_sort([3,2], DESC, NULLS FIRST)
1716+
array_sort([3,2], 'DESC', 'NULLS FIRST')
17171717
---------- Output ---------
1718-
ARRAY_SORT([3, 2]) , DESC , NULLS FIRST)
1718+
ARRAY_SORT([3, 2]) , 'DESC' , 'NULLS FIRST')
17191719
---------- AST ------------
17201720
ArraySort {
17211721
span: Some(
1722-
0..36,
1722+
0..40,
17231723
),
17241724
expr: Array {
17251725
span: Some(

‎src/query/functions/src/scalars/array.rs

+16-17
Original file line numberDiff line numberDiff line change
@@ -768,23 +768,22 @@ fn register_array_aggr(registry: &mut FunctionRegistry) {
768768
FunctionProperty::default(),
769769
|_| FunctionDomain::Full,
770770
vectorize_1_arg::<ArrayType<GenericType<0>>, ArrayType<GenericType<0>>>(|arr, _| {
771-
let data_type = arr.data_type();
772-
let sort_desc = vec![SortColumnDescription {
773-
offset: 0,
774-
asc: sort_desc.0,
775-
nulls_first: sort_desc.1,
776-
}];
777-
let columns = vec![arr];
778-
let len = columns.get(0).map_or(1, |c| c.len());
779-
let columns = columns
780-
.iter()
781-
.map(|col| BlockEntry {
782-
data_type: col.data_type(),
783-
value: Value::Column(col.clone()),
784-
})
785-
.collect();
786-
let sort_block = DataBlock::sort(&DataBlock::new(columns, len), &sort_desc, None).unwrap();
787-
sort_block.columns()[0].value.convert_to_full_column(&data_type, 1)
771+
let len = arr.len();
772+
if arr.len() > 1 {
773+
let sort_desc = vec![SortColumnDescription {
774+
offset: 0,
775+
asc: sort_desc.0,
776+
nulls_first: sort_desc.1,
777+
}];
778+
let columns = vec![BlockEntry{
779+
data_type: arr.data_type(),
780+
value: Value::Column(arr.clone())
781+
}];
782+
let sort_block = DataBlock::sort(&DataBlock::new(columns, len), &sort_desc, None).unwrap();
783+
sort_block.columns()[0].value.clone().into_column().unwrap()
784+
} else {
785+
Value::Column(arr.clone())
786+
}
788787
},
789788
),
790789
);

‎src/query/functions/tests/it/scalars/array.rs

+13-12
Original file line numberDiff line numberDiff line change
@@ -485,34 +485,35 @@ fn test_array_any(file: &mut impl Write) {
485485
fn test_array_sort(file: &mut impl Write) {
486486
run_ast(file, "array_sort([])", &[]);
487487
run_ast(file, "array_sort(NULL)", &[]);
488-
run_ast(file, "array_sort([8, 20, 1, 2, 3, 4, 5, 6, 7], ASC)", &[]);
489-
run_ast(file, "array_sort([], ASC)", &[]);
490-
run_ast(file, "array_sort([], DESC)", &[]);
491-
run_ast(file, "array_sort([8, 20, 1, 2, 3, 4, 5, 6, 7], DESC)", &[]);
488+
run_ast(file, "array_sort([8, 20, 1, 2, 3, 4, 5, 6, 7], 'ASC')", &[]);
489+
run_ast(file, "array_sort([], 'ASC')", &[]);
490+
run_ast(file, "array_sort([], 'DESC')", &[]);
491+
run_ast(file, "array_sort([8, 20, 1, 2, 3, 4, 5, 6, 7], 'DESC')", &[
492+
]);
492493
run_ast(file, "array_sort([9.32, 0, 1.2, 3.4, 5.6, 7.8])", &[]);
493494
run_ast(file, "array_sort(['x', 0, 1.2, 3.4, 5.6, 7.8])", &[]);
494495
run_ast(
495496
file,
496-
"array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], DESC, NULLS FIRST)",
497+
"array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'DESC', 'NULLS FIRST')",
497498
&[],
498499
);
499-
run_ast(file, "array_sort([], DESC, NULLS FIRST)", &[]);
500+
run_ast(file, "array_sort([], 'DESC', 'NULLS FIRST')", &[]);
500501
run_ast(
501502
file,
502-
"array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], DESC, NULLS LAST)",
503+
"array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'DESC', 'NULLS LAST')",
503504
&[],
504505
);
505-
run_ast(file, "array_sort([], DESC, NULLS LAST)", &[]);
506+
run_ast(file, "array_sort([], 'DESC', 'NULLS LAST')", &[]);
506507
run_ast(
507508
file,
508-
"array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], ASC, NULLS FIRST)",
509+
"array_sort([1.2, NULL, 3.4, 5.6, '2.2', NULL], 'ASC', 'NULLS FIRST')",
509510
&[],
510511
);
511-
run_ast(file, "array_sort([], ASC, NULLS FIRST)", &[]);
512+
run_ast(file, "array_sort([], 'ASC', 'NULLS FIRST')", &[]);
512513
run_ast(
513514
file,
514-
"array_sort(['z', 'a', NULL, 'v', 'd', NULL], ASC, NULLS LAST)",
515+
"array_sort(['z', 'a', NULL, 'v', 'd', NULL], 'ASC', 'NULLS LAST')",
515516
&[],
516517
);
517-
run_ast(file, "array_sort([], ASC, NULLS LAST)", &[]);
518+
run_ast(file, "array_sort([], 'ASC', 'NULLS LAST')", &[]);
518519
}

‎src/query/functions/tests/it/scalars/parser.rs

+5-8
Original file line numberDiff line numberDiff line change
@@ -285,14 +285,11 @@ pub fn transform_expr(ast: AExpr, columns: &[(&str, DataType)]) -> RawExpr {
285285
asc,
286286
null_first,
287287
} => {
288-
let name = if asc && null_first {
289-
"array_sort_asc_null_first".to_string()
290-
} else if asc && !null_first {
291-
"array_sort_asc_null_last".to_string()
292-
} else if !asc && null_first {
293-
"array_sort_desc_null_first".to_string()
294-
} else {
295-
"array_sort_desc_null_last".to_string()
288+
let name = match (asc, null_first) {
289+
(true, true) => "array_sort_asc_null_first".to_string(),
290+
(true, false) => "array_sort_asc_null_last".to_string(),
291+
(false, true) => "array_sort_desc_null_first".to_string(),
292+
(false, false) => "array_sort_desc_null_last".to_string(),
296293
};
297294
RawExpr::FunctionCall {
298295
span,

0 commit comments

Comments
 (0)