From 149ddd943ad59646ebc30af770b844481bd12266 Mon Sep 17 00:00:00 2001 From: baishen Date: Tue, 14 Jan 2025 16:41:20 +0800 Subject: [PATCH] fix(query): window function support params --- src/query/ast/src/parser/expr.rs | 30 ++++- src/query/ast/tests/it/parser.rs | 1 + src/query/ast/tests/it/testdata/expr.txt | 114 ++++++++++++++++++ .../query/window_function/window_basic.test | 29 +++++ 4 files changed, 169 insertions(+), 5 deletions(-) diff --git a/src/query/ast/src/parser/expr.rs b/src/query/ast/src/parser/expr.rs index ec3ffe28e109f..854923757dc22 100644 --- a/src/query/ast/src/parser/expr.rs +++ b/src/query/ast/src/parser/expr.rs @@ -1083,18 +1083,36 @@ pub fn expr_element(i: Input) -> IResult> { }, }, ); + let function_call_with_params_window = map( + rule! { + #function_name + ~ "(" ~ #comma_separated_list1(subexpr(0)) ~ ")" + ~ "(" ~ DISTINCT? ~ #comma_separated_list0(subexpr(0))? ~ ")" + ~ #window_function + }, + |(name, _, params, _, _, opt_distinct, opt_args, _, window)| ExprElement::FunctionCall { + func: FunctionCall { + distinct: opt_distinct.is_some(), + name, + args: opt_args.unwrap_or_default(), + params, + window: Some(window), + lambda: None, + }, + }, + ); let function_call_with_params = map( rule! { #function_name - ~ ("(" ~ #comma_separated_list1(subexpr(0)) ~ ")")? + ~ "(" ~ #comma_separated_list1(subexpr(0)) ~ ")" ~ "(" ~ DISTINCT? ~ #comma_separated_list0(subexpr(0))? ~ ")" }, - |(name, params, _, opt_distinct, opt_args, _)| ExprElement::FunctionCall { + |(name, _, params, _, _, opt_distinct, opt_args, _)| ExprElement::FunctionCall { func: FunctionCall { distinct: opt_distinct.is_some(), name, args: opt_args.unwrap_or_default(), - params: params.map(|(_, x, _)| x).unwrap_or_default(), + params, window: None, lambda: None, }, @@ -1376,7 +1394,6 @@ pub fn expr_element(i: Input) -> IResult> { | #interval_expr : "`INTERVAL `" | #extract : "`EXTRACT((YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND | WEEK) FROM ...)`" | #date_part : "`DATE_PART((YEAR | QUARTER | MONTH | DAY | HOUR | MINUTE | SECOND | WEEK), ...)`" - ), rule!( #substring : "`SUBSTRING(... [FROM ...] [FOR ...])`" @@ -1388,9 +1405,12 @@ pub fn expr_element(i: Input) -> IResult> { | #count_all_with_window : "`COUNT(*) OVER ...`" | #function_call_with_lambda : "`function(..., x -> ...)`" | #function_call_with_window : "`function(...) OVER ([ PARTITION BY , ... ] [ ORDER BY , ... ] [ ])`" + | #function_call_with_params_window : "`function(...)(...) OVER ([ PARTITION BY , ... ] [ ORDER BY , ... ] [ ])`" | #function_call_with_params : "`function(...)(...)`" | #function_call : "`function(...)`" - | #case : "`CASE ... END`" + ), + rule!( + #case : "`CASE ... END`" | #tuple : "`( [, ...])`" | #subquery : "`(SELECT ...)`" | #column_ref : "" diff --git a/src/query/ast/tests/it/parser.rs b/src/query/ast/tests/it/parser.rs index 1d8d789906695..052648e507d76 100644 --- a/src/query/ast/tests/it/parser.rs +++ b/src/query/ast/tests/it/parser.rs @@ -1251,6 +1251,7 @@ fn test_expr() { r#"COUNT() OVER (ORDER BY hire_date ROWS UNBOUNDED PRECEDING)"#, r#"COUNT() OVER (ORDER BY hire_date ROWS CURRENT ROW)"#, r#"COUNT() OVER (ORDER BY hire_date ROWS 3 PRECEDING)"#, + r#"QUANTILE_CONT(0.5)(salary) OVER (PARTITION BY department ORDER BY hire_date)"#, r#"ARRAY_APPLY([1,2,3], x -> x + 1)"#, r#"ARRAY_FILTER(col, y -> y % 2 = 0)"#, r#"(current_timestamp, current_timestamp(), now())"#, diff --git a/src/query/ast/tests/it/testdata/expr.txt b/src/query/ast/tests/it/testdata/expr.txt index eb279496999c7..af6d394434354 100644 --- a/src/query/ast/tests/it/testdata/expr.txt +++ b/src/query/ast/tests/it/testdata/expr.txt @@ -4229,6 +4229,120 @@ FunctionCall { } +---------- Input ---------- +QUANTILE_CONT(0.5)(salary) OVER (PARTITION BY department ORDER BY hire_date) +---------- Output --------- +QUANTILE_CONT(0.5)(salary) OVER (PARTITION BY department ORDER BY hire_date) +---------- AST ------------ +FunctionCall { + span: Some( + 0..76, + ), + func: FunctionCall { + distinct: false, + name: Identifier { + span: Some( + 0..13, + ), + name: "QUANTILE_CONT", + quote: None, + ident_type: None, + }, + args: [ + ColumnRef { + span: Some( + 19..25, + ), + column: ColumnRef { + database: None, + table: None, + column: Name( + Identifier { + span: Some( + 19..25, + ), + name: "salary", + quote: None, + ident_type: None, + }, + ), + }, + }, + ], + params: [ + Literal { + span: Some( + 14..17, + ), + value: Decimal256 { + value: 5, + precision: 76, + scale: 1, + }, + }, + ], + window: Some( + WindowDesc { + ignore_nulls: None, + window: WindowSpec( + WindowSpec { + existing_window_name: None, + partition_by: [ + ColumnRef { + span: Some( + 46..56, + ), + column: ColumnRef { + database: None, + table: None, + column: Name( + Identifier { + span: Some( + 46..56, + ), + name: "department", + quote: None, + ident_type: None, + }, + ), + }, + }, + ], + order_by: [ + OrderByExpr { + expr: ColumnRef { + span: Some( + 66..75, + ), + column: ColumnRef { + database: None, + table: None, + column: Name( + Identifier { + span: Some( + 66..75, + ), + name: "hire_date", + quote: None, + ident_type: None, + }, + ), + }, + }, + asc: None, + nulls_first: None, + }, + ], + window_frame: None, + }, + ), + }, + ), + lambda: None, + }, +} + + ---------- Input ---------- ARRAY_APPLY([1,2,3], x -> x + 1) ---------- Output --------- diff --git a/tests/sqllogictests/suites/query/window_function/window_basic.test b/tests/sqllogictests/suites/query/window_function/window_basic.test index 75110235c9de6..694e675219bce 100644 --- a/tests/sqllogictests/suites/query/window_function/window_basic.test +++ b/tests/sqllogictests/suites/query/window_function/window_basic.test @@ -606,6 +606,34 @@ select grouping(salary), grouping(depname), sum(grouping(salary)) over (partitio 1 0 3 1 1 1 +query TII +SELECT depname, empno, quantile_cont(salary) OVER (PARTITION BY depname ORDER BY empno) FROM empsalary ORDER BY depname, empno +---- +develop 7 4200.0 +develop 8 5100.0 +develop 9 4500.0 +develop 10 4850.0 +develop 11 5200.0 +personnel 2 3900.0 +personnel 5 3700.0 +sales 1 5000.0 +sales 3 4900.0 +sales 4 4800.0 + +query TII +SELECT depname, empno, quantile_cont(0.8)(salary) OVER (PARTITION BY depname ORDER BY empno) FROM empsalary ORDER BY depname, empno +---- +develop 7 4200.0 +develop 8 5640.0 +develop 9 5400.0 +develop 10 5520.0 +develop 11 5360.0 +personnel 2 3900.0 +personnel 5 3820.0 +sales 1 5000.0 +sales 3 4960.0 +sales 4 4920.0 + # Window func in subquery query I SELECT * FROM (SELECT row_number() OVER (PARTITION BY depname ORDER BY salary) rn FROM empsalary ORDER BY depname, rn) order by 1; @@ -846,3 +874,4 @@ Product B 1200 NULL statement ok DROP DATABASE test_window_basic; +