From 52bda81d67bfb9c0c5ffe1e9fc811d0932eaec4b Mon Sep 17 00:00:00 2001 From: kould Date: Mon, 13 Jul 2026 21:52:52 +0800 Subject: [PATCH] perf: sort window rows in place --- src/binder/window.rs | 28 ++-- src/execution/dql/sort.rs | 112 ++++++++++------ src/execution/dql/window.rs | 120 ++++++++++-------- src/execution/dql/window/function.rs | 16 +-- src/optimizer/heuristic/optimizer.rs | 4 +- .../rule/implementation/dql/window.rs | 32 ++++- src/planner/operator/mod.rs | 4 +- src/planner/operator/visitor.rs | 11 +- src/planner/operator/visitor_mut.rs | 4 +- src/planner/operator/window.rs | 57 +++++++-- tests/slt/window.slt | 2 +- 11 files changed, 243 insertions(+), 147 deletions(-) diff --git a/src/binder/window.rs b/src/binder/window.rs index 60dac596..792cd61c 100644 --- a/src/binder/window.rs +++ b/src/binder/window.rs @@ -18,7 +18,7 @@ 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::sort::SortField; use crate::planner::operator::window::WindowOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan, PlanArena}; @@ -208,26 +208,16 @@ impl> Binder<'_, '_, T, A> } 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)), - ); - } + let partition_by_len = group.partition_by.len(); children = LogicalPlan::new( Operator::Window(WindowOperator { - partition_by: group.partition_by, - order_by: group.order_by, + sort_fields: group + .partition_by + .into_iter() + .map(SortField::from) + .chain(group.order_by) + .collect(), + partition_by_len, functions: group.functions, output_columns: group.output_columns, }), diff --git a/src/execution/dql/sort.rs b/src/execution/dql/sort.rs index 449ccc4c..c40bbe4f 100644 --- a/src/execution/dql/sort.rs +++ b/src/execution/dql/sort.rs @@ -23,6 +23,7 @@ use crate::types::tuple::Tuple; use bumpalo::Bump; use std::cmp::Ordering; use std::mem::{self, transmute, MaybeUninit}; +use std::ops::{Deref, DerefMut}; pub(crate) type BumpVec<'bump, T> = bumpalo::collections::Vec<'bump, T>; @@ -51,17 +52,35 @@ impl<'a, T> NullableVec<'a, T> { } #[inline] - pub(crate) fn into_iter(self) -> impl Iterator + 'a { - self.0 - .into_iter() - .map(|item| unsafe { item.assume_init_read() }) + pub(crate) fn pop(&mut self) -> Option { + self.0.pop().map(|item| unsafe { item.assume_init() }) + } + + pub(crate) fn truncate(&mut self, len: usize) { + while self.len() > len { + self.pop(); + } } } -pub(crate) fn sort_tuples<'a>( +impl Deref for NullableVec<'_, T> { + type Target = [T]; + + fn deref(&self) -> &Self::Target { + unsafe { std::slice::from_raw_parts(self.0.as_ptr().cast(), self.0.len()) } + } +} + +impl DerefMut for NullableVec<'_, T> { + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { std::slice::from_raw_parts_mut(self.0.as_mut_ptr().cast(), self.0.len()) } + } +} + +pub(crate) fn sort_tuples( sort_fields: &[SortField], - mut tuples: NullableVec<'a, (usize, Tuple)>, -) -> Result + 'a, DatabaseError> { + tuples: &mut NullableVec<'_, (usize, Tuple)>, +) -> Result<(), DatabaseError> { let fn_nulls_first = |nulls_first: bool| { if nulls_first { Ordering::Greater @@ -114,12 +133,12 @@ pub(crate) fn sort_tuples<'a>( }); drop(eval_values); - Ok(tuples.into_iter().map(|(_, tuple)| tuple)) + Ok(()) } pub struct Sort { - output: Option>>, - arena: Box, + rows: NullableVec<'static, (usize, Tuple)>, + _arena: Box, sort_fields: Vec, limit: Option, input: ExecId, @@ -136,9 +155,15 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Sort { transaction: &T, ) -> ExecId { let input = build_read(arena, plan_arena, input, cache, transaction); + let sort_arena = Box::::default(); + let rows = unsafe { + transmute::, NullableVec<'static, (usize, Tuple)>>( + NullableVec::new(&sort_arena), + ) + }; arena.push(ExecNode::Sort(Sort { - output: None, - arena: Box::::default(), + rows, + _arena: sort_arena, sort_fields, limit, input, @@ -152,31 +177,23 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Sort { arena: &mut ExecArena<'a, T>, plan_arena: &mut crate::planner::PlanArena<'a>, ) -> Result<(), DatabaseError> { - if self.output.is_none() { - let mut tuples = NullableVec::new(&self.arena); - + loop { + if let Some((_, tuple)) = self.rows.pop() { + arena.produce_tuple(tuple); + return Ok(()); + } while arena.next_tuple(self.input, plan_arena)? { - let offset = tuples.len(); - tuples.put((offset, mem::take(arena.result_tuple_mut()))); + let offset = self.rows.len(); + self.rows.put((offset, mem::take(arena.result_tuple_mut()))); } - - let limit = self.limit.unwrap_or(tuples.len()); - let rows = sort_tuples(&self.sort_fields, tuples)?; - // The arena lives at a stable boxed address, so we can keep the iterator - // and resume it across executor polls. - self.output = Some(unsafe { - transmute:: + '_>, Box>>( - Box::new(rows.take(limit)), - ) - }); - } - - if let Some(tuple) = self.output.as_mut().and_then(std::iter::Iterator::next) { - arena.produce_tuple(tuple); - } else { - arena.finish(); + if self.rows.is_empty() { + arena.finish(); + return Ok(()); + } + sort_tuples(&self.sort_fields, &mut self.rows)?; + self.rows.truncate(self.limit.unwrap_or(self.rows.len())); + self.rows.reverse(); } - Ok(()) } } @@ -192,6 +209,17 @@ mod test { use crate::types::LogicalType; use bumpalo::Bump; + fn sorted_rows<'a>( + sort_fields: &[SortField], + mut tuples: NullableVec<'a, (usize, Tuple)>, + ) -> Result + 'a, DatabaseError> { + sort_tuples(sort_fields, &mut tuples)?; + Ok(tuples.0.into_iter().map(|item| { + let (_, tuple) = unsafe { item.assume_init() }; + tuple + })) + } + #[test] fn test_single_value_desc_and_null_first() -> Result<(), DatabaseError> { let table_arena = crate::planner::TableArenaCell::default(); @@ -295,19 +323,19 @@ mod test { } }; - fn_asc_and_nulls_first_eq(Box::new(sort_tuples( + fn_asc_and_nulls_first_eq(Box::new(sorted_rows( &fn_sort_fields(true, true), fn_tuples(), )?)); - fn_asc_and_nulls_last_eq(Box::new(sort_tuples( + fn_asc_and_nulls_last_eq(Box::new(sorted_rows( &fn_sort_fields(true, false), fn_tuples(), )?)); - fn_desc_and_nulls_first_eq(Box::new(sort_tuples( + fn_desc_and_nulls_first_eq(Box::new(sorted_rows( &fn_sort_fields(false, true), fn_tuples(), )?)); - fn_desc_and_nulls_last_eq(Box::new(sort_tuples( + fn_desc_and_nulls_last_eq(Box::new(sorted_rows( &fn_sort_fields(false, false), fn_tuples(), )?)); @@ -525,19 +553,19 @@ mod test { } }; - fn_asc_1_and_nulls_first_1_and_asc_2_and_nulls_first_2_eq(Box::new(sort_tuples( + fn_asc_1_and_nulls_first_1_and_asc_2_and_nulls_first_2_eq(Box::new(sorted_rows( &fn_sort_fields(true, true, true, true), fn_tuples(), )?)); - fn_asc_1_and_nulls_last_1_and_asc_2_and_nulls_first_2_eq(Box::new(sort_tuples( + fn_asc_1_and_nulls_last_1_and_asc_2_and_nulls_first_2_eq(Box::new(sorted_rows( &fn_sort_fields(true, false, true, true), fn_tuples(), )?)); - fn_desc_1_and_nulls_first_1_and_asc_2_and_nulls_first_2_eq(Box::new(sort_tuples( + fn_desc_1_and_nulls_first_1_and_asc_2_and_nulls_first_2_eq(Box::new(sorted_rows( &fn_sort_fields(false, true, true, true), fn_tuples(), )?)); - fn_desc_1_and_nulls_last_1_and_asc_2_and_nulls_first_2_eq(Box::new(sort_tuples( + fn_desc_1_and_nulls_last_1_and_asc_2_and_nulls_first_2_eq(Box::new(sorted_rows( &fn_sort_fields(false, false, true, true), fn_tuples(), )?)); diff --git a/src/execution/dql/window.rs b/src/execution/dql/window.rs index f6640eb3..3a34dfbd 100644 --- a/src/execution/dql/window.rs +++ b/src/execution/dql/window.rs @@ -13,29 +13,30 @@ // limitations under the License. use crate::errors::DatabaseError; +use crate::execution::dql::sort::{sort_tuples, NullableVec}; 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; +use bumpalo::Bump; +use std::mem::{self, transmute}; mod function; use function::WindowFunction; pub struct Window { - partition_by: Vec, - order_by: Vec, + rows: NullableVec<'static, (usize, Tuple)>, + _arena: Box, + sort_fields: Vec, + partition_by_len: usize, functions: Vec>, input: ExecId, - pending: Option, - rows: Vec, } impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Window { @@ -50,8 +51,8 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Window { ) -> ExecId { let input = build_read(arena, plan_arena, input, cache, transaction); let WindowOperator { - partition_by, - order_by, + sort_fields, + partition_by_len, functions: window_functions, .. } = operator; @@ -60,27 +61,33 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Window { let crate::expression::window::WindowFunction { kind, args, ty } = function; functions.push(function::new(kind, args, ty)); } + let window_arena = Box::::default(); + let rows = unsafe { + transmute::, NullableVec<'static, (usize, Tuple)>>( + NullableVec::new(&window_arena), + ) + }; arena.push(ExecNode::Window(Window { - partition_by, - order_by, + rows, + _arena: window_arena, + sort_fields, + partition_by_len, functions, input, - pending: None, - rows: Vec::new(), })) } } fn evaluate_partition( - rows: &mut [Tuple], + rows: &mut [(usize, 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() { + let output_offset = first.1.values.len(); + for (_, row) in rows.iter_mut() { row.values .resize(output_offset + functions.len(), DataValue::Null); } @@ -94,8 +101,8 @@ fn evaluate_partition( '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]))? + if field.expr.eval(Some(&rows[peer_end - 1].1))? + != field.expr.eval(Some(&rows[peer_end].1))? { break 'peer; } @@ -118,40 +125,44 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window { plan_arena: &mut crate::planner::PlanArena<'a>, ) -> Result<(), DatabaseError> { loop { - if let Some(tuple) = self.rows.pop() { + 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 { + while arena.next_tuple(self.input, plan_arena)? { + let offset = self.rows.len(); + self.rows.put((offset, mem::take(arena.result_tuple_mut()))); + } + if self.rows.is_empty() { arena.finish(); return Ok(()); - }; - self.rows.push(first); + } + if !self.sort_fields.is_empty() { + sort_tuples(&self.sort_fields, &mut self.rows)?; + } - 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; - break; + let mut partition_start = 0; + while partition_start < self.rows.len() { + let mut partition_end = partition_start + 1; + 'partition: while partition_end < self.rows.len() { + // TODO: Cache evaluated partition keys to avoid recalculating the previous row. + for field in &self.sort_fields[..self.partition_by_len] { + if field.expr.eval(Some(&self.rows[partition_end - 1].1))? + != field.expr.eval(Some(&self.rows[partition_end].1))? + { + break 'partition; + } } + partition_end += 1; } - if same_partition { - self.rows.push(tuple); - } else { - self.pending = Some(tuple); - break; - } + evaluate_partition( + &mut self.rows[partition_start..partition_end], + &self.sort_fields[self.partition_by_len..], + &mut self.functions, + )?; + partition_start = partition_end; } - - evaluate_partition(&mut self.rows, &self.order_by, &mut self.functions)?; self.rows.reverse(); } } @@ -167,6 +178,7 @@ mod tests { use crate::expression::window::{ WindowFunction as WindowExpressionFunction, WindowFunctionKind, }; + use crate::expression::ScalarExpression; use crate::planner::operator::values::ValuesOperator; use crate::planner::operator::Operator; use crate::planner::Childrens; @@ -178,10 +190,11 @@ mod tests { ScalarExpression::column_expr(ColumnRef::new(position + 1), position) } - fn rows(values: &[i32]) -> Vec { + fn rows(values: &[i32]) -> Vec<(usize, Tuple)> { values .iter() - .map(|value| Tuple::new(None, vec![DataValue::Int32(*value)])) + .enumerate() + .map(|(index, value)| (index, Tuple::new(None, vec![DataValue::Int32(*value)]))) .collect() } @@ -212,7 +225,9 @@ mod tests { evaluate_partition(&mut rows, &[column(0).asc()], &mut functions())?; assert_eq!( - rows.into_iter().map(|row| row.values).collect::>(), + rows.into_iter() + .map(|(_, row)| row.values) + .collect::>(), vec![ vec![ 10.into(), @@ -246,7 +261,9 @@ mod tests { evaluate_partition(&mut rows, &[], &mut functions())?; assert_eq!( - rows.into_iter().map(|row| row.values).collect::>(), + rows.into_iter() + .map(|(_, row)| row.values) + .collect::>(), vec![ vec![ 3.into(), @@ -302,19 +319,22 @@ mod tests { let input = LogicalPlan::new( Operator::Values(ValuesOperator { rows: vec![ - vec![1.into(), 10.into()], - vec![1.into(), 10.into()], + vec![2.into(), 7.into()], vec![1.into(), 20.into()], vec![2.into(), 5.into()], - vec![2.into(), 7.into()], + vec![1.into(), 10.into()], + vec![1.into(), 10.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()], + sort_fields: vec![ + ScalarExpression::column_expr(input_columns[0], 0).asc(), + ScalarExpression::column_expr(input_columns[1], 1).asc(), + ], + partition_by_len: 1, functions: vec![ WindowExpressionFunction { kind: WindowFunctionKind::RowNumber, diff --git a/src/execution/dql/window/function.rs b/src/execution/dql/window/function.rs index eef66bef..73bcd530 100644 --- a/src/execution/dql/window/function.rs +++ b/src/execution/dql/window/function.rs @@ -29,7 +29,7 @@ pub(super) trait WindowFunction { fn evaluate( &mut self, - rows: &mut [Tuple], + rows: &mut [(usize, Tuple)], peer: Range, peer_index: usize, output_position: usize, @@ -41,13 +41,13 @@ struct RowNumber; impl WindowFunction for RowNumber { fn evaluate( &mut self, - rows: &mut [Tuple], + rows: &mut [(usize, Tuple)], peer: Range, _peer_index: usize, output_position: usize, ) -> Result<(), DatabaseError> { let start = peer.start; - for (offset, row) in rows[peer].iter_mut().enumerate() { + for (offset, (_, row)) in rows[peer].iter_mut().enumerate() { row.values[output_position] = DataValue::Int64((start + offset + 1) as i64); } Ok(()) @@ -61,7 +61,7 @@ struct Rank { impl WindowFunction for Rank { fn evaluate( &mut self, - rows: &mut [Tuple], + rows: &mut [(usize, Tuple)], peer: Range, peer_index: usize, output_position: usize, @@ -71,7 +71,7 @@ impl WindowFunction for Rank { } else { peer.start + 1 }; - for row in &mut rows[peer] { + for (_, row) in &mut rows[peer] { row.values[output_position] = DataValue::Int64(rank as i64); } Ok(()) @@ -93,7 +93,7 @@ impl WindowFunction for Aggregate { fn evaluate( &mut self, - rows: &mut [Tuple], + rows: &mut [(usize, Tuple)], peer: Range, _peer_index: usize, output_position: usize, @@ -101,12 +101,12 @@ impl WindowFunction for Aggregate { let Some(accumulator) = self.accumulator.as_mut() else { unreachable!() }; - for row in &rows[peer.clone()] { + 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] { + for (_, row) in &mut rows[peer] { row.values[output_position] = result.clone(); } Ok(()) diff --git a/src/optimizer/heuristic/optimizer.rs b/src/optimizer/heuristic/optimizer.rs index 5231adea..a6c5ee8b 100644 --- a/src/optimizer/heuristic/optimizer.rs +++ b/src/optimizer/heuristic/optimizer.rs @@ -557,8 +557,8 @@ 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::Window(op) if self.contains(ImplementationRuleImpl::Window) => { + Some(PhysicalOption::new(PlanImpl::Window, op.sort_option())) } Operator::Analyze(_) if self.contains(ImplementationRuleImpl::Analyze) => { Some(PhysicalOption::new(PlanImpl::Analyze, SortOption::None)) diff --git a/src/optimizer/rule/implementation/dql/window.rs b/src/optimizer/rule/implementation/dql/window.rs index 6d08b9d1..f740b1d7 100644 --- a/src/optimizer/rule/implementation/dql/window.rs +++ b/src/optimizer/rule/implementation/dql/window.rs @@ -16,8 +16,7 @@ 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 crate::planner::operator::{Operator, PhysicalOption, PlanImpl}; use std::sync::LazyLock; static WINDOW_PATTERN: LazyLock = LazyLock::new(|| Pattern { @@ -28,8 +27,27 @@ static WINDOW_PATTERN: LazyLock = LazyLock::new(|| Pattern { #[derive(Clone)] pub struct WindowImplementation; -single_mapping!( - WindowImplementation, - WINDOW_PATTERN, - PhysicalOption::new(PlanImpl::Window, SortOption::Follow) -); +impl MatchPattern for WindowImplementation { + fn pattern(&self) -> &Pattern { + &WINDOW_PATTERN + } +} + +impl ImplementationRule for WindowImplementation { + fn update_best_option( + &self, + op: &Operator, + _: &crate::planner::PlanArena, + _: &StatisticMetaLoader<'_>, + best_physical_option: &mut BestPhysicalOption, + ) -> Result<(), DatabaseError> { + if let Operator::Window(op) = op { + crate::optimizer::core::rule::keep_best_physical_option( + best_physical_option, + PhysicalOption::new(PlanImpl::Window, op.sort_option()), + None, + ); + } + Ok(()) + } +} diff --git a/src/planner/operator/mod.rs b/src/planner/operator/mod.rs index 6a4a9f3f..0296a200 100644 --- a/src/planner/operator/mod.rs +++ b/src/planner/operator/mod.rs @@ -322,9 +322,9 @@ impl Operator { op: &'operator window::WindowOperator, ) -> Result<(), DatabaseError> { for expr in op - .partition_by + .sort_fields .iter() - .chain(op.order_by.iter().map(|field| &field.expr)) + .map(|field| &field.expr) .chain(op.functions.iter().flat_map(|function| &function.args)) { ExprVisitor::visit(self, expr)?; diff --git a/src/planner/operator/visitor.rs b/src/planner/operator/visitor.rs index a25107ce..47d5e6cf 100644 --- a/src/planner/operator/visitor.rs +++ b/src/planner/operator/visitor.rs @@ -261,9 +261,9 @@ impl<'a, V: ExprVisitor<'a>> OperatorVisitor<'a> for OperatorExprVisitor<'_, V> fn visit_window(&mut self, op: &'a window::WindowOperator) -> Result<(), DatabaseError> { for expr in op - .partition_by + .sort_fields .iter() - .chain(op.order_by.iter().map(|field| &field.expr)) + .map(|field| &field.expr) .chain(op.functions.iter().flat_map(|function| &function.args)) { ExprVisitor::visit(self.visitor, expr)?; @@ -462,8 +462,11 @@ pub(crate) mod tests { schema_ref: vec![column_ref], }), Operator::Window(window::WindowOperator { - partition_by: vec![17_i32.into()], - order_by: vec![SortField::from(ScalarExpression::from(18_i32))], + sort_fields: vec![ + SortField::from(ScalarExpression::from(17_i32)), + SortField::from(ScalarExpression::from(18_i32)), + ], + partition_by_len: 1, functions: vec![WindowFunction { kind: WindowFunctionKind::RowNumber, args: Vec::new(), diff --git a/src/planner/operator/visitor_mut.rs b/src/planner/operator/visitor_mut.rs index 388a88f4..b3e3d96e 100644 --- a/src/planner/operator/visitor_mut.rs +++ b/src/planner/operator/visitor_mut.rs @@ -282,9 +282,9 @@ impl<'a, V: ExprVisitorMut<'a>> OperatorVisitorMut<'a> for OperatorExprVisitorMu fn visit_window(&mut self, op: &'a mut window::WindowOperator) -> Result<(), DatabaseError> { for expr in op - .partition_by + .sort_fields .iter_mut() - .chain(op.order_by.iter_mut().map(|field| &mut field.expr)) + .map(|field| &mut field.expr) .chain( op.functions .iter_mut() diff --git a/src/planner/operator/window.rs b/src/planner/operator/window.rs index cb763342..a5506937 100644 --- a/src/planner/operator/window.rs +++ b/src/planner/operator/window.rs @@ -14,22 +14,36 @@ 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 crate::planner::operator::SortOption; 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 sort_fields: Vec, + pub partition_by_len: usize, pub functions: Vec, pub output_columns: Vec, } +impl WindowOperator { + pub(crate) fn sort_option(&self) -> SortOption { + if self.sort_fields.is_empty() { + SortOption::Follow + } else { + SortOption::OrderBy { + fields: self.sort_fields.clone(), + ignore_prefix_len: 0, + } + } + } +} + impl fmt::Display for WindowOperator { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let (partition_by, order_by) = self.sort_fields.split_at(self.partition_by_len); write!( f, "Window [{}]", @@ -38,21 +52,24 @@ impl fmt::Display for WindowOperator { .map(|expr| format!("{expr:?}")) .join(", ") )?; - if !self.partition_by.is_empty() || !self.order_by.is_empty() { + if !self.sort_fields.is_empty() { write!(f, " ->")?; } - if !self.partition_by.is_empty() { + if !partition_by.is_empty() { write!( f, " Partition By [{}]", - self.partition_by.iter().map(ToString::to_string).join(", ") + partition_by + .iter() + .map(|field| field.expr.to_string()) + .join(", ") )?; } - if !self.order_by.is_empty() { + if !order_by.is_empty() { write!( f, " Order By [{}]", - self.order_by.iter().map(ToString::to_string).join(", ") + order_by.iter().map(ToString::to_string).join(", ") )?; } Ok(()) @@ -64,6 +81,7 @@ impl fmt::Display for WindowOperator { mod tests { use super::*; use crate::expression::window::WindowFunctionKind; + use crate::expression::ScalarExpression; use crate::planner::TableArena; use crate::serdes::{ReferenceSerialization, ReferenceTables}; use crate::storage::rocksdb::RocksTransaction; @@ -71,9 +89,14 @@ mod tests { use std::io::{Cursor, Seek, SeekFrom}; fn operator(partition_by: Vec, order_by: Vec) -> WindowOperator { + let partition_by_len = partition_by.len(); WindowOperator { - partition_by, - order_by, + sort_fields: partition_by + .into_iter() + .map(SortField::from) + .chain(order_by) + .collect(), + partition_by_len, functions: vec![WindowFunction { kind: WindowFunctionKind::RowNumber, args: Vec::new(), @@ -99,6 +122,20 @@ mod tests { operator(vec![1.into()], vec![ScalarExpression::from(2).desc()]).to_string(), format!("{function} -> Partition By [1] Order By [2 Desc Nulls Last]") ); + assert_eq!( + operator(Vec::new(), Vec::new()).sort_option(), + SortOption::Follow + ); + assert_eq!( + operator(vec![1.into()], vec![ScalarExpression::from(2).desc()]).sort_option(), + SortOption::OrderBy { + fields: vec![ + ScalarExpression::from(1).asc(), + ScalarExpression::from(2).desc(), + ], + ignore_prefix_len: 0, + } + ); } #[test] diff --git a/tests/slt/window.slt b/tests/slt/window.slt index fee542b0..e1d5488b 100644 --- a/tests/slt/window.slt +++ b/tests/slt/window.slt @@ -45,7 +45,7 @@ explain select 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)] +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: 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: 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,