Skip to content

Commit

Permalink
fix(query): window function support params
Browse files Browse the repository at this point in the history
  • Loading branch information
@b41sh
b41sh committed Jan 14, 2025
1 parent d1a4975 commit 149ddd9
Show file tree
Hide file tree
Showing 4 changed files with 169 additions and 5 deletions.
30 changes: 25 additions & 5 deletions src/query/ast/src/parser/expr.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1083,18 +1083,36 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
},
},
);
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,
},
Expand Down Expand Up @@ -1376,7 +1394,6 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
| #interval_expr : "`INTERVAL <str_literal>`"
| #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 ...])`"
Expand All @@ -1388,9 +1405,12 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
| #count_all_with_window : "`COUNT(*) OVER ...`"
| #function_call_with_lambda : "`function(..., x -> ...)`"
| #function_call_with_window : "`function(...) OVER ([ PARTITION BY <expr>, ... ] [ ORDER BY <expr>, ... ] [ <window frame> ])`"
| #function_call_with_params_window : "`function(...)(...) OVER ([ PARTITION BY <expr>, ... ] [ ORDER BY <expr>, ... ] [ <window frame> ])`"
| #function_call_with_params : "`function(...)(...)`"
| #function_call : "`function(...)`"
| #case : "`CASE ... END`"
),
rule!(
#case : "`CASE ... END`"
| #tuple : "`(<expr> [, ...])`"
| #subquery : "`(SELECT ...)`"
| #column_ref : "<column>"
Expand Down
1 change: 1 addition & 0 deletions src/query/ast/tests/it/parser.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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())"#,
Expand Down
114 changes: 114 additions & 0 deletions src/query/ast/tests/it/testdata/expr.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 ---------
Expand Down
29 changes: 29 additions & 0 deletions tests/sqllogictests/suites/query/window_function/window_basic.test
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -846,3 +874,4 @@ Product B 1200 NULL

statement ok
DROP DATABASE test_window_basic;

0 comments on commit 149ddd9

Please sign in to comment.