Skip to content

Commit 6ae1708

Browse files
KKouldKKould
committed
feat: add window function support
1 parent 362157d commit 6ae1708

40 files changed

Lines changed: 1518 additions & 192 deletions

‎AGENT.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,8 @@ PRs that modify logic but leave obvious test gaps untouched may be rejected.
8181

8282
### 3.1 Prefer Simple Code
8383

84+
- Implement functionality with the smallest reasonable change and the simplest correct design.
85+
- Do not over-engineer or introduce abstractions for hypothetical future requirements.
8486
- Prefer straightforward control flow
8587
- Avoid unnecessary abstractions
8688
- Avoid premature generalization

‎docs/features.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,11 @@ If `unsafe_txdb_checkpoint` is not enabled, `build_rocksdb()` returns an explici
183183
- [x] Exists
184184
- [x] Group By
185185
- [x] Having
186+
- [x] Window functions:
187+
- `row_number()`, `rank()`, `dense_rank()`
188+
- `count()`, `sum()`, `avg()`, `min()`, `max()` with `OVER`
189+
- `PARTITION BY` and window `ORDER BY` with the default frame
190+
- Explicit frames, named windows, and `QUALIFY` are not yet supported
186191
- [x] Order By
187192
- [x] Limit
188193
- [x] Show Tables

‎src/binder/aggregate.rs‎

Lines changed: 28 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ use std::collections::HashSet;
1616

1717
use super::{Binder, QueryBindStep};
1818
use crate::errors::DatabaseError;
19-
use crate::expression::function::scala::ScalarFunction;
2019
use crate::expression::visitor::{walk_expr, ExprVisitor};
2120
use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut};
2221
use crate::planner::LogicalPlan;
@@ -27,6 +26,22 @@ use crate::{
2726
planner::operator::{aggregate::AggregateOperator, sort::SortField},
2827
};
2928

