Skip to content

Commit c19209e

Browse files
KKouldKKould
authored
feat: add window function support (#375)
* feat: add window function support * docs: clarify pull request expectations * refactor: simplify window output binding * test: expand window function coverage * refactor: type aggregate and window binding * test: cover accumulators and window operator
1 parent 362157d commit c19209e

42 files changed

Lines changed: 2057 additions & 259 deletions

Some content is hidden

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

‎AGENT.md‎

Lines changed: 4 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
@@ -149,6 +151,8 @@ KiteSQL aims to stay easy to build, easy to audit, and easy to understand.
149151

150152
A valid PR should:
151153

154+
- Follow the repository pull request template
155+
- Describe test coverage and verified behavior instead of listing commands that were run
152156
- Compile cleanly
153157
- Pass all tests via make
154158
- Include new tests if behavior changes

‎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/expr.rs‎

Lines changed: 30 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -476,76 +476,60 @@ impl<'a, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'a, '_, T
476476
})
477477
}
478478

479-
pub(crate) fn bind_function_call(
479+
pub(crate) fn bind_aggregate_function(
480480
&mut self,
481-
function_name: String,
482-
mut args: Vec<ScalarExpression>,
481+
kind: AggKind,
482+
args: Vec<ScalarExpression>,
483483
is_distinct: bool,
484484
arena: &mut PlanArena,
485485
) -> Result<ScalarExpression, DatabaseError> {
486-
match function_name.as_str() {
487-
"count" => {
486+
let ty = match kind {
487+
AggKind::Count => {
488488
if args.len() != 1 {
489489
return Err(DatabaseError::MisMatch("number of count() parameters", "1"));
490490
}
491-
return Ok(ScalarExpression::AggCall {
492-
distinct: is_distinct,
493-
kind: AggKind::Count,
494-
args,
495-
ty: LogicalType::Integer,
496-
});
491+
LogicalType::Integer
497492
}
498-
"sum" => {
493+
AggKind::Sum => {
499494
if args.len() != 1 {
500495
return Err(DatabaseError::MisMatch("number of sum() parameters", "1"));
501496
}
502-
let ty = args[0].return_type(arena).into_owned();
503-
504-
return Ok(ScalarExpression::AggCall {
505-
distinct: is_distinct,
506-
kind: AggKind::Sum,
507-
args,
508-
ty,
509-
});
497+
args[0].return_type(arena).into_owned()
510498
}
511-
"min" => {
499+
AggKind::Min => {
512500
if args.len() != 1 {
513501
return Err(DatabaseError::MisMatch("number of min() parameters", "1"));
514502
}
515-
let ty = args[0].return_type(arena).into_owned();
516-
517-
return Ok(ScalarExpression::AggCall {
518-
distinct: is_distinct,
519-
kind: AggKind::Min,
520-
args,
521-
ty,
522-
});
503+
args[0].return_type(arena).into_owned()
523504
}
524-
"max" => {
505+
AggKind::Max => {
525506
if args.len() != 1 {
526507
return Err(DatabaseError::MisMatch("number of max() parameters", "1"));
527508
}
528-
let ty = args[0].return_type(arena).into_owned();
529-
530-
return Ok(ScalarExpression::AggCall {
531-
distinct: is_distinct,
532-
kind: AggKind::Max,
533-
args,
534-
ty,
535-
});
509+
args[0].return_type(arena).into_owned()
536510
}
537-
"avg" => {
511+
AggKind::Avg => {
538512
if args.len() != 1 {
539513
return Err(DatabaseError::MisMatch("number of avg() parameters", "1"));
540514
}
541-
542-
return Ok(ScalarExpression::AggCall {
543-
distinct: is_distinct,
544-
kind: AggKind::Avg,
545-
args,
546-
ty: LogicalType::Double,
547-
});
515+
LogicalType::Double
548516
}
517+
};
518+
Ok(ScalarExpression::AggCall {
519+
distinct: is_distinct,
520+
kind,
521+
args,
522+
ty,
523+
})
524+
}
525+
526+
pub(crate) fn bind_function_call(
527+
&mut self,
528+
function_name: String,
529+
mut args: Vec<ScalarExpression>,
530+
arena: &mut PlanArena,
531+
) -> Result<ScalarExpression, DatabaseError> {
532+
match function_name.as_str() {
549533
"if" => {
550534
if args.len() != 3 {
551535
return Err(DatabaseError::MisMatch("number of if() parameters", "3"));

‎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,

0 commit comments

Comments
 (0)