diff --git a/src/binder/window.rs b/src/binder/window.rs index 792cd61c..9627016f 100644 --- a/src/binder/window.rs +++ b/src/binder/window.rs @@ -19,6 +19,7 @@ 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; +use crate::planner::operator::sort::SortOperator; use crate::planner::operator::window::WindowOperator; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan, PlanArena}; @@ -209,14 +210,23 @@ impl> Binder<'_, '_, T, A> for group in groups { let partition_by_len = group.partition_by.len(); + let sort_fields: Vec<_> = group + .partition_by + .into_iter() + .map(SortField::from) + .chain(group.order_by) + .collect(); + if !sort_fields.is_empty() { + children = LogicalPlan::new( + Operator::Sort(SortOperator { + sort_fields: sort_fields.clone(), + }), + Childrens::Only(Box::new(children)), + ); + } children = LogicalPlan::new( Operator::Window(WindowOperator { - sort_fields: group - .partition_by - .into_iter() - .map(SortField::from) - .chain(group.order_by) - .collect(), + sort_fields, partition_by_len, functions: group.functions, output_columns: group.output_columns, diff --git a/src/execution/dql/window.rs b/src/execution/dql/window.rs index 3a34dfbd..851da911 100644 --- a/src/execution/dql/window.rs +++ b/src/execution/dql/window.rs @@ -4,7 +4,7 @@ // 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 +// 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, @@ -13,29 +13,66 @@ // 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::window::WindowFunctionKind; 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 bumpalo::Bump; -use std::mem::{self, transmute}; +use std::mem; mod function; use function::WindowFunction; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Retention { + Row, + Peer, + Partition, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Boundary { + Partition, + Peer, +} + +#[derive(Default)] +struct WindowState { + buffered: Vec<(usize, Tuple)>, + pending: Option<(Tuple, Boundary)>, + sort_values: Vec, + started: bool, + partition_rows: usize, + peer_start: usize, + peer_index: usize, +} + +impl WindowState { + fn begin_partition(&mut self) { + self.partition_rows = 0; + self.peer_start = 0; + self.peer_index = 0; + } + + fn begin_peer(&mut self) { + self.peer_start = self.partition_rows; + self.peer_index += 1; + } +} + pub struct Window { - rows: NullableVec<'static, (usize, Tuple)>, - _arena: Box, + state: WindowState, + retention: Retention, sort_fields: Vec, partition_by_len: usize, functions: Vec>, + input_exhausted: bool, input: ExecId, } @@ -56,66 +93,115 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Window { 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)); - } - let window_arena = Box::::default(); - let rows = unsafe { - transmute::, NullableVec<'static, (usize, Tuple)>>( - NullableVec::new(&window_arena), - ) + let has_aggregate = window_functions + .iter() + .any(|function| matches!(function.kind, WindowFunctionKind::Aggregate(_))); + let retention = if !has_aggregate { + Retention::Row + } else if sort_fields.len() == partition_by_len { + Retention::Partition + } else { + Retention::Peer }; + let functions = window_functions + .into_iter() + .map(|function| function::new(function.kind, function.args, function.ty)) + .collect(); arena.push(ExecNode::Window(Window { - rows, - _arena: window_arena, + state: WindowState::default(), + retention, sort_fields, partition_by_len, functions, + input_exhausted: false, input, })) } } -fn evaluate_partition( - rows: &mut [(usize, Tuple)], - order_by: &[SortField], - functions: &mut [Box], -) -> Result<(), DatabaseError> { - let Some(first) = rows.first() else { - return Ok(()); - }; - let output_offset = first.1.values.len(); - for (_, row) in rows.iter_mut() { - row.values - .resize(output_offset + functions.len(), DataValue::Null); +impl Window { + fn update_keys(&mut self, tuple: &Tuple) -> Result, DatabaseError> { + let mut boundary = (!self.state.started).then_some(Boundary::Partition); + for (index, field) in self.sort_fields.iter().enumerate() { + let value = field.expr.eval(Some(tuple))?; + if self.state.started && self.state.sort_values[index] != value { + if index < self.partition_by_len { + boundary = Some(Boundary::Partition); + } else if boundary.is_none() { + boundary = Some(Boundary::Peer); + } + self.state.sort_values[index] = value; + } else if !self.state.started { + self.state.sort_values.push(value); + } + } + self.state.started = true; + Ok(boundary) + } + + fn reset_functions(&mut self) -> Result<(), DatabaseError> { + for function in &mut self.functions { + function.reset()?; + } + Ok(()) } - for function in functions.iter_mut() { - function.reset()?; + + fn eval_functions(&mut self) -> Result<(), DatabaseError> { + if self.state.buffered.is_empty() { + return Ok(()); + } + let output_offset = self.state.buffered[0].1.values.len(); + for (_, row) in &mut self.state.buffered { + row.values + .resize(output_offset + self.functions.len(), DataValue::Null); + } + let len = self.state.buffered.len(); + for (slot, function) in self.functions.iter_mut().enumerate() { + function.evaluate( + &mut self.state.buffered, + 0..len, + self.state.peer_start, + self.state.peer_index, + output_offset + slot, + )?; + } + self.state.buffered.reverse(); + Ok(()) } - 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() { - // 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].1))? - != field.expr.eval(Some(&rows[peer_end].1))? - { - break 'peer; - } + + fn eval(&mut self, tuple: Tuple, boundary: Option) -> Result { + let boundary = match boundary { + Some(boundary) => Some(boundary), + None => self.update_keys(&tuple)?, + }; + if let Some(boundary) = boundary { + let reached_boundary = boundary == Boundary::Partition + || boundary == Boundary::Peer && self.retention == Retention::Peer; + if reached_boundary && !self.state.buffered.is_empty() { + self.state.pending = Some((tuple, boundary)); + self.eval_functions()?; + return Ok(true); + } + } + + match boundary { + Some(Boundary::Partition) => { + self.state.begin_partition(); + self.reset_functions()?; } - peer_end += 1; + Some(Boundary::Peer) => self.state.begin_peer(), + None => {} } - for (slot, function) in functions.iter_mut().enumerate() { - function.evaluate(rows, peer_start..peer_end, peer_index, output_offset + slot)?; + + let row_index = self.state.partition_rows; + self.state.partition_rows += 1; + self.state.buffered.push((row_index, tuple)); + if self.retention == Retention::Row { + self.eval_functions()?; + return Ok(true); } - peer_start = peer_end; - peer_index += 1; + Ok(false) } - Ok(()) } impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window { @@ -124,46 +210,31 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window { arena: &mut ExecArena<'a, T>, plan_arena: &mut crate::planner::PlanArena<'a>, ) -> Result<(), DatabaseError> { + let mut output_ready = true; loop { - if let Some((_, tuple)) = self.rows.pop() { - arena.produce_tuple(tuple); - return Ok(()); - } - - while arena.next_tuple(self.input, plan_arena)? { - let offset = self.rows.len(); - self.rows.put((offset, mem::take(arena.result_tuple_mut()))); + if output_ready { + if let Some((_, tuple)) = self.state.buffered.pop() { + arena.produce_tuple(tuple); + return Ok(()); + } } - if self.rows.is_empty() { + if self.input_exhausted { arena.finish(); return Ok(()); } - if !self.sort_fields.is_empty() { - sort_tuples(&self.sort_fields, &mut self.rows)?; - } - 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; - } - evaluate_partition( - &mut self.rows[partition_start..partition_end], - &self.sort_fields[self.partition_by_len..], - &mut self.functions, - )?; - partition_start = partition_end; - } - self.rows.reverse(); + let (tuple, boundary) = if let Some((tuple, boundary)) = self.state.pending.take() { + (tuple, Some(boundary)) + } else if arena.next_tuple(self.input, plan_arena)? { + (mem::take(arena.result_tuple_mut()), None) + } else { + self.eval_functions()?; + self.input_exhausted = true; + output_ready = true; + continue; + }; + + output_ready = self.eval(tuple, boundary)?; } } } @@ -172,206 +243,96 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Window { #[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::catalog::ColumnRef; use crate::expression::agg::AggKind; - 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; - 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<(usize, Tuple)> { - values - .iter() - .enumerate() - .map(|(index, value)| (index, Tuple::new(None, vec![DataValue::Int32(*value)]))) - .collect() + fn window( + retention: Retention, + sort_fields: Vec, + partition_by_len: usize, + functions: Vec>, + ) -> Window { + Window { + state: WindowState::default(), + retention, + sort_fields, + partition_by_len, + functions, + input_exhausted: false, + input: 0, + } } - 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, - ), - ] + fn tuple(values: &[i32]) -> Tuple { + Tuple::new(None, values.iter().copied().map(DataValue::from).collect()) } #[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::>(), + fn row_materialization_streams_rows() -> Result<(), DatabaseError> { + let mut window = window( + Retention::Row, + vec![column(0).asc(), column(1).asc()], + 1, 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() - ], - ] + function::new( + WindowFunctionKind::RowNumber, + Vec::new(), + LogicalType::Bigint, + ), + function::new(WindowFunctionKind::Rank, Vec::new(), LogicalType::Bigint), + ], ); + for (value, expected) in [(10, [1_i64, 1]), (10, [2, 1]), (20, [3, 3])] { + window.eval(tuple(&[1, value]), None)?; + let row = window.state.buffered.pop().unwrap().1; + assert_eq!(row.values[2..], expected.map(DataValue::from)); + } 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() - ], - ] + fn peer_materialization_waits_for_peer_boundary() -> Result<(), DatabaseError> { + let mut window = window( + Retention::Peer, + vec![column(0).asc(), column(1).asc()], + 1, + vec![function::new( + WindowFunctionKind::Aggregate(AggKind::Sum), + vec![column(1)], + LogicalType::Integer, + )], ); + window.eval(tuple(&[1, 10]), None)?; + window.eval(tuple(&[1, 10]), None)?; + window.eval(tuple(&[1, 20]), None)?; + assert_eq!(window.state.buffered.len(), 2); + assert!(window.state.pending.is_some()); 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![2.into(), 7.into()], - vec![1.into(), 20.into()], - vec![2.into(), 5.into()], - vec![1.into(), 10.into()], - vec![1.into(), 10.into()], - ], - schema_ref: input_columns.clone(), - }), - Childrens::None, - ); - let operator = WindowOperator { - 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, - 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()], - ] + fn partition_materialization_waits_for_partition_boundary() -> Result<(), DatabaseError> { + let mut window = window( + Retention::Partition, + vec![column(0).asc()], + 1, + vec![function::new( + WindowFunctionKind::Aggregate(AggKind::Sum), + vec![column(1)], + LogicalType::Integer, + )], ); + window.eval(tuple(&[1, 3]), None)?; + window.eval(tuple(&[1, 7]), None)?; + window.eval(tuple(&[2, 5]), None)?; + assert_eq!(window.state.buffered.len(), 2); + assert!(window.state.pending.is_some()); Ok(()) } } diff --git a/src/execution/dql/window/function.rs b/src/execution/dql/window/function.rs index 73bcd530..8484792f 100644 --- a/src/execution/dql/window/function.rs +++ b/src/execution/dql/window/function.rs @@ -31,6 +31,7 @@ pub(super) trait WindowFunction { &mut self, rows: &mut [(usize, Tuple)], peer: Range, + peer_start: usize, peer_index: usize, output_position: usize, ) -> Result<(), DatabaseError>; @@ -43,12 +44,12 @@ impl WindowFunction for RowNumber { &mut self, rows: &mut [(usize, Tuple)], peer: Range, + _peer_start: usize, _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); + for (row_index, row) in &mut rows[peer] { + row.values[output_position] = DataValue::Int64((*row_index + 1) as i64); } Ok(()) } @@ -63,13 +64,14 @@ impl WindowFunction for Rank { &mut self, rows: &mut [(usize, Tuple)], peer: Range, + peer_start: usize, peer_index: usize, output_position: usize, ) -> Result<(), DatabaseError> { let rank = if self.dense { peer_index + 1 } else { - peer.start + 1 + peer_start + 1 }; for (_, row) in &mut rows[peer] { row.values[output_position] = DataValue::Int64(rank as i64); @@ -95,6 +97,7 @@ impl WindowFunction for Aggregate { &mut self, rows: &mut [(usize, Tuple)], peer: Range, + _peer_start: usize, _peer_index: usize, output_position: usize, ) -> Result<(), DatabaseError> { diff --git a/tests/slt/window.slt b/tests/slt/window.slt index e1d5488b..3d3ba6f9 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: 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)] +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)] 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: OrderBy: (#2 Asc Nulls Last, #3 Asc Nulls Last, #1 Asc Nulls Last) ignore_prefix_len: 0)] 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,