29+
struct AggregateCallCollector<'a> {
30+
agg_calls: &'a mut Vec<ScalarExpression>,
31+
}
32+
33+
impl<'expr> ExprVisitor<'expr> for AggregateCallCollector<'_> {
34+
fn visit(&mut self, expr: &'expr ScalarExpression) -> Result<(), DatabaseError> {
35+
match expr {
36+
ScalarExpression::AggCall { .. } => self.agg_calls.push(expr.clone()),
37+
ScalarExpression::Alias { expr, .. } => self.visit(expr)?,
38+
ScalarExpression::Empty | ScalarExpression::TableFunction(_) => unreachable!(),
39+
_ => walk_expr(self, expr)?,
40+
}
41+
Ok(())
42+
}
43+
}
44+
3045
impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A> {
3146
pub fn bind_aggregate(
3247
&mut self,
@@ -48,7 +63,7 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
4863
select_items: &mut [ScalarExpression],
4964
) -> Result<(), DatabaseError> {
5065
for column in select_items {
51-
self.visit_column_agg_expr(column)?;
66+
self.collect_aggregate_calls(column)?;
5267
}
5368
Ok(())
5469
}
@@ -77,14 +92,14 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
7792
F: FnMut(&mut Self, I::Item) -> Result<SortField, DatabaseError>,
7893
{
7994
if let Some(having) = having.as_mut() {
80-
self.visit_column_agg_expr(having)?;
95+
self.collect_aggregate_calls(having)?;
8196
}
8297
let mut return_orderby = None;
8398
if let Some(orderby) = orderby {
8499
let mut fields = Vec::new();
85100
for orderby in orderby {
86-
let mut field = bind_sort_field(self, orderby)?;
87-
self.visit_column_agg_expr(&mut field.expr)?;
101+
let field = bind_sort_field(self, orderby)?;
102+
self.collect_aggregate_calls(&field.expr)?;
88103
fields.push(field);
89104
}
90105
return_orderby = Some(fields);
@@ -119,119 +134,11 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
119134
Ok(())
120135
}
121136

122-
fn visit_column_agg_expr(&mut self, expr: &mut ScalarExpression) -> Result<(), DatabaseError> {
123-
match expr {
124-
ScalarExpression::AggCall { .. } => {
125-
self.context.agg_calls.push(expr.clone());
126-
}
127-
ScalarExpression::TypeCast { expr, .. } => self.visit_column_agg_expr(expr)?,
128-
ScalarExpression::IsNull { expr, .. } => self.visit_column_agg_expr(expr)?,
129-
ScalarExpression::Unary { expr, .. } => self.visit_column_agg_expr(expr)?,
130-
ScalarExpression::Alias { expr, .. } => self.visit_column_agg_expr(expr)?,
131-
ScalarExpression::Binary {
132-
left_expr,
133-
right_expr,
134-
..
135-
} => {
136-
self.visit_column_agg_expr(left_expr)?;
137-
self.visit_column_agg_expr(right_expr)?;
138-
}
139-
ScalarExpression::In { expr, args, .. } => {
140-
self.visit_column_agg_expr(expr)?;
141-
for arg in args {
142-
self.visit_column_agg_expr(arg)?;
143-
}
144-
}
145-
ScalarExpression::Between {
146-
expr,
147-
left_expr,
148-
right_expr,
149-
..
150-
} => {
151-
self.visit_column_agg_expr(expr)?;
152-
self.visit_column_agg_expr(left_expr)?;
153-
self.visit_column_agg_expr(right_expr)?;
154-
}
155-
ScalarExpression::SubString {
156-
expr,
157-
for_expr,
158-
from_expr,
159-
} => {
160-
self.visit_column_agg_expr(expr)?;
161-
if let Some(expr) = for_expr {
162-
self.visit_column_agg_expr(expr)?;
163-
}
164-
if let Some(expr) = from_expr {
165-
self.visit_column_agg_expr(expr)?;
166-
}
167-
}
168-
ScalarExpression::Position { expr, in_expr } => {
169-
self.visit_column_agg_expr(expr)?;
170-
self.visit_column_agg_expr(in_expr)?;
171-
}
172-
ScalarExpression::Trim {
173-
expr,
174-
trim_what_expr,
175-
..
176-
} => {
177-
self.visit_column_agg_expr(expr)?;
178-
if let Some(trim_what_expr) = trim_what_expr {
179-
self.visit_column_agg_expr(trim_what_expr)?;
180-
}
181-
}
182-
ScalarExpression::Constant(_) | ScalarExpression::ColumnRef { .. } => (),
183-
ScalarExpression::Empty => unreachable!(),
184-
ScalarExpression::Tuple(args)
185-
| ScalarExpression::ScalaFunction(ScalarFunction { args, .. })
186-
| ScalarExpression::Coalesce { exprs: args, .. } => {
187-
for expr in args {
188-
self.visit_column_agg_expr(expr)?;
189-
}
190-
}
191-
ScalarExpression::If {
192-
condition,
193-
left_expr,
194-
right_expr,
195-
..
196-
} => {
197-
self.visit_column_agg_expr(condition)?;
198-
self.visit_column_agg_expr(left_expr)?;
199-
self.visit_column_agg_expr(right_expr)?;
200-
}
201-
ScalarExpression::IfNull {
202-
left_expr,
203-
right_expr,
204-
..
205-
}
206-
| ScalarExpression::NullIf {
207-
left_expr,
208-
right_expr,
209-
..
210-
} => {
211-
self.visit_column_agg_expr(left_expr)?;
212-
self.visit_column_agg_expr(right_expr)?;
213-
}
214-
ScalarExpression::CaseWhen {
215-
operand_expr,
216-
expr_pairs,
217-
else_expr,
218-
..
219-
} => {
220-
if let Some(expr) = operand_expr {
221-
self.visit_column_agg_expr(expr)?;
222-
}
223-
for (expr_1, expr_2) in expr_pairs {
224-
self.visit_column_agg_expr(expr_1)?;
225-
self.visit_column_agg_expr(expr_2)?;
226-
}
227-
if let Some(expr) = else_expr {
228-
self.visit_column_agg_expr(expr)?;
229-
}
230-
}
231-
ScalarExpression::TableFunction(_) => unreachable!(),
137+
fn collect_aggregate_calls(&mut self, expr: &ScalarExpression) -> Result<(), DatabaseError> {
138+
AggregateCallCollector {
139+
agg_calls: &mut self.context.agg_calls,
232140
}
233-
234-
Ok(())
141+
.visit(expr)
235142
}
236143

237144
/// Validate select exprs must appear in the GROUP BY clause or be used in
@@ -269,6 +176,10 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
269176
HashSet::from_iter(group_raw_exprs.iter().copied());
270177

271178
for expr in select_items {
179+
if expr.has_window_call()? {
180+
HavingOrderByValidator::new(groupby, &self.context.agg_calls).visit(expr)?;
181+
continue;
182+
}
272183
if expr.has_agg_call()? {
273184
continue;
274185
}

‎src/binder/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ mod show_table;
4848
mod show_view;
4949
mod truncate;
5050
mod update;
51+
mod window;
5152

5253
#[cfg(feature = "parser")]
5354
pub use parser::{command_type, prepare, prepare_all, CommandType, Statement};
@@ -90,6 +91,7 @@ pub enum QueryBindStep {
9091
Where,
9192
Agg,
9293
Having,
94+
Window,
9395
Distinct,
9496
Sort,
9597
Project,

‎src/binder/parser.rs‎

Lines changed: 93 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,7 @@
1414

1515
use super::select::{
1616
BindPlanAggregated, BindPlanComplete, BindPlanDistinct, BindPlanFiltered, BindPlanFrom,
17-
BindPlanHaving, BindPlanProjected, BindPlanSelectList, BindPlanStart, JoinConstraintInput,
18-
TableAliasInput,
17+
BindPlanProjected, BindPlanSelectList, BindPlanStart, JoinConstraintInput, TableAliasInput,
1918
};
2019
use super::{is_valid_identifier, with_query_bind_step, Binder, QueryBindStep, SetOperatorKind};
2120
#[cfg(feature = "copy")]
@@ -47,6 +46,7 @@ pub(super) use sqlparser::ast::{
4746
ObjectType, OrderByExpr, OrderByKind, Query, Select, SelectInto, SelectItem,
4847
SelectItemQualifiedWildcardKind, SetExpr, SetOperator, SetQuantifier, Spanned, TableAlias,
4948
TableConstraint, TableFactor, TableObject, TableWithJoins, TypedString, UnaryOperator, Value,
49+
WindowType,
5050
};
5151
#[cfg(feature = "copy")]
5252
pub(super) use sqlparser::ast::{CopyOption, CopySource, CopyTarget};
@@ -1463,7 +1463,7 @@ where
14631463
}
14641464
}
14651465

1466-
impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanHaving<'s, 'a, 'b, 'arena, T, A>
1466+
impl<'s, 'a: 'b, 'b, 'arena, T, A> super::select::BindPlanWindowed<'s, 'a, 'b, 'arena, T, A>
14671467
where
14681468
T: Transaction,
14691469
A: AsRef<[(&'static str, DataValue)]>,
@@ -2392,7 +2392,15 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
23922392
arena: &mut PlanArena,
23932393
) -> Result<ScalarExpression, DatabaseError> {
23942394
let func_span = func.span();
2395-
let Function { name, args, .. } = func;
2395+
let Function {
2396+
name,
2397+
args,
2398+
over,
2399+
filter,
2400+
null_treatment,
2401+
within_group,
2402+
..
2403+
} = func;
23962404
let (func_args, is_distinct) = match args {
23972405
FunctionArguments::List(args) => (
23982406
args.args.as_slice(),
@@ -2425,10 +2433,78 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
24252433
}
24262434
let function_name = name.to_string().to_lowercase();
24272435

2436+
if let Some(over) = over {
2437+
if filter.is_some() || null_treatment.is_some() || !within_group.is_empty() {
2438+
return Err(DatabaseError::UnsupportedStmt(
2439+
"FILTER, NULL treatment, and WITHIN GROUP are not supported for window functions"
2440+
.to_string(),
2441+
));
2442+
}
2443+
return self
2444+
.bind_window_call(function_name, args, is_distinct, over, arena)
2445+
.map_err(|err| attach_span_if_absent(err, func_span));
2446+
}
2447+
24282448
self.bind_function_call(function_name, args, is_distinct, arena)
24292449
.map_err(|err| attach_span_if_absent(err, func_span))
24302450
}
24312451

2452+
fn bind_window_call(
2453+
&mut self,
2454+
function_name: String,
2455+
args: Vec<ScalarExpression>,
2456+
is_distinct: bool,
2457+
over: &WindowType,
2458+
arena: &mut PlanArena,
2459+
) -> Result<ScalarExpression, DatabaseError> {
2460+
if !matches!(
2461+
self.context.step_now(),
2462+
QueryBindStep::Project | QueryBindStep::Sort
2463+
) {
2464+
return Err(DatabaseError::UnsupportedStmt(
2465+
"window functions are only allowed in SELECT and ORDER BY".to_string(),
2466+
));
2467+
}
2468+
if is_distinct {
2469+
return Err(DatabaseError::UnsupportedStmt(
2470+
"DISTINCT window aggregates are not supported".to_string(),
2471+
));
2472+
}
2473+
let WindowType::WindowSpec(spec) = over else {
2474+
return Err(DatabaseError::UnsupportedStmt(
2475+
"named windows are not supported".to_string(),
2476+
));
2477+
};
2478+
if spec.window_name.is_some() {
2479+
return Err(DatabaseError::UnsupportedStmt(
2480+
"inherited named windows are not supported".to_string(),
2481+
));
2482+
}
2483+
if spec.window_frame.is_some() {
2484+
return Err(DatabaseError::UnsupportedStmt(
2485+
"explicit window frames are not supported".to_string(),
2486+
));
2487+
}
2488+
2489+
let partition_by = spec
2490+
.partition_by
2491+
.iter()
2492+
.map(|expr| self.bind_expr(expr, arena))
2493+
.collect::<Result<Vec<_>, _>>()?;
2494+
let order_by = spec
2495+
.order_by
2496+
.iter()
2497+
.map(|OrderByExpr { expr, options, .. }| {
2498+
Ok(SortField::new(
2499+
self.bind_expr(expr, arena)?,
2500+
options.asc.unwrap_or(true),
2501+
options.nulls_first.unwrap_or(false),
2502+
))
2503+
})
2504+
.collect::<Result<Vec<_>, DatabaseError>>()?;
2505+
self.bind_window_function(function_name, args, partition_by, order_by, arena)
2506+
}
2507+
24322508
pub fn bind_set_expr(
24332509
&mut self,
24342510
set_expr: &SetExpr,
@@ -2511,15 +2587,28 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
25112587
having,
25122588
distinct,
25132589
into,
2590+
named_window,
2591+
qualify,
25142592
..
25152593
} = select;
2594+
if !named_window.is_empty() {
2595+
return Err(DatabaseError::UnsupportedStmt(
2596+
"named windows are not supported".to_string(),
2597+
));
2598+
}
2599+
if qualify.is_some() {
2600+
return Err(DatabaseError::UnsupportedStmt(
2601+
"QUALIFY is not supported".to_string(),
2602+
));
2603+
}
25162604
Ok(self
25172605
.build_plan(arena)
25182606
.from_sql(from)?
25192607
.select_list_from_sql(projection)?
25202608
.where_sql(selection.as_ref())?
25212609
.aggregate_sql(group_by, having.as_ref(), orderby)?
25222610
.having()?
2611+
.window()?
25232612
.distinct_sql(distinct.as_ref())?
25242613
.order_by()?
25252614
.project()?

0 commit comments

Comments
 (0)