From 6ae17088f9f933905f5ac07112646e7e0b3879ac Mon Sep 17 00:00:00 2001 From: kould Date: Sun, 12 Jul 2026 23:26:11 +0800 Subject: [PATCH 1/6] feat: add window function support --- AGENT.md | 2 + docs/features.md | 5 + src/binder/aggregate.rs | 145 ++-------- src/binder/mod.rs | 2 + src/binder/parser.rs | 97 ++++++- src/binder/select.rs | 41 ++- src/binder/window.rs | 255 ++++++++++++++++++ src/db.rs | 1 + src/execution/dql/aggregate/avg.rs | 58 +++- src/execution/dql/aggregate/count.rs | 33 ++- src/execution/dql/aggregate/hash_agg.rs | 5 +- src/execution/dql/aggregate/min_max.rs | 21 +- src/execution/dql/aggregate/mod.rs | 57 ++-- src/execution/dql/aggregate/simple_agg.rs | 5 +- src/execution/dql/aggregate/sum.rs | 20 +- src/execution/dql/mod.rs | 1 + src/execution/dql/window.rs | 156 +++++++++++ src/execution/dql/window/function.rs | 137 ++++++++++ src/execution/mod.rs | 12 + src/expression/evaluator.rs | 3 + src/expression/mod.rs | 76 +++++- src/expression/range_detacher.rs | 6 +- src/expression/visitor.rs | 15 ++ src/expression/visitor_mut.rs | 15 ++ src/expression/window.rs | 46 ++++ src/optimizer/heuristic/optimizer.rs | 3 + src/optimizer/rule/implementation/dql/mod.rs | 1 + .../rule/implementation/dql/window.rs | 35 +++ src/optimizer/rule/implementation/mod.rs | 15 +- .../rule/normalization/column_pruning.rs | 3 +- src/optimizer/rule/normalization/mod.rs | 3 +- src/orm/README.md | 27 ++ src/orm/mod.rs | 122 +++++++++ src/planner/mod.rs | 8 + src/planner/operator/mod.rs | 20 ++ src/planner/operator/visitor.rs | 30 ++- src/planner/operator/visitor_mut.rs | 23 +- src/planner/operator/window.rs | 60 +++++ tests/macros-test/src/main.rs | 32 ++- tests/slt/window.slt | 114 ++++++++ 40 files changed, 1518 insertions(+), 192 deletions(-) create mode 100644 src/binder/window.rs create mode 100644 src/execution/dql/window.rs create mode 100644 src/execution/dql/window/function.rs create mode 100644 src/expression/window.rs create mode 100644 src/optimizer/rule/implementation/dql/window.rs create mode 100644 src/planner/operator/window.rs create mode 100644 tests/slt/window.slt diff --git a/AGENT.md b/AGENT.md index 602d72a7..224c5840 100644 --- a/AGENT.md +++ b/AGENT.md @@ -81,6 +81,8 @@ PRs that modify logic but leave obvious test gaps untouched may be rejected. ### 3.1 Prefer Simple Code +- Implement functionality with the smallest reasonable change and the simplest correct design. +- Do not over-engineer or introduce abstractions for hypothetical future requirements. - Prefer straightforward control flow - Avoid unnecessary abstractions - Avoid premature generalization diff --git a/docs/features.md b/docs/features.md index e366ff00..b66d5ef1 100644 --- a/docs/features.md +++ b/docs/features.md @@ -183,6 +183,11 @@ If `unsafe_txdb_checkpoint` is not enabled, `build_rocksdb()` returns an explici - [x] Exists - [x] Group By - [x] Having +- [x] Window functions: + - `row_number()`, `rank()`, `dense_rank()` + - `count()`, `sum()`, `avg()`, `min()`, `max()` with `OVER` + - `PARTITION BY` and window `ORDER BY` with the default frame + - Explicit frames, named windows, and `QUALIFY` are not yet supported - [x] Order By - [x] Limit - [x] Show Tables diff --git a/src/binder/aggregate.rs b/src/binder/aggregate.rs index 40ad6f21..8216716e 100644 --- a/src/binder/aggregate.rs +++ b/src/binder/aggregate.rs @@ -16,7 +16,6 @@ use std::collections::HashSet; use super::{Binder, QueryBindStep}; use crate::errors::DatabaseError; -use crate::expression::function::scala::ScalarFunction; use crate::expression::visitor::{walk_expr, ExprVisitor}; use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut}; use crate::planner::LogicalPlan; @@ -27,6 +26,22 @@ use crate::{ planner::operator::{aggregate::AggregateOperator, sort::SortField}, }; +struct AggregateCallCollector<'a> { + agg_calls: &'a mut Vec, +} + +impl<'expr> ExprVisitor<'expr> for AggregateCallCollector<'_> { + fn visit(&mut self, expr: &'expr ScalarExpression) -> Result<(), DatabaseError> { + match expr { + ScalarExpression::AggCall { .. } => self.agg_calls.push(expr.clone()), + ScalarExpression::Alias { expr, .. } => self.visit(expr)?, + ScalarExpression::Empty | ScalarExpression::TableFunction(_) => unreachable!(), + _ => walk_expr(self, expr)?, + } + Ok(()) + } +} + impl> Binder<'_, '_, T, A> { pub fn bind_aggregate( &mut self, @@ -48,7 +63,7 @@ impl> Binder<'_, '_, T, A> select_items: &mut [ScalarExpression], ) -> Result<(), DatabaseError> { for column in select_items { - self.visit_column_agg_expr(column)?; + self.collect_aggregate_calls(column)?; } Ok(()) } @@ -77,14 +92,14 @@ impl> Binder<'_, '_, T, A> F: FnMut(&mut Self, I::Item) -> Result, { if let Some(having) = having.as_mut() { - self.visit_column_agg_expr(having)?; + self.collect_aggregate_calls(having)?; } let mut return_orderby = None; if let Some(orderby) = orderby { let mut fields = Vec::new(); for orderby in orderby { - let mut field = bind_sort_field(self, orderby)?; - self.visit_column_agg_expr(&mut field.expr)?; + let field = bind_sort_field(self, orderby)?; + self.collect_aggregate_calls(&field.expr)?; fields.push(field); } return_orderby = Some(fields); @@ -119,119 +134,11 @@ impl> Binder<'_, '_, T, A> Ok(()) } - fn visit_column_agg_expr(&mut self, expr: &mut ScalarExpression) -> Result<(), DatabaseError> { - match expr { - ScalarExpression::AggCall { .. } => { - self.context.agg_calls.push(expr.clone()); - } - ScalarExpression::TypeCast { expr, .. } => self.visit_column_agg_expr(expr)?, - ScalarExpression::IsNull { expr, .. } => self.visit_column_agg_expr(expr)?, - ScalarExpression::Unary { expr, .. } => self.visit_column_agg_expr(expr)?, - ScalarExpression::Alias { expr, .. } => self.visit_column_agg_expr(expr)?, - ScalarExpression::Binary { - left_expr, - right_expr, - .. - } => { - self.visit_column_agg_expr(left_expr)?; - self.visit_column_agg_expr(right_expr)?; - } - ScalarExpression::In { expr, args, .. } => { - self.visit_column_agg_expr(expr)?; - for arg in args { - self.visit_column_agg_expr(arg)?; - } - } - ScalarExpression::Between { - expr, - left_expr, - right_expr, - .. - } => { - self.visit_column_agg_expr(expr)?; - self.visit_column_agg_expr(left_expr)?; - self.visit_column_agg_expr(right_expr)?; - } - ScalarExpression::SubString { - expr, - for_expr, - from_expr, - } => { - self.visit_column_agg_expr(expr)?; - if let Some(expr) = for_expr { - self.visit_column_agg_expr(expr)?; - } - if let Some(expr) = from_expr { - self.visit_column_agg_expr(expr)?; - } - } - ScalarExpression::Position { expr, in_expr } => { - self.visit_column_agg_expr(expr)?; - self.visit_column_agg_expr(in_expr)?; - } - ScalarExpression::Trim { - expr, - trim_what_expr, - .. - } => { - self.visit_column_agg_expr(expr)?; - if let Some(trim_what_expr) = trim_what_expr { - self.visit_column_agg_expr(trim_what_expr)?; - } - } - ScalarExpression::Constant(_) | ScalarExpression::ColumnRef { .. } => (), - ScalarExpression::Empty => unreachable!(), - ScalarExpression::Tuple(args) - | ScalarExpression::ScalaFunction(ScalarFunction { args, .. }) - | ScalarExpression::Coalesce { exprs: args, .. } => { - for expr in args { - self.visit_column_agg_expr(expr)?; - } - } - ScalarExpression::If { - condition, - left_expr, - right_expr, - .. - } => { - self.visit_column_agg_expr(condition)?; - self.visit_column_agg_expr(left_expr)?; - self.visit_column_agg_expr(right_expr)?; - } - ScalarExpression::IfNull { - left_expr, - right_expr, - .. - } - | ScalarExpression::NullIf { - left_expr, - right_expr, - .. - } => { - self.visit_column_agg_expr(left_expr)?; - self.visit_column_agg_expr(right_expr)?; - } - ScalarExpression::CaseWhen { - operand_expr, - expr_pairs, - else_expr, - .. - } => { - if let Some(expr) = operand_expr { - self.visit_column_agg_expr(expr)?; - } - for (expr_1, expr_2) in expr_pairs { - self.visit_column_agg_expr(expr_1)?; - self.visit_column_agg_expr(expr_2)?; - } - if let Some(expr) = else_expr { - self.visit_column_agg_expr(expr)?; - } - } - ScalarExpression::TableFunction(_) => unreachable!(), + fn collect_aggregate_calls(&mut self, expr: &ScalarExpression) -> Result<(), DatabaseError> { + AggregateCallCollector { + agg_calls: &mut self.context.agg_calls, } - - Ok(()) + .visit(expr) } /// Validate select exprs must appear in the GROUP BY clause or be used in @@ -269,6 +176,10 @@ impl> Binder<'_, '_, T, A> HashSet::from_iter(group_raw_exprs.iter().copied()); for expr in select_items { + if expr.has_window_call()? { + HavingOrderByValidator::new(groupby, &self.context.agg_calls).visit(expr)?; + continue; + } if expr.has_agg_call()? { continue; } diff --git a/src/binder/mod.rs b/src/binder/mod.rs index 4323ae6d..15efb0fc 100644 --- a/src/binder/mod.rs +++ b/src/binder/mod.rs @@ -48,6 +48,7 @@ mod show_table; mod show_view; mod truncate; mod update; +mod window; #[cfg(feature = "parser")] pub use parser::{command_type, prepare, prepare_all, CommandType, Statement}; @@ -90,6 +91,7 @@ pub enum QueryBindStep { Where, Agg, Having, + Window, Distinct, Sort, Project, diff --git a/src/binder/parser.rs b/src/binder/parser.rs index 62b8499a..41cbe304 100644 --- a/src/binder/parser.rs +++ b/src/binder/parser.rs @@ -14,8 +14,7 @@ use super::select::{ BindPlanAggregated, BindPlanComplete, BindPlanDistinct, BindPlanFiltered, BindPlanFrom, - BindPlanHaving, BindPlanProjected, BindPlanSelectList, BindPlanStart, JoinConstraintInput, - TableAliasInput, + BindPlanProjected, BindPlanSelectList, BindPlanStart, JoinConstraintInput, TableAliasInput, }; use super::{is_valid_identifier, with_query_bind_step, Binder, QueryBindStep, SetOperatorKind}; #[cfg(feature = "copy")] @@ -47,6 +46,7 @@ pub(super) use sqlparser::ast::{ ObjectType, OrderByExpr, OrderByKind, Query, Select, SelectInto, SelectItem, SelectItemQualifiedWildcardKind, SetExpr, SetOperator, SetQuantifier, Spanned, TableAlias, TableConstraint, TableFactor, TableObject, TableWithJoins, TypedString, UnaryOperator, Value, + WindowType, }; #[cfg(feature = "copy")] pub(super) use sqlparser::ast::{CopyOption, CopySource, CopyTarget}; @@ -1463,7 +1463,7 @@ where } } -impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanHaving<'s, 'a, 'b, 'arena, T, A> +impl<'s, 'a: 'b, 'b, 'arena, T, A> super::select::BindPlanWindowed<'s, 'a, 'b, 'arena, T, A> where T: Transaction, A: AsRef<[(&'static str, DataValue)]>, @@ -2392,7 +2392,15 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder< arena: &mut PlanArena, ) -> Result { let func_span = func.span(); - let Function { name, args, .. } = func; + let Function { + name, + args, + over, + filter, + null_treatment, + within_group, + .. + } = func; let (func_args, is_distinct) = match args { FunctionArguments::List(args) => ( args.args.as_slice(), @@ -2425,10 +2433,78 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder< } let function_name = name.to_string().to_lowercase(); + if let Some(over) = over { + if filter.is_some() || null_treatment.is_some() || !within_group.is_empty() { + return Err(DatabaseError::UnsupportedStmt( + "FILTER, NULL treatment, and WITHIN GROUP are not supported for window functions" + .to_string(), + )); + } + return self + .bind_window_call(function_name, args, is_distinct, over, arena) + .map_err(|err| attach_span_if_absent(err, func_span)); + } + self.bind_function_call(function_name, args, is_distinct, arena) .map_err(|err| attach_span_if_absent(err, func_span)) } + fn bind_window_call( + &mut self, + function_name: String, + args: Vec, + is_distinct: bool, + over: &WindowType, + arena: &mut PlanArena, + ) -> Result { + if !matches!( + self.context.step_now(), + QueryBindStep::Project | QueryBindStep::Sort + ) { + return Err(DatabaseError::UnsupportedStmt( + "window functions are only allowed in SELECT and ORDER BY".to_string(), + )); + } + if is_distinct { + return Err(DatabaseError::UnsupportedStmt( + "DISTINCT window aggregates are not supported".to_string(), + )); + } + let WindowType::WindowSpec(spec) = over else { + return Err(DatabaseError::UnsupportedStmt( + "named windows are not supported".to_string(), + )); + }; + if spec.window_name.is_some() { + return Err(DatabaseError::UnsupportedStmt( + "inherited named windows are not supported".to_string(), + )); + } + if spec.window_frame.is_some() { + return Err(DatabaseError::UnsupportedStmt( + "explicit window frames are not supported".to_string(), + )); + } + + let partition_by = spec + .partition_by + .iter() + .map(|expr| self.bind_expr(expr, arena)) + .collect::, _>>()?; + let order_by = spec + .order_by + .iter() + .map(|OrderByExpr { expr, options, .. }| { + Ok(SortField::new( + self.bind_expr(expr, arena)?, + options.asc.unwrap_or(true), + options.nulls_first.unwrap_or(false), + )) + }) + .collect::, DatabaseError>>()?; + self.bind_window_function(function_name, args, partition_by, order_by, arena) + } + pub fn bind_set_expr( &mut self, set_expr: &SetExpr, @@ -2511,8 +2587,20 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder< having, distinct, into, + named_window, + qualify, .. } = select; + if !named_window.is_empty() { + return Err(DatabaseError::UnsupportedStmt( + "named windows are not supported".to_string(), + )); + } + if qualify.is_some() { + return Err(DatabaseError::UnsupportedStmt( + "QUALIFY is not supported".to_string(), + )); + } Ok(self .build_plan(arena) .from_sql(from)? @@ -2520,6 +2608,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder< .where_sql(selection.as_ref())? .aggregate_sql(group_by, having.as_ref(), orderby)? .having()? + .window()? .distinct_sql(distinct.as_ref())? .order_by()? .project()? diff --git a/src/binder/select.rs b/src/binder/select.rs index ce8d69a7..eace5f7b 100644 --- a/src/binder/select.rs +++ b/src/binder/select.rs @@ -232,6 +232,18 @@ where orderby: Option>, } +pub(crate) struct BindPlanWindowed<'s, 'a, 'b, 'arena, T, A> +where + T: Transaction, + A: AsRef<[(&'static str, DataValue)]>, +{ + binder: &'s mut Binder<'a, 'b, T, A>, + arena: &'s mut crate::planner::PlanArena<'arena>, + plan: LogicalPlan, + select_list: Vec, + orderby: Option>, +} + pub(crate) struct BindPlanDistinct<'s, 'a, 'b, 'arena, T, A> where T: Transaction, @@ -357,6 +369,7 @@ where |_binder, _arena, order| Ok(order), )? .having()? + .window()? .distinct(false)? .order_by()?; Ok(BindPlanSelectList { @@ -379,6 +392,7 @@ where |_binder, _arena, order| Ok(order), )? .having()? + .window()? .distinct(false)? .order_by()?; Ok(BindPlanSelectList { @@ -433,7 +447,7 @@ where #[cfg(feature = "orm")] pub fn finish(self) -> Result { for expr in &self.select_list { - if expr.has_agg_call()? { + if expr.has_agg_call()? || expr.has_window_call()? { return self.aggregate_without_group()?.finish(); } } @@ -579,6 +593,31 @@ where } impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanHaving<'s, 'a, 'b, 'arena, T, A> +where + T: Transaction, + A: AsRef<[(&'static str, DataValue)]>, +{ + pub(crate) fn window( + mut self, + ) -> Result, DatabaseError> { + self.plan = self.binder.bind_window( + self.plan, + &mut self.select_list, + &mut self.orderby, + self.arena, + )?; + + Ok(BindPlanWindowed { + binder: self.binder, + arena: self.arena, + plan: self.plan, + select_list: self.select_list, + orderby: self.orderby, + }) + } +} + +impl<'s, 'a: 'b, 'b, 'arena, T, A> BindPlanWindowed<'s, 'a, 'b, 'arena, T, A> where T: Transaction, A: AsRef<[(&'static str, DataValue)]>, diff --git a/src/binder/window.rs b/src/binder/window.rs new file mode 100644 index 00000000..9a48531c --- /dev/null +++ b/src/binder/window.rs @@ -0,0 +1,255 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{Binder, QueryBindStep}; +use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef}; +use crate::errors::DatabaseError; +use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut}; +use crate::expression::window::{WindowCall, WindowFunction, WindowFunctionKind, WindowSpec}; +use crate::expression::ScalarExpression; +use crate::planner::operator::sort::{SortField, SortOperator}; +use crate::planner::operator::window::WindowOperator; +use crate::planner::operator::Operator; +use crate::planner::{Childrens, LogicalPlan, PlanArena}; +use crate::storage::Transaction; +use crate::types::value::DataValue; +use crate::types::LogicalType; + +struct WindowCollector<'a, 'p> { + arena: &'a mut PlanArena<'p>, + windows: Vec<(WindowCall, ColumnRef)>, +} + +impl ExprVisitorMut<'_> for WindowCollector<'_, '_> { + fn visit(&mut self, expr: &mut ScalarExpression) -> Result<(), DatabaseError> { + let ScalarExpression::WindowCall(window) = expr else { + return walk_mut_expr(self, expr); + }; + if let Some((_, output_column)) = self + .windows + .iter() + .find(|(candidate, _)| candidate == window) + { + *expr = ScalarExpression::column_expr(*output_column, 0); + return Ok(()); + } + + let output_name = expr.output_name(self.arena); + let ScalarExpression::WindowCall(window) = std::mem::replace(expr, ScalarExpression::Empty) + else { + unreachable!() + }; + let output_column = self.arena.alloc_column(ColumnCatalog::new( + output_name, + true, + ColumnDesc::new(window.function.ty.clone(), None, false, None)?, + )); + self.windows.push((window, output_column)); + *expr = ScalarExpression::column_expr(output_column, 0); + Ok(()) + } +} + +struct WindowOutputBinder { + outputs: Vec<(ColumnRef, usize)>, +} + +impl ExprVisitorMut<'_> for WindowOutputBinder { + fn visit_column_ref( + &mut self, + column: &mut ColumnRef, + position: &mut usize, + ) -> Result<(), DatabaseError> { + if let Some((_, output_position)) = self + .outputs + .iter() + .find(|(output_column, _)| output_column == column) + { + *position = *output_position; + } + Ok(()) + } +} + +struct WindowGroup { + partition_by: Vec, + order_by: Vec, + functions: Vec, + output_columns: Vec, +} + +impl> Binder<'_, '_, T, A> { + pub(crate) fn bind_window_function( + &mut self, + function_name: String, + args: Vec, + partition_by: Vec, + order_by: Vec, + arena: &mut PlanArena, + ) -> Result { + if !matches!( + self.context.step_now(), + QueryBindStep::Project | QueryBindStep::Sort + ) { + return Err(DatabaseError::UnsupportedStmt( + "window functions are only allowed in SELECT and ORDER BY".to_string(), + )); + } + for expr in args + .iter() + .chain(&partition_by) + .chain(order_by.iter().map(|field| &field.expr)) + { + if expr.has_window_call()? { + return Err(DatabaseError::UnsupportedStmt( + "window functions cannot be nested".to_string(), + )); + } + } + + let (kind, args, ty) = match function_name.as_str() { + "row_number" | "rank" | "dense_rank" => { + if !args.is_empty() { + return Err(DatabaseError::MisMatch( + "number of ranking function parameters", + "0", + )); + } + let kind = match function_name.as_str() { + "row_number" => WindowFunctionKind::RowNumber, + "rank" => WindowFunctionKind::Rank, + "dense_rank" => WindowFunctionKind::DenseRank, + _ => unreachable!(), + }; + (kind, args, LogicalType::Bigint) + } + "count" | "sum" | "avg" | "min" | "max" => { + let ScalarExpression::AggCall { kind, args, ty, .. } = + self.bind_function_call(function_name, args, false, arena)? + else { + unreachable!() + }; + (WindowFunctionKind::Aggregate(kind), args, ty) + } + _ => { + return Err(DatabaseError::UnsupportedStmt(format!( + "window function `{function_name}` is not supported" + ))) + } + }; + + Ok(ScalarExpression::WindowCall(WindowCall { + function: WindowFunction { kind, args, ty }, + spec: WindowSpec { + partition_by, + order_by, + }, + })) + } + + pub(crate) fn bind_window( + &mut self, + mut children: LogicalPlan, + select_list: &mut [ScalarExpression], + order_by: &mut Option>, + arena: &mut PlanArena, + ) -> Result { + let mut collector = WindowCollector { + arena, + windows: Vec::new(), + }; + for expr in select_list.iter_mut() { + collector.visit(expr)?; + } + if let Some(order_by) = order_by.as_mut() { + for field in order_by { + collector.visit(&mut field.expr)?; + } + } + if collector.windows.is_empty() { + return Ok(children); + } + let windows = collector.windows; + + let base_position = children.output_schema(arena).len(); + let mut groups: Vec = Vec::new(); + for (window, output_column) in windows { + let WindowCall { function, spec } = window; + let group_index = groups.iter().position(|group| { + group.partition_by == spec.partition_by && group.order_by == spec.order_by + }); + + if let Some(index) = group_index { + groups[index].functions.push(function); + groups[index].output_columns.push(output_column); + } else { + groups.push(WindowGroup { + partition_by: spec.partition_by, + order_by: spec.order_by, + functions: vec![function], + output_columns: vec![output_column], + }); + } + } + + let mut outputs = Vec::new(); + let mut position = base_position; + for group in &groups { + for column in &group.output_columns { + outputs.push((*column, position)); + position += 1; + } + } + + let mut output_binder = WindowOutputBinder { outputs }; + for expr in select_list { + output_binder.visit(expr)?; + } + if let Some(order_by) = order_by { + for field in order_by { + output_binder.visit(&mut field.expr)?; + } + } + + for group in groups { + let sort_fields = group + .partition_by + .iter() + .cloned() + .map(SortField::from) + .chain(group.order_by.iter().cloned()) + .collect::>(); + if !sort_fields.is_empty() { + children = LogicalPlan::new( + Operator::Sort(SortOperator { + sort_fields, + limit: None, + }), + Childrens::Only(Box::new(children)), + ); + } + children = LogicalPlan::new( + Operator::Window(WindowOperator { + partition_by: group.partition_by, + order_by: group.order_by, + functions: group.functions, + output_columns: group.output_columns, + }), + Childrens::Only(Box::new(children)), + ); + } + self.context.step(QueryBindStep::Window); + Ok(children) + } +} diff --git a/src/db.rs b/src/db.rs index 65ad0f8d..ebec405f 100644 --- a/src/db.rs +++ b/src/db.rs @@ -390,6 +390,7 @@ fn default_optimizer_pipeline() -> HepOptimizerPipeline { ImplementationRuleImpl::Sort, ImplementationRuleImpl::TopK, ImplementationRuleImpl::Values, + ImplementationRuleImpl::Window, // DML ImplementationRuleImpl::Analyze, #[cfg(feature = "copy")] diff --git a/src/execution/dql/aggregate/avg.rs b/src/execution/dql/aggregate/avg.rs index a169fd20..061edec5 100644 --- a/src/execution/dql/aggregate/avg.rs +++ b/src/execution/dql/aggregate/avg.rs @@ -23,6 +23,7 @@ use std::borrow::Cow; pub struct AvgAccumulator { inner: Option, count: usize, + result: DataValue, } impl AvgAccumulator { @@ -30,6 +31,7 @@ impl AvgAccumulator { Self { inner: None, count: 0, + result: DataValue::Null, } } } @@ -45,33 +47,67 @@ impl Accumulator for AvgAccumulator { }; acc.update_value(value)?; self.count += 1; + self.result = DataValue::Null; } Ok(()) } - fn evaluate(self: Box) -> Result { - let Self { inner, count } = *self; - let Some(acc) = inner else { - return Ok(DataValue::Null); + fn evaluate(&mut self) -> Result<(), DatabaseError> { + if !self.result.is_null() { + return Ok(()); + } + let Some(acc) = &self.inner else { + return Ok(()); }; - let mut value = acc.into_result(); + let mut value = Cow::Borrowed(acc.result()); let value_ty = value.logical_type(); - if count == 0 { - return Ok(DataValue::Null); + if self.count == 0 { + return Ok(()); } let quantity = if value_ty.is_signed_numeric() { - DataValue::Int64(count as i64) + DataValue::Int64(self.count as i64) } else { - DataValue::UInt32(count as u32) + DataValue::UInt32(self.count as u32) }; let quantity_ty = quantity.logical_type(); if value_ty != quantity_ty { - value = value.cast(&quantity_ty)? + value = Cow::Owned(value.into_owned().cast(&quantity_ty)?) } let evaluator = binary_create(Cow::Owned(quantity_ty), BinaryOperator::Divide)?; - evaluator.binary_eval(&value, &quantity) + self.result = evaluator.binary_eval(value.as_ref(), &quantity)?; + Ok(()) + } + + fn result(&self) -> &DataValue { + &self.result + } + + fn result_owned(self: Box) -> DataValue { + self.result + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn invalidate_cached_result_after_update() -> Result<(), DatabaseError> { + let mut accumulator = AvgAccumulator::new(); + accumulator.update_value(&DataValue::Int32(2))?; + accumulator.evaluate()?; + let first = accumulator.result().clone(); + + accumulator.update_value(&DataValue::Int32(4))?; + accumulator.evaluate()?; + let second = accumulator.result().clone(); + + assert_ne!(first, second); + accumulator.evaluate()?; + assert_eq!(&second, accumulator.result()); + Ok(()) } } diff --git a/src/execution/dql/aggregate/count.rs b/src/execution/dql/aggregate/count.rs index e9dcc0ef..08c20c72 100644 --- a/src/execution/dql/aggregate/count.rs +++ b/src/execution/dql/aggregate/count.rs @@ -18,37 +18,48 @@ use crate::types::value::DataValue; use std::collections::HashSet; pub struct CountAccumulator { - result: i32, + result: DataValue, } impl CountAccumulator { pub fn new() -> Self { - Self { result: 0 } + Self { + result: DataValue::Int32(0), + } } } impl Accumulator for CountAccumulator { fn update_value(&mut self, value: &DataValue) -> Result<(), DatabaseError> { if !value.is_null() { - self.result += 1; + let DataValue::Int32(result) = &mut self.result else { + unreachable!() + }; + *result += 1; } Ok(()) } - fn evaluate(self: Box) -> Result { - Ok(DataValue::Int32(self.result)) + fn result(&self) -> &DataValue { + &self.result + } + + fn result_owned(self: Box) -> DataValue { + self.result } } pub struct DistinctCountAccumulator { distinct_values: HashSet, + result: DataValue, } impl DistinctCountAccumulator { pub fn new() -> Self { Self { distinct_values: HashSet::default(), + result: DataValue::Int32(0), } } } @@ -56,13 +67,19 @@ impl DistinctCountAccumulator { impl Accumulator for DistinctCountAccumulator { fn update_value(&mut self, value: &DataValue) -> Result<(), DatabaseError> { if !value.is_null() { - self.distinct_values.insert(value.clone()); + if self.distinct_values.insert(value.clone()) { + self.result = DataValue::Int32(self.distinct_values.len() as i32); + } } Ok(()) } - fn evaluate(self: Box) -> Result { - Ok(DataValue::Int32(self.distinct_values.len() as i32)) + fn result(&self) -> &DataValue { + &self.result + } + + fn result_owned(self: Box) -> DataValue { + self.result } } diff --git a/src/execution/dql/aggregate/hash_agg.rs b/src/execution/dql/aggregate/hash_agg.rs index c8f54979..fd5b38cf 100644 --- a/src/execution/dql/aggregate/hash_agg.rs +++ b/src/execution/dql/aggregate/hash_agg.rs @@ -124,8 +124,9 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for HashAggExecutor { output.values.clear(); output.values.reserve(accs.len() + group_keys.len()); - for acc in accs { - output.values.push(acc.evaluate()?); + for mut acc in accs { + acc.evaluate()?; + output.values.push(acc.result_owned()); } output.values.extend(group_keys); arena.resume(); diff --git a/src/execution/dql/aggregate/min_max.rs b/src/execution/dql/aggregate/min_max.rs index 812f5b63..186dc665 100644 --- a/src/execution/dql/aggregate/min_max.rs +++ b/src/execution/dql/aggregate/min_max.rs @@ -20,7 +20,7 @@ use crate::types::value::DataValue; use std::borrow::Cow; pub struct MinMaxAccumulator { - inner: Option, + result: DataValue, op: BinaryOperator, } @@ -32,16 +32,19 @@ impl MinMaxAccumulator { BinaryOperator::Gt }; - Self { inner: None, op } + Self { + result: DataValue::Null, + op, + } } } impl Accumulator for MinMaxAccumulator { fn update_value(&mut self, value: &DataValue) -> Result<(), DatabaseError> { if !value.is_null() { - if let Some(inner_value) = &self.inner { + if !self.result.is_null() { let evaluator = binary_create(Cow::Owned(value.logical_type()), self.op)?; - if let DataValue::Boolean(result) = evaluator.binary_eval(inner_value, value)? { + if let DataValue::Boolean(result) = evaluator.binary_eval(&self.result, value)? { result } else { return Err(DatabaseError::InvalidType); @@ -49,13 +52,17 @@ impl Accumulator for MinMaxAccumulator { } else { true } - .then(|| self.inner = Some(value.clone())); + .then(|| self.result = value.clone()); } Ok(()) } - fn evaluate(self: Box) -> Result { - Ok(self.inner.unwrap_or(DataValue::Null)) + fn result(&self) -> &DataValue { + &self.result + } + + fn result_owned(self: Box) -> DataValue { + self.result } } diff --git a/src/execution/dql/aggregate/mod.rs b/src/execution/dql/aggregate/mod.rs index 917654b3..3e125c1e 100644 --- a/src/execution/dql/aggregate/mod.rs +++ b/src/execution/dql/aggregate/mod.rs @@ -38,35 +38,46 @@ pub trait Accumulator { /// updates the accumulator's state from a vector of arrays. fn update_value(&mut self, value: &DataValue) -> Result<(), DatabaseError>; - /// returns its value based on its current state. - fn evaluate(self: Box) -> Result; + /// evaluates its result based on its current state. + fn evaluate(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn result(&self) -> &DataValue; + + fn result_owned(self: Box) -> DataValue; } -fn create_accumulator(expr: &ScalarExpression) -> Result, DatabaseError> { - if let ScalarExpression::AggCall { - kind, ty, distinct, .. - } = expr - { - Ok(match (kind, distinct) { - (AggKind::Count, false) => Box::new(CountAccumulator::new()), - (AggKind::Count, true) => Box::new(DistinctCountAccumulator::new()), - (AggKind::Sum, false) => Box::new(SumAccumulator::new(Cow::Borrowed(ty))?), - (AggKind::Sum, true) => Box::new(DistinctSumAccumulator::new(ty)?), - (AggKind::Min, _) => Box::new(MinMaxAccumulator::new(false)), - (AggKind::Max, _) => Box::new(MinMaxAccumulator::new(true)), - (AggKind::Avg, _) => Box::new(AvgAccumulator::new()), - }) - } else { - unreachable!( - "create_accumulator called with non-aggregate expression {}", - expr - ); - } +pub(crate) fn create_accumulator( + kind: AggKind, + ty: &crate::types::LogicalType, + distinct: bool, +) -> Result, DatabaseError> { + Ok(match (kind, distinct) { + (AggKind::Count, false) => Box::new(CountAccumulator::new()), + (AggKind::Count, true) => Box::new(DistinctCountAccumulator::new()), + (AggKind::Sum, false) => Box::new(SumAccumulator::new(Cow::Borrowed(ty))?), + (AggKind::Sum, true) => Box::new(DistinctSumAccumulator::new(ty)?), + (AggKind::Min, _) => Box::new(MinMaxAccumulator::new(false)), + (AggKind::Max, _) => Box::new(MinMaxAccumulator::new(true)), + (AggKind::Avg, _) => Box::new(AvgAccumulator::new()), + }) } #[inline] pub(crate) fn create_accumulators( exprs: &[ScalarExpression], ) -> Result>, DatabaseError> { - exprs.iter().map(create_accumulator).try_collect() + exprs + .iter() + .map(|expr| { + let ScalarExpression::AggCall { + kind, ty, distinct, .. + } = expr + else { + unreachable!("create_accumulators called with non-aggregate expression {expr}") + }; + create_accumulator(*kind, ty, *distinct) + }) + .try_collect() } diff --git a/src/execution/dql/aggregate/simple_agg.rs b/src/execution/dql/aggregate/simple_agg.rs index 6986e826..74a5f15d 100644 --- a/src/execution/dql/aggregate/simple_agg.rs +++ b/src/execution/dql/aggregate/simple_agg.rs @@ -81,8 +81,9 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for SimpleAggExecutor { output.pk = None; output.values.clear(); output.values.reserve(accs.len()); - for acc in accs { - output.values.push(acc.evaluate()?); + for mut acc in accs { + acc.evaluate()?; + output.values.push(acc.result_owned()); } self.returned = true; arena.resume(); diff --git a/src/execution/dql/aggregate/sum.rs b/src/execution/dql/aggregate/sum.rs index dd721a5c..d0ced0a8 100644 --- a/src/execution/dql/aggregate/sum.rs +++ b/src/execution/dql/aggregate/sum.rs @@ -35,10 +35,6 @@ impl SumAccumulator { evaluator: binary_create(ty, BinaryOperator::Plus)?, }) } - - pub(super) fn into_result(self) -> DataValue { - self.result - } } impl Accumulator for SumAccumulator { @@ -54,8 +50,12 @@ impl Accumulator for SumAccumulator { Ok(()) } - fn evaluate(self: Box) -> Result { - Ok(self.into_result()) + fn result(&self) -> &DataValue { + &self.result + } + + fn result_owned(self: Box) -> DataValue { + self.result } } @@ -83,7 +83,11 @@ impl Accumulator for DistinctSumAccumulator { Ok(()) } - fn evaluate(self: Box) -> Result { - Ok(self.inner.into_result()) + fn result(&self) -> &DataValue { + self.inner.result() + } + + fn result_owned(self: Box) -> DataValue { + self.inner.result } } diff --git a/src/execution/dql/mod.rs b/src/execution/dql/mod.rs index 774c44e5..fbe27e87 100644 --- a/src/execution/dql/mod.rs +++ b/src/execution/dql/mod.rs @@ -33,6 +33,7 @@ pub(crate) mod sort; pub(crate) mod top_k; pub(crate) mod union; pub(crate) mod values; +pub(crate) mod window; #[cfg(all(test, not(target_arch = "wasm32")))] pub(crate) mod test { diff --git a/src/execution/dql/window.rs b/src/execution/dql/window.rs new file mode 100644 index 00000000..1884ac9c --- /dev/null +++ b/src/execution/dql/window.rs @@ -0,0 +1,156 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::execution::{ + build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor, +}; +use crate::expression::ScalarExpression; +use crate::planner::operator::sort::SortField; +use crate::planner::operator::window::WindowOperator; +use crate::planner::LogicalPlan; +use crate::storage::Transaction; +use crate::types::tuple::Tuple; +use crate::types::value::DataValue; +use std::mem; + +mod function; + +use function::WindowFunction; + +pub struct Window { + partition_by: Vec, + order_by: Vec, + functions: Vec>, + input: ExecId, + pending: Option, + rows: Vec, +} + +impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Window { + type Input = (WindowOperator, LogicalPlan); + + fn into_executor( + (operator, input): Self::Input, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut crate::planner::PlanArena<'a>, + cache: ExecutionContext<'_>, + transaction: &T, + ) -> ExecId { + let input = build_read(arena, plan_arena, input, cache, transaction); + let WindowOperator { + partition_by, + order_by, + functions: window_functions, + .. + } = operator; + let mut functions = Vec::with_capacity(window_functions.len()); + for function in window_functions { + let crate::expression::window::WindowFunction { kind, args, ty } = function; + functions.push(function::new(kind, args, ty)); + } + arena.push(ExecNode::Window(Window { + partition_by, + order_by, + functions, + input, + pending: None, + rows: Vec::new(), + })) + } +} + +fn evaluate_partition( + rows: &mut [Tuple], + order_by: &[SortField], + functions: &mut [Box], +) -> Result<(), DatabaseError> { + let Some(first) = rows.first() else { + return Ok(()); + }; + let output_offset = first.values.len(); + for row in rows.iter_mut() { + row.values + .resize(output_offset + functions.len(), DataValue::Null); + } + for function in functions.iter_mut() { + function.reset()?; + } + let mut peer_start = 0; + let mut peer_index = 0; + while peer_start < rows.len() { + let mut peer_end = peer_start + 1; + 'peer: while peer_end < rows.len() { + for field in order_by { + if field.expr.eval(Some(&rows[peer_end - 1]))? + != field.expr.eval(Some(&rows[peer_end]))? + { + break 'peer; + } + } + peer_end += 1; + } + for (slot, function) in functions.iter_mut().enumerate() { + function.evaluate(rows, peer_start..peer_end, peer_index, output_offset + slot)?; + } + peer_start = peer_end; + peer_index += 1; + } + Ok(()) +} + +impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window { + fn next_tuple( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut crate::planner::PlanArena<'a>, + ) -> Result<(), DatabaseError> { + loop { + if let Some(tuple) = self.rows.pop() { + arena.produce_tuple(tuple); + return Ok(()); + } + + let first = if let Some(tuple) = self.pending.take() { + tuple + } else if arena.next_tuple(self.input, plan_arena)? { + mem::take(arena.result_tuple_mut()) + } else { + arena.finish(); + return Ok(()); + }; + self.rows.push(first); + + while arena.next_tuple(self.input, plan_arena)? { + let tuple = mem::take(arena.result_tuple_mut()); + let mut same_partition = true; + for expr in &self.partition_by { + if expr.eval(self.rows.last())? != expr.eval(Some(&tuple))? { + same_partition = false; + break; + } + } + if same_partition { + self.rows.push(tuple); + } else { + self.pending = Some(tuple); + break; + } + } + + evaluate_partition(&mut self.rows, &self.order_by, &mut self.functions)?; + self.rows.reverse(); + } + } +} diff --git a/src/execution/dql/window/function.rs b/src/execution/dql/window/function.rs new file mode 100644 index 00000000..eef66bef --- /dev/null +++ b/src/execution/dql/window/function.rs @@ -0,0 +1,137 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::execution::dql::aggregate::{create_accumulator, Accumulator}; +use crate::expression::agg::AggKind; +use crate::expression::window::WindowFunctionKind; +use crate::expression::ScalarExpression; +use crate::types::tuple::Tuple; +use crate::types::value::DataValue; +use crate::types::LogicalType; +use std::ops::Range; + +pub(super) trait WindowFunction { + fn reset(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn evaluate( + &mut self, + rows: &mut [Tuple], + peer: Range, + peer_index: usize, + output_position: usize, + ) -> Result<(), DatabaseError>; +} + +struct RowNumber; + +impl WindowFunction for RowNumber { + fn evaluate( + &mut self, + rows: &mut [Tuple], + peer: Range, + _peer_index: usize, + output_position: usize, + ) -> Result<(), DatabaseError> { + let start = peer.start; + for (offset, row) in rows[peer].iter_mut().enumerate() { + row.values[output_position] = DataValue::Int64((start + offset + 1) as i64); + } + Ok(()) + } +} + +struct Rank { + dense: bool, +} + +impl WindowFunction for Rank { + fn evaluate( + &mut self, + rows: &mut [Tuple], + peer: Range, + peer_index: usize, + output_position: usize, + ) -> Result<(), DatabaseError> { + let rank = if self.dense { + peer_index + 1 + } else { + peer.start + 1 + }; + for row in &mut rows[peer] { + row.values[output_position] = DataValue::Int64(rank as i64); + } + Ok(()) + } +} + +struct Aggregate { + kind: AggKind, + ty: LogicalType, + arg: ScalarExpression, + accumulator: Option>, +} + +impl WindowFunction for Aggregate { + fn reset(&mut self) -> Result<(), DatabaseError> { + self.accumulator = Some(create_accumulator(self.kind, &self.ty, false)?); + Ok(()) + } + + fn evaluate( + &mut self, + rows: &mut [Tuple], + peer: Range, + _peer_index: usize, + output_position: usize, + ) -> Result<(), DatabaseError> { + let Some(accumulator) = self.accumulator.as_mut() else { + unreachable!() + }; + for row in &rows[peer.clone()] { + accumulator.update_value(&self.arg.eval(Some(row))?)?; + } + accumulator.evaluate()?; + let result = accumulator.result(); + for row in &mut rows[peer] { + row.values[output_position] = result.clone(); + } + Ok(()) + } +} + +pub(super) fn new( + kind: WindowFunctionKind, + args: Vec, + ty: LogicalType, +) -> Box { + match kind { + WindowFunctionKind::RowNumber => Box::new(RowNumber), + WindowFunctionKind::Rank => Box::new(Rank { dense: false }), + WindowFunctionKind::DenseRank => Box::new(Rank { dense: true }), + WindowFunctionKind::Aggregate(kind) => { + let Some(arg) = args.into_iter().next() else { + unreachable!() + }; + Box::new(Aggregate { + kind, + ty, + arg, + accumulator: None, + }) + } + } +} diff --git a/src/execution/mod.rs b/src/execution/mod.rs index 6438caa0..97cadb6c 100644 --- a/src/execution/mod.rs +++ b/src/execution/mod.rs @@ -63,6 +63,7 @@ use crate::execution::dql::sort::Sort; use crate::execution::dql::top_k::TopK; use crate::execution::dql::union::Union; use crate::execution::dql::values::Values; +use crate::execution::dql::window::Window; use crate::expression::ScalarExpression; use crate::planner::operator::join::JoinCondition; use crate::planner::operator::{Operator, PhysicalOption, PlanImpl}; @@ -203,6 +204,7 @@ pub(crate) enum ExecNode<'a, T: Transaction + 'a> { Union(Union), Update(Update), Values(Values), + Window(Window), Empty, } @@ -343,6 +345,9 @@ impl<'a, T: Transaction + 'a> ExecNode<'a, T> { ExecNode::Values(exec) => { >::next_tuple(exec, arena, plan_arena) } + ExecNode::Window(exec) => { + >::next_tuple(exec, arena, plan_arena) + } ExecNode::Empty => unreachable!("executor node re-entered while active"), } } @@ -800,6 +805,13 @@ where cache, transaction, ), + Operator::Window(op) => >::into_executor( + (op, childrens.pop_only()), + arena, + plan_arena, + cache, + transaction, + ), Operator::ShowTable => as ReadExecutor<'a, T>>::into_executor( ShowTables { metas: None }, arena, diff --git a/src/expression/evaluator.rs b/src/expression/evaluator.rs index e2e04548..6bec4bf4 100644 --- a/src/expression/evaluator.rs +++ b/src/expression/evaluator.rs @@ -337,6 +337,9 @@ impl ScalarExpression { result.unwrap_or(DataValue::Null).cast(ty) } ScalarExpression::TableFunction(_) => unreachable!(), + ScalarExpression::WindowCall(_) => Err(DatabaseError::UnsupportedStmt( + "window calls must be evaluated by the window executor".to_string(), + )), } } } diff --git a/src/expression/mod.rs b/src/expression/mod.rs index 48340e2a..6a902335 100644 --- a/src/expression/mod.rs +++ b/src/expression/mod.rs @@ -44,6 +44,7 @@ pub mod range_detacher; pub mod simplify; pub mod visitor; pub mod visitor_mut; +pub mod window; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum TrimWhereField { @@ -157,6 +158,7 @@ pub enum ScalarExpression { else_expr: Option>, ty: LogicalType, }, + WindowCall(window::WindowCall), } impl From for ScalarExpression { @@ -498,7 +500,14 @@ impl ScalarExpression { } | ScalarExpression::CaseWhen { ty: return_type, .. - } => Cow::Borrowed(return_type), + } + | ScalarExpression::WindowCall(window::WindowCall { + function: + window::WindowFunction { + ty: return_type, .. + }, + .. + }) => Cow::Borrowed(return_type), ScalarExpression::IsNull { .. } | ScalarExpression::In { .. } | ScalarExpression::Between { .. } => Cow::Owned(LogicalType::Boolean), @@ -683,6 +692,31 @@ impl ScalarExpression { Ok(checker.has_agg) } + pub fn has_window_call(&self) -> Result { + struct WindowCallChecker(bool); + + impl<'a> ExprVisitor<'a> for WindowCallChecker { + fn visit(&mut self, expr: &'a ScalarExpression) -> Result<(), DatabaseError> { + if !self.0 { + walk_expr(self, expr)?; + } + Ok(()) + } + + fn visit_window( + &mut self, + _window: &'a window::WindowCall, + ) -> Result<(), DatabaseError> { + self.0 = true; + Ok(()) + } + } + + let mut checker = WindowCallChecker(false); + checker.visit(self)?; + Ok(checker.0) + } + fn output_name_by(&self, fn_display: &impl Fn(ColumnRef) -> N) -> String { match self { ScalarExpression::Constant(value) => format!("{value}"), @@ -743,6 +777,46 @@ impl ScalarExpression { args_str ) } + ScalarExpression::WindowCall(window) => { + let args = window + .function + .args + .iter() + .map(|expr| expr.output_name_by(fn_display)) + .join(", "); + let function = match window.function.kind { + window::WindowFunctionKind::RowNumber => "row_number".to_string(), + window::WindowFunctionKind::Rank => "rank".to_string(), + window::WindowFunctionKind::DenseRank => "dense_rank".to_string(), + window::WindowFunctionKind::Aggregate(kind) => { + format!("{kind:?}").to_lowercase() + } + }; + let mut spec = Vec::new(); + if !window.spec.partition_by.is_empty() { + spec.push(format!( + "partition by {}", + window + .spec + .partition_by + .iter() + .map(|expr| expr.output_name_by(fn_display)) + .join(", ") + )); + } + if !window.spec.order_by.is_empty() { + spec.push(format!( + "order by {}", + window + .spec + .order_by + .iter() + .map(ToString::to_string) + .join(", ") + )); + } + format!("{function}({args}) over ({})", spec.join(" ")) + } ScalarExpression::In { args, negated, diff --git a/src/expression/range_detacher.rs b/src/expression/range_detacher.rs index a5b67cac..8dc26b01 100644 --- a/src/expression/range_detacher.rs +++ b/src/expression/range_detacher.rs @@ -348,7 +348,8 @@ impl<'a, 'p> RangeDetacher<'a, 'p> { | ScalarExpression::IfNull { .. } | ScalarExpression::NullIf { .. } | ScalarExpression::Coalesce { .. } - | ScalarExpression::CaseWhen { .. } => None, + | ScalarExpression::CaseWhen { .. } + | ScalarExpression::WindowCall(_) => None, ScalarExpression::Tuple(_) | ScalarExpression::TableFunction(_) | ScalarExpression::Empty => unreachable!(), @@ -368,7 +369,8 @@ impl<'a, 'p> RangeDetacher<'a, 'p> { | ScalarExpression::IfNull { .. } | ScalarExpression::NullIf { .. } | ScalarExpression::Coalesce { .. } - | ScalarExpression::CaseWhen { .. } => None, + | ScalarExpression::CaseWhen { .. } + | ScalarExpression::WindowCall(_) => None, ScalarExpression::TableFunction(_) | ScalarExpression::Empty => unreachable!(), }) } diff --git a/src/expression/visitor.rs b/src/expression/visitor.rs index c8ca08d4..ecd85a3f 100644 --- a/src/expression/visitor.rs +++ b/src/expression/visitor.rs @@ -17,6 +17,7 @@ use crate::errors::DatabaseError; use crate::expression::agg::AggKind; use crate::expression::function::scala::ScalarFunction; use crate::expression::function::table::TableFunction; +use crate::expression::window::WindowCall; use crate::expression::TrimWhereField; use crate::expression::{AliasType, BinaryOperator, ScalarExpression, UnaryOperator}; use crate::types::evaluator::{BinaryEvaluatorRef, CastEvaluatorRef, UnaryEvaluatorRef}; @@ -99,6 +100,19 @@ pub trait ExprVisitor<'a>: Sized { Ok(()) } + fn visit_window(&mut self, window: &'a WindowCall) -> Result<(), DatabaseError> { + for expr in window + .function + .args + .iter() + .chain(&window.spec.partition_by) + .chain(window.spec.order_by.iter().map(|field| &field.expr)) + { + self.visit(expr)?; + } + Ok(()) + } + fn visit_in( &mut self, _negated: bool, @@ -356,5 +370,6 @@ pub fn walk_expr<'a, V: ExprVisitor<'a>>( else_expr.as_deref(), ty, ), + ScalarExpression::WindowCall(window) => visitor.visit_window(window), } } diff --git a/src/expression/visitor_mut.rs b/src/expression/visitor_mut.rs index dbddafd4..1ea298ce 100644 --- a/src/expression/visitor_mut.rs +++ b/src/expression/visitor_mut.rs @@ -17,6 +17,7 @@ use crate::errors::DatabaseError; use crate::expression::agg::AggKind; use crate::expression::function::scala::ScalarFunction; use crate::expression::function::table::TableFunction; +use crate::expression::window::WindowCall; use crate::expression::TrimWhereField; use crate::expression::{AliasType, BinaryOperator, ScalarExpression, UnaryOperator}; use crate::types::evaluator::{BinaryEvaluatorRef, CastEvaluatorRef, UnaryEvaluatorRef}; @@ -122,6 +123,19 @@ pub trait ExprVisitorMut<'a>: Sized { Ok(()) } + fn visit_window(&mut self, window: &'a mut WindowCall) -> Result<(), DatabaseError> { + for expr in window + .function + .args + .iter_mut() + .chain(&mut window.spec.partition_by) + .chain(window.spec.order_by.iter_mut().map(|field| &mut field.expr)) + { + self.visit(expr)?; + } + Ok(()) + } + fn visit_in( &mut self, _negated: bool, @@ -376,5 +390,6 @@ pub fn walk_mut_expr<'a, V: ExprVisitorMut<'a>>( else_expr, ty, } => visitor.visit_case_when(operand_expr, expr_pairs, else_expr, ty), + ScalarExpression::WindowCall(window) => visitor.visit_window(window), } } diff --git a/src/expression/window.rs b/src/expression/window.rs new file mode 100644 index 00000000..ddc70a91 --- /dev/null +++ b/src/expression/window.rs @@ -0,0 +1,46 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::expression::agg::AggKind; +use crate::expression::ScalarExpression; +use crate::planner::operator::sort::SortField; +use crate::types::LogicalType; +use kite_sql_serde_macros::ReferenceSerialization; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, ReferenceSerialization)] +pub enum WindowFunctionKind { + RowNumber, + Rank, + DenseRank, + Aggregate(AggKind), +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, ReferenceSerialization)] +pub struct WindowFunction { + pub kind: WindowFunctionKind, + pub args: Vec, + pub ty: LogicalType, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, ReferenceSerialization)] +pub struct WindowSpec { + pub partition_by: Vec, + pub order_by: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, ReferenceSerialization)] +pub struct WindowCall { + pub function: WindowFunction, + pub spec: WindowSpec, +} diff --git a/src/optimizer/heuristic/optimizer.rs b/src/optimizer/heuristic/optimizer.rs index 84b2e3d6..5231adea 100644 --- a/src/optimizer/heuristic/optimizer.rs +++ b/src/optimizer/heuristic/optimizer.rs @@ -557,6 +557,9 @@ impl ImplementationRuleIndex { Operator::Values(_) if self.contains(ImplementationRuleImpl::Values) => { Some(PhysicalOption::new(PlanImpl::Values, SortOption::None)) } + Operator::Window(_) if self.contains(ImplementationRuleImpl::Window) => { + Some(PhysicalOption::new(PlanImpl::Window, SortOption::Follow)) + } Operator::Analyze(_) if self.contains(ImplementationRuleImpl::Analyze) => { Some(PhysicalOption::new(PlanImpl::Analyze, SortOption::None)) } diff --git a/src/optimizer/rule/implementation/dql/mod.rs b/src/optimizer/rule/implementation/dql/mod.rs index 400de6cc..9258b248 100644 --- a/src/optimizer/rule/implementation/dql/mod.rs +++ b/src/optimizer/rule/implementation/dql/mod.rs @@ -26,3 +26,4 @@ pub(crate) mod sort; pub(crate) mod table_scan; pub(crate) mod top_k; pub(crate) mod values; +pub(crate) mod window; diff --git a/src/optimizer/rule/implementation/dql/window.rs b/src/optimizer/rule/implementation/dql/window.rs new file mode 100644 index 00000000..6d08b9d1 --- /dev/null +++ b/src/optimizer/rule/implementation/dql/window.rs @@ -0,0 +1,35 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::optimizer::core::pattern::{Pattern, PatternChildrenPredicate}; +use crate::optimizer::core::rule::{BestPhysicalOption, ImplementationRule, MatchPattern}; +use crate::optimizer::core::statistics_meta::StatisticMetaLoader; +use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::single_mapping; +use std::sync::LazyLock; + +static WINDOW_PATTERN: LazyLock = LazyLock::new(|| Pattern { + predicate: |op| matches!(op, Operator::Window(_)), + children: PatternChildrenPredicate::None, +}); + +#[derive(Clone)] +pub struct WindowImplementation; + +single_mapping!( + WindowImplementation, + WINDOW_PATTERN, + PhysicalOption::new(PlanImpl::Window, SortOption::Follow) +); diff --git a/src/optimizer/rule/implementation/mod.rs b/src/optimizer/rule/implementation/mod.rs index d358bd22..a93e6dca 100644 --- a/src/optimizer/rule/implementation/mod.rs +++ b/src/optimizer/rule/implementation/mod.rs @@ -53,6 +53,7 @@ use crate::optimizer::rule::implementation::dql::table_scan::{ }; use crate::optimizer::rule::implementation::dql::top_k::TopKImplementation; use crate::optimizer::rule::implementation::dql::values::ValuesImplementation; +use crate::optimizer::rule::implementation::dql::window::WindowImplementation; use crate::planner::operator::Operator; #[repr(usize)] @@ -86,10 +87,11 @@ pub enum ImplementationRuleRootTag { DropColumn, DropTable, Truncate, + Window, } impl ImplementationRuleRootTag { - pub const COUNT: usize = Self::Truncate as usize + 1; + pub const COUNT: usize = Self::Window as usize + 1; pub fn from_operator(operator: &Operator) -> Option { match operator { @@ -121,6 +123,7 @@ impl ImplementationRuleRootTag { Operator::DropColumn(_) => Some(Self::DropColumn), Operator::DropTable(_) => Some(Self::DropTable), Operator::Truncate(_) => Some(Self::Truncate), + Operator::Window(_) => Some(Self::Window), Operator::ShowTable | Operator::ShowView | Operator::Explain @@ -170,6 +173,7 @@ pub enum ImplementationRuleImpl { DropColumn, DropTable, Truncate, + Window, } impl MatchPattern for ImplementationRuleImpl { @@ -205,6 +209,7 @@ impl MatchPattern for ImplementationRuleImpl { ImplementationRuleImpl::DropTable => DropTableImplementation.pattern(), ImplementationRuleImpl::Truncate => TruncateImplementation.pattern(), ImplementationRuleImpl::Analyze => AnalyzeImplementation.pattern(), + ImplementationRuleImpl::Window => WindowImplementation.pattern(), } } } @@ -244,6 +249,7 @@ impl ImplementationRuleImpl { ImplementationRuleImpl::DropColumn => ImplementationRuleRootTag::DropColumn, ImplementationRuleImpl::DropTable => ImplementationRuleRootTag::DropTable, ImplementationRuleImpl::Truncate => ImplementationRuleRootTag::Truncate, + ImplementationRuleImpl::Window => ImplementationRuleRootTag::Window, } } } @@ -293,6 +299,7 @@ impl ImplementationRule for ImplementationRuleImpl { ImplementationRuleImpl::DropTable => update!(DropTableImplementation), ImplementationRuleImpl::Truncate => update!(TruncateImplementation), ImplementationRuleImpl::Analyze => update!(AnalyzeImplementation), + ImplementationRuleImpl::Window => update!(WindowImplementation), } Ok(()) @@ -450,6 +457,12 @@ mod tests { |op| matches!(op, Operator::Values(_)), PlanImpl::Values, )?; + assert_sql_rule( + "select row_number() over (order by c1) from t1", + ImplementationRuleImpl::Window, + |op| matches!(op, Operator::Window(_)), + PlanImpl::Window, + )?; assert_sql_rule( "select * from t1", ImplementationRuleImpl::SeqScan, diff --git a/src/optimizer/rule/normalization/column_pruning.rs b/src/optimizer/rule/normalization/column_pruning.rs index bdfacf63..41882c92 100644 --- a/src/optimizer/rule/normalization/column_pruning.rs +++ b/src/optimizer/rule/normalization/column_pruning.rs @@ -501,7 +501,8 @@ impl ColumnPruning { | Operator::Filter(_) | Operator::Union(_) | Operator::SetMembership(_) - | Operator::TopK(_) => { + | Operator::TopK(_) + | Operator::Window(_) => { if matches!(operator, Operator::ScalarApply(_) | Operator::MarkApply(_)) { let mut child_required = required_columns; Self::extend_operator_referenced_columns(operator, &mut child_required, arena)?; diff --git a/src/optimizer/rule/normalization/mod.rs b/src/optimizer/rule/normalization/mod.rs index 0df08ea5..f8b41f1f 100644 --- a/src/optimizer/rule/normalization/mod.rs +++ b/src/optimizer/rule/normalization/mod.rs @@ -133,7 +133,8 @@ impl NormalizationRuleRootTag { | Operator::FunctionScan(_) | Operator::Update(_) | Operator::Union(_) - | Operator::SetMembership(_) => None, + | Operator::SetMembership(_) + | Operator::Window(_) => None, #[cfg(feature = "copy")] Operator::CopyFromFile(_) | Operator::CopyToFile(_) => None, } diff --git a/src/orm/README.md b/src/orm/README.md index bfff6424..032b74c1 100644 --- a/src/orm/README.md +++ b/src/orm/README.md @@ -97,6 +97,33 @@ fields, or `order_by_expr` for computed sort expressions. Call `.asc()`, `.desc()`, `.nulls_first()`, or `.nulls_last()` when the default ascending/nulls-last order is not enough. +Window expressions use `WindowSpec` inside a projection closure: + +```rust,ignore +use kite_sql::orm::WindowSpec; + +let rows = database.bind(|ctx| { + ctx.from::()? + .project_tuple(|e| { + let category = e.column(EventLog::category())?; + let score = e.column(EventLog::score())?; + let spec = WindowSpec::new() + .partition_by(category.clone()) + .order_by(score.clone().desc()); + Ok(vec![ + category, + e.row_number(spec.clone())?, + e.sum_over(score, spec)?, + ]) + })? + .finish() +})?; +``` + +Window expressions use the explicit `row_number`, `rank`, `dense_rank`, +`count_over`, `count_all_over`, `sum_over`, `avg_over`, `min_over`, and +`max_over` methods. + Joins and set operations use the same binder-backed style: ```rust,ignore diff --git a/src/orm/mod.rs b/src/orm/mod.rs index e491d64f..f8cca8ee 100644 --- a/src/orm/mod.rs +++ b/src/orm/mod.rs @@ -101,6 +101,13 @@ pub struct FieldSort { nulls_first: bool, } +/// Partitioning and ordering for an ORM window expression. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct WindowSpec { + partition_by: Vec, + order_by: Vec, +} + #[derive(Debug, Clone, PartialEq, Eq)] struct QuerySource { table_name: String, @@ -325,6 +332,22 @@ pub trait IntoOrmScalarExpression { fn into_orm_scalar(self) -> ScalarExpression; } +impl WindowSpec { + pub fn new() -> Self { + Self::default() + } + + pub fn partition_by(mut self, expr: impl IntoOrmScalarExpression) -> Self { + self.partition_by.push(expr.into_orm_scalar()); + self + } + + pub fn order_by(mut self, field: SortField) -> Self { + self.order_by.push(field); + self + } +} + impl IntoOrmScalarExpression for E where E: Into, @@ -494,6 +517,23 @@ where .map(|expr| self.wrap(expr)) } + fn window( + self, + name: &'static str, + args: Vec, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.binder() + .bind_window_function( + name.to_string(), + args, + spec.partition_by, + spec.order_by, + self.arena(), + ) + .map(|expr| self.wrap(expr)) + } + fn scalar_subquery( self, build: F, @@ -1485,6 +1525,88 @@ where self.function(name, args) } + fn aggregate_window( + &self, + name: &'static str, + expr: impl IntoOrmScalarExpression, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.handle() + .window(name, vec![expr.into_orm_scalar()], spec) + } + + pub fn row_number( + &self, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.handle().window("row_number", Vec::new(), spec) + } + + pub fn rank( + &self, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.handle().window("rank", Vec::new(), spec) + } + + pub fn dense_rank( + &self, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.handle().window("dense_rank", Vec::new(), spec) + } + + pub fn count_over( + &self, + expr: impl IntoOrmScalarExpression, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.aggregate_window("count", expr, spec) + } + + pub fn sum_over( + &self, + expr: impl IntoOrmScalarExpression, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.aggregate_window("sum", expr, spec) + } + + pub fn avg_over( + &self, + expr: impl IntoOrmScalarExpression, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.aggregate_window("avg", expr, spec) + } + + pub fn min_over( + &self, + expr: impl IntoOrmScalarExpression, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.aggregate_window("min", expr, spec) + } + + pub fn max_over( + &self, + expr: impl IntoOrmScalarExpression, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.aggregate_window("max", expr, spec) + } + + pub fn count_all_over( + &self, + spec: WindowSpec, + ) -> Result, DatabaseError> { + self.handle().window( + "count", + vec![Binder::<'bind, 'parent, T, A>::wildcard_expr()], + spec, + ) + } + pub fn count_all(&self) -> Result, DatabaseError> { self.function( "count", diff --git a/src/planner/mod.rs b/src/planner/mod.rs index 559831fe..99e20338 100644 --- a/src/planner/mod.rs +++ b/src/planner/mod.rs @@ -188,6 +188,14 @@ impl LogicalPlan { Childrens::Only(child) => child.output_schema(arena).clone(), _ => unreachable!(), }, + Operator::Window(op) => match childrens { + Childrens::Only(child) => { + let mut schema = child.output_schema(arena).clone(); + schema.extend_from_slice(&op.output_columns); + schema + } + _ => unreachable!(), + }, Operator::ScalarApply(_) | Operator::Join(_) => match childrens { Childrens::Twins { left, right } => { let mut schema = left.output_schema(arena).clone(); diff --git a/src/planner/operator/mod.rs b/src/planner/operator/mod.rs index e3b03ac4..6a4a9f3f 100644 --- a/src/planner/operator/mod.rs +++ b/src/planner/operator/mod.rs @@ -46,6 +46,7 @@ pub mod update; pub mod values; pub mod visitor; pub mod visitor_mut; +pub mod window; use self::{ aggregate::AggregateOperator, alter_table::add_column::AddColumnOperator, @@ -134,6 +135,7 @@ pub enum Operator { CopyFromFile(CopyFromFileOperator), #[cfg(feature = "copy")] CopyToFile(CopyToFileOperator), + Window(window::WindowOperator), } #[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] @@ -200,6 +202,7 @@ pub enum PlanImpl { #[cfg(feature = "copy")] CopyToFile, Analyze, + Window, } impl Operator { @@ -314,6 +317,21 @@ impl Operator { Ok(()) } + fn visit_window( + &mut self, + op: &'operator window::WindowOperator, + ) -> Result<(), DatabaseError> { + for expr in op + .partition_by + .iter() + .chain(op.order_by.iter().map(|field| &field.expr)) + .chain(op.functions.iter().flat_map(|function| &function.args)) + { + ExprVisitor::visit(self, expr)?; + } + Ok(()) + } + fn visit_top_k(&mut self, op: &'operator TopKOperator) -> Result<(), DatabaseError> { for field in &op.sort_fields { ExprVisitor::visit(self, &field.expr)?; @@ -469,6 +487,7 @@ impl fmt::Display for Operator { Operator::CopyToFile(op) => write!(f, "{op}"), Operator::Union(op) => write!(f, "{op}"), Operator::SetMembership(op) => write!(f, "{op}"), + Operator::Window(op) => write!(f, "{op}"), } } } @@ -538,6 +557,7 @@ impl fmt::Display for PlanImpl { #[cfg(feature = "copy")] PlanImpl::CopyToFile => write!(f, "CopyToFile"), PlanImpl::Analyze => write!(f, "Analyze"), + PlanImpl::Window => write!(f, "Window"), } } } diff --git a/src/planner/operator/visitor.rs b/src/planner/operator/visitor.rs index 06911e8c..a25107ce 100644 --- a/src/planner/operator/visitor.rs +++ b/src/planner/operator/visitor.rs @@ -69,6 +69,10 @@ pub trait OperatorVisitor<'a>: Sized { Ok(()) } + fn visit_window(&mut self, _op: &'a window::WindowOperator) -> Result<(), DatabaseError> { + Ok(()) + } + fn visit_limit(&mut self, _op: &'a LimitOperator) -> Result<(), DatabaseError> { Ok(()) } @@ -255,6 +259,18 @@ impl<'a, V: ExprVisitor<'a>> OperatorVisitor<'a> for OperatorExprVisitor<'_, V> Ok(()) } + fn visit_window(&mut self, op: &'a window::WindowOperator) -> Result<(), DatabaseError> { + for expr in op + .partition_by + .iter() + .chain(op.order_by.iter().map(|field| &field.expr)) + .chain(op.functions.iter().flat_map(|function| &function.args)) + { + ExprVisitor::visit(self.visitor, expr)?; + } + Ok(()) + } + fn visit_top_k(&mut self, op: &'a TopKOperator) -> Result<(), DatabaseError> { for field in &op.sort_fields { ExprVisitor::visit(self.visitor, &field.expr)?; @@ -336,6 +352,7 @@ pub fn walk_operator<'a, V: OperatorVisitor<'a>>( Operator::CopyFromFile(op) => visitor.visit_copy_from_file(op), #[cfg(feature = "copy")] Operator::CopyToFile(op) => visitor.visit_copy_to_file(op), + Operator::Window(op) => visitor.visit_window(op), } } @@ -350,6 +367,7 @@ pub(crate) mod tests { ArcTableFunctionImpl, TableFunction, TableFunctionCatalog, }; use crate::expression::visitor::{walk_expr, ExprVisitor}; + use crate::expression::window::{WindowFunction, WindowFunctionKind}; use crate::expression::ScalarExpression; use crate::function::numbers::Numbers; use crate::planner::operator::alter_table::change_column::NotNullChange; @@ -443,6 +461,16 @@ pub(crate) mod tests { rows: vec![vec![DataValue::Int32(1)]], schema_ref: vec![column_ref], }), + Operator::Window(window::WindowOperator { + partition_by: vec![17_i32.into()], + order_by: vec![SortField::from(ScalarExpression::from(18_i32))], + functions: vec![WindowFunction { + kind: WindowFunctionKind::RowNumber, + args: Vec::new(), + ty: LogicalType::Bigint, + }], + output_columns: vec![column_ref], + }), Operator::ShowTable, Operator::ShowView, Operator::Explain, @@ -579,7 +607,7 @@ pub(crate) mod tests { for operator in &operators { visitor.visit_operator(operator)?; } - assert_eq!(counter.0, 18); + assert_eq!(counter.0, 20); Ok(()) } diff --git a/src/planner/operator/visitor_mut.rs b/src/planner/operator/visitor_mut.rs index c5b740cb..388a88f4 100644 --- a/src/planner/operator/visitor_mut.rs +++ b/src/planner/operator/visitor_mut.rs @@ -75,6 +75,10 @@ pub trait OperatorVisitorMut<'a>: Sized { Ok(()) } + fn visit_window(&mut self, _op: &'a mut window::WindowOperator) -> Result<(), DatabaseError> { + Ok(()) + } + fn visit_limit(&mut self, _op: &'a mut LimitOperator) -> Result<(), DatabaseError> { Ok(()) } @@ -276,6 +280,22 @@ impl<'a, V: ExprVisitorMut<'a>> OperatorVisitorMut<'a> for OperatorExprVisitorMu Ok(()) } + fn visit_window(&mut self, op: &'a mut window::WindowOperator) -> Result<(), DatabaseError> { + for expr in op + .partition_by + .iter_mut() + .chain(op.order_by.iter_mut().map(|field| &mut field.expr)) + .chain( + op.functions + .iter_mut() + .flat_map(|function| &mut function.args), + ) + { + ExprVisitorMut::visit(self.visitor, expr)?; + } + Ok(()) + } + fn visit_top_k(&mut self, op: &'a mut TopKOperator) -> Result<(), DatabaseError> { for field in &mut op.sort_fields { ExprVisitorMut::visit(self.visitor, &mut field.expr)?; @@ -360,6 +380,7 @@ pub fn walk_mut_operator<'a, V: OperatorVisitorMut<'a>>( Operator::CopyFromFile(op) => visitor.visit_copy_from_file(op), #[cfg(feature = "copy")] Operator::CopyToFile(op) => visitor.visit_copy_to_file(op), + Operator::Window(op) => visitor.visit_window(op), } } @@ -398,7 +419,7 @@ mod tests { visitor.visit_operator(operator)?; } } - assert_eq!(counter.0, 18); + assert_eq!(counter.0, 20); Ok(()) } diff --git a/src/planner/operator/window.rs b/src/planner/operator/window.rs new file mode 100644 index 00000000..03197501 --- /dev/null +++ b/src/planner/operator/window.rs @@ -0,0 +1,60 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::catalog::ColumnRef; +use crate::expression::window::WindowFunction; +use crate::expression::ScalarExpression; +use crate::iter_ext::Itertools; +use crate::planner::operator::sort::SortField; +use kite_sql_serde_macros::ReferenceSerialization; +use std::fmt; + +#[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] +pub struct WindowOperator { + pub partition_by: Vec, + pub order_by: Vec, + pub functions: Vec, + pub output_columns: Vec, +} + +impl fmt::Display for WindowOperator { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "Window [{}]", + self.functions + .iter() + .map(|expr| format!("{expr:?}")) + .join(", ") + )?; + if !self.partition_by.is_empty() || !self.order_by.is_empty() { + write!(f, " ->")?; + } + if !self.partition_by.is_empty() { + write!( + f, + " Partition By [{}]", + self.partition_by.iter().map(ToString::to_string).join(", ") + )?; + } + if !self.order_by.is_empty() { + write!( + f, + " Order By [{}]", + self.order_by.iter().map(ToString::to_string).join(", ") + )?; + } + Ok(()) + } +} diff --git a/tests/macros-test/src/main.rs b/tests/macros-test/src/main.rs index b91cba13..b686f17f 100644 --- a/tests/macros-test/src/main.rs +++ b/tests/macros-test/src/main.rs @@ -24,7 +24,7 @@ mod test { use kite_sql::expression::function::FunctionSummary; use kite_sql::expression::BinaryOperator; use kite_sql::expression::ScalarExpression; - use kite_sql::orm::OrmQueryResultExt; + use kite_sql::orm::{OrmQueryResultExt, WindowSpec}; use kite_sql::planner::{MetaArena, PlanArena, TableArena, TableArenaCell}; use kite_sql::storage::rocksdb::RocksStorage; use kite_sql::types::evaluator::binary_create; @@ -2095,6 +2095,36 @@ mod test { score: 5, })?; + let mut windowed_scores = database + .bind(|ctx| { + ctx.from::()? + .project_tuple(|e| { + let id = e.column(EventLog::id())?; + let category = e.column(EventLog::category())?; + let score = e.column(EventLog::score())?; + let ordered = WindowSpec::new() + .partition_by(category.clone()) + .order_by(score.clone().asc()); + let partition = WindowSpec::new().partition_by(category); + Ok(vec![ + id, + e.row_number(ordered.clone())?, + e.rank(ordered.clone())?, + e.dense_rank(ordered)?, + e.sum_over(score, partition.clone())?, + e.count_all_over(partition)?, + ]) + })? + .finish() + })? + .project_tuple::<(i32, i64, i64, i64, i32, i32)>() + .collect::, _>>()?; + windowed_scores.sort_by_key(|row| row.0); + assert_eq!( + windowed_scores, + vec![(1, 1, 1, 1, 30, 2), (2, 2, 2, 2, 30, 2), (3, 1, 1, 1, 5, 1)] + ); + let mut grouped_categories = database .bind(|ctx| { ctx.from::()? diff --git a/tests/slt/window.slt b/tests/slt/window.slt new file mode 100644 index 00000000..6d56904b --- /dev/null +++ b/tests/slt/window.slt @@ -0,0 +1,114 @@ +statement ok +create table window_test(id int primary key, k int, v int null) + +statement ok +insert into window_test values + (1, 1, 10), + (2, 1, 10), + (3, 1, 20), + (4, 2, 5), + (5, 2, null) + +query IIIIII +select id, + row_number() over (partition by k order by v, id), + rank() over (partition by k order by v), + dense_rank() over (partition by k order by v), + sum(v) over (partition by k order by v), + sum(v) over (partition by k) +from window_test +order by id +---- +1 1 1 1 20 40 +2 2 1 1 20 40 +3 3 3 2 40 40 +4 1 1 1 5 5 +5 2 2 2 5 5 + +query III +select id, + row_number() over (partition by k order by v, id), + row_number() over (order by v desc, id) +from window_test +order by id +---- +1 1 2 +2 2 3 +3 3 1 +4 1 4 +5 2 5 + +query T +explain select + row_number() over (partition by k order by v, id), + rank() over (partition by k order by v, id), + row_number() over (order by v desc, id) +from window_test +---- +Projection [#4, #5, #6] [Project => (Sort Option: Follow)] Window [WindowFunction { kind: RowNumber, args: [], ty: Bigint }] -> Order By [#3 Desc Nulls Last, #1 Asc Nulls Last] [Window => (Sort Option: Follow)] Sort By #3 Desc Nulls Last, #1 Asc Nulls Last [Sort => (Sort Option: OrderBy: (#3 Desc Nulls Last, #1 Asc Nulls Last) ignore_prefix_len: 0)] Window [WindowFunction { kind: RowNumber, args: [], ty: Bigint }, WindowFunction { kind: Rank, args: [], ty: Bigint }] -> Partition By [#2] Order By [#3 Asc Nulls Last, #1 Asc Nulls Last] [Window => (Sort Option: Follow)] Sort By #2 Asc Nulls Last, #3 Asc Nulls Last, #1 Asc Nulls Last [Sort => (Sort Option: OrderBy: (#2 Asc Nulls Last, #3 Asc Nulls Last, #1 Asc Nulls Last) ignore_prefix_len: 0)] TableScan window_test -> [#1, #2, #3] [SeqScan => (Sort Option: None)] + +query III +select id, + row_number() over (partition by k order by v, id), + row_number() over (partition by k order by v, id) +from window_test +order by id +---- +1 1 1 +2 2 2 +3 3 3 +4 1 1 +5 2 2 + +query IIII +select k, + sum(v), + sum(sum(v)) over (order by k), + row_number() over (order by k) +from window_test +group by k +order by k +---- +1 40 40 1 +2 5 45 2 + +query III +select id, + count(*) over (partition by k), + count(v) over (partition by k) +from window_test +order by id +---- +1 3 3 +2 3 3 +3 3 3 +4 2 1 +5 2 1 + +statement ok +create view window_view as +select id, k, row_number() over (partition by k order by v, id) as rn +from window_test + +query III +select id, k, rn from window_view order by id +---- +1 1 1 +2 1 2 +3 1 3 +4 2 1 +5 2 2 + +statement ok +drop view window_view + +statement error +select sum(v) over (order by id rows between 1 preceding and current row) +from window_test + +statement error +select row_number() over named_window from window_test +window named_window as (order by id) + +statement error +select id from window_test qualify id = 1 From bdd3de799272b85714de8a1ee89153cb4d564a8f Mon Sep 17 00:00:00 2001 From: kould Date: Sun, 12 Jul 2026 23:28:18 +0800 Subject: [PATCH 2/6] docs: clarify pull request expectations --- AGENT.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/AGENT.md b/AGENT.md index 224c5840..7e2cd95c 100644 --- a/AGENT.md +++ b/AGENT.md @@ -151,6 +151,8 @@ KiteSQL aims to stay easy to build, easy to audit, and easy to understand. A valid PR should: +- Follow the repository pull request template +- Describe test coverage and verified behavior instead of listing commands that were run - Compile cleanly - Pass all tests via make - Include new tests if behavior changes From 4798c61ee8934da9ffc1867eb43e982862e4d9ac Mon Sep 17 00:00:00 2001 From: kould Date: Sun, 12 Jul 2026 23:53:25 +0800 Subject: [PATCH 3/6] refactor: simplify window output binding --- src/binder/window.rs | 40 ++++++++++++++++--------------------- src/execution/dql/window.rs | 2 ++ 2 files changed, 19 insertions(+), 23 deletions(-) diff --git a/src/binder/window.rs b/src/binder/window.rs index 9a48531c..1ba3a442 100644 --- a/src/binder/window.rs +++ b/src/binder/window.rs @@ -61,22 +61,24 @@ impl ExprVisitorMut<'_> for WindowCollector<'_, '_> { } } -struct WindowOutputBinder { - outputs: Vec<(ColumnRef, usize)>, +struct WindowOutputBinder<'a> { + groups: &'a [WindowGroup], + base_position: usize, } -impl ExprVisitorMut<'_> for WindowOutputBinder { +impl ExprVisitorMut<'_> for WindowOutputBinder<'_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, position: &mut usize, ) -> Result<(), DatabaseError> { - if let Some((_, output_position)) = self - .outputs + if let Some(output_position) = self + .groups .iter() - .find(|(output_column, _)| output_column == column) + .flat_map(|group| &group.output_columns) + .position(|output_column| output_column == column) { - *position = *output_position; + *position = self.base_position + output_position; } Ok(()) } @@ -203,24 +205,16 @@ impl> Binder<'_, '_, T, A> } } - let mut outputs = Vec::new(); - let mut position = base_position; - for group in &groups { - for column in &group.output_columns { - outputs.push((*column, position)); - position += 1; - } - } - - let mut output_binder = WindowOutputBinder { outputs }; - for expr in select_list { + let mut output_binder = WindowOutputBinder { + groups: &groups, + base_position, + }; + for expr in select_list + .iter_mut() + .chain(order_by.iter_mut().flatten().map(|field| &mut field.expr)) + { output_binder.visit(expr)?; } - if let Some(order_by) = order_by { - for field in order_by { - output_binder.visit(&mut field.expr)?; - } - } for group in groups { let sort_fields = group diff --git a/src/execution/dql/window.rs b/src/execution/dql/window.rs index 1884ac9c..aee6a88e 100644 --- a/src/execution/dql/window.rs +++ b/src/execution/dql/window.rs @@ -92,6 +92,7 @@ fn evaluate_partition( while peer_start < rows.len() { let mut peer_end = peer_start + 1; 'peer: while peer_end < rows.len() { + // TODO: Cache evaluated order keys to avoid recalculating the previous row. for field in order_by { if field.expr.eval(Some(&rows[peer_end - 1]))? != field.expr.eval(Some(&rows[peer_end]))? @@ -135,6 +136,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window { while arena.next_tuple(self.input, plan_arena)? { let tuple = mem::take(arena.result_tuple_mut()); let mut same_partition = true; + // TODO: Cache evaluated partition keys to avoid recalculating the previous row. for expr in &self.partition_by { if expr.eval(self.rows.last())? != expr.eval(Some(&tuple))? { same_partition = false; From 70cb927e2b06f4bf27ec302b833e062f0409e900 Mon Sep 17 00:00:00 2001 From: kould Date: Mon, 13 Jul 2026 00:12:29 +0800 Subject: [PATCH 4/6] test: expand window function coverage --- src/execution/dql/window.rs | 200 ++++++++++++++++++++++++++++++++++ tests/macros-test/src/main.rs | 35 +++++- tests/slt/window.slt | 27 +++++ 3 files changed, 258 insertions(+), 4 deletions(-) diff --git a/src/execution/dql/window.rs b/src/execution/dql/window.rs index aee6a88e..f6640eb3 100644 --- a/src/execution/dql/window.rs +++ b/src/execution/dql/window.rs @@ -156,3 +156,203 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window { } } } + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef}; + use crate::execution::{empty_context, execute_input, try_collect}; + use crate::expression::agg::AggKind; + use crate::expression::window::{ + WindowFunction as WindowExpressionFunction, WindowFunctionKind, + }; + use crate::planner::operator::values::ValuesOperator; + use crate::planner::operator::Operator; + use crate::planner::Childrens; + use crate::storage::memory::MemoryStorage; + use crate::storage::Storage; + use crate::types::LogicalType; + + fn column(position: usize) -> ScalarExpression { + ScalarExpression::column_expr(ColumnRef::new(position + 1), position) + } + + fn rows(values: &[i32]) -> Vec { + values + .iter() + .map(|value| Tuple::new(None, vec![DataValue::Int32(*value)])) + .collect() + } + + fn functions() -> Vec> { + vec![ + function::new( + WindowFunctionKind::RowNumber, + Vec::new(), + LogicalType::Bigint, + ), + function::new(WindowFunctionKind::Rank, Vec::new(), LogicalType::Bigint), + function::new( + WindowFunctionKind::DenseRank, + Vec::new(), + LogicalType::Bigint, + ), + function::new( + WindowFunctionKind::Aggregate(AggKind::Sum), + vec![column(0)], + LogicalType::Integer, + ), + ] + } + + #[test] + fn evaluate_peer_groups() -> Result<(), DatabaseError> { + let mut rows = rows(&[10, 10, 20]); + evaluate_partition(&mut rows, &[column(0).asc()], &mut functions())?; + + assert_eq!( + rows.into_iter().map(|row| row.values).collect::>(), + vec![ + vec![ + 10.into(), + 1_i64.into(), + 1_i64.into(), + 1_i64.into(), + 20.into() + ], + vec![ + 10.into(), + 2_i64.into(), + 1_i64.into(), + 1_i64.into(), + 20.into() + ], + vec![ + 20.into(), + 3_i64.into(), + 3_i64.into(), + 2_i64.into(), + 40.into() + ], + ] + ); + Ok(()) + } + + #[test] + fn evaluate_without_order_by() -> Result<(), DatabaseError> { + let mut rows = rows(&[3, 7]); + evaluate_partition(&mut rows, &[], &mut functions())?; + + assert_eq!( + rows.into_iter().map(|row| row.values).collect::>(), + vec![ + vec![ + 3.into(), + 1_i64.into(), + 1_i64.into(), + 1_i64.into(), + 10.into() + ], + vec![ + 7.into(), + 2_i64.into(), + 1_i64.into(), + 1_i64.into(), + 10.into() + ], + ] + ); + Ok(()) + } + + #[test] + fn evaluate_empty_partition() -> Result<(), DatabaseError> { + let mut rows = Vec::new(); + evaluate_partition(&mut rows, &[], &mut functions())?; + assert!(rows.is_empty()); + Ok(()) + } + + #[test] + fn execute_partitions() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut plan_arena = crate::planner::PlanArena::new(&table_arena); + let input_desc = ColumnDesc::new(LogicalType::Integer, None, false, None)?; + let input_columns = ["partition", "value"] + .map(|name| { + plan_arena.alloc_column(ColumnCatalog::new( + name.to_string(), + true, + input_desc.clone(), + )) + }) + .to_vec(); + let output_desc = ColumnDesc::new(LogicalType::Bigint, None, false, None)?; + let output_columns = ["row_number", "rank"] + .map(|name| { + plan_arena.alloc_column(ColumnCatalog::new( + name.to_string(), + true, + output_desc.clone(), + )) + }) + .to_vec(); + let input = LogicalPlan::new( + Operator::Values(ValuesOperator { + rows: vec![ + vec![1.into(), 10.into()], + vec![1.into(), 10.into()], + vec![1.into(), 20.into()], + vec![2.into(), 5.into()], + vec![2.into(), 7.into()], + ], + schema_ref: input_columns.clone(), + }), + Childrens::None, + ); + let operator = WindowOperator { + partition_by: vec![ScalarExpression::column_expr(input_columns[0], 0)], + order_by: vec![ScalarExpression::column_expr(input_columns[1], 1).asc()], + functions: vec![ + WindowExpressionFunction { + kind: WindowFunctionKind::RowNumber, + args: Vec::new(), + ty: LogicalType::Bigint, + }, + WindowExpressionFunction { + kind: WindowFunctionKind::Rank, + args: Vec::new(), + ty: LogicalType::Bigint, + }, + ], + output_columns, + }; + let table_cache = crate::storage::TableCache::default(); + let view_cache = crate::storage::ViewCache::default(); + let meta_cache = crate::storage::StatisticsMetaCache::default(); + let storage = MemoryStorage::new(); + let transaction = storage.transaction()?; + + let tuples = try_collect(execute_input::<_, Window>( + (operator, input), + empty_context(&table_cache, &view_cache, &meta_cache), + plan_arena, + &transaction, + ))?; + + assert_eq!( + tuples.into_iter().map(|row| row.values).collect::>(), + vec![ + vec![1.into(), 10.into(), 1_i64.into(), 1_i64.into()], + vec![1.into(), 10.into(), 2_i64.into(), 1_i64.into()], + vec![1.into(), 20.into(), 3_i64.into(), 3_i64.into()], + vec![2.into(), 5.into(), 1_i64.into(), 1_i64.into()], + vec![2.into(), 7.into(), 2_i64.into(), 2_i64.into()], + ] + ); + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/tests/macros-test/src/main.rs b/tests/macros-test/src/main.rs index b686f17f..1f0a05e4 100644 --- a/tests/macros-test/src/main.rs +++ b/tests/macros-test/src/main.rs @@ -2111,20 +2111,47 @@ mod test { e.row_number(ordered.clone())?, e.rank(ordered.clone())?, e.dense_rank(ordered)?, - e.sum_over(score, partition.clone())?, - e.count_all_over(partition)?, + e.sum_over(score.clone(), partition.clone())?, + e.count_all_over(partition.clone())?, + e.count_over(score.clone(), partition.clone())?, + e.avg_over(score, partition)?, ]) })? .finish() })? - .project_tuple::<(i32, i64, i64, i64, i32, i32)>() + .project_tuple::<(i32, i64, i64, i64, i32, i32, i32, f64)>() .collect::, _>>()?; windowed_scores.sort_by_key(|row| row.0); assert_eq!( windowed_scores, - vec![(1, 1, 1, 1, 30, 2), (2, 2, 2, 2, 30, 2), (3, 1, 1, 1, 5, 1)] + vec![ + (1, 1, 1, 1, 30, 2, 2, 15.0), + (2, 2, 2, 2, 30, 2, 2, 15.0), + (3, 1, 1, 1, 5, 1, 1, 5.0), + ] ); + let mut windowed_min_max = database + .bind(|ctx| { + ctx.from::()? + .project_tuple(|e| { + let id = e.column(EventLog::id())?; + let category = e.column(EventLog::category())?; + let score = e.column(EventLog::score())?; + let partition = WindowSpec::new().partition_by(category); + Ok(vec![ + id, + e.min_over(score.clone(), partition.clone())?, + e.max_over(score, partition)?, + ]) + })? + .finish() + })? + .project_tuple::<(i32, i32, i32)>() + .collect::, _>>()?; + windowed_min_max.sort_by_key(|row| row.0); + assert_eq!(windowed_min_max, vec![(1, 10, 20), (2, 10, 20), (3, 5, 5)]); + let mut grouped_categories = database .bind(|ctx| { ctx.from::()? diff --git a/tests/slt/window.slt b/tests/slt/window.slt index 6d56904b..fee542b0 100644 --- a/tests/slt/window.slt +++ b/tests/slt/window.slt @@ -85,6 +85,11 @@ order by id 4 2 1 5 2 1 +query I +select row_number() over () from window_test where id = 1 +---- +1 + statement ok create view window_view as select id, k, row_number() over (partition by k order by v, id) as rn @@ -112,3 +117,25 @@ window named_window as (order by id) statement error select id from window_test qualify id = 1 + +statement error +select count(distinct v) over () from window_test + +statement error +select sum(v) filter (where v > 0) over () from window_test + +statement error +select row_number(id) over () from window_test + +statement error +select lower('x') over () from window_test + +statement error +select sum(row_number() over ()) over () from window_test + +statement error +select id from window_test where row_number() over () = 1 + +statement error +select row_number() over (named_window order by id) from window_test +window named_window as (partition by k) From 04ac467cb266944a7c0e92d5fe22872093705c4b Mon Sep 17 00:00:00 2001 From: kould Date: Mon, 13 Jul 2026 00:36:10 +0800 Subject: [PATCH 5/6] refactor: type aggregate and window binding --- src/binder/expr.rs | 76 ++++++++++++++--------------------- src/binder/parser.rs | 24 ++++++++--- src/binder/window.rs | 29 +++++-------- src/expression/agg.rs | 21 ++++++++++ src/expression/mod.rs | 9 +---- src/expression/window.rs | 20 +++++++++ src/orm/mod.rs | 70 +++++++++++++++++++------------- tests/macros-test/src/main.rs | 24 +++++------ 8 files changed, 156 insertions(+), 117 deletions(-) diff --git a/src/binder/expr.rs b/src/binder/expr.rs index 1b215240..9a9b6cfb 100644 --- a/src/binder/expr.rs +++ b/src/binder/expr.rs @@ -476,76 +476,60 @@ impl<'a, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'a, '_, T }) } - pub(crate) fn bind_function_call( + pub(crate) fn bind_aggregate_function( &mut self, - function_name: String, - mut args: Vec, + kind: AggKind, + args: Vec, is_distinct: bool, arena: &mut PlanArena, ) -> Result { - match function_name.as_str() { - "count" => { + let ty = match kind { + AggKind::Count => { if args.len() != 1 { return Err(DatabaseError::MisMatch("number of count() parameters", "1")); } - return Ok(ScalarExpression::AggCall { - distinct: is_distinct, - kind: AggKind::Count, - args, - ty: LogicalType::Integer, - }); + LogicalType::Integer } - "sum" => { + AggKind::Sum => { if args.len() != 1 { return Err(DatabaseError::MisMatch("number of sum() parameters", "1")); } - let ty = args[0].return_type(arena).into_owned(); - - return Ok(ScalarExpression::AggCall { - distinct: is_distinct, - kind: AggKind::Sum, - args, - ty, - }); + args[0].return_type(arena).into_owned() } - "min" => { + AggKind::Min => { if args.len() != 1 { return Err(DatabaseError::MisMatch("number of min() parameters", "1")); } - let ty = args[0].return_type(arena).into_owned(); - - return Ok(ScalarExpression::AggCall { - distinct: is_distinct, - kind: AggKind::Min, - args, - ty, - }); + args[0].return_type(arena).into_owned() } - "max" => { + AggKind::Max => { if args.len() != 1 { return Err(DatabaseError::MisMatch("number of max() parameters", "1")); } - let ty = args[0].return_type(arena).into_owned(); - - return Ok(ScalarExpression::AggCall { - distinct: is_distinct, - kind: AggKind::Max, - args, - ty, - }); + args[0].return_type(arena).into_owned() } - "avg" => { + AggKind::Avg => { if args.len() != 1 { return Err(DatabaseError::MisMatch("number of avg() parameters", "1")); } - - return Ok(ScalarExpression::AggCall { - distinct: is_distinct, - kind: AggKind::Avg, - args, - ty: LogicalType::Double, - }); + LogicalType::Double } + }; + Ok(ScalarExpression::AggCall { + distinct: is_distinct, + kind, + args, + ty, + }) + } + + pub(crate) fn bind_function_call( + &mut self, + function_name: String, + mut args: Vec, + arena: &mut PlanArena, + ) -> Result { + match function_name.as_str() { "if" => { if args.len() != 3 { return Err(DatabaseError::MisMatch("number of if() parameters", "3")); diff --git a/src/binder/parser.rs b/src/binder/parser.rs index 41cbe304..390a56fe 100644 --- a/src/binder/parser.rs +++ b/src/binder/parser.rs @@ -23,8 +23,10 @@ use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef, TableName}; use crate::db::{BindSource, DBTransaction, Database, DatabaseIter, TransactionIter}; use crate::errors::{DatabaseError, SqlErrorSpan}; use crate::expression; +use crate::expression::agg::AggKind; use crate::expression::simplify::ConstantCalculator; use crate::expression::visitor_mut::ExprVisitorMut; +use crate::expression::window::WindowFunctionKind; use crate::expression::{AliasType, ScalarExpression}; use crate::iter_ext::Itertools; use crate::parser::parse_sql; @@ -2440,18 +2442,30 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder< .to_string(), )); } + let Some(kind) = WindowFunctionKind::from_name(&function_name) else { + return Err(attach_span_if_absent( + DatabaseError::UnsupportedStmt(format!( + "window function `{function_name}` is not supported" + )), + func_span, + )); + }; return self - .bind_window_call(function_name, args, is_distinct, over, arena) + .bind_window_call(kind, args, is_distinct, over, arena) .map_err(|err| attach_span_if_absent(err, func_span)); } - self.bind_function_call(function_name, args, is_distinct, arena) - .map_err(|err| attach_span_if_absent(err, func_span)) + let result = if let Some(kind) = AggKind::from_name(&function_name) { + self.bind_aggregate_function(kind, args, is_distinct, arena) + } else { + self.bind_function_call(function_name, args, arena) + }; + result.map_err(|err| attach_span_if_absent(err, func_span)) } fn bind_window_call( &mut self, - function_name: String, + kind: WindowFunctionKind, args: Vec, is_distinct: bool, over: &WindowType, @@ -2502,7 +2516,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder< )) }) .collect::, DatabaseError>>()?; - self.bind_window_function(function_name, args, partition_by, order_by, arena) + self.bind_window_function(kind, args, partition_by, order_by, arena) } pub fn bind_set_expr( diff --git a/src/binder/window.rs b/src/binder/window.rs index 1ba3a442..60dac596 100644 --- a/src/binder/window.rs +++ b/src/binder/window.rs @@ -94,7 +94,7 @@ struct WindowGroup { impl> Binder<'_, '_, T, A> { pub(crate) fn bind_window_function( &mut self, - function_name: String, + kind: WindowFunctionKind, args: Vec, partition_by: Vec, order_by: Vec, @@ -120,34 +120,25 @@ impl> Binder<'_, '_, T, A> } } - let (kind, args, ty) = match function_name.as_str() { - "row_number" | "rank" | "dense_rank" => { + let (args, ty) = match kind { + WindowFunctionKind::RowNumber + | WindowFunctionKind::Rank + | WindowFunctionKind::DenseRank => { if !args.is_empty() { return Err(DatabaseError::MisMatch( "number of ranking function parameters", "0", )); } - let kind = match function_name.as_str() { - "row_number" => WindowFunctionKind::RowNumber, - "rank" => WindowFunctionKind::Rank, - "dense_rank" => WindowFunctionKind::DenseRank, - _ => unreachable!(), - }; - (kind, args, LogicalType::Bigint) + (args, LogicalType::Bigint) } - "count" | "sum" | "avg" | "min" | "max" => { - let ScalarExpression::AggCall { kind, args, ty, .. } = - self.bind_function_call(function_name, args, false, arena)? + WindowFunctionKind::Aggregate(agg_kind) => { + let ScalarExpression::AggCall { args, ty, .. } = + self.bind_aggregate_function(agg_kind, args, false, arena)? else { unreachable!() }; - (WindowFunctionKind::Aggregate(kind), args, ty) - } - _ => { - return Err(DatabaseError::UnsupportedStmt(format!( - "window function `{function_name}` is not supported" - ))) + (args, ty) } }; diff --git a/src/expression/agg.rs b/src/expression/agg.rs index fb45df4e..1761d03b 100644 --- a/src/expression/agg.rs +++ b/src/expression/agg.rs @@ -24,6 +24,27 @@ pub enum AggKind { } impl AggKind { + pub(crate) fn from_name(name: &str) -> Option { + match name { + "avg" => Some(Self::Avg), + "max" => Some(Self::Max), + "min" => Some(Self::Min), + "sum" => Some(Self::Sum), + "count" => Some(Self::Count), + _ => None, + } + } + + pub(crate) fn name(self) -> &'static str { + match self { + Self::Avg => "avg", + Self::Max => "max", + Self::Min => "min", + Self::Sum => "sum", + Self::Count => "count", + } + } + pub fn allow_distinct(&self) -> bool { match self { AggKind::Avg => false, diff --git a/src/expression/mod.rs b/src/expression/mod.rs index 6a902335..16fe1ad2 100644 --- a/src/expression/mod.rs +++ b/src/expression/mod.rs @@ -784,14 +784,7 @@ impl ScalarExpression { .iter() .map(|expr| expr.output_name_by(fn_display)) .join(", "); - let function = match window.function.kind { - window::WindowFunctionKind::RowNumber => "row_number".to_string(), - window::WindowFunctionKind::Rank => "rank".to_string(), - window::WindowFunctionKind::DenseRank => "dense_rank".to_string(), - window::WindowFunctionKind::Aggregate(kind) => { - format!("{kind:?}").to_lowercase() - } - }; + let function = window.function.kind.name(); let mut spec = Vec::new(); if !window.spec.partition_by.is_empty() { spec.push(format!( diff --git a/src/expression/window.rs b/src/expression/window.rs index ddc70a91..8c8eed5e 100644 --- a/src/expression/window.rs +++ b/src/expression/window.rs @@ -26,6 +26,26 @@ pub enum WindowFunctionKind { Aggregate(AggKind), } +impl WindowFunctionKind { + pub(crate) fn from_name(name: &str) -> Option { + match name { + "row_number" => Some(Self::RowNumber), + "rank" => Some(Self::Rank), + "dense_rank" => Some(Self::DenseRank), + name => AggKind::from_name(name).map(Self::Aggregate), + } + } + + pub(crate) fn name(self) -> &'static str { + match self { + Self::RowNumber => "row_number", + Self::Rank => "rank", + Self::DenseRank => "dense_rank", + Self::Aggregate(kind) => kind.name(), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Hash, ReferenceSerialization)] pub struct WindowFunction { pub kind: WindowFunctionKind, diff --git a/src/orm/mod.rs b/src/orm/mod.rs index f8cca8ee..c2f21582 100644 --- a/src/orm/mod.rs +++ b/src/orm/mod.rs @@ -9,6 +9,8 @@ use crate::db::{ BindSource, DBTransaction, Database, DatabaseIter, OrmIter, ResultIter, TransactionIter, }; use crate::errors::DatabaseError; +pub use crate::expression::agg::AggKind; +use crate::expression::window::WindowFunctionKind; use crate::expression::{self, AliasType, ScalarExpression}; use crate::planner::operator::alter_table::change_column::{DefaultChange, NotNullChange}; use crate::planner::operator::join::JoinType; @@ -513,24 +515,28 @@ where args: Vec, ) -> Result, DatabaseError> { self.binder() - .bind_function_call(name.into(), args, false, self.arena()) + .bind_function_call(name.into(), args, self.arena()) + .map(|expr| self.wrap(expr)) + } + + fn aggregate( + self, + kind: AggKind, + args: Vec, + ) -> Result, DatabaseError> { + self.binder() + .bind_aggregate_function(kind, args, false, self.arena()) .map(|expr| self.wrap(expr)) } fn window( self, - name: &'static str, + kind: WindowFunctionKind, args: Vec, spec: WindowSpec, ) -> Result, DatabaseError> { self.binder() - .bind_window_function( - name.to_string(), - args, - spec.partition_by, - spec.order_by, - self.arena(), - ) + .bind_window_function(kind, args, spec.partition_by, spec.order_by, self.arena()) .map(|expr| self.wrap(expr)) } @@ -1516,44 +1522,54 @@ where pub fn aggregate( &self, - name: impl Into, + kind: AggKind, args: impl IntoIterator, ) -> Result, DatabaseError> where E: IntoOrmScalarExpression, { - self.function(name, args) + let args = args + .into_iter() + .map(IntoOrmScalarExpression::into_orm_scalar) + .collect(); + self.handle().aggregate(kind, args) } fn aggregate_window( &self, - name: &'static str, + kind: AggKind, expr: impl IntoOrmScalarExpression, spec: WindowSpec, ) -> Result, DatabaseError> { - self.handle() - .window(name, vec![expr.into_orm_scalar()], spec) + self.handle().window( + WindowFunctionKind::Aggregate(kind), + vec![expr.into_orm_scalar()], + spec, + ) } pub fn row_number( &self, spec: WindowSpec, ) -> Result, DatabaseError> { - self.handle().window("row_number", Vec::new(), spec) + self.handle() + .window(WindowFunctionKind::RowNumber, Vec::new(), spec) } pub fn rank( &self, spec: WindowSpec, ) -> Result, DatabaseError> { - self.handle().window("rank", Vec::new(), spec) + self.handle() + .window(WindowFunctionKind::Rank, Vec::new(), spec) } pub fn dense_rank( &self, spec: WindowSpec, ) -> Result, DatabaseError> { - self.handle().window("dense_rank", Vec::new(), spec) + self.handle() + .window(WindowFunctionKind::DenseRank, Vec::new(), spec) } pub fn count_over( @@ -1561,7 +1577,7 @@ where expr: impl IntoOrmScalarExpression, spec: WindowSpec, ) -> Result, DatabaseError> { - self.aggregate_window("count", expr, spec) + self.aggregate_window(AggKind::Count, expr, spec) } pub fn sum_over( @@ -1569,7 +1585,7 @@ where expr: impl IntoOrmScalarExpression, spec: WindowSpec, ) -> Result, DatabaseError> { - self.aggregate_window("sum", expr, spec) + self.aggregate_window(AggKind::Sum, expr, spec) } pub fn avg_over( @@ -1577,7 +1593,7 @@ where expr: impl IntoOrmScalarExpression, spec: WindowSpec, ) -> Result, DatabaseError> { - self.aggregate_window("avg", expr, spec) + self.aggregate_window(AggKind::Avg, expr, spec) } pub fn min_over( @@ -1585,7 +1601,7 @@ where expr: impl IntoOrmScalarExpression, spec: WindowSpec, ) -> Result, DatabaseError> { - self.aggregate_window("min", expr, spec) + self.aggregate_window(AggKind::Min, expr, spec) } pub fn max_over( @@ -1593,7 +1609,7 @@ where expr: impl IntoOrmScalarExpression, spec: WindowSpec, ) -> Result, DatabaseError> { - self.aggregate_window("max", expr, spec) + self.aggregate_window(AggKind::Max, expr, spec) } pub fn count_all_over( @@ -1601,15 +1617,15 @@ where spec: WindowSpec, ) -> Result, DatabaseError> { self.handle().window( - "count", + WindowFunctionKind::Aggregate(AggKind::Count), vec![Binder::<'bind, 'parent, T, A>::wildcard_expr()], spec, ) } pub fn count_all(&self) -> Result, DatabaseError> { - self.function( - "count", + self.aggregate( + AggKind::Count, vec![Binder::<'bind, 'parent, T, A>::wildcard_expr()], ) } @@ -3831,12 +3847,12 @@ mod tests { .project_tuple(|e| { Ok(vec![ e.column(OrmUnitOrder::user_id())?, - e.aggregate("sum", [e.column(OrmUnitOrder::amount())?])?, + e.aggregate(AggKind::Sum, [e.column(OrmUnitOrder::amount())?])?, ]) })? .group_by(|e| e.column(OrmUnitOrder::user_id()))? .having(|e| { - e.aggregate("sum", [e.column(OrmUnitOrder::amount())?])? + e.aggregate(AggKind::Sum, [e.column(OrmUnitOrder::amount())?])? .gte(200) })? .order_by_expr(|e| Ok(e.column(OrmUnitOrder::user_id())?.asc()))? diff --git a/tests/macros-test/src/main.rs b/tests/macros-test/src/main.rs index 1f0a05e4..37b84c66 100644 --- a/tests/macros-test/src/main.rs +++ b/tests/macros-test/src/main.rs @@ -24,7 +24,7 @@ mod test { use kite_sql::expression::function::FunctionSummary; use kite_sql::expression::BinaryOperator; use kite_sql::expression::ScalarExpression; - use kite_sql::orm::{OrmQueryResultExt, WindowSpec}; + use kite_sql::orm::{AggKind, OrmQueryResultExt, WindowSpec}; use kite_sql::planner::{MetaArena, PlanArena, TableArena, TableArenaCell}; use kite_sql::storage::rocksdb::RocksStorage; use kite_sql::types::evaluator::binary_create; @@ -982,7 +982,7 @@ mod test { ctx.from::()? .project_value(|e| { let id = e.column(User::id())?; - e.aggregate("sum", vec![id]) + e.aggregate(AggKind::Sum, vec![id]) })? .finish() })? @@ -1010,7 +1010,7 @@ mod test { ctx.from::()? .project_value(|e| { let id = e.column(User::id())?; - let min_id = e.aggregate("min", vec![id])?; + let min_id = e.aggregate(AggKind::Min, vec![id])?; Ok(e.alias(min_id, "min_user_id")) })? .finish() @@ -1025,7 +1025,7 @@ mod test { ctx.from::()? .project_value(|e| { let id = e.column(User::id())?; - let max_id = e.aggregate("max", vec![id])?; + let max_id = e.aggregate(AggKind::Max, vec![id])?; Ok(e.alias(max_id, "max_user_id")) })? .finish() @@ -1200,7 +1200,7 @@ mod test { ctx.from::()? .project_value(|e| { let id = e.column(User::id())?; - e.aggregate("max", vec![id]) + e.aggregate(AggKind::Max, vec![id]) })? .finish() }) @@ -1860,7 +1860,7 @@ mod test { ctx.from::()? .project_value(|e| { let id = e.column(User::id())?; - e.aggregate("max", vec![id]) + e.aggregate(AggKind::Max, vec![id]) })? .finish() })?; @@ -2212,7 +2212,7 @@ mod test { .project_tuple(|e| { let category = e.column(EventLog::category())?; let score = e.column(EventLog::score())?; - let total_score = e.aggregate("sum", vec![score])?; + let total_score = e.aggregate(AggKind::Sum, vec![score])?; Ok(vec![category, e.alias(total_score, "total_score")]) })? .group_by_scalar(EventLog::category())? @@ -2232,7 +2232,7 @@ mod test { .project_tuple(|e| { let category = e.column(EventLog::category())?; let score = e.column(EventLog::score())?; - let total_score = e.aggregate("sum", vec![score])?; + let total_score = e.aggregate(AggKind::Sum, vec![score])?; let total_count = e.count_all()?; Ok(vec![ category, @@ -2256,7 +2256,7 @@ mod test { .project_tuple(|e| { let category = e.column(EventLog::category())?; let score = e.column(EventLog::score())?; - let total_score = e.aggregate("sum", vec![score])?; + let total_score = e.aggregate(AggKind::Sum, vec![score])?; let total_count = e.count_all()?; Ok(vec![ category, @@ -2333,7 +2333,7 @@ mod test { ctx.from::()? .project_value(|e| { let id = e.column(User::id())?; - e.aggregate("max", vec![id]) + e.aggregate(AggKind::Max, vec![id]) })? .finish() })?; @@ -2401,7 +2401,7 @@ mod test { ctx.from::()? .project_value(|e| { let id = e.column(User::id())?; - e.aggregate("max", vec![id]) + e.aggregate(AggKind::Max, vec![id]) })? .finish() }) @@ -2447,7 +2447,7 @@ mod test { ctx.from::()? .project_value(|e| { let id = e.column(User::id())?; - e.aggregate("max", vec![id]) + e.aggregate(AggKind::Max, vec![id]) })? .having(|e| { let max_id = e.scalar_subquery(|ctx| { From 692aef60c409ffdf7641b5e551a832050ee362cb Mon Sep 17 00:00:00 2001 From: kould Date: Mon, 13 Jul 2026 00:49:58 +0800 Subject: [PATCH 6/6] test: cover accumulators and window operator --- src/execution/dql/aggregate/avg.rs | 19 ++++++++ src/execution/dql/aggregate/count.rs | 35 ++++++++++++-- src/execution/dql/aggregate/min_max.rs | 23 +++++++++ src/execution/dql/aggregate/sum.rs | 29 ++++++++++++ src/expression/mod.rs | 18 +++++++ src/planner/operator/window.rs | 65 ++++++++++++++++++++++++++ 6 files changed, 185 insertions(+), 4 deletions(-) diff --git a/src/execution/dql/aggregate/avg.rs b/src/execution/dql/aggregate/avg.rs index 061edec5..c47461de 100644 --- a/src/execution/dql/aggregate/avg.rs +++ b/src/execution/dql/aggregate/avg.rs @@ -108,6 +108,25 @@ mod tests { assert_ne!(first, second); accumulator.evaluate()?; assert_eq!(&second, accumulator.result()); + assert_eq!(Box::new(accumulator).result_owned(), second); + Ok(()) + } + + #[test] + fn empty_and_unsigned_results() -> Result<(), DatabaseError> { + let mut empty = AvgAccumulator::new(); + empty.evaluate()?; + assert_eq!(empty.result(), &DataValue::Null); + assert_eq!(Box::new(empty).result_owned(), DataValue::Null); + + let mut accumulator = AvgAccumulator::new(); + accumulator.update_value(&DataValue::UInt32(2))?; + accumulator.update_value(&DataValue::UInt32(4))?; + accumulator.evaluate()?; + assert_eq!( + Box::new(accumulator).result_owned(), + DataValue::Float64(3.0.into()) + ); Ok(()) } } diff --git a/src/execution/dql/aggregate/count.rs b/src/execution/dql/aggregate/count.rs index 08c20c72..df3cbbbb 100644 --- a/src/execution/dql/aggregate/count.rs +++ b/src/execution/dql/aggregate/count.rs @@ -66,10 +66,8 @@ impl DistinctCountAccumulator { impl Accumulator for DistinctCountAccumulator { fn update_value(&mut self, value: &DataValue) -> Result<(), DatabaseError> { - if !value.is_null() { - if self.distinct_values.insert(value.clone()) { - self.result = DataValue::Int32(self.distinct_values.len() as i32); - } + if !value.is_null() && self.distinct_values.insert(value.clone()) { + self.result = DataValue::Int32(self.distinct_values.len() as i32); } Ok(()) @@ -83,3 +81,32 @@ impl Accumulator for DistinctCountAccumulator { self.result } } + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + + #[test] + fn count_results() -> Result<(), DatabaseError> { + let mut accumulator = CountAccumulator::new(); + for value in [DataValue::Null, 1.into(), 1.into()] { + accumulator.update_value(&value)?; + } + assert_eq!(accumulator.result(), &DataValue::Int32(2)); + assert_eq!(Box::new(accumulator).result_owned(), DataValue::Int32(2)); + Ok(()) + } + + #[test] + fn distinct_count_results() -> Result<(), DatabaseError> { + let mut accumulator = DistinctCountAccumulator::new(); + for value in [DataValue::Null, 1.into(), 1.into(), 2.into()] { + accumulator.update_value(&value)?; + } + assert_eq!(accumulator.result(), &DataValue::Int32(2)); + assert_eq!(Box::new(accumulator).result_owned(), DataValue::Int32(2)); + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/execution/dql/aggregate/min_max.rs b/src/execution/dql/aggregate/min_max.rs index 186dc665..85de9240 100644 --- a/src/execution/dql/aggregate/min_max.rs +++ b/src/execution/dql/aggregate/min_max.rs @@ -66,3 +66,26 @@ impl Accumulator for MinMaxAccumulator { self.result } } + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + + #[test] + fn min_max_results() -> Result<(), DatabaseError> { + for (is_max, expected) in [(false, 1), (true, 3)] { + let mut accumulator = MinMaxAccumulator::new(is_max); + for value in [DataValue::Null, 3.into(), 1.into(), 2.into()] { + accumulator.update_value(&value)?; + } + assert_eq!(accumulator.result(), &DataValue::Int32(expected)); + assert_eq!( + Box::new(accumulator).result_owned(), + DataValue::Int32(expected) + ); + } + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/execution/dql/aggregate/sum.rs b/src/execution/dql/aggregate/sum.rs index d0ced0a8..f85c50c3 100644 --- a/src/execution/dql/aggregate/sum.rs +++ b/src/execution/dql/aggregate/sum.rs @@ -91,3 +91,32 @@ impl Accumulator for DistinctSumAccumulator { self.inner.result } } + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + + #[test] + fn sum_results() -> Result<(), DatabaseError> { + let mut accumulator = SumAccumulator::new(Cow::Borrowed(&LogicalType::Integer))?; + for value in [DataValue::Null, 2.into(), 3.into()] { + accumulator.update_value(&value)?; + } + assert_eq!(accumulator.result(), &DataValue::Int32(5)); + assert_eq!(Box::new(accumulator).result_owned(), DataValue::Int32(5)); + Ok(()) + } + + #[test] + fn distinct_sum_results() -> Result<(), DatabaseError> { + let mut accumulator = DistinctSumAccumulator::new(&LogicalType::Integer)?; + for value in [DataValue::Null, 2.into(), 2.into(), 3.into()] { + accumulator.update_value(&value)?; + } + assert_eq!(accumulator.result(), &DataValue::Int32(5)); + assert_eq!(Box::new(accumulator).result_owned(), DataValue::Int32(5)); + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/expression/mod.rs b/src/expression/mod.rs index 16fe1ad2..eec778c2 100644 --- a/src/expression/mod.rs +++ b/src/expression/mod.rs @@ -1116,6 +1116,7 @@ mod test { use crate::expression::function::table::{ ArcTableFunctionImpl, TableFunction, TableFunctionCatalog, TableFunctionImpl, }; + use crate::expression::window::{WindowCall, WindowFunction, WindowFunctionKind, WindowSpec}; use crate::expression::TrimWhereField; use crate::expression::{AliasType, BinaryOperator, ScalarExpression, UnaryOperator}; use crate::function::current_date::CurrentDate; @@ -1599,6 +1600,23 @@ mod test { &mut reference_tables, &mut plan_arena, )?; + fn_assert( + &mut cursor, + ScalarExpression::WindowCall(WindowCall { + function: WindowFunction { + kind: WindowFunctionKind::Aggregate(AggKind::Sum), + args: vec![ScalarExpression::Constant(1.into())], + ty: LogicalType::Integer, + }, + spec: WindowSpec { + partition_by: vec![ScalarExpression::Constant(2.into())], + order_by: vec![ScalarExpression::Constant(3.into()).desc()], + }, + }), + Some(&context), + &mut reference_tables, + &mut plan_arena, + )?; Ok(()) } diff --git a/src/planner/operator/window.rs b/src/planner/operator/window.rs index 03197501..cb763342 100644 --- a/src/planner/operator/window.rs +++ b/src/planner/operator/window.rs @@ -58,3 +58,68 @@ impl fmt::Display for WindowOperator { Ok(()) } } + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::expression::window::WindowFunctionKind; + use crate::planner::TableArena; + use crate::serdes::{ReferenceSerialization, ReferenceTables}; + use crate::storage::rocksdb::RocksTransaction; + use crate::types::LogicalType; + use std::io::{Cursor, Seek, SeekFrom}; + + fn operator(partition_by: Vec, order_by: Vec) -> WindowOperator { + WindowOperator { + partition_by, + order_by, + functions: vec![WindowFunction { + kind: WindowFunctionKind::RowNumber, + args: Vec::new(), + ty: LogicalType::Bigint, + }], + output_columns: Vec::new(), + } + } + + #[test] + fn display_window_spec() { + let function = "Window [WindowFunction { kind: RowNumber, args: [], ty: Bigint }]"; + assert_eq!(operator(Vec::new(), Vec::new()).to_string(), function); + assert_eq!( + operator(vec![1.into()], Vec::new()).to_string(), + format!("{function} -> Partition By [1]") + ); + assert_eq!( + operator(Vec::new(), vec![ScalarExpression::from(2).desc()]).to_string(), + format!("{function} -> Order By [2 Desc Nulls Last]") + ); + assert_eq!( + operator(vec![1.into()], vec![ScalarExpression::from(2).desc()]).to_string(), + format!("{function} -> Partition By [1] Order By [2 Desc Nulls Last]") + ); + } + + #[test] + fn serialization_roundtrip() -> Result<(), crate::errors::DatabaseError> { + let source = operator(vec![1.into()], vec![ScalarExpression::from(2).desc()]); + let mut cursor = Cursor::new(Vec::new()); + let mut reference_tables = ReferenceTables::new(); + let arena = TableArena::default(); + source.encode(&mut cursor, false, &mut reference_tables, &arena)?; + cursor.seek(SeekFrom::Start(0))?; + + assert_eq!( + WindowOperator::decode::( + &mut cursor, + None, + &reference_tables, + &mut TableArena::default(), + )?, + source + ); + Ok(()) + } +} +// GRCOV_EXCL_STOP