Skip to content

Commit 791e3a8

Browse files
KKouldKKould
authored
feat: support WITHIN GROUP clause for AggregateFunction (#17519)
* feat: support `WITHIN GROUP` clause for `AggregateFunction` Signed-off-by: Kould <kould2333@gmail.com> * chore: fix e2e tests Signed-off-by: Kould <kould2333@gmail.com> * feat: support order by with expression Signed-off-by: Kould <kould2333@gmail.com> * refactor: add aggregate_sort_adaptor.rs to sort `InputColumns` before aggregate_function Signed-off-by: Kould <kould2333@gmail.com> * perf: `sort_desc` will try reuse column Signed-off-by: Kould <kould2333@gmail.com> * test: add e2e test for `null first` Signed-off-by: Kould <kould2333@gmail.com> --------- Signed-off-by: Kould <kould2333@gmail.com>
1 parent 812e9c3 commit 791e3a8

File tree

95 files changed

+2035
-204
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

95 files changed

+2035
-204
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -961,6 +961,7 @@ pub struct FunctionCall {
961961
pub name: Identifier,
962962
pub args: Vec<Expr>,
963963
pub params: Vec<Expr>,
964+
pub order_by: Vec<OrderByExpr>,
964965
pub window: Option<WindowDesc>,
965966
pub lambda: Option<Lambda>,
966967
}
@@ -972,6 +973,7 @@ impl Display for FunctionCall {
972973
name,
973974
args,
974975
params,
976+
order_by,
975977
window,
976978
lambda,
977979
} = self;
@@ -991,6 +993,11 @@ impl Display for FunctionCall {
991993
}
992994
write!(f, ")")?;
993995

996+
if !order_by.is_empty() {
997+
write!(f, " WITHIN GROUP ( ORDER BY ")?;
998+
write_comma_separated_list(f, &self.order_by)?;
999+
write!(f, " )")?;
1000+
}
9941001
if let Some(window) = window {
9951002
if let Some(ignore_null) = window.ignore_nulls {
9961003
if ignore_null {

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -612,6 +612,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
612612
),
613613
args: vec![source],
614614
params: vec![],
615+
order_by: vec![],
615616
window: None,
616617
lambda: Some(Lambda {
617618
params: vec![param.clone()],
@@ -628,6 +629,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
628629
name: Identifier::from_name(transform_span(elem.span.tokens), "array_map"),
629630
args: vec![source],
630631
params: vec![],
632+
order_by: vec![],
631633
window: None,
632634
lambda: Some(Lambda {
633635
params: vec![param.clone()],
@@ -794,6 +796,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
794796
name,
795797
args: [vec![lhs], args].concat(),
796798
params: vec![],
799+
order_by: vec![],
797800
window: None,
798801
lambda,
799802
},
@@ -1028,6 +1031,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
10281031
name,
10291032
args: opt_args.unwrap_or_default(),
10301033
params: vec![],
1034+
order_by: vec![],
10311035
window: None,
10321036
lambda: None,
10331037
},
@@ -1044,6 +1048,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
10441048
name,
10451049
args: vec![arg],
10461050
params: vec![],
1051+
order_by: vec![],
10471052
window: None,
10481053
lambda: Some(Lambda {
10491054
params,
@@ -1064,11 +1069,31 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
10641069
name,
10651070
args: opt_args.unwrap_or_default(),
10661071
params: vec![],
1072+
order_by: vec![],
10671073
window: Some(window),
10681074
lambda: None,
10691075
},
10701076
},
10711077
);
1078+
let function_call_with_within_group_window = map(
1079+
rule! {
1080+
#function_name
1081+
~ "(" ~ DISTINCT? ~ #comma_separated_list0(subexpr(0))? ~ ")"
1082+
~ #within_group
1083+
~ #window_function?
1084+
},
1085+
|(name, _, opt_distinct, opt_args, _, order_by, window)| ExprElement::FunctionCall {
1086+
func: FunctionCall {
1087+
distinct: opt_distinct.is_some(),
1088+
name,
1089+
args: opt_args.unwrap_or_default(),
1090+
params: vec![],
1091+
order_by,
1092+
window,
1093+
lambda: None,
1094+
},
1095+
},
1096+
);
10721097
let function_call_with_params_window = map(
10731098
rule! {
10741099
#function_name
@@ -1082,6 +1107,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
10821107
name,
10831108
args: opt_args.unwrap_or_default(),
10841109
params,
1110+
order_by: vec![],
10851111
window,
10861112
lambda: None,
10871113
},
@@ -1326,6 +1352,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
13261352
name: Identifier::from_name(transform_span(span.tokens), "current_timestamp"),
13271353
args: vec![],
13281354
params: vec![],
1355+
order_by: vec![],
13291356
window: None,
13301357
lambda: None,
13311358
},
@@ -1373,6 +1400,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
13731400
| #count_all_with_window : "`COUNT(*) OVER ...`"
13741401
| #function_call_with_lambda : "`function(..., x -> ...)`"
13751402
| #function_call_with_window : "`function(...) OVER ([ PARTITION BY <expr>, ... ] [ ORDER BY <expr>, ... ] [ <window frame> ])`"
1403+
| #function_call_with_within_group_window: "`function(...) [ WITHIN GROUP ( ORDER BY <expr>, ... ) ] OVER ([ PARTITION BY <expr>, ... ] [ ORDER BY <expr>, ... ] [ <window frame> ])`"
13761404
| #function_call_with_params_window : "`function(...)(...) OVER ([ PARTITION BY <expr>, ... ] [ ORDER BY <expr>, ... ] [ <window frame> ])`"
13771405
| #function_call : "`function(...)`"
13781406
),
@@ -2029,6 +2057,7 @@ pub(crate) fn make_func_get_variable(span: Span, name: String) -> Expr {
20292057
value: Literal::String(name),
20302058
}],
20312059
params: vec![],
2060+
order_by: vec![],
20322061
window: None,
20332062
lambda: None,
20342063
},

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1232,6 +1232,15 @@ pub fn window_spec_ident(i: Input) -> IResult<Window> {
12321232
))(i)
12331233
}
12341234

1235+
pub fn within_group(i: Input) -> IResult<Vec<OrderByExpr>> {
1236+
map(
1237+
rule! {
1238+
WITHIN ~ GROUP ~ "(" ~ ORDER ~ ^BY ~ ^#comma_separated_list1(order_by_expr) ~ ")"
1239+
},
1240+
|(_, _, _, _, _, order_by, _)| order_by,
1241+
)(i)
1242+
}
1243+
12351244
pub fn window_function(i: Input) -> IResult<WindowDesc> {
12361245
map(
12371246
rule! {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1256,6 +1256,8 @@ pub enum TokenKind {
12561256
WINDOW,
12571257
#[token("WITH", ignore(ascii_case))]
12581258
WITH,
1259+
#[token("WITHIN", ignore(ascii_case))]
1260+
WITHIN,
12591261
#[token("XML", ignore(ascii_case))]
12601262
XML,
12611263
#[token("XOR", ignore(ascii_case))]

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1255,6 +1255,8 @@ fn test_expr() {
12551255
r#"a is distinct from b"#,
12561256
r#"1 is not distinct from null"#,
12571257
r#"{'k1':1,'k2':2}"#,
1258+
// within group
1259+
r#"LISTAGG(salary, '|') WITHIN GROUP (ORDER BY salary DESC NULLS LAST)"#,
12581260
// window expr
12591261
r#"ROW_NUMBER() OVER (ORDER BY salary DESC)"#,
12601262
r#"SUM(salary) OVER ()"#,

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,7 @@ FunctionCall {
563563
},
564564
],
565565
params: [],
566+
order_by: [],
566567
window: None,
567568
lambda: None,
568569
},
@@ -644,6 +645,7 @@ FunctionCall {
644645
},
645646
],
646647
params: [],
648+
order_by: [],
647649
window: None,
648650
lambda: None,
649651
},
@@ -669,6 +671,7 @@ FunctionCall {
669671
},
670672
],
671673
params: [],
674+
order_by: [],
672675
window: None,
673676
lambda: None,
674677
},
@@ -750,6 +753,7 @@ FunctionCall {
750753
},
751754
],
752755
params: [],
756+
order_by: [],
753757
window: None,
754758
lambda: None,
755759
},
@@ -775,6 +779,7 @@ FunctionCall {
775779
},
776780
],
777781
params: [],
782+
order_by: [],
778783
window: None,
779784
lambda: None,
780785
},
@@ -875,6 +880,7 @@ FunctionCall {
875880
},
876881
],
877882
params: [],
883+
order_by: [],
878884
window: None,
879885
lambda: None,
880886
},
@@ -946,6 +952,7 @@ FunctionCall {
946952
},
947953
],
948954
params: [],
955+
order_by: [],
949956
window: None,
950957
lambda: None,
951958
},
@@ -1020,6 +1027,7 @@ FunctionCall {
10201027
},
10211028
],
10221029
params: [],
1030+
order_by: [],
10231031
window: None,
10241032
lambda: Some(
10251033
Lambda {
@@ -1086,6 +1094,7 @@ FunctionCall {
10861094
},
10871095
],
10881096
params: [],
1097+
order_by: [],
10891098
window: None,
10901099
lambda: Some(
10911100
Lambda {

0 commit comments

Comments
 (0)