From 73041800272346f4941aed37e546de025086bf49 Mon Sep 17 00:00:00 2001 From: doki Date: Fri, 22 Apr 2022 12:06:29 +0800 Subject: [PATCH 01/18] window func expression --- common/ast/src/parser/expr/expr_visitor.rs | 9 + common/functions/src/lib.rs | 1 + common/functions/src/window/mod.rs | 3 + common/functions/src/window/window_frame.rs | 228 ++++++++++++++++++ common/planners/src/plan_expression.rs | 38 +++ common/planners/src/plan_expression_chain.rs | 6 + common/planners/src/plan_expression_common.rs | 28 +++ .../planners/src/plan_expression_rewriter.rs | 62 +++++ .../planners/src/plan_expression_visitor.rs | 14 ++ common/planners/src/plan_node_rewriter.rs | 29 +++ query/src/sql/statements/analyzer_expr.rs | 87 ++++++- .../sql/statements/query/query_normalizer.rs | 17 +- query/src/sql/statements/statement_select.rs | 15 +- 13 files changed, 524 insertions(+), 13 deletions(-) create mode 100644 common/functions/src/window/mod.rs create mode 100644 common/functions/src/window/window_frame.rs diff --git a/common/ast/src/parser/expr/expr_visitor.rs b/common/ast/src/parser/expr/expr_visitor.rs index 4eeb86467f28..6a4cddee229b 100644 --- a/common/ast/src/parser/expr/expr_visitor.rs +++ b/common/ast/src/parser/expr/expr_visitor.rs @@ -163,6 +163,15 @@ pub trait ExprVisitor: Sized + Send { }; } + if let Some(over) = &function.over { + for partition_by in &over.partition_by { + ExprTraverser::accept(partition_by, self).await?; + } + for order_by in &over.order_by { + ExprTraverser::accept(&order_by.expr, self).await?; + } + } + Ok(()) } diff --git a/common/functions/src/lib.rs b/common/functions/src/lib.rs index 9340791ea80c..a17897503d1a 100644 --- a/common/functions/src/lib.rs +++ b/common/functions/src/lib.rs @@ -19,6 +19,7 @@ pub mod aggregates; pub mod rdoc; pub mod scalars; +pub mod window; use aggregates::AggregateFunctionFactory; use scalars::FunctionFactory; diff --git a/common/functions/src/window/mod.rs b/common/functions/src/window/mod.rs new file mode 100644 index 000000000000..21e32b7615d1 --- /dev/null +++ b/common/functions/src/window/mod.rs @@ -0,0 +1,3 @@ +mod window_frame; + +pub use window_frame::WindowFrame; diff --git a/common/functions/src/window/window_frame.rs b/common/functions/src/window/window_frame.rs new file mode 100644 index 000000000000..3a9e75b4fe93 --- /dev/null +++ b/common/functions/src/window/window_frame.rs @@ -0,0 +1,228 @@ +// Copyright 2021 Datafuse Labs. +// +// 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 std::cmp::Ordering; +use std::fmt; +use std::hash::Hash; +use std::hash::Hasher; + +use common_exception::ErrorCode; +use sqlparser::ast; + +/// The frame-spec determines which output rows are read by an aggregate window function. +/// +/// The ending frame boundary can be omitted (if the BETWEEN and AND keywords that surround the +/// starting frame boundary are also omitted), in which case the ending frame boundary defaults to +/// CURRENT ROW. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash, serde::Serialize, serde::Deserialize, +)] +pub struct WindowFrame { + /// A frame type - either ROWS, RANGE or GROUPS + pub units: WindowFrameUnits, + /// A starting frame boundary + pub start_bound: WindowFrameBound, + /// An ending frame boundary + pub end_bound: WindowFrameBound, +} + +impl TryFrom for WindowFrame { + type Error = ErrorCode; + + fn try_from(value: ast::WindowFrame) -> Result { + let start_bound = value.start_bound.into(); + let end_bound = value + .end_bound + .map(WindowFrameBound::from) + .unwrap_or(WindowFrameBound::CurrentRow); + + if let WindowFrameBound::Following(None) = start_bound { + Err(ErrorCode::LogicalError( + "Invalid window frame: start bound cannot be unbounded following".to_owned(), + )) + } else if let WindowFrameBound::Preceding(None) = end_bound { + Err(ErrorCode::LogicalError( + "Invalid window frame: end bound cannot be unbounded preceding".to_owned(), + )) + } else if start_bound > end_bound { + Err(ErrorCode::LogicalError(format!( + "Invalid window frame: start bound ({}) cannot be larger than end bound ({})", + start_bound, end_bound + ))) + } else { + let units = value.units.into(); + if units == WindowFrameUnits::Range { + for bound in &[start_bound, end_bound] { + match bound { + WindowFrameBound::Preceding(Some(v)) + | WindowFrameBound::Following(Some(v)) + if *v > 0 => + { + Err(ErrorCode::UnImplement(format!( + "With WindowFrameUnits={}, the bound cannot be {} PRECEDING or FOLLOWING at the moment", + units, v + ))) + } + _ => Ok(()), + }?; + } + } + Ok(Self { + units, + start_bound, + end_bound, + }) + } + } +} + +impl Default for WindowFrame { + fn default() -> Self { + WindowFrame { + units: WindowFrameUnits::Range, + start_bound: WindowFrameBound::Preceding(None), + end_bound: WindowFrameBound::CurrentRow, + } + } +} + +/// There are three frame types: ROWS, GROUPS, and RANGE. The frame type determines how the +/// starting and ending boundaries of the frame are measured. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash, serde::Serialize, serde::Deserialize, +)] +pub enum WindowFrameUnits { + /// The ROWS frame type means that the starting and ending boundaries for the frame are + /// determined by counting individual rows relative to the current row. + Rows, + /// The RANGE frame type requires that the ORDER BY clause of the window have exactly one + /// term. Call that term "X". With the RANGE frame type, the elements of the frame are + /// determined by computing the value of expression X for all rows in the partition and framing + /// those rows for which the value of X is within a certain range of the value of X for the + /// current row. + Range, + /// The GROUPS frame type means that the starting and ending boundaries are determine + /// by counting "groups" relative to the current group. A "group" is a set of rows that all have + /// equivalent values for all all terms of the window ORDER BY clause. + Groups, +} + +impl fmt::Display for WindowFrameUnits { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str(match self { + WindowFrameUnits::Rows => "ROWS", + WindowFrameUnits::Range => "RANGE", + WindowFrameUnits::Groups => "GROUPS", + }) + } +} + +impl From for WindowFrameUnits { + fn from(value: ast::WindowFrameUnits) -> Self { + match value { + ast::WindowFrameUnits::Range => Self::Range, + ast::WindowFrameUnits::Groups => Self::Groups, + ast::WindowFrameUnits::Rows => Self::Rows, + } + } +} + +#[derive(Debug, Clone, Copy, Eq, serde::Serialize, serde::Deserialize)] +pub enum WindowFrameBound { + /// 1. UNBOUNDED PRECEDING + /// The frame boundary is the first row in the partition. + /// + /// 2. PRECEDING + /// must be a non-negative constant numeric expression. The boundary is a row that + /// is "units" prior to the current row. + Preceding(Option), + /// 3. The current row. + /// + /// For RANGE and GROUPS frame types, peers of the current row are also + /// included in the frame, unless specifically excluded by the EXCLUDE clause. + /// This is true regardless of whether CURRENT ROW is used as the starting or ending frame + /// boundary. + CurrentRow, + /// 4. This is the same as " PRECEDING" except that the boundary is units after the + /// current rather than before the current row. + /// + /// 5. UNBOUNDED FOLLOWING + /// The frame boundary is the last row in the partition. + Following(Option), +} + +impl WindowFrameBound { + /// get the rank of this window frame bound. + /// + /// the rank is a tuple of (u8, u64) because we'll firstly compare the kind and then the value + /// which requires special handling e.g. with preceding the larger the value the smaller the + /// rank and also for 0 preceding / following it is the same as current row + fn get_rank(&self) -> (u8, u64) { + match self { + WindowFrameBound::Preceding(None) => (0, 0), + WindowFrameBound::Following(None) => (4, 0), + WindowFrameBound::Preceding(Some(0)) + | WindowFrameBound::CurrentRow + | WindowFrameBound::Following(Some(0)) => (2, 0), + WindowFrameBound::Preceding(Some(v)) => (1, u64::MAX - *v), + WindowFrameBound::Following(Some(v)) => (3, *v), + } + } +} + +impl From for WindowFrameBound { + fn from(value: ast::WindowFrameBound) -> Self { + match value { + ast::WindowFrameBound::Preceding(v) => Self::Preceding(v), + ast::WindowFrameBound::Following(v) => Self::Following(v), + ast::WindowFrameBound::CurrentRow => Self::CurrentRow, + } + } +} + +impl fmt::Display for WindowFrameBound { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + WindowFrameBound::CurrentRow => f.write_str("CURRENT ROW"), + WindowFrameBound::Preceding(None) => f.write_str("UNBOUNDED PRECEDING"), + WindowFrameBound::Following(None) => f.write_str("UNBOUNDED FOLLOWING"), + WindowFrameBound::Preceding(Some(n)) => write!(f, "{} PRECEDING", n), + WindowFrameBound::Following(Some(n)) => write!(f, "{} FOLLOWING", n), + } + } +} + +impl PartialEq for WindowFrameBound { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl PartialOrd for WindowFrameBound { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for WindowFrameBound { + fn cmp(&self, other: &Self) -> Ordering { + self.get_rank().cmp(&other.get_rank()) + } +} + +impl Hash for WindowFrameBound { + fn hash(&self, state: &mut H) { + self.get_rank().hash(state) + } +} diff --git a/common/planners/src/plan_expression.rs b/common/planners/src/plan_expression.rs index 78bcc2acf092..1cc19a411931 100644 --- a/common/planners/src/plan_expression.rs +++ b/common/planners/src/plan_expression.rs @@ -21,6 +21,7 @@ use common_exception::ErrorCode; use common_exception::Result; use common_functions::aggregates::AggregateFunctionFactory; use common_functions::aggregates::AggregateFunctionRef; +use common_functions::window::WindowFrame; use once_cell::sync::Lazy; use crate::plan_expression_common::ExpressionDataTypeVisitor; @@ -94,6 +95,18 @@ pub enum Expression { args: Vec, }, + /// WindowFunction + WindowFunction { + /// operation performed + func: Box, + /// partition by + partition_by: Vec, + /// order by + order_by: Vec, + /// Window frame + window_frame: Option, + }, + /// A sort expression, that can be used to sort values. Sort { /// The expression to sort on @@ -407,6 +420,31 @@ impl fmt::Debug for Expression { Ok(()) } + Expression::WindowFunction { + func, + partition_by, + order_by, + window_frame, + } => { + write!(f, "{:?}", func)?; + write!(f, " OVER( ")?; + if !partition_by.is_empty() { + write!(f, "PARTITION BY {:?} ", partition_by)?; + } + if !order_by.is_empty() { + write!(f, "ORDER BY {:?} ", order_by)?; + } + if let Some(window_frame) = window_frame { + write!( + f, + "{} BETWEEN {} AND {} ", + window_frame.units, window_frame.start_bound, window_frame.end_bound + )?; + } + write!(f, ")")?; + Ok(()) + } + Expression::Sort { expr, .. } => write!(f, "{:?}", expr), Expression::Wildcard => write!(f, "*"), Expression::Cast { diff --git a/common/planners/src/plan_expression_chain.rs b/common/planners/src/plan_expression_chain.rs index 1e0fe3e20351..a1d05a6c0f48 100644 --- a/common/planners/src/plan_expression_chain.rs +++ b/common/planners/src/plan_expression_chain.rs @@ -191,7 +191,13 @@ impl ExpressionChain { "Action must be a non-aggregated function.", )); } + + Expression::WindowFunction { .. } => { + todo!("not figure out yet @doki") + } + Expression::Wildcard | Expression::Sort { .. } => {} + Expression::Cast { expr: sub_expr, data_type, diff --git a/common/planners/src/plan_expression_common.rs b/common/planners/src/plan_expression_common.rs index ed3c2293f462..100ef60088ec 100644 --- a/common/planners/src/plan_expression_common.rs +++ b/common/planners/src/plan_expression_common.rs @@ -303,6 +303,24 @@ where F: Fn(&Expression) -> Result> { .collect::>>()?, }), + Expression::WindowFunction { + func, + partition_by, + order_by, + window_frame, + } => Ok(Expression::WindowFunction { + func: Box::new(clone_with_replacement(func, replacement_fn)?), + partition_by: partition_by + .iter() + .map(|e| clone_with_replacement(e, replacement_fn)) + .collect::>>()?, + order_by: order_by + .iter() + .map(|e| clone_with_replacement(e, replacement_fn)) + .collect::>>()?, + window_frame: window_frame.to_owned(), + }), + Expression::Sort { expr: nested_expr, asc, @@ -468,6 +486,16 @@ impl ExpressionVisitor for ExpressionDataTypeVisitor { self.stack.push(return_type); Ok(self) } + Expression::WindowFunction { + partition_by, + order_by, + .. + } => { + for _ in 0..partition_by.len() + order_by.len() { + self.stack.remove(0); + } + Ok(self) + } Expression::Cast { data_type, .. } => { let inner_type = match self.stack.pop() { None => Err(ErrorCode::LogicalError( diff --git a/common/planners/src/plan_expression_rewriter.rs b/common/planners/src/plan_expression_rewriter.rs index 5cfdf4940493..3233bd587b17 100644 --- a/common/planners/src/plan_expression_rewriter.rs +++ b/common/planners/src/plan_expression_rewriter.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use common_datavalues::prelude::*; use common_exception::ErrorCode; use common_exception::Result; +use common_functions::window::WindowFrame; use crate::Expression; use crate::ExpressionVisitor; @@ -124,6 +125,21 @@ pub trait ExpressionRewriter: Sized { }) } + fn mutate_window_function( + &mut self, + func: Expression, + partition_by: Vec, + order_by: Vec, + window_frame: Option, + ) -> Result { + Ok(Expression::WindowFunction { + func: Box::new(func), + partition_by, + order_by, + window_frame, + }) + } + fn mutate_cast( &mut self, typ: &DataTypePtr, @@ -286,6 +302,52 @@ impl ExpressionVisitor for ExpressionRewriteVisitor { self.stack.push(new_expr); Ok(self) } + Expression::WindowFunction { + func: _, + partition_by, + order_by, + window_frame, + } => { + let new_func = self.stack.pop().unwrap(); + + let mut new_partition_by = Vec::with_capacity(partition_by.len()); + for (i, _) in partition_by.iter().enumerate() { + match self.stack.pop() { + Some(new_partition_by_expr) => new_partition_by.push(new_partition_by_expr), + None => { + return Err(ErrorCode::LogicalError(format!( + "WindowFunction expects {} partition by arguments, actual {}", + partition_by.len(), + i + ))) + } + } + } + + let mut new_order_by = Vec::with_capacity(order_by.len()); + for (i, _) in order_by.iter().enumerate() { + match self.stack.pop() { + Some(new_order_by_expr) => new_order_by.push(new_order_by_expr), + None => { + return Err(ErrorCode::LogicalError(format!( + "WindowFunction expects {} order by arguments, actual {}", + order_by.len(), + i + ))) + } + } + } + + let new_expr = self.inner.mutate_window_function( + new_func, + new_partition_by, + new_order_by, + window_frame.to_owned(), + )?; + + self.stack.push(new_expr); + Ok(self) + } Expression::Cast { data_type, .. } => match self.stack.pop() { None => Err(ErrorCode::LogicalError( "Cast expr expected 1 parameters, actual 0.", diff --git a/common/planners/src/plan_expression_visitor.rs b/common/planners/src/plan_expression_visitor.rs index 7a4aaf62bbdd..6d0599c5c795 100644 --- a/common/planners/src/plan_expression_visitor.rs +++ b/common/planners/src/plan_expression_visitor.rs @@ -66,6 +66,20 @@ pub trait ExpressionVisitor: Sized { stack.push(RecursionProcessing::Call(arg)); } } + Expression::WindowFunction { + func, + partition_by, + order_by, + .. + } => { + stack.push(RecursionProcessing::Call(func)); + for part_by_expr in partition_by { + stack.push(RecursionProcessing::Call(part_by_expr)); + } + for order_by_expr in order_by { + stack.push(RecursionProcessing::Call(order_by_expr)); + } + } Expression::Cast { expr, .. } => { stack.push(RecursionProcessing::Call(expr)); } diff --git a/common/planners/src/plan_node_rewriter.rs b/common/planners/src/plan_node_rewriter.rs index 5d520bdb021d..d2f2265d7a3f 100644 --- a/common/planners/src/plan_node_rewriter.rs +++ b/common/planners/src/plan_node_rewriter.rs @@ -639,6 +639,8 @@ impl RewriteHelper { } } + Expression::WindowFunction { .. } => todo!("not figure out yet @doki"), + Expression::Alias(alias, plan) => { if data.inside_aliases.contains(alias) { return Result::Err(ErrorCode::SyntaxException(format!( @@ -730,6 +732,17 @@ impl RewriteHelper { } Expression::ScalarFunction { args, .. } => args.clone(), Expression::AggregateFunction { args, .. } => args.clone(), + Expression::WindowFunction { + func, + partition_by, + order_by, + .. + } => { + let mut v = vec![func.as_ref().clone()]; + v.extend(partition_by.clone()); + v.extend(order_by.clone()); + v + } Expression::Wildcard => vec![], Expression::Sort { expr, .. } => vec![expr.as_ref().clone()], Expression::Cast { expr, .. } => vec![expr.as_ref().clone()], @@ -768,6 +781,22 @@ impl RewriteHelper { } v } + Expression::WindowFunction { + func, + partition_by, + order_by, + .. + } => { + let mut v = vec![]; + v.append(&mut Self::expression_plan_columns(func)?); + for part_by_expr in partition_by { + v.append(&mut Self::expression_plan_columns(part_by_expr)?) + } + for order_by_expr in order_by { + v.append(&mut Self::expression_plan_columns(order_by_expr)?) + } + v + } Expression::Wildcard => vec![], Expression::Sort { expr, .. } => Self::expression_plan_columns(expr)?, Expression::Cast { expr, .. } => Self::expression_plan_columns(expr)?, diff --git a/query/src/sql/statements/analyzer_expr.rs b/query/src/sql/statements/analyzer_expr.rs index 38947d655dfa..07bfd0468283 100644 --- a/query/src/sql/statements/analyzer_expr.rs +++ b/query/src/sql/statements/analyzer_expr.rs @@ -35,6 +35,7 @@ use sqlparser::ast::Ident; use sqlparser::ast::Query; use sqlparser::ast::UnaryOperator; use sqlparser::ast::Value; +use sqlparser::ast::WindowSpec; use crate::procedures::ContextFunction; use crate::sessions::QueryContext; @@ -128,6 +129,29 @@ impl ExpressionAnalyzer { } fn analyze_function(&self, info: &FunctionExprInfo, args: &mut Vec) -> Result<()> { + // window function partition by and order by args + let mut partition_by_args = vec![]; + let mut order_by_args = vec![]; + if let Some(window_spec) = &info.over { + let order_by_args_count = window_spec.order_by.len(); + let partition_by_args_count = window_spec.partition_by.len(); + + for i in 0..partition_by_args_count + order_by_args_count { + match args.pop() { + None => { + return Err(ErrorCode::LogicalError("It's a bug.")); + } + Some(arg) => { + if i < order_by_args_count { + order_by_args.insert(0, arg); + } else { + partition_by_args.insert(0, arg); + } + } + } + } + } + let mut arguments = Vec::with_capacity(info.args_count); for _ in 0..info.args_count { match args.pop() { @@ -140,16 +164,23 @@ impl ExpressionAnalyzer { } } - args.push( - match AggregateFunctionFactory::instance().check(&info.name) { - true => self.aggr_function(info, &arguments), - false => match info.kind { - OperatorKind::Unary => Self::unary_function(info, &arguments), - OperatorKind::Binary => Self::binary_function(info, &arguments), - OperatorKind::Other => self.other_function(info, &arguments), - }, - }?, - ); + let func_expr = match AggregateFunctionFactory::instance().check(&info.name) { + true => self.aggr_function(info, &arguments), + false => match info.kind { + OperatorKind::Unary => Self::unary_function(info, &arguments), + OperatorKind::Binary => Self::binary_function(info, &arguments), + OperatorKind::Other => self.other_function(info, &arguments), + }, + }?; + + let func_expr = match &info.over { + Some(window_spec) => { + self.window_function(func_expr, partition_by_args, order_by_args, window_spec)? + } + None => func_expr, + }; + + args.push(func_expr); Ok(()) } @@ -242,6 +273,37 @@ impl ExpressionAnalyzer { } } + fn window_function( + &self, + func: Expression, + partition_by: Vec, + order_by_exprs: Vec, + window_spec: &WindowSpec, + ) -> Result { + let order_by: Vec = order_by_exprs + .into_iter() + .zip(window_spec.order_by.clone()) + .map(|(expr_sort_on, parser_sort_expr)| Expression::Sort { + expr: Box::new(expr_sort_on.clone()), + asc: parser_sort_expr.asc.unwrap_or(true), + nulls_first: parser_sort_expr.nulls_first.unwrap_or(true), + origin_expr: Box::new(expr_sort_on), + }) + .collect(); + + let window_frame = window_spec + .window_frame + .clone() + .map(|wf| wf.try_into().unwrap()); + + Ok(Expression::WindowFunction { + func: Box::new(func), + partition_by, + order_by, + window_frame, + }) + } + fn analyze_identifier(&self, ident: &Ident, arguments: &mut Vec) -> Result<()> { let column_name = ident.clone().value; arguments.push(Expression::Column(column_name)); @@ -418,6 +480,7 @@ struct FunctionExprInfo { args_count: usize, kind: OperatorKind, parameters: Vec, + over: Option, } struct InListInfo { @@ -447,6 +510,7 @@ impl ExprRPNItem { args_count, kind: OperatorKind::Other, parameters: Vec::new(), + over: None, }) } @@ -457,6 +521,7 @@ impl ExprRPNItem { args_count: 2, kind: OperatorKind::Binary, parameters: Vec::new(), + over: None, }) } @@ -467,6 +532,7 @@ impl ExprRPNItem { args_count: 1, kind: OperatorKind::Unary, parameters: Vec::new(), + over: None, }) } } @@ -532,6 +598,7 @@ impl ExprRPNBuilder { args_count: function.args.len(), kind: OperatorKind::Other, parameters: function.params.to_owned(), + over: function.over.clone(), })); } Expr::Cast { data_type, .. } => { diff --git a/query/src/sql/statements/query/query_normalizer.rs b/query/src/sql/statements/query/query_normalizer.rs index 487a93f2062e..e0271a30ec13 100644 --- a/query/src/sql/statements/query/query_normalizer.rs +++ b/query/src/sql/statements/query/query_normalizer.rs @@ -21,6 +21,7 @@ use common_planners::extract_aliases; use common_planners::find_aggregate_exprs_in_expr; use common_planners::resolve_aliases_to_exprs; use common_planners::Expression; +use common_tracing::tracing::log::debug; use sqlparser::ast::Expr; use sqlparser::ast::OffsetRows; use sqlparser::ast::SelectItem; @@ -64,26 +65,32 @@ impl QueryNormalizer { if let Err(cause) = self.visit_filter(query).await { return Err(cause.add_message_back(" (while in analyze select filter)")); } + debug!("ir after analyzing filter:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_projection(query).await { return Err(cause.add_message_back(" (while in analyze select projection)")); } + debug!("ir after analyzing projection:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_group_by(query).await { return Err(cause.add_message_back(" (while in analyze select group by)")); } + debug!("ir after analyzing group by:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_having(query).await { return Err(cause.add_message_back(" (while in analyze select having)")); } + debug!("ir after analyzing having:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_order_by(query).await { return Err(cause.add_message_back(" (while in analyze select order by)")); } + debug!("ir after analyzing order by:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_limit(query).await { return Err(cause.add_message_back(" (while in analyze select limit)")); } + debug!("ir after analyzing limit:\n{:?}", self.query_ast_ir); Ok(self.query_ast_ir) } @@ -102,7 +109,15 @@ impl QueryNormalizer { self.aliases_map = extract_aliases(&projection_expressions); for projection_expression in &projection_expressions { - self.add_aggregate_function(projection_expression)?; + match projection_expression { + Expression::WindowFunction { .. } => continue, + Expression::Alias(_, expr) + if matches!(expr.as_ref(), Expression::WindowFunction { .. }) => + { + continue + } + _ => self.add_aggregate_function(projection_expression)?, + } } self.query_ast_ir.projection_expressions = projection_expressions; diff --git a/query/src/sql/statements/statement_select.rs b/query/src/sql/statements/statement_select.rs index 617f21592dac..d7d34f1fbdff 100644 --- a/query/src/sql/statements/statement_select.rs +++ b/query/src/sql/statements/statement_select.rs @@ -24,6 +24,7 @@ use common_planners::find_aggregate_exprs_in_expr; use common_planners::rebase_expr; use common_planners::Expression; use common_tracing::tracing; +use common_tracing::tracing::log::debug; use sqlparser::ast::Expr; use sqlparser::ast::Offset; use sqlparser::ast::OrderByExpr; @@ -64,12 +65,18 @@ impl AnalyzableStatement for DfQueryStatement { let mut joined_schema = analyzer.analyze(self).await?; let mut ir = QueryNormalizer::normalize(ctx.clone(), self).await?; - - let has_aggregation = !find_aggregate_exprs(&ir.projection_expressions).is_empty(); + debug!("QueryASTIR after normalize:\n{:?}", ir); QualifiedRewriter::rewrite(&joined_schema, ctx.clone(), &mut ir)?; + debug!("QueryASTIR after qualified rewrite:\n{:?}", ir); + + let has_aggregation = !find_aggregate_exprs(&ir.projection_expressions).is_empty(); QueryCollectPushDowns::collect_extras(&mut ir, &mut joined_schema, has_aggregation)?; + debug!("QueryASTIR after push downs:\n{:?}", ir); + let analyze_state = self.analyze_query(ir).await?; + debug!("QueryAnalyzeState:\n{:?}", analyze_state); + self.check_and_finalize(joined_schema, analyze_state, ctx) .await } @@ -200,6 +207,10 @@ impl DfQueryStatement { ) -> Result { let dry_run_res = Self::verify_with_dry_run(&schema, &state)?; state.finalize_schema = dry_run_res.schema().clone(); + debug!( + "QueryAnalyzeState finalized schema:\n{}", + state.finalize_schema + ); let mut tables_desc = schema.take_tables_desc(); From 0d576220d53d0af7930e84d1f6ba808d4fa1e096 Mon Sep 17 00:00:00 2001 From: doki Date: Fri, 22 Apr 2022 17:04:05 +0800 Subject: [PATCH 02/18] update structure of window function expression --- common/planners/src/plan_expression.rs | 21 +++- common/planners/src/plan_expression_common.rs | 15 ++- .../planners/src/plan_expression_rewriter.rs | 35 ++++-- .../planners/src/plan_expression_visitor.rs | 6 +- common/planners/src/plan_node_rewriter.rs | 10 +- query/src/sql/statements/analyzer_expr.rs | 113 ++++++++++-------- .../src/sql/statements/analyzer_statement.rs | 1 + .../sql/statements/query/query_normalizer.rs | 17 +-- query/src/sql/statements/statement_select.rs | 8 +- 9 files changed, 130 insertions(+), 96 deletions(-) diff --git a/common/planners/src/plan_expression.rs b/common/planners/src/plan_expression.rs index 1cc19a411931..5fd2ca6e1449 100644 --- a/common/planners/src/plan_expression.rs +++ b/common/planners/src/plan_expression.rs @@ -98,12 +98,14 @@ pub enum Expression { /// WindowFunction WindowFunction { /// operation performed - func: Box, + op: String, + /// arguments + args: Vec, /// partition by partition_by: Vec, /// order by order_by: Vec, - /// Window frame + /// window frame window_frame: Option, }, @@ -421,12 +423,22 @@ impl fmt::Debug for Expression { } Expression::WindowFunction { - func, + op, + args, partition_by, order_by, window_frame, } => { - write!(f, "{:?}", func)?; + write!(f, "{}(", op)?; + for (i, arg) in args.iter().enumerate() { + if i < args.len() { + write!(f, "{},", arg.column_name())?; + } else { + write!(f, "{}", arg.column_name())?; + } + } + write!(f, ")")?; + write!(f, " OVER( ")?; if !partition_by.is_empty() { write!(f, "PARTITION BY {:?} ", partition_by)?; @@ -442,6 +454,7 @@ impl fmt::Debug for Expression { )?; } write!(f, ")")?; + Ok(()) } diff --git a/common/planners/src/plan_expression_common.rs b/common/planners/src/plan_expression_common.rs index 100ef60088ec..7b4b32d98e33 100644 --- a/common/planners/src/plan_expression_common.rs +++ b/common/planners/src/plan_expression_common.rs @@ -75,6 +75,12 @@ pub fn expand_aggregate_arg_exprs(exprs: &[Expression]) -> Vec { res } +pub fn find_window_exprs(exprs: &[Expression]) -> Vec { + find_exprs_in_exprs(exprs, &|nested_expr| { + matches!(nested_expr, Expression::WindowFunction { .. }) + }) +} + /// Collect all deeply nested `Expression::Column`'s. They are returned in order of /// appearance (depth first), with duplicates omitted. pub fn find_column_exprs(exprs: &[Expression]) -> Vec { @@ -304,12 +310,17 @@ where F: Fn(&Expression) -> Result> { }), Expression::WindowFunction { - func, + op, + args, partition_by, order_by, window_frame, } => Ok(Expression::WindowFunction { - func: Box::new(clone_with_replacement(func, replacement_fn)?), + op: op.clone(), + args: args + .iter() + .map(|e| clone_with_replacement(e, replacement_fn)) + .collect::>>()?, partition_by: partition_by .iter() .map(|e| clone_with_replacement(e, replacement_fn)) diff --git a/common/planners/src/plan_expression_rewriter.rs b/common/planners/src/plan_expression_rewriter.rs index 3233bd587b17..4b236883591e 100644 --- a/common/planners/src/plan_expression_rewriter.rs +++ b/common/planners/src/plan_expression_rewriter.rs @@ -127,13 +127,15 @@ pub trait ExpressionRewriter: Sized { fn mutate_window_function( &mut self, - func: Expression, + op: String, + args: Vec, partition_by: Vec, order_by: Vec, window_frame: Option, ) -> Result { Ok(Expression::WindowFunction { - func: Box::new(func), + op, + args, partition_by, order_by, window_frame, @@ -303,17 +305,23 @@ impl ExpressionVisitor for ExpressionRewriteVisitor { Ok(self) } Expression::WindowFunction { - func: _, + op, + args, partition_by, order_by, window_frame, } => { - let new_func = self.stack.pop().unwrap(); - let mut new_partition_by = Vec::with_capacity(partition_by.len()); - for (i, _) in partition_by.iter().enumerate() { + let mut new_order_by = Vec::with_capacity(order_by.len()); + for i in 0..partition_by.len() + order_by.len() { match self.stack.pop() { - Some(new_partition_by_expr) => new_partition_by.push(new_partition_by_expr), + Some(expr) => { + if i < order_by.len() { + new_order_by.push(expr); + } else { + new_partition_by.push(expr); + } + } None => { return Err(ErrorCode::LogicalError(format!( "WindowFunction expects {} partition by arguments, actual {}", @@ -324,14 +332,14 @@ impl ExpressionVisitor for ExpressionRewriteVisitor { } } - let mut new_order_by = Vec::with_capacity(order_by.len()); - for (i, _) in order_by.iter().enumerate() { + let mut new_args = Vec::with_capacity(args.len()); + for i in 0..args.len() { match self.stack.pop() { - Some(new_order_by_expr) => new_order_by.push(new_order_by_expr), + Some(expr) => new_args.push(expr), None => { return Err(ErrorCode::LogicalError(format!( - "WindowFunction expects {} order by arguments, actual {}", - order_by.len(), + "WindowFunction expects {} partition by arguments, actual {}", + partition_by.len(), i ))) } @@ -339,7 +347,8 @@ impl ExpressionVisitor for ExpressionRewriteVisitor { } let new_expr = self.inner.mutate_window_function( - new_func, + op.clone(), + new_args, new_partition_by, new_order_by, window_frame.to_owned(), diff --git a/common/planners/src/plan_expression_visitor.rs b/common/planners/src/plan_expression_visitor.rs index 6d0599c5c795..60b079aa6316 100644 --- a/common/planners/src/plan_expression_visitor.rs +++ b/common/planners/src/plan_expression_visitor.rs @@ -67,12 +67,14 @@ pub trait ExpressionVisitor: Sized { } } Expression::WindowFunction { - func, + args, partition_by, order_by, .. } => { - stack.push(RecursionProcessing::Call(func)); + for arg_exppr in args { + stack.push(RecursionProcessing::Call(arg_exppr)); + } for part_by_expr in partition_by { stack.push(RecursionProcessing::Call(part_by_expr)); } diff --git a/common/planners/src/plan_node_rewriter.rs b/common/planners/src/plan_node_rewriter.rs index d2f2265d7a3f..a0c31362aad1 100644 --- a/common/planners/src/plan_node_rewriter.rs +++ b/common/planners/src/plan_node_rewriter.rs @@ -733,12 +733,12 @@ impl RewriteHelper { Expression::ScalarFunction { args, .. } => args.clone(), Expression::AggregateFunction { args, .. } => args.clone(), Expression::WindowFunction { - func, + args, partition_by, order_by, .. } => { - let mut v = vec![func.as_ref().clone()]; + let mut v = args.clone(); v.extend(partition_by.clone()); v.extend(order_by.clone()); v @@ -782,13 +782,15 @@ impl RewriteHelper { v } Expression::WindowFunction { - func, + args, partition_by, order_by, .. } => { let mut v = vec![]; - v.append(&mut Self::expression_plan_columns(func)?); + for arg_expr in args { + v.append(&mut Self::expression_plan_columns(arg_expr)?) + } for part_by_expr in partition_by { v.append(&mut Self::expression_plan_columns(part_by_expr)?) } diff --git a/query/src/sql/statements/analyzer_expr.rs b/query/src/sql/statements/analyzer_expr.rs index 07bfd0468283..379eea8a71ad 100644 --- a/query/src/sql/statements/analyzer_expr.rs +++ b/query/src/sql/statements/analyzer_expr.rs @@ -129,58 +129,33 @@ impl ExpressionAnalyzer { } fn analyze_function(&self, info: &FunctionExprInfo, args: &mut Vec) -> Result<()> { - // window function partition by and order by args - let mut partition_by_args = vec![]; - let mut order_by_args = vec![]; - if let Some(window_spec) = &info.over { - let order_by_args_count = window_spec.order_by.len(); - let partition_by_args_count = window_spec.partition_by.len(); - - for i in 0..partition_by_args_count + order_by_args_count { - match args.pop() { - None => { - return Err(ErrorCode::LogicalError("It's a bug.")); - } - Some(arg) => { - if i < order_by_args_count { - order_by_args.insert(0, arg); - } else { - partition_by_args.insert(0, arg); + let func = match &info.over { + Some(_) => self.window_function(info, args)?, + None => { + let mut arguments = Vec::with_capacity(info.args_count); + for _ in 0..info.args_count { + match args.pop() { + None => { + return Err(ErrorCode::LogicalError("It's a bug.")); + } + Some(arg) => { + arguments.insert(0, arg); } } } - } - } - - let mut arguments = Vec::with_capacity(info.args_count); - for _ in 0..info.args_count { - match args.pop() { - None => { - return Err(ErrorCode::LogicalError("It's a bug.")); - } - Some(arg) => { - arguments.insert(0, arg); - } - } - } - let func_expr = match AggregateFunctionFactory::instance().check(&info.name) { - true => self.aggr_function(info, &arguments), - false => match info.kind { - OperatorKind::Unary => Self::unary_function(info, &arguments), - OperatorKind::Binary => Self::binary_function(info, &arguments), - OperatorKind::Other => self.other_function(info, &arguments), - }, - }?; - - let func_expr = match &info.over { - Some(window_spec) => { - self.window_function(func_expr, partition_by_args, order_by_args, window_spec)? + match AggregateFunctionFactory::instance().check(&info.name) { + true => self.aggr_function(info, &arguments), + false => match info.kind { + OperatorKind::Unary => Self::unary_function(info, &arguments), + OperatorKind::Binary => Self::binary_function(info, &arguments), + OperatorKind::Other => self.other_function(info, &arguments), + }, + }? } - None => func_expr, }; - args.push(func_expr); + args.push(func); Ok(()) } @@ -275,12 +250,47 @@ impl ExpressionAnalyzer { fn window_function( &self, - func: Expression, - partition_by: Vec, - order_by_exprs: Vec, - window_spec: &WindowSpec, + info: &FunctionExprInfo, + args: &mut Vec, ) -> Result { - let order_by: Vec = order_by_exprs + // window function partition by and order by args + let mut partition_by = vec![]; + let mut order_by_expr = vec![]; + if let Some(window_spec) = &info.over { + let order_by_args_count = window_spec.order_by.len(); + let partition_by_args_count = window_spec.partition_by.len(); + + for i in 0..partition_by_args_count + order_by_args_count { + match args.pop() { + None => { + return Err(ErrorCode::LogicalError("It's a bug.")); + } + Some(arg) => { + if i < order_by_args_count { + order_by_expr.insert(0, arg); + } else { + partition_by.insert(0, arg); + } + } + } + } + } + + let mut arguments = Vec::with_capacity(info.args_count); + for _ in 0..info.args_count { + match args.pop() { + None => { + return Err(ErrorCode::LogicalError("It's a bug.")); + } + Some(arg) => { + arguments.insert(0, arg); + } + } + } + + let window_spec = info.over.as_ref().unwrap(); + + let order_by: Vec = order_by_expr .into_iter() .zip(window_spec.order_by.clone()) .map(|(expr_sort_on, parser_sort_expr)| Expression::Sort { @@ -297,7 +307,8 @@ impl ExpressionAnalyzer { .map(|wf| wf.try_into().unwrap()); Ok(Expression::WindowFunction { - func: Box::new(func), + op: info.name.clone(), + args: arguments, partition_by, order_by, window_frame, diff --git a/query/src/sql/statements/analyzer_statement.rs b/query/src/sql/statements/analyzer_statement.rs index d998f730c5aa..fbf8f141a67e 100644 --- a/query/src/sql/statements/analyzer_statement.rs +++ b/query/src/sql/statements/analyzer_statement.rs @@ -54,6 +54,7 @@ pub struct QueryAnalyzeState { pub aggregate_expressions: Vec, pub before_group_by_expressions: Vec, + // pub window_expressions: Vec, pub limit: Option, pub offset: Option, diff --git a/query/src/sql/statements/query/query_normalizer.rs b/query/src/sql/statements/query/query_normalizer.rs index e0271a30ec13..487a93f2062e 100644 --- a/query/src/sql/statements/query/query_normalizer.rs +++ b/query/src/sql/statements/query/query_normalizer.rs @@ -21,7 +21,6 @@ use common_planners::extract_aliases; use common_planners::find_aggregate_exprs_in_expr; use common_planners::resolve_aliases_to_exprs; use common_planners::Expression; -use common_tracing::tracing::log::debug; use sqlparser::ast::Expr; use sqlparser::ast::OffsetRows; use sqlparser::ast::SelectItem; @@ -65,32 +64,26 @@ impl QueryNormalizer { if let Err(cause) = self.visit_filter(query).await { return Err(cause.add_message_back(" (while in analyze select filter)")); } - debug!("ir after analyzing filter:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_projection(query).await { return Err(cause.add_message_back(" (while in analyze select projection)")); } - debug!("ir after analyzing projection:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_group_by(query).await { return Err(cause.add_message_back(" (while in analyze select group by)")); } - debug!("ir after analyzing group by:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_having(query).await { return Err(cause.add_message_back(" (while in analyze select having)")); } - debug!("ir after analyzing having:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_order_by(query).await { return Err(cause.add_message_back(" (while in analyze select order by)")); } - debug!("ir after analyzing order by:\n{:?}", self.query_ast_ir); if let Err(cause) = self.analyze_limit(query).await { return Err(cause.add_message_back(" (while in analyze select limit)")); } - debug!("ir after analyzing limit:\n{:?}", self.query_ast_ir); Ok(self.query_ast_ir) } @@ -109,15 +102,7 @@ impl QueryNormalizer { self.aliases_map = extract_aliases(&projection_expressions); for projection_expression in &projection_expressions { - match projection_expression { - Expression::WindowFunction { .. } => continue, - Expression::Alias(_, expr) - if matches!(expr.as_ref(), Expression::WindowFunction { .. }) => - { - continue - } - _ => self.add_aggregate_function(projection_expression)?, - } + self.add_aggregate_function(projection_expression)?; } self.query_ast_ir.projection_expressions = projection_expressions; diff --git a/query/src/sql/statements/statement_select.rs b/query/src/sql/statements/statement_select.rs index d7d34f1fbdff..97a623e2b4dc 100644 --- a/query/src/sql/statements/statement_select.rs +++ b/query/src/sql/statements/statement_select.rs @@ -65,17 +65,17 @@ impl AnalyzableStatement for DfQueryStatement { let mut joined_schema = analyzer.analyze(self).await?; let mut ir = QueryNormalizer::normalize(ctx.clone(), self).await?; - debug!("QueryASTIR after normalize:\n{:?}", ir); + tracing::debug!("\nQueryASTIR after normalize:\n{:?}", ir); QualifiedRewriter::rewrite(&joined_schema, ctx.clone(), &mut ir)?; - debug!("QueryASTIR after qualified rewrite:\n{:?}", ir); + tracing::debug!("\nQueryASTIR after qualified rewrite:\n{:?}", ir); let has_aggregation = !find_aggregate_exprs(&ir.projection_expressions).is_empty(); QueryCollectPushDowns::collect_extras(&mut ir, &mut joined_schema, has_aggregation)?; - debug!("QueryASTIR after push downs:\n{:?}", ir); + tracing::debug!("\nQueryASTIR after push downs:\n{:?}", ir); let analyze_state = self.analyze_query(ir).await?; - debug!("QueryAnalyzeState:\n{:?}", analyze_state); + tracing::debug!("\nQueryAnalyzeState:\n{:?}", analyze_state); self.check_and_finalize(joined_schema, analyze_state, ctx) .await From 962f50f4d768bca2ed377a7470d26112861af8fa Mon Sep 17 00:00:00 2001 From: doki Date: Sat, 23 Apr 2022 14:29:27 +0800 Subject: [PATCH 03/18] window plan --- common/planners/src/lib.rs | 3 ++ common/planners/src/plan_expression.rs | 2 +- common/planners/src/plan_expression_chain.rs | 4 +- common/planners/src/plan_expression_common.rs | 19 ++++++--- common/planners/src/plan_node.rs | 5 +++ common/planners/src/plan_node_builder.rs | 17 ++++++++ .../planners/src/plan_node_display_indent.rs | 1 + common/planners/src/plan_node_rewriter.rs | 10 +++++ common/planners/src/plan_node_visitor.rs | 7 ++++ common/planners/src/plan_window_aggr.rs | 40 +++++++++++++++++++ query/src/interpreters/interpreter_select.rs | 2 + .../servers/mysql/mysql_interactive_worker.rs | 7 +++- query/src/sql/plan_parser.rs | 19 ++++++++- .../src/sql/statements/analyzer_statement.rs | 8 +++- query/src/sql/statements/statement_select.rs | 6 +++ 15 files changed, 136 insertions(+), 14 deletions(-) create mode 100644 common/planners/src/plan_window_aggr.rs diff --git a/common/planners/src/lib.rs b/common/planners/src/lib.rs index 1293c17a230c..a119b105aeca 100644 --- a/common/planners/src/lib.rs +++ b/common/planners/src/lib.rs @@ -98,6 +98,7 @@ mod plan_user_udf_drop; mod plan_view_alter; mod plan_view_create; mod plan_view_drop; +mod plan_window_aggr; pub use plan_aggregator_final::AggregatorFinalPlan; pub use plan_aggregator_partial::AggregatorPartialPlan; @@ -125,6 +126,8 @@ pub use plan_expression_common::extract_aliases; pub use plan_expression_common::find_aggregate_exprs; pub use plan_expression_common::find_aggregate_exprs_in_expr; pub use plan_expression_common::find_columns_not_satisfy_exprs; +pub use plan_expression_common::find_window_exprs; +pub use plan_expression_common::find_window_exprs_in_expr; pub use plan_expression_common::rebase_expr; pub use plan_expression_common::rebase_expr_from_input; pub use plan_expression_common::resolve_aliases_to_exprs; diff --git a/common/planners/src/plan_expression.rs b/common/planners/src/plan_expression.rs index 5fd2ca6e1449..a9c0e31c65d6 100644 --- a/common/planners/src/plan_expression.rs +++ b/common/planners/src/plan_expression.rs @@ -431,7 +431,7 @@ impl fmt::Debug for Expression { } => { write!(f, "{}(", op)?; for (i, arg) in args.iter().enumerate() { - if i < args.len() { + if i < args.len() - 1 { write!(f, "{},", arg.column_name())?; } else { write!(f, "{}", arg.column_name())?; diff --git a/common/planners/src/plan_expression_chain.rs b/common/planners/src/plan_expression_chain.rs index a1d05a6c0f48..acfc69565ecb 100644 --- a/common/planners/src/plan_expression_chain.rs +++ b/common/planners/src/plan_expression_chain.rs @@ -192,9 +192,7 @@ impl ExpressionChain { )); } - Expression::WindowFunction { .. } => { - todo!("not figure out yet @doki") - } + Expression::WindowFunction { .. } => {} Expression::Wildcard | Expression::Sort { .. } => {} diff --git a/common/planners/src/plan_expression_common.rs b/common/planners/src/plan_expression_common.rs index 7b4b32d98e33..2619c1649772 100644 --- a/common/planners/src/plan_expression_common.rs +++ b/common/planners/src/plan_expression_common.rs @@ -51,6 +51,19 @@ pub fn find_aggregate_exprs_in_expr(expr: &Expression) -> Vec { }) } +/// Collect all deeply nested `Expression::WindowFunction`. +pub fn find_window_exprs(exprs: &[Expression]) -> Vec { + find_exprs_in_exprs(exprs, &|nested_expr| { + matches!(nested_expr, Expression::WindowFunction { .. }) + }) +} + +pub fn find_window_exprs_in_expr(expr: &Expression) -> Vec { + find_exprs_in_expr(expr, &|nest_exprs| { + matches!(nest_exprs, Expression::WindowFunction { .. }) + }) +} + /// Collect all arguments from aggregation function and append to this exprs /// [ColumnExpr(b), Aggr(sum(a, b))] ---> [ColumnExpr(b), ColumnExpr(a)] @@ -75,12 +88,6 @@ pub fn expand_aggregate_arg_exprs(exprs: &[Expression]) -> Vec { res } -pub fn find_window_exprs(exprs: &[Expression]) -> Vec { - find_exprs_in_exprs(exprs, &|nested_expr| { - matches!(nested_expr, Expression::WindowFunction { .. }) - }) -} - /// Collect all deeply nested `Expression::Column`'s. They are returned in order of /// appearance (depth first), with duplicates omitted. pub fn find_column_exprs(exprs: &[Expression]) -> Vec { diff --git a/common/planners/src/plan_node.rs b/common/planners/src/plan_node.rs index 8d2b19963012..384bf00ecc75 100644 --- a/common/planners/src/plan_node.rs +++ b/common/planners/src/plan_node.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use common_datavalues::DataSchemaRef; +use crate::plan_window_aggr::WindowAggrPlan; use crate::AggregatorFinalPlan; use crate::AggregatorPartialPlan; use crate::AlterUserPlan; @@ -85,6 +86,7 @@ pub enum PlanNode { AggregatorFinal(AggregatorFinalPlan), Filter(FilterPlan), Having(HavingPlan), + WindowAggr(WindowAggrPlan), Sort(SortPlan), Limit(LimitPlan), LimitBy(LimitByPlan), @@ -184,6 +186,7 @@ impl PlanNode { PlanNode::AggregatorFinal(v) => v.schema(), PlanNode::Filter(v) => v.schema(), PlanNode::Having(v) => v.schema(), + PlanNode::WindowAggr(v) => v.schema(), PlanNode::Limit(v) => v.schema(), PlanNode::LimitBy(v) => v.schema(), PlanNode::ReadSource(v) => v.schema(), @@ -282,6 +285,7 @@ impl PlanNode { PlanNode::AggregatorFinal(_) => "AggregatorFinalPlan", PlanNode::Filter(_) => "FilterPlan", PlanNode::Having(_) => "HavingPlan", + PlanNode::WindowAggr(_) => "WindowAggrPlan", PlanNode::Limit(_) => "LimitPlan", PlanNode::LimitBy(_) => "LimitByPlan", PlanNode::ReadSource(_) => "ReadSourcePlan", @@ -377,6 +381,7 @@ impl PlanNode { PlanNode::AggregatorFinal(v) => vec![v.input.clone()], PlanNode::Filter(v) => vec![v.input.clone()], PlanNode::Having(v) => vec![v.input.clone()], + PlanNode::WindowAggr(v) => vec![v.input.clone()], PlanNode::Limit(v) => vec![v.input.clone()], PlanNode::Explain(v) => vec![v.input.clone()], PlanNode::Select(v) => vec![v.input.clone()], diff --git a/common/planners/src/plan_node_builder.rs b/common/planners/src/plan_node_builder.rs index fbdea7b9ae1d..6420f18f1513 100644 --- a/common/planners/src/plan_node_builder.rs +++ b/common/planners/src/plan_node_builder.rs @@ -20,6 +20,7 @@ use common_exception::Result; use crate::col; use crate::plan_subqueries_set::SubQueriesSetPlan; +use crate::plan_window_aggr::WindowAggrPlan; use crate::validate_expression; use crate::AggregatorFinalPlan; use crate::AggregatorPartialPlan; @@ -214,6 +215,22 @@ impl PlanBuilder { }))) } + /// Apply a window function + pub fn window_aggr(&self, expr: Expression) -> Result { + let window_func = expr.clone(); + let input = self.wrap_subquery_plan(&[expr])?; + let input_schema = input.schema(); + let mut input_fields = input_schema.fields().to_owned(); + let window_field = window_func.to_data_field(&input_schema).unwrap(); + input_fields.push(window_field); + let schema = Arc::new(DataSchema::new(input_fields)); + Ok(Self::from(&PlanNode::WindowAggr(WindowAggrPlan { + window_func, + input, + schema, + }))) + } + pub fn sort(&self, exprs: &[Expression]) -> Result { Ok(Self::from(&PlanNode::Sort(SortPlan { order_by: exprs.to_vec(), diff --git a/common/planners/src/plan_node_display_indent.rs b/common/planners/src/plan_node_display_indent.rs index 3df392627a95..fa83a2c791d0 100644 --- a/common/planners/src/plan_node_display_indent.rs +++ b/common/planners/src/plan_node_display_indent.rs @@ -68,6 +68,7 @@ impl<'a> fmt::Display for PlanNodeIndentFormatDisplay<'a> { PlanNode::AggregatorFinal(plan) => Self::format_aggregator_final(f, plan), PlanNode::Filter(plan) => write!(f, "Filter: {:?}", plan.predicate), PlanNode::Having(plan) => write!(f, "Having: {:?}", plan.predicate), + PlanNode::WindowAggr(plan) => write!(f, "WindowAggr: {:?}", plan.window_func), PlanNode::Sort(plan) => Self::format_sort(f, plan), PlanNode::Limit(plan) => Self::format_limit(f, plan), PlanNode::SubQueryExpression(plan) => Self::format_subquery_expr(f, plan), diff --git a/common/planners/src/plan_node_rewriter.rs b/common/planners/src/plan_node_rewriter.rs index a0c31362aad1..33cfd68530e5 100644 --- a/common/planners/src/plan_node_rewriter.rs +++ b/common/planners/src/plan_node_rewriter.rs @@ -23,6 +23,7 @@ use common_exception::Result; use crate::plan_broadcast::BroadcastPlan; use crate::plan_subqueries_set::SubQueriesSetPlan; +use crate::plan_window_aggr::WindowAggrPlan; use crate::AggregatorFinalPlan; use crate::AggregatorPartialPlan; use crate::AlterUserPlan; @@ -113,6 +114,7 @@ pub trait PlanRewriter: Sized { PlanNode::Broadcast(plan) => self.rewrite_broadcast(plan), PlanNode::Remote(plan) => self.rewrite_remote(plan), PlanNode::Having(plan) => self.rewrite_having(plan), + PlanNode::WindowAggr(plan) => self.rewrite_window_aggr(plan), PlanNode::Expression(plan) => self.rewrite_expression(plan), PlanNode::Sort(plan) => self.rewrite_sort(plan), PlanNode::Limit(plan) => self.rewrite_limit(plan), @@ -311,6 +313,14 @@ pub trait PlanRewriter: Sized { PlanBuilder::from(&new_input).having(new_predicate)?.build() } + fn rewrite_window_aggr(&mut self, plan: &WindowAggrPlan) -> Result { + let new_input = self.rewrite_plan_node(plan.input.as_ref())?; + let new_window_func = self.rewrite_expr(&new_input.schema(), &plan.window_func)?; + PlanBuilder::from(&new_input) + .window_aggr(new_window_func)? + .build() + } + fn rewrite_sort(&mut self, plan: &SortPlan) -> Result { let new_input = self.rewrite_plan_node(plan.input.as_ref())?; let new_order_by = self.rewrite_exprs(&new_input.schema(), &plan.order_by)?; diff --git a/common/planners/src/plan_node_visitor.rs b/common/planners/src/plan_node_visitor.rs index 5db3148669be..62ee4cee96a7 100644 --- a/common/planners/src/plan_node_visitor.rs +++ b/common/planners/src/plan_node_visitor.rs @@ -16,6 +16,7 @@ use common_exception::Result; use crate::plan_broadcast::BroadcastPlan; use crate::plan_subqueries_set::SubQueriesSetPlan; +use crate::plan_window_aggr::WindowAggrPlan; use crate::AggregatorFinalPlan; use crate::AggregatorPartialPlan; use crate::AlterUserPlan; @@ -126,6 +127,7 @@ pub trait PlanVisitor { PlanNode::Broadcast(plan) => self.visit_broadcast(plan), PlanNode::Remote(plan) => self.visit_remote(plan), PlanNode::Having(plan) => self.visit_having(plan), + PlanNode::WindowAggr(plan) => self.visit_window_aggr(plan), PlanNode::Expression(plan) => self.visit_expression(plan), PlanNode::Limit(plan) => self.visit_limit(plan), PlanNode::LimitBy(plan) => self.visit_limit_by(plan), @@ -288,6 +290,11 @@ pub trait PlanVisitor { self.visit_expr(&plan.predicate) } + fn visit_window_aggr(&mut self, plan: &WindowAggrPlan) -> Result<()> { + self.visit_plan_node(plan.input.as_ref())?; + self.visit_expr(&plan.window_func) + } + fn visit_sort(&mut self, plan: &SortPlan) -> Result<()> { self.visit_plan_node(plan.input.as_ref())?; self.visit_exprs(&plan.order_by) diff --git a/common/planners/src/plan_window_aggr.rs b/common/planners/src/plan_window_aggr.rs new file mode 100644 index 000000000000..96b16adadee4 --- /dev/null +++ b/common/planners/src/plan_window_aggr.rs @@ -0,0 +1,40 @@ +// Copyright 2022 Datafuse Labs. +// +// 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 std::sync::Arc; + +use common_datavalues::DataSchemaRef; + +use crate::Expression; +use crate::PlanNode; + +#[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq)] +pub struct WindowAggrPlan { + /// The window function expression + pub window_func: Expression, + /// The incoming logical plan + pub input: Arc, + /// output schema + pub schema: DataSchemaRef, +} + +impl WindowAggrPlan { + pub fn schema(&self) -> DataSchemaRef { + self.schema.clone() + } + + pub fn set_input(&mut self, node: &PlanNode) { + self.input = Arc::new(node.clone()); + } +} diff --git a/query/src/interpreters/interpreter_select.rs b/query/src/interpreters/interpreter_select.rs index af5e3b9b4a14..2e50cb65991e 100644 --- a/query/src/interpreters/interpreter_select.rs +++ b/query/src/interpreters/interpreter_select.rs @@ -83,6 +83,8 @@ impl Interpreter for SelectInterpreter { return Ok(Box::pin(self.ctx.try_create_abortable(executor_stream)?)); } let optimized_plan = self.rewrite_plan()?; + tracing::debug!("\nget optimized plan:\n{:?}", optimized_plan); + plan_schedulers::schedule_query(&self.ctx, &optimized_plan).await } diff --git a/query/src/servers/mysql/mysql_interactive_worker.rs b/query/src/servers/mysql/mysql_interactive_worker.rs index 74e6e45429c5..8b664641d168 100644 --- a/query/src/servers/mysql/mysql_interactive_worker.rs +++ b/query/src/servers/mysql/mysql_interactive_worker.rs @@ -24,6 +24,7 @@ use common_exception::ToErrorCode; use common_io::prelude::*; use common_planners::PlanNode; use common_tracing::tracing; +use common_tracing::tracing::debug; use common_tracing::tracing::Instrument; use metrics::histogram; use opensrv_mysql::AsyncMysqlShim; @@ -274,6 +275,8 @@ impl InteractiveWorkerBase { let context = self.session.create_query_context().await?; context.attach_query_str(query); let (plan, hints) = PlanParser::parse_with_hint(query, context.clone()).await; + let plan = plan?; + debug!("\nget query plan:\n{:?}", plan); match hints .iter() @@ -305,11 +308,11 @@ impl InteractiveWorkerBase { #[tracing::instrument(level = "debug", skip(plan, context))] async fn exec_query( - plan: Result, + plan: PlanNode, context: &Arc, ) -> Result<(Vec, String)> { let instant = Instant::now(); - let interpreter = InterpreterFactory::get(context.clone(), plan?)?; + let interpreter = InterpreterFactory::get(context.clone(), plan)?; let query_result = context.try_spawn( async move { diff --git a/query/src/sql/plan_parser.rs b/query/src/sql/plan_parser.rs index 81ea7119cf67..2b6d90bfa799 100644 --- a/query/src/sql/plan_parser.rs +++ b/query/src/sql/plan_parser.rs @@ -78,7 +78,8 @@ impl PlanParser { let group_by = Self::build_group_by_plan(filter, data)?; let before_order = Self::build_before_order(group_by, data)?; let having = Self::build_having_plan(before_order, data)?; - let order_by = Self::build_order_by_plan(having, data)?; + let window = Self::build_window_plan(having, data)?; + let order_by = Self::build_order_by_plan(window, data)?; let projection = Self::build_projection_plan(order_by, data)?; let limit = Self::build_limit_plan(projection, data)?; @@ -173,6 +174,22 @@ impl PlanParser { } } + fn build_window_plan(plan: PlanNode, data: &QueryAnalyzeState) -> Result { + match data.window_expressions.is_empty() { + true => Ok(plan), + false => { + let exprs = data.window_expressions.to_vec(); + Ok(exprs.into_iter().fold(plan, |input, window_func| { + PlanBuilder::from(&input) + .window_aggr(window_func) + .unwrap() + .build() + .unwrap() + })) + } + } + } + fn build_order_by_plan(plan: PlanNode, data: &QueryAnalyzeState) -> Result { match data.order_by_expressions.is_empty() { true => Ok(plan), diff --git a/query/src/sql/statements/analyzer_statement.rs b/query/src/sql/statements/analyzer_statement.rs index fbf8f141a67e..bd04993564c1 100644 --- a/query/src/sql/statements/analyzer_statement.rs +++ b/query/src/sql/statements/analyzer_statement.rs @@ -54,7 +54,8 @@ pub struct QueryAnalyzeState { pub aggregate_expressions: Vec, pub before_group_by_expressions: Vec, - // pub window_expressions: Vec, + pub window_expressions: Vec, + pub limit: Option, pub offset: Option, @@ -87,6 +88,7 @@ impl Default for QueryAnalyzeState { group_by_expressions: vec![], aggregate_expressions: vec![], before_group_by_expressions: vec![], + window_expressions: vec![], limit: None, offset: None, relation: QueryRelation::None, @@ -115,6 +117,10 @@ impl Debug for QueryAnalyzeState { debug_struct.field("aggregate", &self.aggregate_expressions); } + if !self.window_expressions.is_empty() { + debug_struct.field("window_aggr", &self.window_expressions); + } + if !self.expressions.is_empty() { match self.order_by_expressions.is_empty() { true => debug_struct.field("before_projection", &self.expressions), diff --git a/query/src/sql/statements/statement_select.rs b/query/src/sql/statements/statement_select.rs index 97a623e2b4dc..5d9221117ea9 100644 --- a/query/src/sql/statements/statement_select.rs +++ b/query/src/sql/statements/statement_select.rs @@ -21,6 +21,7 @@ use common_exception::Result; use common_planners::expand_aggregate_arg_exprs; use common_planners::find_aggregate_exprs; use common_planners::find_aggregate_exprs_in_expr; +use common_planners::find_window_exprs_in_expr; use common_planners::rebase_expr; use common_planners::Expression; use common_tracing::tracing; @@ -180,6 +181,11 @@ impl DfQueryStatement { _ => state.add_expression(item), } + let window_exprs = find_window_exprs_in_expr(item); + if !window_exprs.is_empty() { + state.window_expressions.extend(window_exprs); + } + let rebased_expr = rebase_expr(item, &state.expressions)?; state.projection_expressions.push(rebased_expr); } From e180ef14c2cde213a0fb1a87ab4f37e837d1b2e7 Mon Sep 17 00:00:00 2001 From: doki Date: Sat, 23 Apr 2022 21:00:00 +0800 Subject: [PATCH 04/18] local commit --- common/planners/src/plan_expression.rs | 14 +++++--- common/planners/src/plan_expression_common.rs | 1 - query/src/sql/plan_parser.rs | 18 +++++++++++ .../src/sql/statements/query/query_ast_ir.rs | 32 +++++++++++++++++++ .../sql/statements/query/query_normalizer.rs | 13 ++++++++ query/src/sql/statements/statement_select.rs | 25 ++++++++++++--- 6 files changed, 93 insertions(+), 10 deletions(-) diff --git a/common/planners/src/plan_expression.rs b/common/planners/src/plan_expression.rs index a9c0e31c65d6..d82817b35101 100644 --- a/common/planners/src/plan_expression.rs +++ b/common/planners/src/plan_expression.rs @@ -439,17 +439,23 @@ impl fmt::Debug for Expression { } write!(f, ")")?; - write!(f, " OVER( ")?; + write!(f, " OVER(")?; if !partition_by.is_empty() { - write!(f, "PARTITION BY {:?} ", partition_by)?; + write!(f, "PARTITION BY {:?}", partition_by)?; } if !order_by.is_empty() { - write!(f, "ORDER BY {:?} ", order_by)?; + if !partition_by.is_empty() { + write!(f, " ")?; + } + write!(f, "ORDER BY {:?}", order_by)?; } if let Some(window_frame) = window_frame { + if !partition_by.is_empty() || !order_by.is_empty() { + write!(f, " ")?; + } write!( f, - "{} BETWEEN {} AND {} ", + "{} BETWEEN {} AND {}", window_frame.units, window_frame.start_bound, window_frame.end_bound )?; } diff --git a/common/planners/src/plan_expression_common.rs b/common/planners/src/plan_expression_common.rs index 2619c1649772..4b77cba4a8d1 100644 --- a/common/planners/src/plan_expression_common.rs +++ b/common/planners/src/plan_expression_common.rs @@ -66,7 +66,6 @@ pub fn find_window_exprs_in_expr(expr: &Expression) -> Vec { /// Collect all arguments from aggregation function and append to this exprs /// [ColumnExpr(b), Aggr(sum(a, b))] ---> [ColumnExpr(b), ColumnExpr(a)] - pub fn expand_aggregate_arg_exprs(exprs: &[Expression]) -> Vec { let mut res = vec![]; for expr in exprs { diff --git a/query/src/sql/plan_parser.rs b/query/src/sql/plan_parser.rs index 2b6d90bfa799..a397b4d8c0d4 100644 --- a/query/src/sql/plan_parser.rs +++ b/query/src/sql/plan_parser.rs @@ -21,6 +21,7 @@ use common_planners::Expression; use common_planners::PlanBuilder; use common_planners::PlanNode; use common_planners::SelectPlan; +use common_tracing::tracing::debug; use crate::sessions::QueryContext; use crate::sql::statements::AnalyzableStatement; @@ -74,14 +75,31 @@ impl PlanParser { pub fn build_query_plan(data: &QueryAnalyzeState) -> Result { let from = Self::build_from_plan(data)?; + debug!("\nfrom plan node:\n{:?}", from); + let filter = Self::build_filter_plan(from, data)?; + debug!("\nfilter plan node:\n{:?}", filter); + let group_by = Self::build_group_by_plan(filter, data)?; + debug!("\ngroup_by plan node:\n{:?}", group_by); + let before_order = Self::build_before_order(group_by, data)?; + debug!("\nbefore_order plan node:\n{:?}", before_order); + let having = Self::build_having_plan(before_order, data)?; + debug!("\nhaving plan node:\n{:?}", having); + let window = Self::build_window_plan(having, data)?; + debug!("\nwindow plan node:\n{:?}", window); + let order_by = Self::build_order_by_plan(window, data)?; + debug!("\norder_by plan node:\n{:?}", order_by); + let projection = Self::build_projection_plan(order_by, data)?; + debug!("\nprojection plan node:\n{:?}", projection); + let limit = Self::build_limit_plan(projection, data)?; + debug!("\nlimit plan node:\n{:?}", limit); Ok(PlanNode::Select(SelectPlan { input: Arc::new(limit), diff --git a/query/src/sql/statements/query/query_ast_ir.rs b/query/src/sql/statements/query/query_ast_ir.rs index ac1408af26c7..d98b908642f3 100644 --- a/query/src/sql/statements/query/query_ast_ir.rs +++ b/query/src/sql/statements/query/query_ast_ir.rs @@ -24,6 +24,7 @@ pub struct QueryASTIR { pub group_by_expressions: Vec, pub having_predicate: Option, pub aggregate_expressions: Vec, + pub window_expressions: Vec, pub order_by_expressions: Vec, pub projection_expressions: Vec, pub limit: Option, @@ -43,6 +44,7 @@ pub trait QueryASTIRVisitor { Self::visit_group_by(&mut ir.group_by_expressions, data)?; Self::visit_order_by(&mut ir.order_by_expressions, data)?; Self::visit_aggregates(&mut ir.aggregate_expressions, data)?; + Self::visit_window(&mut ir.window_expressions, data)?; Self::visit_projection(&mut ir.projection_expressions, data)?; Ok(()) } @@ -73,6 +75,24 @@ pub trait QueryASTIRVisitor { Ok(()) } + Expression::WindowFunction { + args, + partition_by, + order_by, + .. + } => { + for expr in args { + Self::visit_recursive_expr(expr, data)?; + } + for expr in partition_by { + Self::visit_recursive_expr(expr, data)?; + } + for expr in order_by { + Self::visit_recursive_expr(expr, data)?; + } + + Ok(()) + } Expression::Sort { expr, origin_expr, .. } => { @@ -108,6 +128,14 @@ pub trait QueryASTIRVisitor { Ok(()) } + fn visit_window(exprs: &mut Vec, data: &mut Data) -> Result<()> { + for expr in exprs { + Self::visit_recursive_expr(expr, data)?; + } + + Ok(()) + } + fn visit_order_by(exprs: &mut Vec, data: &mut Data) -> Result<()> { for expr in exprs { Self::visit_recursive_expr(expr, data)?; @@ -145,6 +173,10 @@ impl Debug for QueryASTIR { debug_struct.field("aggregate", &self.aggregate_expressions); } + if !self.window_expressions.is_empty() { + debug_struct.field("window", &self.window_expressions); + } + if !self.order_by_expressions.is_empty() { debug_struct.field("order by", &self.order_by_expressions); } diff --git a/query/src/sql/statements/query/query_normalizer.rs b/query/src/sql/statements/query/query_normalizer.rs index 487a93f2062e..c176049be4d3 100644 --- a/query/src/sql/statements/query/query_normalizer.rs +++ b/query/src/sql/statements/query/query_normalizer.rs @@ -19,6 +19,7 @@ use common_exception::ErrorCode; use common_exception::Result; use common_planners::extract_aliases; use common_planners::find_aggregate_exprs_in_expr; +use common_planners::find_window_exprs_in_expr; use common_planners::resolve_aliases_to_exprs; use common_planners::Expression; use sqlparser::ast::Expr; @@ -47,6 +48,7 @@ impl QueryNormalizer { group_by_expressions: vec![], having_predicate: None, aggregate_expressions: vec![], + window_expressions: vec![], order_by_expressions: vec![], projection_expressions: vec![], limit: None, @@ -103,6 +105,7 @@ impl QueryNormalizer { for projection_expression in &projection_expressions { self.add_aggregate_function(projection_expression)?; + self.add_window_function(projection_expression)?; } self.query_ast_ir.projection_expressions = projection_expressions; @@ -231,4 +234,14 @@ impl QueryNormalizer { Ok(()) } + + fn add_window_function(&mut self, expr: &Expression) -> Result<()> { + for window_expr in find_window_exprs_in_expr(expr) { + if !self.query_ast_ir.window_expressions.contains(&window_expr) { + self.query_ast_ir.window_expressions.push(window_expr); + } + } + + Ok(()) + } } diff --git a/query/src/sql/statements/statement_select.rs b/query/src/sql/statements/statement_select.rs index 5d9221117ea9..1f3722e9211b 100644 --- a/query/src/sql/statements/statement_select.rs +++ b/query/src/sql/statements/statement_select.rs @@ -21,6 +21,7 @@ use common_exception::Result; use common_planners::expand_aggregate_arg_exprs; use common_planners::find_aggregate_exprs; use common_planners::find_aggregate_exprs_in_expr; +use common_planners::find_window_exprs; use common_planners::find_window_exprs_in_expr; use common_planners::rebase_expr; use common_planners::Expression; @@ -75,6 +76,8 @@ impl AnalyzableStatement for DfQueryStatement { QueryCollectPushDowns::collect_extras(&mut ir, &mut joined_schema, has_aggregation)?; tracing::debug!("\nQueryASTIR after push downs:\n{:?}", ir); + // todo collect window_functions @doki + let analyze_state = self.analyze_query(ir).await?; tracing::debug!("\nQueryAnalyzeState:\n{:?}", analyze_state); @@ -153,6 +156,18 @@ impl DfQueryStatement { Self::analyze_aggregate(&ir.aggregate_expressions, &mut analyze_state)?; } + if !ir.window_expressions.is_empty() { + let mut expressions = Vec::with_capacity(analyze_state.expressions.len()); + for expression in &analyze_state.expressions { + let expression = rebase_expr(expression, &ir.window_expressions)?; + expressions.push(expression); + } + + analyze_state.expressions = expressions; + + Self::analyze_window(&ir.window_expressions, &mut analyze_state)?; + } + Ok(analyze_state) } @@ -174,6 +189,11 @@ impl DfQueryStatement { Ok(()) } + fn analyze_window(exprs: &[Expression], state: &mut QueryAnalyzeState) -> Result<()> { + let window_functions = find_window_exprs(exprs); + Ok(()) + } + fn analyze_projection(exprs: &[Expression], state: &mut QueryAnalyzeState) -> Result<()> { for item in exprs { match item { @@ -181,11 +201,6 @@ impl DfQueryStatement { _ => state.add_expression(item), } - let window_exprs = find_window_exprs_in_expr(item); - if !window_exprs.is_empty() { - state.window_expressions.extend(window_exprs); - } - let rebased_expr = rebase_expr(item, &state.expressions)?; state.projection_expressions.push(rebased_expr); } From f455a6a636b584ebadd33f6d5161f1dc7ec4780e Mon Sep 17 00:00:00 2001 From: doki Date: Sat, 30 Apr 2022 20:37:48 +0800 Subject: [PATCH 05/18] window logical plan done --- .../src/sql/statements/analyzer_statement.rs | 1 - query/src/sql/statements/statement_select.rs | 93 +++++++++++++++---- 2 files changed, 74 insertions(+), 20 deletions(-) diff --git a/query/src/sql/statements/analyzer_statement.rs b/query/src/sql/statements/analyzer_statement.rs index 589497f12ea8..1ec87d27f862 100644 --- a/query/src/sql/statements/analyzer_statement.rs +++ b/query/src/sql/statements/analyzer_statement.rs @@ -55,7 +55,6 @@ pub struct QueryAnalyzeState { pub group_by_expressions: Vec, pub aggregate_expressions: Vec, - // window functions pub window_expressions: Vec, // rebase on projection expressions without aliases, aggregate and group by expressions diff --git a/query/src/sql/statements/statement_select.rs b/query/src/sql/statements/statement_select.rs index 3cae7e20814e..583e31c5bfe5 100644 --- a/query/src/sql/statements/statement_select.rs +++ b/query/src/sql/statements/statement_select.rs @@ -21,8 +21,6 @@ use common_exception::Result; use common_planners::expand_aggregate_arg_exprs; use common_planners::find_aggregate_exprs; use common_planners::find_aggregate_exprs_in_expr; -use common_planners::find_window_exprs; -use common_planners::find_window_exprs_in_expr; use common_planners::rebase_expr; use common_planners::Expression; use common_tracing::tracing; @@ -74,9 +72,9 @@ impl AnalyzableStatement for DfQueryStatement { let has_aggregation = !find_aggregate_exprs(&ir.projection_expressions).is_empty(); QueryCollectPushDowns::collect_extras(&mut ir, &mut joined_schema, has_aggregation)?; - // todo collect window_functions @doki - let analyze_state = self.analyze_query(ir).await?; + tracing::debug!("analyze state is:\n{:?}", analyze_state); + self.check_and_finalize(joined_schema, analyze_state, ctx) .await } @@ -153,14 +151,6 @@ impl DfQueryStatement { } if !ir.window_expressions.is_empty() { - let mut expressions = Vec::with_capacity(analyze_state.expressions.len()); - for expression in &analyze_state.expressions { - let expression = rebase_expr(expression, &ir.window_expressions)?; - expressions.push(expression); - } - - analyze_state.expressions = expressions; - Self::analyze_window(&ir.window_expressions, &mut analyze_state)?; } @@ -199,8 +189,7 @@ impl DfQueryStatement { _ => item.clone(), }; - // support select distinct aggr_func()... - // todo may need before_distinct_expressions @doki + let distinct_expr = rebase_expr(&distinct_expr, &state.expressions)?; let distinct_expr = rebase_expr(&distinct_expr, &state.group_by_expressions)?; let distinct_expr = rebase_expr(&distinct_expr, &state.aggregate_expressions)?; let distinct_expr = rebase_expr(&distinct_expr, &state.window_expressions)?; @@ -210,16 +199,60 @@ impl DfQueryStatement { Ok(()) } - fn analyze_window(exprs: &[Expression], state: &mut QueryAnalyzeState) -> Result<()> { - let window_functions = find_window_exprs(exprs); + fn analyze_window(window_exprs: &[Expression], state: &mut QueryAnalyzeState) -> Result<()> { + for expr in window_exprs { + match expr { + Expression::WindowFunction { + args, + partition_by, + order_by, + .. + } => { + for arg in args { + state.add_expression(arg); + } + for partition_by_expr in partition_by { + state.add_expression(partition_by_expr); + } + for order_by_expr in order_by { + match order_by_expr { + Expression::Sort { expr, .. } => state.add_expression(expr), + _ => { + return Err(ErrorCode::LogicalError(format!( + "Found non-sort expression {:?} while analyzing order by expressions of window expressions", + order_by_expr + ))) + } + } + } + } + _ => { + return Err(ErrorCode::LogicalError(format!( + "Found non-window expression {:?} while analyzing window expressions!", + expr + ))) + } + } + } + + for expr in window_exprs { + let base_exprs = &state.expressions; + state + .window_expressions + .push(rebase_expr(expr, base_exprs)?); + } + Ok(()) } fn analyze_projection(exprs: &[Expression], state: &mut QueryAnalyzeState) -> Result<()> { for item in exprs { - match item { - Expression::Alias(_, expr) => state.add_expression(expr), - _ => state.add_expression(item), + let expr = match item { + Expression::Alias(_, expr) => expr, + _ => item, + }; + if !matches!(expr, Expression::WindowFunction { .. }) { + state.add_expression(expr); } let rebased_expr = rebase_expr(item, &state.expressions)?; @@ -341,6 +374,28 @@ impl DfQueryStatement { } } + if !state.window_expressions.is_empty() { + let new_len = state.window_expressions.len() + state.expressions.len(); + let mut new_expression = Vec::with_capacity(new_len); + + for expr in &state.window_expressions { + new_expression.push(expr); + } + + for expr in &state.expressions { + new_expression.push(expr); + } + + match Self::dry_run_exprs_ref(&new_expression, &data_block) { + Ok(res) => { + data_block = res; + } + Err(cause) => { + return Err(cause.add_message_back(" (while in select window aggr)")); + } + } + } + if !state.order_by_expressions.is_empty() { if let Err(cause) = Self::dry_run_exprs(&state.order_by_expressions, &data_block) { return Err(cause.add_message_back(" (while in select order by)")); From a121073ff33910b0de6d0cabf97d67f25930efae Mon Sep 17 00:00:00 2001 From: doki Date: Sat, 30 Apr 2022 21:02:07 +0800 Subject: [PATCH 06/18] fix --- query/src/common/expression_evaluator.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/query/src/common/expression_evaluator.rs b/query/src/common/expression_evaluator.rs index 8fb76f83c142..277fd3578a23 100644 --- a/query/src/common/expression_evaluator.rs +++ b/query/src/common/expression_evaluator.rs @@ -105,6 +105,10 @@ impl ExpressionEvaluator { "Unsupported AggregateFunction scalar expression", )), + Expression::WindowFunction { .. } => Err(ErrorCode::LogicalError( + "Unsupported WindowFunction scalar expression", + )), + Expression::Sort { .. } => Err(ErrorCode::LogicalError( "Unsupported Sort scalar expression", )), From 89a707a637b9483bb8cfe99d30631498326de36c Mon Sep 17 00:00:00 2001 From: doki Date: Thu, 12 May 2022 16:53:30 +0800 Subject: [PATCH 07/18] temporary commit --- Cargo.lock | 7 + common/functions/src/window/mod.rs | 2 + common/functions/src/window/window_frame.rs | 15 +- common/planners/src/lib.rs | 3 +- common/planners/src/plan_expression.rs | 23 ++- common/planners/src/plan_expression_common.rs | 2 + .../planners/src/plan_expression_rewriter.rs | 30 +-- .../planners/src/plan_expression_visitor.rs | 4 +- common/planners/src/plan_node.rs | 10 +- common/planners/src/plan_node_builder.rs | 6 +- .../planners/src/plan_node_display_indent.rs | 4 +- common/planners/src/plan_node_rewriter.rs | 8 +- common/planners/src/plan_node_visitor.rs | 6 +- ...lan_window_aggr.rs => plan_window_func.rs} | 4 +- query/Cargo.toml | 1 + .../plan_schedulers/plan_scheduler_query.rs | 1 + query/src/optimizers/optimizer_scatters.rs | 20 ++ .../pipelines/processors/pipeline_builder.rs | 15 ++ query/src/pipelines/transforms/mod.rs | 3 + .../transforms/transform_window_func.rs | 189 ++++++++++++++++++ .../pipelines/transforms/window_func/mod.rs | 13 ++ query/src/sql/plan_parser.rs | 39 +++- query/src/sql/statements/analyzer_expr.rs | 20 ++ .../src/sql/statements/analyzer_statement.rs | 2 +- 24 files changed, 367 insertions(+), 60 deletions(-) rename common/planners/src/{plan_window_aggr.rs => plan_window_func.rs} (95%) create mode 100644 query/src/pipelines/transforms/transform_window_func.rs create mode 100644 query/src/pipelines/transforms/window_func/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 0a1523d8d3b8..7be802cec8f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1785,6 +1785,7 @@ dependencies = [ "criterion", "dyn-clone", "enum_dispatch", + "enum_extract", "futures", "headers", "http", @@ -2086,6 +2087,12 @@ dependencies = [ "syn", ] +[[package]] +name = "enum_extract" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11578f9e8496eeb626c549e1b340a57e479840ff7ae13fe12d4dbd97c18779d1" + [[package]] name = "enumflags2" version = "0.7.4" diff --git a/common/functions/src/window/mod.rs b/common/functions/src/window/mod.rs index 21e32b7615d1..074748ca3e34 100644 --- a/common/functions/src/window/mod.rs +++ b/common/functions/src/window/mod.rs @@ -1,3 +1,5 @@ mod window_frame; pub use window_frame::WindowFrame; +pub use window_frame::WindowFrameBound; +pub use window_frame::WindowFrameUnits; diff --git a/common/functions/src/window/window_frame.rs b/common/functions/src/window/window_frame.rs index 3a9e75b4fe93..87d0251a32e8 100644 --- a/common/functions/src/window/window_frame.rs +++ b/common/functions/src/window/window_frame.rs @@ -103,27 +103,22 @@ impl Default for WindowFrame { Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash, serde::Serialize, serde::Deserialize, )] pub enum WindowFrameUnits { - /// The ROWS frame type means that the starting and ending boundaries for the frame are - /// determined by counting individual rows relative to the current row. - Rows, /// The RANGE frame type requires that the ORDER BY clause of the window have exactly one /// term. Call that term "X". With the RANGE frame type, the elements of the frame are /// determined by computing the value of expression X for all rows in the partition and framing /// those rows for which the value of X is within a certain range of the value of X for the /// current row. Range, - /// The GROUPS frame type means that the starting and ending boundaries are determine - /// by counting "groups" relative to the current group. A "group" is a set of rows that all have - /// equivalent values for all all terms of the window ORDER BY clause. - Groups, + /// The ROWS frame type means that the starting and ending boundaries for the frame are + /// determined by counting individual rows relative to the current row. + Rows, } impl fmt::Display for WindowFrameUnits { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.write_str(match self { - WindowFrameUnits::Rows => "ROWS", WindowFrameUnits::Range => "RANGE", - WindowFrameUnits::Groups => "GROUPS", + WindowFrameUnits::Rows => "ROWS", }) } } @@ -132,8 +127,8 @@ impl From for WindowFrameUnits { fn from(value: ast::WindowFrameUnits) -> Self { match value { ast::WindowFrameUnits::Range => Self::Range, - ast::WindowFrameUnits::Groups => Self::Groups, ast::WindowFrameUnits::Rows => Self::Rows, + _ => unimplemented!(), } } } diff --git a/common/planners/src/lib.rs b/common/planners/src/lib.rs index 583f794e7df2..860e5a72f7b7 100644 --- a/common/planners/src/lib.rs +++ b/common/planners/src/lib.rs @@ -98,7 +98,7 @@ mod plan_user_udf_drop; mod plan_view_alter; mod plan_view_create; mod plan_view_drop; -mod plan_window_aggr; +mod plan_window_func; pub use plan_aggregator_final::AggregatorFinalPlan; pub use plan_aggregator_partial::AggregatorPartialPlan; @@ -224,3 +224,4 @@ pub use plan_user_udf_drop::DropUserUDFPlan; pub use plan_view_alter::AlterViewPlan; pub use plan_view_create::CreateViewPlan; pub use plan_view_drop::DropViewPlan; +pub use plan_window_func::WindowFuncPlan; diff --git a/common/planners/src/plan_expression.rs b/common/planners/src/plan_expression.rs index e4b6c6660fac..785d95f1b888 100644 --- a/common/planners/src/plan_expression.rs +++ b/common/planners/src/plan_expression.rs @@ -99,6 +99,8 @@ pub enum Expression { WindowFunction { /// operation performed op: String, + /// params + params: Vec, /// arguments args: Vec, /// partition by @@ -438,20 +440,25 @@ impl fmt::Debug for Expression { Expression::WindowFunction { op, + params, args, partition_by, order_by, window_frame, } => { - write!(f, "{}(", op)?; - for (i, arg) in args.iter().enumerate() { - if i < args.len() - 1 { - write!(f, "{},", arg.column_name())?; - } else { - write!(f, "{}", arg.column_name())?; - } + let args_column_name = args.iter().map(Expression::column_name).collect::>(); + let params_name = params + .iter() + .map(|v| DataValue::custom_display(v, true)) + .collect::>(); + + if params.is_empty() { + write!(f, "{}", op)?; + } else { + write!(f, "{}({})", op, params_name.join(", "))?; } - write!(f, ")")?; + + write!(f, "({})", args_column_name.join(","))?; write!(f, " OVER(")?; if !partition_by.is_empty() { diff --git a/common/planners/src/plan_expression_common.rs b/common/planners/src/plan_expression_common.rs index f04ad63b7e0f..d573c1976930 100644 --- a/common/planners/src/plan_expression_common.rs +++ b/common/planners/src/plan_expression_common.rs @@ -318,12 +318,14 @@ where F: Fn(&Expression) -> Result> { Expression::WindowFunction { op, + params, args, partition_by, order_by, window_frame, } => Ok(Expression::WindowFunction { op: op.clone(), + params: params.clone(), args: args .iter() .map(|e| clone_with_replacement(e, replacement_fn)) diff --git a/common/planners/src/plan_expression_rewriter.rs b/common/planners/src/plan_expression_rewriter.rs index 773e716c1559..17066997abdc 100644 --- a/common/planners/src/plan_expression_rewriter.rs +++ b/common/planners/src/plan_expression_rewriter.rs @@ -128,6 +128,7 @@ pub trait ExpressionRewriter: Sized { fn mutate_window_function( &mut self, op: String, + params: Vec, args: Vec, partition_by: Vec, order_by: Vec, @@ -135,6 +136,7 @@ pub trait ExpressionRewriter: Sized { ) -> Result { Ok(Expression::WindowFunction { op, + params, args, partition_by, order_by, @@ -320,22 +322,16 @@ impl ExpressionVisitor for ExpressionRewriteVisitor { } Expression::WindowFunction { op, + params, args, partition_by, order_by, window_frame, } => { - let mut new_partition_by = Vec::with_capacity(partition_by.len()); - let mut new_order_by = Vec::with_capacity(order_by.len()); - for i in 0..partition_by.len() + order_by.len() { + let mut new_args = Vec::with_capacity(args.len()); + for i in 0..args.len() { match self.stack.pop() { - Some(expr) => { - if i < order_by.len() { - new_order_by.push(expr); - } else { - new_partition_by.push(expr); - } - } + Some(expr) => new_args.push(expr), None => { return Err(ErrorCode::LogicalError(format!( "WindowFunction expects {} partition by arguments, actual {}", @@ -346,10 +342,17 @@ impl ExpressionVisitor for ExpressionRewriteVisitor { } } - let mut new_args = Vec::with_capacity(args.len()); - for i in 0..args.len() { + let mut new_partition_by = Vec::with_capacity(partition_by.len()); + let mut new_order_by = Vec::with_capacity(order_by.len()); + for i in 0..partition_by.len() + order_by.len() { match self.stack.pop() { - Some(expr) => new_args.push(expr), + Some(expr) => { + if i < partition_by.len() { + new_partition_by.push(expr); + } else { + new_order_by.push(expr); + } + } None => { return Err(ErrorCode::LogicalError(format!( "WindowFunction expects {} partition by arguments, actual {}", @@ -362,6 +365,7 @@ impl ExpressionVisitor for ExpressionRewriteVisitor { let new_expr = self.inner.mutate_window_function( op.clone(), + params.to_owned(), new_args, new_partition_by, new_order_by, diff --git a/common/planners/src/plan_expression_visitor.rs b/common/planners/src/plan_expression_visitor.rs index c65f44d10e66..2c73df11867c 100644 --- a/common/planners/src/plan_expression_visitor.rs +++ b/common/planners/src/plan_expression_visitor.rs @@ -72,8 +72,8 @@ pub trait ExpressionVisitor: Sized { order_by, .. } => { - for arg_exppr in args { - stack.push(RecursionProcessing::Call(arg_exppr)); + for arg_expr in args { + stack.push(RecursionProcessing::Call(arg_expr)); } for part_by_expr in partition_by { stack.push(RecursionProcessing::Call(part_by_expr)); diff --git a/common/planners/src/plan_node.rs b/common/planners/src/plan_node.rs index 384bf00ecc75..2f5fe3162abf 100644 --- a/common/planners/src/plan_node.rs +++ b/common/planners/src/plan_node.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use common_datavalues::DataSchemaRef; -use crate::plan_window_aggr::WindowAggrPlan; +use crate::plan_window_func::WindowFuncPlan; use crate::AggregatorFinalPlan; use crate::AggregatorPartialPlan; use crate::AlterUserPlan; @@ -86,7 +86,7 @@ pub enum PlanNode { AggregatorFinal(AggregatorFinalPlan), Filter(FilterPlan), Having(HavingPlan), - WindowAggr(WindowAggrPlan), + WindowFunc(WindowFuncPlan), Sort(SortPlan), Limit(LimitPlan), LimitBy(LimitByPlan), @@ -186,7 +186,7 @@ impl PlanNode { PlanNode::AggregatorFinal(v) => v.schema(), PlanNode::Filter(v) => v.schema(), PlanNode::Having(v) => v.schema(), - PlanNode::WindowAggr(v) => v.schema(), + PlanNode::WindowFunc(v) => v.schema(), PlanNode::Limit(v) => v.schema(), PlanNode::LimitBy(v) => v.schema(), PlanNode::ReadSource(v) => v.schema(), @@ -285,7 +285,7 @@ impl PlanNode { PlanNode::AggregatorFinal(_) => "AggregatorFinalPlan", PlanNode::Filter(_) => "FilterPlan", PlanNode::Having(_) => "HavingPlan", - PlanNode::WindowAggr(_) => "WindowAggrPlan", + PlanNode::WindowFunc(_) => "WindowFuncPlan", PlanNode::Limit(_) => "LimitPlan", PlanNode::LimitBy(_) => "LimitByPlan", PlanNode::ReadSource(_) => "ReadSourcePlan", @@ -381,7 +381,7 @@ impl PlanNode { PlanNode::AggregatorFinal(v) => vec![v.input.clone()], PlanNode::Filter(v) => vec![v.input.clone()], PlanNode::Having(v) => vec![v.input.clone()], - PlanNode::WindowAggr(v) => vec![v.input.clone()], + PlanNode::WindowFunc(v) => vec![v.input.clone()], PlanNode::Limit(v) => vec![v.input.clone()], PlanNode::Explain(v) => vec![v.input.clone()], PlanNode::Select(v) => vec![v.input.clone()], diff --git a/common/planners/src/plan_node_builder.rs b/common/planners/src/plan_node_builder.rs index 6420f18f1513..a47fe0ac13d1 100644 --- a/common/planners/src/plan_node_builder.rs +++ b/common/planners/src/plan_node_builder.rs @@ -20,7 +20,7 @@ use common_exception::Result; use crate::col; use crate::plan_subqueries_set::SubQueriesSetPlan; -use crate::plan_window_aggr::WindowAggrPlan; +use crate::plan_window_func::WindowFuncPlan; use crate::validate_expression; use crate::AggregatorFinalPlan; use crate::AggregatorPartialPlan; @@ -216,7 +216,7 @@ impl PlanBuilder { } /// Apply a window function - pub fn window_aggr(&self, expr: Expression) -> Result { + pub fn window_func(&self, expr: Expression) -> Result { let window_func = expr.clone(); let input = self.wrap_subquery_plan(&[expr])?; let input_schema = input.schema(); @@ -224,7 +224,7 @@ impl PlanBuilder { let window_field = window_func.to_data_field(&input_schema).unwrap(); input_fields.push(window_field); let schema = Arc::new(DataSchema::new(input_fields)); - Ok(Self::from(&PlanNode::WindowAggr(WindowAggrPlan { + Ok(Self::from(&PlanNode::WindowFunc(WindowFuncPlan { window_func, input, schema, diff --git a/common/planners/src/plan_node_display_indent.rs b/common/planners/src/plan_node_display_indent.rs index 0aa2baad16c6..75791883590f 100644 --- a/common/planners/src/plan_node_display_indent.rs +++ b/common/planners/src/plan_node_display_indent.rs @@ -70,7 +70,9 @@ impl<'a> fmt::Display for PlanNodeIndentFormatDisplay<'a> { PlanNode::AggregatorFinal(plan) => Self::format_aggregator_final(f, plan), PlanNode::Filter(plan) => write!(f, "Filter: {:?}", plan.predicate), PlanNode::Having(plan) => write!(f, "Having: {:?}", plan.predicate), - PlanNode::WindowAggr(plan) => write!(f, "WindowAggr: {:?}", plan.window_func), + PlanNode::WindowFunc(plan) => { + write!(f, "WindowFunc: {:?}", plan.window_func) + } PlanNode::Sort(plan) => Self::format_sort(f, plan), PlanNode::Limit(plan) => Self::format_limit(f, plan), PlanNode::SubQueryExpression(plan) => Self::format_subquery_expr(f, plan), diff --git a/common/planners/src/plan_node_rewriter.rs b/common/planners/src/plan_node_rewriter.rs index cbe99484a5c6..cc1dc18ab305 100644 --- a/common/planners/src/plan_node_rewriter.rs +++ b/common/planners/src/plan_node_rewriter.rs @@ -23,7 +23,7 @@ use common_exception::Result; use crate::plan_broadcast::BroadcastPlan; use crate::plan_subqueries_set::SubQueriesSetPlan; -use crate::plan_window_aggr::WindowAggrPlan; +use crate::plan_window_func::WindowFuncPlan; use crate::AggregatorFinalPlan; use crate::AggregatorPartialPlan; use crate::AlterUserPlan; @@ -114,7 +114,7 @@ pub trait PlanRewriter: Sized { PlanNode::Broadcast(plan) => self.rewrite_broadcast(plan), PlanNode::Remote(plan) => self.rewrite_remote(plan), PlanNode::Having(plan) => self.rewrite_having(plan), - PlanNode::WindowAggr(plan) => self.rewrite_window_aggr(plan), + PlanNode::WindowFunc(plan) => self.rewrite_window_func(plan), PlanNode::Expression(plan) => self.rewrite_expression(plan), PlanNode::Sort(plan) => self.rewrite_sort(plan), PlanNode::Limit(plan) => self.rewrite_limit(plan), @@ -313,11 +313,11 @@ pub trait PlanRewriter: Sized { PlanBuilder::from(&new_input).having(new_predicate)?.build() } - fn rewrite_window_aggr(&mut self, plan: &WindowAggrPlan) -> Result { + fn rewrite_window_func(&mut self, plan: &WindowFuncPlan) -> Result { let new_input = self.rewrite_plan_node(plan.input.as_ref())?; let new_window_func = self.rewrite_expr(&new_input.schema(), &plan.window_func)?; PlanBuilder::from(&new_input) - .window_aggr(new_window_func)? + .window_func(new_window_func)? .build() } diff --git a/common/planners/src/plan_node_visitor.rs b/common/planners/src/plan_node_visitor.rs index 62ee4cee96a7..7d92960cdd20 100644 --- a/common/planners/src/plan_node_visitor.rs +++ b/common/planners/src/plan_node_visitor.rs @@ -16,7 +16,7 @@ use common_exception::Result; use crate::plan_broadcast::BroadcastPlan; use crate::plan_subqueries_set::SubQueriesSetPlan; -use crate::plan_window_aggr::WindowAggrPlan; +use crate::plan_window_func::WindowFuncPlan; use crate::AggregatorFinalPlan; use crate::AggregatorPartialPlan; use crate::AlterUserPlan; @@ -127,7 +127,7 @@ pub trait PlanVisitor { PlanNode::Broadcast(plan) => self.visit_broadcast(plan), PlanNode::Remote(plan) => self.visit_remote(plan), PlanNode::Having(plan) => self.visit_having(plan), - PlanNode::WindowAggr(plan) => self.visit_window_aggr(plan), + PlanNode::WindowFunc(plan) => self.visit_window_func(plan), PlanNode::Expression(plan) => self.visit_expression(plan), PlanNode::Limit(plan) => self.visit_limit(plan), PlanNode::LimitBy(plan) => self.visit_limit_by(plan), @@ -290,7 +290,7 @@ pub trait PlanVisitor { self.visit_expr(&plan.predicate) } - fn visit_window_aggr(&mut self, plan: &WindowAggrPlan) -> Result<()> { + fn visit_window_func(&mut self, plan: &WindowFuncPlan) -> Result<()> { self.visit_plan_node(plan.input.as_ref())?; self.visit_expr(&plan.window_func) } diff --git a/common/planners/src/plan_window_aggr.rs b/common/planners/src/plan_window_func.rs similarity index 95% rename from common/planners/src/plan_window_aggr.rs rename to common/planners/src/plan_window_func.rs index 96b16adadee4..24ce6a8b3b11 100644 --- a/common/planners/src/plan_window_aggr.rs +++ b/common/planners/src/plan_window_func.rs @@ -20,7 +20,7 @@ use crate::Expression; use crate::PlanNode; #[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq)] -pub struct WindowAggrPlan { +pub struct WindowFuncPlan { /// The window function expression pub window_func: Expression, /// The incoming logical plan @@ -29,7 +29,7 @@ pub struct WindowAggrPlan { pub schema: DataSchemaRef, } -impl WindowAggrPlan { +impl WindowFuncPlan { pub fn schema(&self) -> DataSchemaRef { self.schema.clone() } diff --git a/query/Cargo.toml b/query/Cargo.toml index a7d7460efa04..01fb6b02eaa5 100644 --- a/query/Cargo.toml +++ b/query/Cargo.toml @@ -73,6 +73,7 @@ chrono-tz = "0.6.1" clap = { version = "3.1.8", features = ["derive", "env"] } dyn-clone = "1.0.5" enum_dispatch = "0.3.8" +enum_extract = "0.1.1" futures = "0.3.21" headers = "0.3.7" http = "0.2.6" diff --git a/query/src/interpreters/plan_schedulers/plan_scheduler_query.rs b/query/src/interpreters/plan_schedulers/plan_scheduler_query.rs index 2ee2930b5aa2..570c059f2f80 100644 --- a/query/src/interpreters/plan_schedulers/plan_scheduler_query.rs +++ b/query/src/interpreters/plan_schedulers/plan_scheduler_query.rs @@ -50,6 +50,7 @@ pub async fn schedule_query( let pipeline_builder = PipelineBuilder::create(ctx.clone()); let mut in_local_pipeline = pipeline_builder.build(&scheduled_tasks.get_local_task())?; + tracing::debug!("local_pipeline:\n{:?}", in_local_pipeline); match in_local_pipeline.execute().await { Ok(stream) => Ok(ScheduledStream::create(ctx.clone(), scheduled, stream)), diff --git a/query/src/optimizers/optimizer_scatters.rs b/query/src/optimizers/optimizer_scatters.rs index 373da3caae1e..7a88ebd6b982 100644 --- a/query/src/optimizers/optimizer_scatters.rs +++ b/query/src/optimizers/optimizer_scatters.rs @@ -31,6 +31,7 @@ use common_planners::ReadDataSourcePlan; use common_planners::SortPlan; use common_planners::StageKind; use common_planners::StagePlan; +use common_planners::WindowFuncPlan; use crate::optimizers::Optimizer; use crate::sessions::QueryContext; @@ -109,6 +110,14 @@ impl ScattersOptimizerImpl { } } + fn cluster_window(&mut self, plan: &WindowFuncPlan) -> Result { + todo!() + } + + fn standalone_window(&mut self, plan: &WindowFuncPlan) -> Result { + todo!() + } + fn cluster_sort(&mut self, plan: &SortPlan) -> Result { // Order by we convergent it in local node self.running_mode = RunningMode::Standalone; @@ -251,6 +260,17 @@ impl PlanRewriter for ScattersOptimizerImpl { } } + fn rewrite_window_func(&mut self, plan: &WindowFuncPlan) -> Result { + let new_input = Arc::new(self.rewrite_plan_node(&plan.input)?); + + self.input = Some(new_input.clone()); + + match self.running_mode { + RunningMode::Cluster => self.cluster_window(plan), + RunningMode::Standalone => self.standalone_window(plan), + } + } + fn rewrite_sort(&mut self, plan: &SortPlan) -> Result { self.input = Some(Arc::new(self.rewrite_plan_node(plan.input.as_ref())?)); diff --git a/query/src/pipelines/processors/pipeline_builder.rs b/query/src/pipelines/processors/pipeline_builder.rs index a72b6336cc98..6699aaf544cd 100644 --- a/query/src/pipelines/processors/pipeline_builder.rs +++ b/query/src/pipelines/processors/pipeline_builder.rs @@ -32,6 +32,7 @@ use common_planners::SinkPlan; use common_planners::SortPlan; use common_planners::StagePlan; use common_planners::SubQueriesSetPlan; +use common_planners::WindowFuncPlan; use common_tracing::tracing; use crate::api::FlightTicket; @@ -53,6 +54,7 @@ use crate::pipelines::transforms::SortPartialTransform; use crate::pipelines::transforms::SourceTransform; use crate::pipelines::transforms::SubQueriesPuller; use crate::pipelines::transforms::WhereTransform; +use crate::pipelines::transforms::WindowFuncTransform; use crate::sessions::QueryContext; pub struct PipelineBuilder { @@ -89,6 +91,7 @@ impl PipelineBuilder { PlanNode::Projection(node) => self.visit_projection(node), PlanNode::AggregatorPartial(node) => self.visit_aggregator_partial(node), PlanNode::AggregatorFinal(node) => self.visit_aggregator_final(node), + PlanNode::WindowFunc(node) => self.visit_window_func(node), PlanNode::Filter(node) => self.visit_filter(node), PlanNode::Having(node) => self.visit_having(node), PlanNode::Sort(node) => self.visit_sort(node), @@ -216,6 +219,18 @@ impl PipelineBuilder { Ok(pipeline) } + fn visit_window_func(&mut self, node: &WindowFuncPlan) -> Result { + let mut pipeline = self.visit(&*node.input)?; + pipeline.add_simple_transform(|| { + Ok(Box::new(WindowFuncTransform::create( + node.window_func.to_owned(), + node.schema.to_owned(), + node.input.schema(), + ))) + })?; + Ok(pipeline) + } + fn visit_filter(&mut self, node: &FilterPlan) -> Result { let mut pipeline = self.visit(&*node.input)?; pipeline.add_simple_transform(|| { diff --git a/query/src/pipelines/transforms/mod.rs b/query/src/pipelines/transforms/mod.rs index b32dc65bfa2e..363bd04f206f 100644 --- a/query/src/pipelines/transforms/mod.rs +++ b/query/src/pipelines/transforms/mod.rs @@ -27,10 +27,12 @@ mod transform_remote; mod transform_sort_merge; mod transform_sort_partial; mod transform_source; +mod transform_window_func; pub mod group_by; mod streams; mod transform_sink; +pub mod window_func; pub use streams::AddOnStream; pub use transform_aggregator_final::AggregatorFinalTransform; @@ -52,3 +54,4 @@ pub use transform_sort_merge::SortMergeTransform; pub use transform_sort_partial::get_sort_descriptions; pub use transform_sort_partial::SortPartialTransform; pub use transform_source::SourceTransform; +pub use transform_window_func::WindowFuncTransform; diff --git a/query/src/pipelines/transforms/transform_window_func.rs b/query/src/pipelines/transforms/transform_window_func.rs new file mode 100644 index 000000000000..4137cd0c39eb --- /dev/null +++ b/query/src/pipelines/transforms/transform_window_func.rs @@ -0,0 +1,189 @@ +// Copyright 2022 Datafuse Labs. +// +// 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 std::any::Any; +use std::sync::Arc; + +use common_datablocks::DataBlock; +use common_datavalues::BooleanType; +use common_datavalues::ColumnRef; +use common_datavalues::DataField; +use common_datavalues::DataSchemaRef; +use common_datavalues::DataTypeImpl; +use common_functions::aggregates::AggregateFunctionFactory; +use common_functions::aggregates::AggregateFunctionRef; +use common_functions::window::WindowFrameBound; +use common_functions::window::WindowFrameUnits; +use common_planners::Expression; +use common_streams::DataBlockStream; +use common_streams::SendableDataBlockStream; +use common_tracing::tracing; +use enum_extract::let_extract; +use futures::StreamExt; + +use crate::pipelines::processors::EmptyProcessor; +use crate::pipelines::processors::Processor; + +pub struct WindowFuncTransform { + window_func: Expression, + schema: DataSchemaRef, + input: Arc, + input_schema: DataSchemaRef, +} + +impl WindowFuncTransform { + pub fn create( + window_func: Expression, + schema: DataSchemaRef, + input_schema: DataSchemaRef, + ) -> Self { + WindowFuncTransform { + window_func, + schema, + input: Arc::new(EmptyProcessor::create()), + input_schema, + } + } +} + +#[async_trait::async_trait] +impl Processor for WindowFuncTransform { + fn name(&self) -> &str { + "WindowAggrFuncTransform" + } + + fn connect_to(&mut self, input: Arc) -> common_exception::Result<()> { + self.input = input; + Ok(()) + } + + fn inputs(&self) -> Vec> { + vec![self.input.clone()] + } + + fn as_any(&self) -> &dyn Any { + self + } + + #[tracing::instrument(level = "debug", name = "window_aggr_func_execute", skip(self))] + async fn execute(&self) -> common_exception::Result { + let mut stream: SendableDataBlockStream = self.input.execute().await?; + let mut data_blocks: Vec = vec![]; + while let Some(block) = stream.next().await { + let block = block?; + data_blocks.push(block); + } + + let_extract!( + Expression::WindowFunction { + op, + params, + args, + partition_by, + order_by, + window_frame + }, + &self.window_func, + panic!() + ); + + let mut arguments = Vec::with_capacity(args.len()); + for arg in args { + arguments.push(arg.to_data_field(&self.input_schema)?); + } + + let is_aggr_func = AggregateFunctionFactory::instance().check(&op); + + let window_cols = match &window_frame { + Some(window_frame) if is_aggr_func => match &window_frame.units { + WindowFrameUnits::Range => { + evaluate_range_with_aggr_func( + AggregateFunctionFactory::instance().get( + op, + params.to_owned(), + arguments, + )?, + &data_blocks, + window_frame.start_bound, + window_frame.end_bound, + ) + .await? + } + WindowFrameUnits::Rows => { + evaluate_rows_with_aggr_func( + AggregateFunctionFactory::instance().get( + op, + params.to_owned(), + arguments, + )?, + &data_blocks, + window_frame.start_bound, + window_frame.end_bound, + ) + .await? + } + }, + None if is_aggr_func => { + evaluate_range_with_aggr_func( + AggregateFunctionFactory::instance().get(op, params.to_owned(), arguments)?, + &data_blocks, + WindowFrameBound::Preceding(None), + WindowFrameBound::CurrentRow, + ) + .await? + } + _ => unimplemented!(), + }; + + // add window func result column to blocks + let data_blocks = data_blocks + .into_iter() + .zip(window_cols) + .map(|(block, window_col)| { + block + .add_column( + window_col, + DataField::new("test", DataTypeImpl::Boolean(BooleanType {})), + ) + .unwrap() + }) + .collect(); + + Ok(Box::pin(DataBlockStream::create( + self.schema.clone(), + None, + data_blocks, + ))) + } +} + +/// evaluate range frame +async fn evaluate_range_with_aggr_func( + func: AggregateFunctionRef, + blocks: &[DataBlock], + start_bound: WindowFrameBound, + end_bound: WindowFrameBound, +) -> common_exception::Result> { + let window_cols = Vec::with_capacity(blocks.len()); +} + +/// evaluate rows frame +async fn evaluate_rows_with_aggr_func( + func: AggregateFunctionRef, + blocks: &[DataBlock], + start_bound: WindowFrameBound, + end_bound: WindowFrameBound, +) -> common_exception::Result> { + todo!() +} diff --git a/query/src/pipelines/transforms/window_func/mod.rs b/query/src/pipelines/transforms/window_func/mod.rs new file mode 100644 index 000000000000..ea0ed57e60ed --- /dev/null +++ b/query/src/pipelines/transforms/window_func/mod.rs @@ -0,0 +1,13 @@ +// Copyright 2022 Datafuse Labs. +// +// 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. diff --git a/query/src/sql/plan_parser.rs b/query/src/sql/plan_parser.rs index a9cbf24bd2e6..b56cb910d897 100644 --- a/query/src/sql/plan_parser.rs +++ b/query/src/sql/plan_parser.rs @@ -224,13 +224,38 @@ impl PlanParser { true => Ok(plan), false => { let exprs = data.window_expressions.to_vec(); - Ok(exprs.into_iter().fold(plan, |input, window_func| { - PlanBuilder::from(&input) - .window_aggr(window_func) - .unwrap() - .build() - .unwrap() - })) + exprs.into_iter().try_fold(plan, |input, window_func| { + let mut new_plan_builder = PlanBuilder::from(&input); + let mut sort_columns = vec![]; + match &window_func { + Expression::WindowFunction { + partition_by, + order_by, + .. + } => { + for expr in partition_by { + sort_columns.push(Expression::Sort { + expr: Box::new(expr.to_owned()), + asc: true, + nulls_first: false, + origin_expr: Box::new(expr.to_owned()), + }); + } + for expr in order_by { + sort_columns.push(expr.clone()); + } + } + _ => { + return Err(ErrorCode::LogicalError( + "It should be a WindowFunction expression!", + )) + } + } + if !sort_columns.is_empty() { + new_plan_builder = new_plan_builder.sort(&sort_columns)?; + } + new_plan_builder.window_func(window_func.clone())?.build() + }) } } } diff --git a/query/src/sql/statements/analyzer_expr.rs b/query/src/sql/statements/analyzer_expr.rs index d82ce9f24000..c66841cede4d 100644 --- a/query/src/sql/statements/analyzer_expr.rs +++ b/query/src/sql/statements/analyzer_expr.rs @@ -285,6 +285,25 @@ impl ExpressionAnalyzer { } } + let mut parameters = Vec::with_capacity(info.parameters.len()); + + for parameter in &info.parameters { + match ValueExprAnalyzer::analyze( + parameter, + self.context.get_current_session().get_type(), + )? { + Expression::Literal { value, .. } => { + parameters.push(value); + } + expr => { + return Err(ErrorCode::SyntaxException(format!( + "Unsupported value expression: {:?}, must be datavalue", + expr + ))); + } + }; + } + let mut arguments = Vec::with_capacity(info.args_count); for _ in 0..info.args_count { match args.pop() { @@ -317,6 +336,7 @@ impl ExpressionAnalyzer { Ok(Expression::WindowFunction { op: info.name.clone(), + params: parameters, args: arguments, partition_by, order_by, diff --git a/query/src/sql/statements/analyzer_statement.rs b/query/src/sql/statements/analyzer_statement.rs index 1ec87d27f862..d8a3c8d65a77 100644 --- a/query/src/sql/statements/analyzer_statement.rs +++ b/query/src/sql/statements/analyzer_statement.rs @@ -107,7 +107,7 @@ impl Debug for QueryAnalyzeState { } if !self.window_expressions.is_empty() { - debug_struct.field("window_aggr", &self.window_expressions); + debug_struct.field("window_func", &self.window_expressions); } if !self.expressions.is_empty() { From 65c3563adfc165102d2042bcba6fefca295e5b43 Mon Sep 17 00:00:00 2001 From: doki Date: Mon, 16 May 2022 16:03:58 +0800 Subject: [PATCH 08/18] basic func --- .../transforms/transform_window_func.rs | 224 ++++++++++-------- 1 file changed, 125 insertions(+), 99 deletions(-) diff --git a/query/src/pipelines/transforms/transform_window_func.rs b/query/src/pipelines/transforms/transform_window_func.rs index 4137cd0c39eb..e2409898396e 100644 --- a/query/src/pipelines/transforms/transform_window_func.rs +++ b/query/src/pipelines/transforms/transform_window_func.rs @@ -15,16 +15,18 @@ use std::any::Any; use std::sync::Arc; +use common_arrow::arrow::array::Array; +use common_arrow::arrow::array::ArrayRef; +use common_arrow::arrow::compute::partition::lexicographical_partition_ranges; +use common_arrow::arrow::compute::sort::SortColumn; use common_datablocks::DataBlock; -use common_datavalues::BooleanType; use common_datavalues::ColumnRef; +use common_datavalues::ColumnWithField; use common_datavalues::DataField; use common_datavalues::DataSchemaRef; -use common_datavalues::DataTypeImpl; +use common_datavalues::Series; +use common_functions::aggregates::eval_aggr; use common_functions::aggregates::AggregateFunctionFactory; -use common_functions::aggregates::AggregateFunctionRef; -use common_functions::window::WindowFrameBound; -use common_functions::window::WindowFrameUnits; use common_planners::Expression; use common_streams::DataBlockStream; use common_streams::SendableDataBlockStream; @@ -55,6 +57,91 @@ impl WindowFuncTransform { input_schema, } } + + /// evaluate range frame + /// requires that the block's already sorted by expressions of partition by and order by sub stmts + async fn evaluate_window_func_col( + &self, + block: &DataBlock, + ) -> common_exception::Result { + // extract the window function + let_extract!( + Expression::WindowFunction { + op, + params, + args, + partition_by, + order_by, + window_frame + }, + &self.window_func, + panic!() + ); + + match window_frame { + None => { + // at the moment, only supports aggr function + let is_aggr_func = AggregateFunctionFactory::instance().check(&op); + match is_aggr_func { + true => { + let mut arguments: Vec = Vec::with_capacity(args.len()); + for arg in args { + arguments.push(arg.to_data_field(&self.input_schema)?); + } + + let mut sort_cols: Vec = + Vec::with_capacity(partition_by.len() + order_by.len()); + let partition_by_arrow_array = partition_by + .iter() + .map(|expr| { + block + .try_column_by_name(&expr.column_name()) + .unwrap() + .as_arrow_array() + }) + .collect::>(); + let partition_by_column = partition_by_arrow_array + .iter() + .map(|array| SortColumn { + values: array.as_ref(), + options: None, + }) + .collect::>(); + + sort_cols.extend(partition_by_column); + // todo process sort expr + + let peer_rows = lexicographical_partition_ranges(&sort_cols).unwrap(); + + let window_cols = peer_rows + .map(|range| { + let offset = range.start; + let length = range.end - range.start; + let peered = block.slice(offset, length); + let args = arguments + .iter() + .map(|f| { + let arg = peered.try_column_by_name(f.name()).unwrap(); + ColumnWithField::new(arg.clone(), f.to_owned()) + }) + .collect::>(); + let agg_result = + eval_aggr(op, params.to_owned(), &args, peered.num_rows()) + .unwrap(); + Series::concat( + &(0..length).map(|_| agg_result.clone()).collect::>(), + ) + .unwrap() + }) + .collect::>(); + Ok(Series::concat(&window_cols).unwrap()) + } + false => unimplemented!(), + } + } + Some(window_frame) => unimplemented!("not yet support window frame"), + } + } } #[async_trait::async_trait] @@ -79,111 +166,50 @@ impl Processor for WindowFuncTransform { #[tracing::instrument(level = "debug", name = "window_aggr_func_execute", skip(self))] async fn execute(&self) -> common_exception::Result { let mut stream: SendableDataBlockStream = self.input.execute().await?; - let mut data_blocks: Vec = vec![]; + let mut blocks: Vec = vec![]; while let Some(block) = stream.next().await { let block = block?; - data_blocks.push(block); + blocks.push(block); } - let_extract!( - Expression::WindowFunction { - op, - params, - args, - partition_by, - order_by, - window_frame - }, - &self.window_func, - panic!() - ); - - let mut arguments = Vec::with_capacity(args.len()); - for arg in args { - arguments.push(arg.to_data_field(&self.input_schema)?); + if blocks.is_empty() { + return Ok(Box::pin(DataBlockStream::create( + self.schema.clone(), + None, + vec![], + ))); } - let is_aggr_func = AggregateFunctionFactory::instance().check(&op); - - let window_cols = match &window_frame { - Some(window_frame) if is_aggr_func => match &window_frame.units { - WindowFrameUnits::Range => { - evaluate_range_with_aggr_func( - AggregateFunctionFactory::instance().get( - op, - params.to_owned(), - arguments, - )?, - &data_blocks, - window_frame.start_bound, - window_frame.end_bound, - ) - .await? - } - WindowFrameUnits::Rows => { - evaluate_rows_with_aggr_func( - AggregateFunctionFactory::instance().get( - op, - params.to_owned(), - arguments, - )?, - &data_blocks, - window_frame.start_bound, - window_frame.end_bound, - ) - .await? - } - }, - None if is_aggr_func => { - evaluate_range_with_aggr_func( - AggregateFunctionFactory::instance().get(op, params.to_owned(), arguments)?, - &data_blocks, - WindowFrameBound::Preceding(None), - WindowFrameBound::CurrentRow, - ) - .await? - } - _ => unimplemented!(), - }; - - // add window func result column to blocks - let data_blocks = data_blocks - .into_iter() - .zip(window_cols) - .map(|(block, window_col)| { - block - .add_column( - window_col, - DataField::new("test", DataTypeImpl::Boolean(BooleanType {})), - ) - .unwrap() + // combine blocks + let schema = blocks[0].schema(); + + let combined_columns = (0..schema.num_fields()) + .map(|i| { + blocks + .iter() + .map(|block| block.column(i).clone()) + .collect::>() }) - .collect(); + .map(|columns| Series::concat(&columns).unwrap()) + .collect::>(); + + let block = DataBlock::create(schema.clone(), combined_columns); + + // evaluate the window function column + let window_col = self.evaluate_window_func_col(&block).await.unwrap(); + + // add window func result column to the block + let block = block + .add_column( + window_col, + self.window_func.to_data_field(&self.input_schema).unwrap(), + ) + .unwrap(); Ok(Box::pin(DataBlockStream::create( self.schema.clone(), None, - data_blocks, + vec![block], ))) } } - -/// evaluate range frame -async fn evaluate_range_with_aggr_func( - func: AggregateFunctionRef, - blocks: &[DataBlock], - start_bound: WindowFrameBound, - end_bound: WindowFrameBound, -) -> common_exception::Result> { - let window_cols = Vec::with_capacity(blocks.len()); -} - -/// evaluate rows frame -async fn evaluate_rows_with_aggr_func( - func: AggregateFunctionRef, - blocks: &[DataBlock], - start_bound: WindowFrameBound, - end_bound: WindowFrameBound, -) -> common_exception::Result> { - todo!() -} From e032aa7391a1cc6888b9509b12cbd189d460fd5f Mon Sep 17 00:00:00 2001 From: doki Date: Fri, 27 May 2022 23:24:52 +0800 Subject: [PATCH 09/18] naive evaluation and rows frame --- common/ast/src/udfs/udf_expr_visitor.rs | 8 +- .../src/columns/column_with_field.rs | 6 + common/functions/src/window/mod.rs | 22 +- .../src/window/{window_frame.rs => window.rs} | 42 +-- .../transforms/transform_window_func.rs | 278 +++++++++++++----- query/src/sql/plan_parser.rs | 33 +-- query/src/sql/statements/statement_select.rs | 2 +- 7 files changed, 239 insertions(+), 152 deletions(-) rename common/functions/src/window/{window_frame.rs => window.rs} (71%) diff --git a/common/ast/src/udfs/udf_expr_visitor.rs b/common/ast/src/udfs/udf_expr_visitor.rs index 1097a9d26acd..0cf5bfc4b62b 100644 --- a/common/ast/src/udfs/udf_expr_visitor.rs +++ b/common/ast/src/udfs/udf_expr_visitor.rs @@ -158,17 +158,17 @@ pub trait UDFExprVisitor: Sized + Send { fn visit_function(&mut self, function: &Function) -> Result<()> { for function_arg in &function.args { match function_arg { - FunctionArg::Named { arg, .. } => self.visit_function_arg(arg)?, - FunctionArg::Unnamed(arg) => self.visit_function_arg(arg)?, + FunctionArg::Named { arg, .. } => self.visit_function_arg(&arg)?, + FunctionArg::Unnamed(arg) => self.visit_function_arg(&arg)?, }; } if let Some(over) = &function.over { for partition_by in &over.partition_by { - UDFExprTraverser::accept(partition_by, self).await?; + UDFExprTraverser::accept(partition_by, self)?; } for order_by in &over.order_by { - UDFExprTraverser::accept(&order_by.expr, self).await?; + UDFExprTraverser::accept(&order_by.expr, self)?; } } diff --git a/common/datavalues/src/columns/column_with_field.rs b/common/datavalues/src/columns/column_with_field.rs index 762d65eecf2d..4a4aaa5ebe28 100644 --- a/common/datavalues/src/columns/column_with_field.rs +++ b/common/datavalues/src/columns/column_with_field.rs @@ -34,6 +34,12 @@ impl ColumnWithField { pub fn data_type(&self) -> &DataTypeImpl { self.field.data_type() } + pub fn slice(&self, offset: usize, length: usize) -> Self { + Self { + column: self.column.slice(offset, length), + field: self.field.to_owned(), + } + } } pub type ColumnsWithField = [ColumnWithField]; diff --git a/common/functions/src/window/mod.rs b/common/functions/src/window/mod.rs index 074748ca3e34..cd12f0784763 100644 --- a/common/functions/src/window/mod.rs +++ b/common/functions/src/window/mod.rs @@ -1,5 +1,19 @@ -mod window_frame; +// Copyright 2022 Datafuse Labs. +// +// 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. -pub use window_frame::WindowFrame; -pub use window_frame::WindowFrameBound; -pub use window_frame::WindowFrameUnits; +mod window; + +pub use window::WindowFrame; +pub use window::WindowFrameBound; +pub use window::WindowFrameUnits; diff --git a/common/functions/src/window/window_frame.rs b/common/functions/src/window/window.rs similarity index 71% rename from common/functions/src/window/window_frame.rs rename to common/functions/src/window/window.rs index 87d0251a32e8..c72c254b58fc 100644 --- a/common/functions/src/window/window_frame.rs +++ b/common/functions/src/window/window.rs @@ -1,4 +1,4 @@ -// Copyright 2021 Datafuse Labs. +// Copyright 2022 Datafuse Labs. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -20,20 +20,12 @@ use std::hash::Hasher; use common_exception::ErrorCode; use sqlparser::ast; -/// The frame-spec determines which output rows are read by an aggregate window function. -/// -/// The ending frame boundary can be omitted (if the BETWEEN and AND keywords that surround the -/// starting frame boundary are also omitted), in which case the ending frame boundary defaults to -/// CURRENT ROW. #[derive( Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash, serde::Serialize, serde::Deserialize, )] pub struct WindowFrame { - /// A frame type - either ROWS, RANGE or GROUPS pub units: WindowFrameUnits, - /// A starting frame boundary pub start_bound: WindowFrameBound, - /// An ending frame boundary pub end_bound: WindowFrameBound, } @@ -97,20 +89,11 @@ impl Default for WindowFrame { } } -/// There are three frame types: ROWS, GROUPS, and RANGE. The frame type determines how the -/// starting and ending boundaries of the frame are measured. #[derive( Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash, serde::Serialize, serde::Deserialize, )] pub enum WindowFrameUnits { - /// The RANGE frame type requires that the ORDER BY clause of the window have exactly one - /// term. Call that term "X". With the RANGE frame type, the elements of the frame are - /// determined by computing the value of expression X for all rows in the partition and framing - /// those rows for which the value of X is within a certain range of the value of X for the - /// current row. Range, - /// The ROWS frame type means that the starting and ending boundaries for the frame are - /// determined by counting individual rows relative to the current row. Rows, } @@ -133,36 +116,15 @@ impl From for WindowFrameUnits { } } +// todo u64 -> Expression #[derive(Debug, Clone, Copy, Eq, serde::Serialize, serde::Deserialize)] pub enum WindowFrameBound { - /// 1. UNBOUNDED PRECEDING - /// The frame boundary is the first row in the partition. - /// - /// 2. PRECEDING - /// must be a non-negative constant numeric expression. The boundary is a row that - /// is "units" prior to the current row. Preceding(Option), - /// 3. The current row. - /// - /// For RANGE and GROUPS frame types, peers of the current row are also - /// included in the frame, unless specifically excluded by the EXCLUDE clause. - /// This is true regardless of whether CURRENT ROW is used as the starting or ending frame - /// boundary. CurrentRow, - /// 4. This is the same as " PRECEDING" except that the boundary is units after the - /// current rather than before the current row. - /// - /// 5. UNBOUNDED FOLLOWING - /// The frame boundary is the last row in the partition. Following(Option), } impl WindowFrameBound { - /// get the rank of this window frame bound. - /// - /// the rank is a tuple of (u8, u64) because we'll firstly compare the kind and then the value - /// which requires special handling e.g. with preceding the larger the value the smaller the - /// rank and also for 0 preceding / following it is the same as current row fn get_rank(&self) -> (u8, u64) { match self { WindowFrameBound::Preceding(None) => (0, 0), diff --git a/query/src/pipelines/transforms/transform_window_func.rs b/query/src/pipelines/transforms/transform_window_func.rs index e2409898396e..90005a4cc0b7 100644 --- a/query/src/pipelines/transforms/transform_window_func.rs +++ b/query/src/pipelines/transforms/transform_window_func.rs @@ -13,20 +13,22 @@ // limitations under the License. use std::any::Any; +use std::ops::Range; use std::sync::Arc; -use common_arrow::arrow::array::Array; use common_arrow::arrow::array::ArrayRef; use common_arrow::arrow::compute::partition::lexicographical_partition_ranges; use common_arrow::arrow::compute::sort::SortColumn; use common_datablocks::DataBlock; -use common_datavalues::ColumnRef; use common_datavalues::ColumnWithField; use common_datavalues::DataField; use common_datavalues::DataSchemaRef; use common_datavalues::Series; use common_functions::aggregates::eval_aggr; use common_functions::aggregates::AggregateFunctionFactory; +use common_functions::window::WindowFrame; +use common_functions::window::WindowFrameBound; +use common_functions::window::WindowFrameUnits; use common_planners::Expression; use common_streams::DataBlockStream; use common_streams::SendableDataBlockStream; @@ -36,6 +38,7 @@ use futures::StreamExt; use crate::pipelines::processors::EmptyProcessor; use crate::pipelines::processors::Processor; +use crate::pipelines::transforms::get_sort_descriptions; pub struct WindowFuncTransform { window_func: Expression, @@ -58,12 +61,8 @@ impl WindowFuncTransform { } } - /// evaluate range frame - /// requires that the block's already sorted by expressions of partition by and order by sub stmts - async fn evaluate_window_func_col( - &self, - block: &DataBlock, - ) -> common_exception::Result { + /// evaluate window function for each frame and return the result column + async fn evaluate_window_func(&self, block: &DataBlock) -> common_exception::Result { // extract the window function let_extract!( Expression::WindowFunction { @@ -78,68 +77,209 @@ impl WindowFuncTransform { panic!() ); - match window_frame { + // sort block by partition_by and order_by exprs + let mut sort_exprs: Vec = + Vec::with_capacity(partition_by.len() + order_by.len()); + sort_exprs.extend( + partition_by + .iter() + .map(|part_by_expr| Expression::Sort { + expr: Box::new(part_by_expr.to_owned()), + asc: true, + nulls_first: false, + origin_expr: Box::new(part_by_expr.to_owned()), + }) + .collect::>(), + ); + sort_exprs.extend(order_by.to_owned()); + let sort_column_desc = get_sort_descriptions(block.schema(), &sort_exprs)?; + let block = DataBlock::sort_block(&block, &sort_column_desc, None)?; + + // set default window frame + let window_frame = match window_frame { None => { + // compute the whole partition + match order_by.is_empty() { + // RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + false => WindowFrame::default(), + // RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + true => WindowFrame { + units: WindowFrameUnits::Range, + start_bound: WindowFrameBound::Preceding(None), + end_bound: WindowFrameBound::Following(None), + }, + } + } + Some(window_frame) => window_frame.to_owned(), + }; + + // determine window frame bounds of each tuple + let frame_bounds = Self::frame_bounds_per_tuple( + &block, + partition_by, + order_by, + window_frame.units, + window_frame.start_bound, + window_frame.end_bound, + ); + + // function calculate + let mut arguments: Vec = Vec::with_capacity(args.len()); + for arg in args { + let arg_field: DataField = arg.to_data_field(&self.input_schema)?; + let arg_column = block.try_column_by_name(arg_field.name()).unwrap(); + arguments.push(ColumnWithField::new(Arc::clone(arg_column), arg_field)); + } + + let window_col_per_tuple = (0..block.num_rows()) + .map(|i| { + let frame = &frame_bounds[i]; + let frame_start = frame.start; + let frame_end = frame.end; + let frame_size = frame_end - frame_start; + let args_slice = arguments + .iter() + .map(|c| c.slice(frame_start, frame_size)) + .collect::>(); // at the moment, only supports aggr function - let is_aggr_func = AggregateFunctionFactory::instance().check(&op); - match is_aggr_func { - true => { - let mut arguments: Vec = Vec::with_capacity(args.len()); - for arg in args { - arguments.push(arg.to_data_field(&self.input_schema)?); + if !AggregateFunctionFactory::instance().check(op) { + unimplemented!("not yet impl built-in window func"); + } + eval_aggr(op, params.to_owned(), &args_slice, frame_size).unwrap() + }) + .collect::>(); + + let window_col = Series::concat(&window_col_per_tuple).unwrap(); + + block.add_column( + window_col, + self.window_func.to_data_field(&self.input_schema).unwrap(), + ) + } + + /// compute frame range for each tuple + fn frame_bounds_per_tuple( + block: &DataBlock, + partition_by: &[Expression], + order_by: &[Expression], + frame_units: WindowFrameUnits, + start: WindowFrameBound, + end: WindowFrameBound, + ) -> Vec> { + match (frame_units, start, end) { + (_, WindowFrameBound::Preceding(None), WindowFrameBound::Following(None)) => { + let partition_by_arrow_array = partition_by + .iter() + .map(|expr| { + block + .try_column_by_name(&expr.column_name()) + .unwrap() + .as_arrow_array() + }) + .collect::>(); + let partition_by_sort_column = partition_by_arrow_array + .iter() + .map(|array| SortColumn { + values: array.as_ref(), + options: None, + }) + .collect::>(); + let mut partition_boundaries = + lexicographical_partition_ranges(&partition_by_sort_column).unwrap(); + let mut partition = partition_boundaries.next().unwrap(); + (0..block.num_rows()) + .map(|i| { + if i >= partition.end && i < block.num_rows() { + partition = partition_boundaries.next().unwrap(); } + partition.clone() + }) + .collect::>() + } + (frame_unit, frame_start, frame_end) => { + let partition_by_arrow_array = partition_by + .iter() + .map(|expr| { + block + .try_column_by_name(&expr.column_name()) + .unwrap() + .as_arrow_array() + }) + .collect::>(); + let mut partition_by_sort_column = partition_by_arrow_array + .iter() + .map(|array| SortColumn { + values: array.as_ref(), + options: None, + }) + .collect::>(); - let mut sort_cols: Vec = - Vec::with_capacity(partition_by.len() + order_by.len()); - let partition_by_arrow_array = partition_by - .iter() - .map(|expr| { - block - .try_column_by_name(&expr.column_name()) - .unwrap() - .as_arrow_array() - }) - .collect::>(); - let partition_by_column = partition_by_arrow_array - .iter() - .map(|array| SortColumn { - values: array.as_ref(), - options: None, - }) - .collect::>(); - - sort_cols.extend(partition_by_column); - // todo process sort expr - - let peer_rows = lexicographical_partition_ranges(&sort_cols).unwrap(); - - let window_cols = peer_rows - .map(|range| { - let offset = range.start; - let length = range.end - range.start; - let peered = block.slice(offset, length); - let args = arguments - .iter() - .map(|f| { - let arg = peered.try_column_by_name(f.name()).unwrap(); - ColumnWithField::new(arg.clone(), f.to_owned()) - }) - .collect::>(); - let agg_result = - eval_aggr(op, params.to_owned(), &args, peered.num_rows()) - .unwrap(); - Series::concat( - &(0..length).map(|_| agg_result.clone()).collect::>(), - ) - .unwrap() + let order_by_arrow_array = order_by + .iter() + .map(|expr| { + block + .try_column_by_name(&expr.column_name()) + .unwrap() + .as_arrow_array() + }) + .collect::>(); + let order_by_sort_column = order_by_arrow_array + .iter() + .map(|array| SortColumn { + values: array.as_ref(), + options: None, + }) + .collect::>(); + partition_by_sort_column.extend(order_by_sort_column); + + let mut frame_boundaries = + lexicographical_partition_ranges(&partition_by_sort_column).unwrap(); + + match frame_unit { + WindowFrameUnits::Rows => { + let mut frame_bound = frame_boundaries.next().unwrap(); + + (0..block.num_rows()) + .map(|i| { + if i >= frame_bound.end && i < block.num_rows() { + frame_bound = frame_boundaries.next().unwrap(); + } + let mut start = frame_bound.start; + let mut end = frame_bound.end; + match frame_start { + WindowFrameBound::Preceding(Some(preceding)) => { + start = std::cmp::max( + start, + if i < preceding as usize { + 0 + } else { + i - preceding as usize + }, + ); + } + WindowFrameBound::CurrentRow => { + start = i; + } + _ => (), + } + match frame_end { + WindowFrameBound::CurrentRow => { + end = i; + } + WindowFrameBound::Following(Some(following)) => { + end = std::cmp::min(end, i + 1 + following as usize); + } + _ => (), + } + start..end }) - .collect::>(); - Ok(Series::concat(&window_cols).unwrap()) + .collect::>() + } + WindowFrameUnits::Range => { + todo!() } - false => unimplemented!(), } } - Some(window_frame) => unimplemented!("not yet support window frame"), } } } @@ -147,7 +287,7 @@ impl WindowFuncTransform { #[async_trait::async_trait] impl Processor for WindowFuncTransform { fn name(&self) -> &str { - "WindowAggrFuncTransform" + "WindowFuncTransform" } fn connect_to(&mut self, input: Arc) -> common_exception::Result<()> { @@ -163,7 +303,7 @@ impl Processor for WindowFuncTransform { self } - #[tracing::instrument(level = "debug", name = "window_aggr_func_execute", skip(self))] + #[tracing::instrument(level = "debug", name = "window_func_execute", skip(self))] async fn execute(&self) -> common_exception::Result { let mut stream: SendableDataBlockStream = self.input.execute().await?; let mut blocks: Vec = vec![]; @@ -196,15 +336,7 @@ impl Processor for WindowFuncTransform { let block = DataBlock::create(schema.clone(), combined_columns); // evaluate the window function column - let window_col = self.evaluate_window_func_col(&block).await.unwrap(); - - // add window func result column to the block - let block = block - .add_column( - window_col, - self.window_func.to_data_field(&self.input_schema).unwrap(), - ) - .unwrap(); + let block = self.evaluate_window_func(&block).await.unwrap(); Ok(Box::pin(DataBlockStream::create( self.schema.clone(), diff --git a/query/src/sql/plan_parser.rs b/query/src/sql/plan_parser.rs index fc77ec327d8f..67e3a10d6288 100644 --- a/query/src/sql/plan_parser.rs +++ b/query/src/sql/plan_parser.rs @@ -264,36 +264,9 @@ impl PlanParser { false => { let exprs = data.window_expressions.to_vec(); exprs.into_iter().try_fold(plan, |input, window_func| { - let mut new_plan_builder = PlanBuilder::from(&input); - let mut sort_columns = vec![]; - match &window_func { - Expression::WindowFunction { - partition_by, - order_by, - .. - } => { - for expr in partition_by { - sort_columns.push(Expression::Sort { - expr: Box::new(expr.to_owned()), - asc: true, - nulls_first: false, - origin_expr: Box::new(expr.to_owned()), - }); - } - for expr in order_by { - sort_columns.push(expr.clone()); - } - } - _ => { - return Err(ErrorCode::LogicalError( - "It should be a WindowFunction expression!", - )) - } - } - if !sort_columns.is_empty() { - new_plan_builder = new_plan_builder.sort(&sort_columns)?; - } - new_plan_builder.window_func(window_func.clone())?.build() + PlanBuilder::from(&input) + .window_func(window_func.clone())? + .build() }) } } diff --git a/query/src/sql/statements/statement_select.rs b/query/src/sql/statements/statement_select.rs index 3382f7349dcb..e26ff7e51940 100644 --- a/query/src/sql/statements/statement_select.rs +++ b/query/src/sql/statements/statement_select.rs @@ -412,7 +412,7 @@ impl DfQueryStatement { data_block = res; } Err(cause) => { - return Err(cause.add_message_back(" (while in select window aggr)")); + return Err(cause.add_message_back(" (while in select window func)")); } } } From 9701d7c7388bfa1287d5ee330951aae228a45221 Mon Sep 17 00:00:00 2001 From: doki Date: Tue, 31 May 2022 22:41:17 +0800 Subject: [PATCH 10/18] stage but with a mock hash function --- common/planners/src/plan_expression_common.rs | 22 ++++++++ .../plan_schedulers/plan_scheduler.rs | 29 ++++++++++ .../plan_schedulers/plan_scheduler_query.rs | 2 +- query/src/optimizers/optimizer_scatters.rs | 53 ++++++++++++++++++- 4 files changed, 103 insertions(+), 3 deletions(-) diff --git a/common/planners/src/plan_expression_common.rs b/common/planners/src/plan_expression_common.rs index 3927e8de323c..a612ed67dec4 100644 --- a/common/planners/src/plan_expression_common.rs +++ b/common/planners/src/plan_expression_common.rs @@ -18,6 +18,7 @@ use std::collections::HashSet; use common_datavalues::prelude::*; use common_exception::ErrorCode; use common_exception::Result; +use common_functions::aggregates::AggregateFunctionFactory; use common_functions::scalars::FunctionFactory; use crate::validate_function_arg; @@ -523,6 +524,9 @@ impl ExpressionVisitor for ExpressionDataTypeVisitor { Ok(self) } Expression::WindowFunction { + op, + params, + args, partition_by, order_by, .. @@ -530,6 +534,24 @@ impl ExpressionVisitor for ExpressionDataTypeVisitor { for _ in 0..partition_by.len() + order_by.len() { self.stack.remove(0); } + + if !AggregateFunctionFactory::instance().check(op) { + return Err(ErrorCode::LogicalError( + "not yet support non aggr window function", + )); + } + + let mut fields = Vec::with_capacity(args.len()); + for arg in args.iter() { + let arg_type = self.stack.remove(0); + fields.push(DataField::new(&arg.column_name(), arg_type)); + } + + let aggregate_window_function = + AggregateFunctionFactory::instance().get(op, params.clone(), fields); + let return_type = aggregate_window_function.unwrap().return_type().unwrap(); + + self.stack.push(return_type); Ok(self) } Expression::Cast { data_type, .. } => { diff --git a/query/src/interpreters/plan_schedulers/plan_scheduler.rs b/query/src/interpreters/plan_schedulers/plan_scheduler.rs index 064e4b412565..0a07af5c03c9 100644 --- a/query/src/interpreters/plan_schedulers/plan_scheduler.rs +++ b/query/src/interpreters/plan_schedulers/plan_scheduler.rs @@ -42,6 +42,7 @@ use common_planners::SortPlan; use common_planners::StageKind; use common_planners::StagePlan; use common_planners::SubQueriesSetPlan; +use common_planners::WindowFuncPlan; use common_tracing::tracing; use crate::api::BroadcastAction; @@ -296,6 +297,7 @@ impl PlanScheduler { match node { PlanNode::AggregatorPartial(plan) => self.visit_aggr_part(plan, tasks), PlanNode::AggregatorFinal(plan) => self.visit_aggr_final(plan, tasks), + PlanNode::WindowFunc(plan) => self.visit_window_func(plan, tasks), PlanNode::Empty(plan) => self.visit_empty(plan, tasks), PlanNode::Projection(plan) => self.visit_projection(plan, tasks), PlanNode::Filter(plan) => self.visit_filter(plan, tasks), @@ -353,6 +355,33 @@ impl PlanScheduler { Ok(()) } + fn visit_window_func(&mut self, plan: &WindowFuncPlan, tasks: &mut Tasks) -> Result<()> { + self.visit_plan_node(plan.input.as_ref(), tasks)?; + match self.running_mode { + RunningMode::Cluster => self.visit_cluster_window_func(plan), + RunningMode::Standalone => self.visit_local_window_func(plan), + } + Ok(()) + } + + fn visit_cluster_window_func(&mut self, plan: &WindowFuncPlan) { + for index in 0..self.nodes_plan.len() { + self.nodes_plan[index] = PlanNode::WindowFunc(WindowFuncPlan { + window_func: plan.window_func.clone(), + input: Arc::new(self.nodes_plan[index].clone()), + schema: plan.schema.clone(), + }) + } + } + + fn visit_local_window_func(&mut self, plan: &WindowFuncPlan) { + self.nodes_plan[self.local_pos] = PlanNode::WindowFunc(WindowFuncPlan { + window_func: plan.window_func.clone(), + input: Arc::new(self.nodes_plan[self.local_pos].clone()), + schema: plan.schema(), + }); + } + fn visit_local_aggr_final(&mut self, plan: &AggregatorFinalPlan) { self.nodes_plan[self.local_pos] = PlanNode::AggregatorFinal(AggregatorFinalPlan { schema: plan.schema.clone(), diff --git a/query/src/interpreters/plan_schedulers/plan_scheduler_query.rs b/query/src/interpreters/plan_schedulers/plan_scheduler_query.rs index 570c059f2f80..f355822ad9a5 100644 --- a/query/src/interpreters/plan_schedulers/plan_scheduler_query.rs +++ b/query/src/interpreters/plan_schedulers/plan_scheduler_query.rs @@ -50,7 +50,7 @@ pub async fn schedule_query( let pipeline_builder = PipelineBuilder::create(ctx.clone()); let mut in_local_pipeline = pipeline_builder.build(&scheduled_tasks.get_local_task())?; - tracing::debug!("local_pipeline:\n{:?}", in_local_pipeline); + tracing::log::debug!("local_pipeline:\n{:?}", in_local_pipeline); match in_local_pipeline.execute().await { Ok(stream) => Ok(ScheduledStream::create(ctx.clone(), scheduled, stream)), diff --git a/query/src/optimizers/optimizer_scatters.rs b/query/src/optimizers/optimizer_scatters.rs index 7a88ebd6b982..45d1217c4458 100644 --- a/query/src/optimizers/optimizer_scatters.rs +++ b/query/src/optimizers/optimizer_scatters.rs @@ -14,6 +14,7 @@ use std::sync::Arc; +use common_datablocks::DataBlock; use common_datavalues::DataSchemaRef; use common_datavalues::DataValue; use common_exception::ErrorCode; @@ -32,6 +33,7 @@ use common_planners::SortPlan; use common_planners::StageKind; use common_planners::StagePlan; use common_planners::WindowFuncPlan; +use enum_extract::let_extract; use crate::optimizers::Optimizer; use crate::sessions::QueryContext; @@ -111,11 +113,58 @@ impl ScattersOptimizerImpl { } fn cluster_window(&mut self, plan: &WindowFuncPlan) -> Result { - todo!() + self.running_mode = RunningMode::Cluster; + + match self.input.take() { + None => Err(ErrorCode::LogicalError("Cluster window input is None")), + Some(input) => { + let_extract!( + Expression::WindowFunction { + op: _op, + params: _params, + args: _args, + partition_by: partition_by, + order_by: _order_by, + window_frame: _window_frame + }, + &plan.window_func, + panic!() + ); + // let sample_block = DataBlock::empty_with_schema(input.schema()); + // let method = DataBlock::choose_hash_method( + // &sample_block, + // &partition_by + // .iter() + // .map(|expr| expr.column_name()) + // .collect::>(), + // )?; + let scatters_expr = Expression::ScalarFunction { + op: String::from("sipHash"), + args: partition_by.to_owned(), + }; + + let scatter_plan = PlanNode::Stage(StagePlan { + scatters_expr, + kind: StageKind::Normal, + input, + }); + + Ok(PlanNode::WindowFunc(WindowFuncPlan { + window_func: plan.window_func.to_owned(), + input: Arc::new(scatter_plan), + schema: plan.schema.to_owned(), + })) + } + } } fn standalone_window(&mut self, plan: &WindowFuncPlan) -> Result { - todo!() + match self.input.take() { + None => Err(ErrorCode::LogicalError("Standalone window input is None")), + Some(input) => PlanBuilder::from(input.as_ref()) + .window_func(plan.window_func.to_owned())? + .build(), + } } fn cluster_sort(&mut self, plan: &SortPlan) -> Result { From ebd377fa166d0f9ff22e13dc2d4abe717ec4a94a Mon Sep 17 00:00:00 2001 From: doki Date: Wed, 1 Jun 2022 10:41:28 +0800 Subject: [PATCH 11/18] support multi partition by --- query/src/optimizers/optimizer_scatters.rs | 23 ++++++++++------------ 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/query/src/optimizers/optimizer_scatters.rs b/query/src/optimizers/optimizer_scatters.rs index 45d1217c4458..26876e503ab6 100644 --- a/query/src/optimizers/optimizer_scatters.rs +++ b/query/src/optimizers/optimizer_scatters.rs @@ -14,7 +14,6 @@ use std::sync::Arc; -use common_datablocks::DataBlock; use common_datavalues::DataSchemaRef; use common_datavalues::DataValue; use common_exception::ErrorCode; @@ -130,18 +129,16 @@ impl ScattersOptimizerImpl { &plan.window_func, panic!() ); - // let sample_block = DataBlock::empty_with_schema(input.schema()); - // let method = DataBlock::choose_hash_method( - // &sample_block, - // &partition_by - // .iter() - // .map(|expr| expr.column_name()) - // .collect::>(), - // )?; - let scatters_expr = Expression::ScalarFunction { - op: String::from("sipHash"), - args: partition_by.to_owned(), - }; + + let mut concat_ws_args = vec![Expression::create_literal(DataValue::String( + "#".as_bytes().to_vec(), + ))]; + concat_ws_args.extend(partition_by.to_owned()); + let concat_partition_by = + Expression::create_scalar_function("concat_ws", concat_ws_args); + + let scatters_expr = + Expression::create_scalar_function("sipHash", vec![concat_partition_by]); let scatter_plan = PlanNode::Stage(StagePlan { scatters_expr, From f6fd7234955b8df634ed31107bcfc71afa66ee2b Mon Sep 17 00:00:00 2001 From: doki Date: Wed, 1 Jun 2022 18:45:58 +0800 Subject: [PATCH 12/18] support range frame --- common/functions/src/window/window.rs | 17 - .../transforms/transform_window_func.rs | 324 ++++++++++++------ 2 files changed, 220 insertions(+), 121 deletions(-) diff --git a/common/functions/src/window/window.rs b/common/functions/src/window/window.rs index c72c254b58fc..57e56298799a 100644 --- a/common/functions/src/window/window.rs +++ b/common/functions/src/window/window.rs @@ -54,22 +54,6 @@ impl TryFrom for WindowFrame { ))) } else { let units = value.units.into(); - if units == WindowFrameUnits::Range { - for bound in &[start_bound, end_bound] { - match bound { - WindowFrameBound::Preceding(Some(v)) - | WindowFrameBound::Following(Some(v)) - if *v > 0 => - { - Err(ErrorCode::UnImplement(format!( - "With WindowFrameUnits={}, the bound cannot be {} PRECEDING or FOLLOWING at the moment", - units, v - ))) - } - _ => Ok(()), - }?; - } - } Ok(Self { units, start_bound, @@ -116,7 +100,6 @@ impl From for WindowFrameUnits { } } -// todo u64 -> Expression #[derive(Debug, Clone, Copy, Eq, serde::Serialize, serde::Deserialize)] pub enum WindowFrameBound { Preceding(Option), diff --git a/query/src/pipelines/transforms/transform_window_func.rs b/query/src/pipelines/transforms/transform_window_func.rs index 90005a4cc0b7..d090bced3c3e 100644 --- a/query/src/pipelines/transforms/transform_window_func.rs +++ b/query/src/pipelines/transforms/transform_window_func.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::any::Any; +use std::cmp::Ordering; use std::ops::Range; use std::sync::Arc; @@ -23,12 +24,15 @@ use common_datablocks::DataBlock; use common_datavalues::ColumnWithField; use common_datavalues::DataField; use common_datavalues::DataSchemaRef; +use common_datavalues::DataValue; use common_datavalues::Series; use common_functions::aggregates::eval_aggr; use common_functions::aggregates::AggregateFunctionFactory; +use common_functions::scalars::assert_numeric; use common_functions::window::WindowFrame; use common_functions::window::WindowFrameBound; use common_functions::window::WindowFrameUnits; +use common_planners::sort; use common_planners::Expression; use common_streams::DataBlockStream; use common_streams::SendableDataBlockStream; @@ -61,7 +65,7 @@ impl WindowFuncTransform { } } - /// evaluate window function for each frame and return the result column + /// evaluate window function for each frame and return the result block async fn evaluate_window_func(&self, block: &DataBlock) -> common_exception::Result { // extract the window function let_extract!( @@ -166,120 +170,232 @@ impl WindowFuncTransform { start: WindowFrameBound, end: WindowFrameBound, ) -> Vec> { + let partition_by_arrow_array = partition_by + .iter() + .map(|expr| { + block + .try_column_by_name(&expr.column_name()) + .unwrap() + .as_arrow_array() + }) + .collect::>(); + let partition_by_sort_column = partition_by_arrow_array + .iter() + .map(|array| SortColumn { + values: array.as_ref(), + options: None, + }) + .collect::>(); + + let mut partition_boundaries = + lexicographical_partition_ranges(&partition_by_sort_column).unwrap(); + let mut partition = partition_boundaries.next().unwrap(); + match (frame_units, start, end) { - (_, WindowFrameBound::Preceding(None), WindowFrameBound::Following(None)) => { - let partition_by_arrow_array = partition_by - .iter() - .map(|expr| { - block - .try_column_by_name(&expr.column_name()) - .unwrap() - .as_arrow_array() - }) - .collect::>(); - let partition_by_sort_column = partition_by_arrow_array - .iter() - .map(|array| SortColumn { - values: array.as_ref(), - options: None, - }) - .collect::>(); - let mut partition_boundaries = - lexicographical_partition_ranges(&partition_by_sort_column).unwrap(); - let mut partition = partition_boundaries.next().unwrap(); - (0..block.num_rows()) - .map(|i| { - if i >= partition.end && i < block.num_rows() { - partition = partition_boundaries.next().unwrap(); + (_, WindowFrameBound::Preceding(None), WindowFrameBound::Following(None)) => (0..block + .num_rows()) + .map(|i| { + if i >= partition.end && i < block.num_rows() { + partition = partition_boundaries.next().unwrap(); + } + partition.clone() + }) + .collect::>(), + (WindowFrameUnits::Rows, frame_start, frame_end) => (0..block.num_rows()) + .map(|i| { + if i >= partition.end && i < block.num_rows() { + partition = partition_boundaries.next().unwrap(); + } + let mut start = partition.start; + let mut end = partition.end; + match frame_start { + WindowFrameBound::Preceding(Some(preceding)) => { + start = std::cmp::max( + start, + if i < preceding as usize { + 0 + } else { + i - preceding as usize + }, + ); } - partition.clone() - }) - .collect::>() - } - (frame_unit, frame_start, frame_end) => { - let partition_by_arrow_array = partition_by - .iter() - .map(|expr| { - block - .try_column_by_name(&expr.column_name()) - .unwrap() - .as_arrow_array() - }) - .collect::>(); - let mut partition_by_sort_column = partition_by_arrow_array - .iter() - .map(|array| SortColumn { - values: array.as_ref(), - options: None, - }) - .collect::>(); + WindowFrameBound::CurrentRow => { + start = i; + } + WindowFrameBound::Following(Some(following)) => { + start = std::cmp::min(end, i + following as usize) + } + _ => (), + } + match frame_end { + WindowFrameBound::Preceding(Some(preceding)) => { + end = std::cmp::max(start, i + 1 - preceding as usize); + } + WindowFrameBound::CurrentRow => { + end = i + 1; + } + WindowFrameBound::Following(Some(following)) => { + end = std::cmp::min(end, i + 1 + following as usize); + } + _ => (), + } + start..end + }) + .collect::>(), + (WindowFrameUnits::Range, frame_start, frame_end) => match (frame_start, frame_end) { + (WindowFrameBound::Preceding(None), WindowFrameBound::CurrentRow) => { + let mut partition_by_sort_column = partition_by_sort_column; + let order_by_arrow_array = order_by + .iter() + .map(|expr| { + block + .try_column_by_name(&expr.column_name()) + .unwrap() + .as_arrow_array() + }) + .collect::>(); + let order_by_sort_column = order_by_arrow_array + .iter() + .map(|array| SortColumn { + values: array.as_ref(), + options: None, + }) + .collect::>(); + partition_by_sort_column.extend(order_by_sort_column); - let order_by_arrow_array = order_by - .iter() - .map(|expr| { - block - .try_column_by_name(&expr.column_name()) - .unwrap() - .as_arrow_array() - }) - .collect::>(); - let order_by_sort_column = order_by_arrow_array - .iter() - .map(|array| SortColumn { - values: array.as_ref(), - options: None, - }) - .collect::>(); - partition_by_sort_column.extend(order_by_sort_column); + let mut peered_boundaries = + lexicographical_partition_ranges(&partition_by_sort_column).unwrap(); + let mut peered = peered_boundaries.next().unwrap(); + (0..block.num_rows()) + .map(|i| { + if i >= partition.end && i < block.num_rows() { + partition = partition_boundaries.next().unwrap(); + } + if i >= peered.end && i < block.num_rows() { + peered = peered_boundaries.next().unwrap(); + } + partition.start..peered.end + }) + .collect::>() + } + (WindowFrameBound::CurrentRow, WindowFrameBound::Following(None)) => { + let mut partition_by_sort_column = partition_by_sort_column; + let order_by_arrow_array = order_by + .iter() + .map(|expr| { + block + .try_column_by_name(&expr.column_name()) + .unwrap() + .as_arrow_array() + }) + .collect::>(); + let order_by_sort_column = order_by_arrow_array + .iter() + .map(|array| SortColumn { + values: array.as_ref(), + options: None, + }) + .collect::>(); + partition_by_sort_column.extend(order_by_sort_column); - let mut frame_boundaries = - lexicographical_partition_ranges(&partition_by_sort_column).unwrap(); + let mut peered_boundaries = + lexicographical_partition_ranges(&partition_by_sort_column).unwrap(); + let mut peered = peered_boundaries.next().unwrap(); + (0..block.num_rows()) + .map(|i| { + if i >= partition.end && i < block.num_rows() { + partition = partition_boundaries.next().unwrap(); + } + if i >= peered.end && i < block.num_rows() { + peered = peered_boundaries.next().unwrap(); + } + peered.start..partition.end + }) + .collect::>() + } + (frame_start, frame_end) => { + assert_eq!( + order_by.len(), + 1, + "Range mode is only possible if the query has exactly one numeric order by expression." + ); + assert_numeric( + block + .schema() + .field_with_name(&order_by[0].column_name()) + .unwrap() + .data_type(), + ) + .expect( + "Range mode is only possible if the query has exactly one numeric order by expression.", + ); - match frame_unit { - WindowFrameUnits::Rows => { - let mut frame_bound = frame_boundaries.next().unwrap(); + let order_by_values = block + .try_column_by_name(&order_by[0].column_name()) + .unwrap() + .to_values(); - (0..block.num_rows()) - .map(|i| { - if i >= frame_bound.end && i < block.num_rows() { - frame_bound = frame_boundaries.next().unwrap(); + (0..block.num_rows()) + .map(|i| { + let current_value = &order_by_values[i]; + if i >= partition.end && i < block.num_rows() { + partition = partition_boundaries.next().unwrap(); + } + let mut start = partition.start; + let mut end = partition.end; + match frame_start { + WindowFrameBound::Preceding(Some(preceding)) => { + let idx = order_by_values.partition_point(|x| { + let min = DataValue::from( + current_value.as_f64().unwrap() - preceding as f64, + ); + x.cmp(&min) == Ordering::Less + }); + start = std::cmp::max(start, idx); } - let mut start = frame_bound.start; - let mut end = frame_bound.end; - match frame_start { - WindowFrameBound::Preceding(Some(preceding)) => { - start = std::cmp::max( - start, - if i < preceding as usize { - 0 - } else { - i - preceding as usize - }, + WindowFrameBound::CurrentRow => { + start = i; + } + WindowFrameBound::Following(Some(following)) => { + let idx = order_by_values.partition_point(|x| { + let max = DataValue::from( + current_value.as_f64().unwrap() + following as f64, ); - } - WindowFrameBound::CurrentRow => { - start = i; - } - _ => (), + x.cmp(&max) == Ordering::Less + }); + start = std::cmp::min(end, idx); } - match frame_end { - WindowFrameBound::CurrentRow => { - end = i; - } - WindowFrameBound::Following(Some(following)) => { - end = std::cmp::min(end, i + 1 + following as usize); - } - _ => (), + _ => (), + } + match frame_end { + WindowFrameBound::Preceding(Some(preceding)) => { + let idx = order_by_values.partition_point(|x| { + let min = DataValue::from( + current_value.as_f64().unwrap() - preceding as f64, + ); + x.cmp(&min) != Ordering::Greater + }); + end = std::cmp::max(start, idx); } - start..end - }) - .collect::>() - } - WindowFrameUnits::Range => { - todo!() - } + WindowFrameBound::CurrentRow => { + end = i + 1; + } + WindowFrameBound::Following(Some(following)) => { + let idx = order_by_values.partition_point(|x| { + let max = DataValue::from( + current_value.as_f64().unwrap() + following as f64, + ); + x.cmp(&max) != Ordering::Greater + }); + end = std::cmp::min(end, idx); + } + _ => (), + } + start..end + }) + .collect::>() } - } + }, } } } From 70b2ed54d2b67d90621b6c4af3f6356e5f46d9f4 Mon Sep 17 00:00:00 2001 From: doki Date: Tue, 7 Jun 2022 00:40:16 +0800 Subject: [PATCH 13/18] evaluation with segment tree and clean the code --- Cargo.lock | 7 ++ common/ast/src/udfs/udf_expr_visitor.rs | 4 +- .../aggregates/aggregate_function_state.rs | 2 +- common/functions/src/window/function.rs | 18 +++ common/functions/src/window/mod.rs | 8 +- .../src/window/{window.rs => window_frame.rs} | 0 query/Cargo.toml | 1 + query/src/optimizers/optimizer_scatters.rs | 2 +- .../transforms/transform_window_func.rs | 114 ++++++++++++++++-- query/src/sql/plan_parser.rs | 4 +- 10 files changed, 136 insertions(+), 24 deletions(-) create mode 100644 common/functions/src/window/function.rs rename common/functions/src/window/{window.rs => window_frame.rs} (100%) diff --git a/Cargo.lock b/Cargo.lock index 575ea328c8db..8901370a8244 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1874,6 +1874,7 @@ dependencies = [ "regex", "reqwest", "rsa", + "segment-tree", "serde", "serde-bridge", "serde_json", @@ -5629,6 +5630,12 @@ dependencies = [ "libc", ] +[[package]] +name = "segment-tree" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f7dbd0d32cabaa6c7c3286d756268247538d613b621227bfe59237d7bbb271a" + [[package]] name = "semver" version = "1.0.4" diff --git a/common/ast/src/udfs/udf_expr_visitor.rs b/common/ast/src/udfs/udf_expr_visitor.rs index 0cf5bfc4b62b..24b1faacc95d 100644 --- a/common/ast/src/udfs/udf_expr_visitor.rs +++ b/common/ast/src/udfs/udf_expr_visitor.rs @@ -158,8 +158,8 @@ pub trait UDFExprVisitor: Sized + Send { fn visit_function(&mut self, function: &Function) -> Result<()> { for function_arg in &function.args { match function_arg { - FunctionArg::Named { arg, .. } => self.visit_function_arg(&arg)?, - FunctionArg::Unnamed(arg) => self.visit_function_arg(&arg)?, + FunctionArg::Named { arg, .. } => self.visit_function_arg(arg)?, + FunctionArg::Unnamed(arg) => self.visit_function_arg(arg)?, }; } diff --git a/common/functions/src/aggregates/aggregate_function_state.rs b/common/functions/src/aggregates/aggregate_function_state.rs index d8b13847bce7..d3176facb805 100644 --- a/common/functions/src/aggregates/aggregate_function_state.rs +++ b/common/functions/src/aggregates/aggregate_function_state.rs @@ -20,7 +20,7 @@ use common_exception::Result; use crate::aggregates::AggregateFunctionRef; -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Default)] pub struct StateAddr { addr: usize, } diff --git a/common/functions/src/window/function.rs b/common/functions/src/window/function.rs new file mode 100644 index 000000000000..d614d587efa3 --- /dev/null +++ b/common/functions/src/window/function.rs @@ -0,0 +1,18 @@ +// Copyright 2022 Datafuse Labs. +// +// 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. + +pub enum WindowFunction { + AggregateFunction, + BuiltInFunction, +} diff --git a/common/functions/src/window/mod.rs b/common/functions/src/window/mod.rs index cd12f0784763..ecb24c836d55 100644 --- a/common/functions/src/window/mod.rs +++ b/common/functions/src/window/mod.rs @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod window; +mod function; +mod window_frame; -pub use window::WindowFrame; -pub use window::WindowFrameBound; -pub use window::WindowFrameUnits; +pub use function::*; +pub use window_frame::*; diff --git a/common/functions/src/window/window.rs b/common/functions/src/window/window_frame.rs similarity index 100% rename from common/functions/src/window/window.rs rename to common/functions/src/window/window_frame.rs diff --git a/query/Cargo.toml b/query/Cargo.toml index c0a903104685..ff79acf96747 100644 --- a/query/Cargo.toml +++ b/query/Cargo.toml @@ -99,6 +99,7 @@ rand = "0.8.5" regex = "1.5.5" reqwest = "0.11.10" rsa = "0.5.0" +segment-tree = "2.0.0" serde = { version = "1.0.136", features = ["derive"] } serde-bridge = "0.0.3" serde_json = "1.0.79" diff --git a/query/src/optimizers/optimizer_scatters.rs b/query/src/optimizers/optimizer_scatters.rs index 26876e503ab6..332b4d576940 100644 --- a/query/src/optimizers/optimizer_scatters.rs +++ b/query/src/optimizers/optimizer_scatters.rs @@ -309,7 +309,7 @@ impl PlanRewriter for ScattersOptimizerImpl { fn rewrite_window_func(&mut self, plan: &WindowFuncPlan) -> Result { let new_input = Arc::new(self.rewrite_plan_node(&plan.input)?); - self.input = Some(new_input.clone()); + self.input = Some(new_input); match self.running_mode { RunningMode::Cluster => self.cluster_window(plan), diff --git a/query/src/pipelines/transforms/transform_window_func.rs b/query/src/pipelines/transforms/transform_window_func.rs index d090bced3c3e..7a256d3162a7 100644 --- a/query/src/pipelines/transforms/transform_window_func.rs +++ b/query/src/pipelines/transforms/transform_window_func.rs @@ -17,28 +17,35 @@ use std::cmp::Ordering; use std::ops::Range; use std::sync::Arc; +use bumpalo::Bump; use common_arrow::arrow::array::ArrayRef; use common_arrow::arrow::compute::partition::lexicographical_partition_ranges; use common_arrow::arrow::compute::sort::SortColumn; use common_datablocks::DataBlock; +use common_datavalues::ColumnRef; use common_datavalues::ColumnWithField; use common_datavalues::DataField; use common_datavalues::DataSchemaRef; +use common_datavalues::DataType; use common_datavalues::DataValue; use common_datavalues::Series; -use common_functions::aggregates::eval_aggr; use common_functions::aggregates::AggregateFunctionFactory; +use common_functions::aggregates::AggregateFunctionRef; +use common_functions::aggregates::StateAddr; use common_functions::scalars::assert_numeric; use common_functions::window::WindowFrame; use common_functions::window::WindowFrameBound; use common_functions::window::WindowFrameUnits; -use common_planners::sort; use common_planners::Expression; use common_streams::DataBlockStream; use common_streams::SendableDataBlockStream; use common_tracing::tracing; use enum_extract::let_extract; use futures::StreamExt; +use segment_tree::ops::Commutative; +use segment_tree::ops::Identity; +use segment_tree::ops::Operation; +use segment_tree::SegmentPoint; use crate::pipelines::processors::EmptyProcessor; use crate::pipelines::processors::Processor; @@ -97,7 +104,9 @@ impl WindowFuncTransform { ); sort_exprs.extend(order_by.to_owned()); let sort_column_desc = get_sort_descriptions(block.schema(), &sort_exprs)?; - let block = DataBlock::sort_block(&block, &sort_column_desc, None)?; + + // todo Fixme: empty partition by + let block = DataBlock::sort_block(block, &sort_column_desc, None)?; // set default window frame let window_frame = match window_frame { @@ -129,27 +138,41 @@ impl WindowFuncTransform { // function calculate let mut arguments: Vec = Vec::with_capacity(args.len()); + let mut arg_fields = Vec::with_capacity(args.len()); for arg in args { let arg_field: DataField = arg.to_data_field(&self.input_schema)?; + arg_fields.push(arg_field.clone()); let arg_column = block.try_column_by_name(arg_field.name()).unwrap(); arguments.push(ColumnWithField::new(Arc::clone(arg_column), arg_field)); } + let function = if !AggregateFunctionFactory::instance().check(op) { + unimplemented!("not yet impl built-in window func"); + } else { + AggregateFunctionFactory::instance().get(op, params.clone(), arg_fields)? + }; + + let arena = Arc::new(Bump::with_capacity( + 2 * block.num_rows() * function.state_layout().size(), + )); + let segment_tree_state = Self::new_partition_aggr_func( + function.clone(), + &arguments + .iter() + .map(|a| a.column().clone()) + .collect::>(), + arena, + ); + let window_col_per_tuple = (0..block.num_rows()) .map(|i| { let frame = &frame_bounds[i]; let frame_start = frame.start; let frame_end = frame.end; - let frame_size = frame_end - frame_start; - let args_slice = arguments - .iter() - .map(|c| c.slice(frame_start, frame_size)) - .collect::>(); - // at the moment, only supports aggr function - if !AggregateFunctionFactory::instance().check(op) { - unimplemented!("not yet impl built-in window func"); - } - eval_aggr(op, params.to_owned(), &args_slice, frame_size).unwrap() + let state = segment_tree_state.query(frame_start, frame_end); + let mut builder = function.return_type().unwrap().create_mutable(1); + function.merge_result(state, builder.as_mut()).unwrap(); + builder.to_column() }) .collect::>(); @@ -398,6 +421,32 @@ impl WindowFuncTransform { }, } } + + /// Actually, the computation is row based + fn new_partition_aggr_func( + func: AggregateFunctionRef, + arguments: &[ColumnRef], + arena: Arc, + ) -> SegmentPoint { + let rows = arguments[0].len(); + let state_per_tuple = (0..rows) + .map(|i| { + arguments + .iter() + .map(|c| c.slice(i, 1)) + .collect::>() + }) + .map(|args| { + let place = arena.alloc_layout(func.state_layout()); + let state_addr = place.into(); + func.init_state(state_addr); + func.accumulate(state_addr, &args, None, 1) + .expect("Failed to initialize the state"); + state_addr + }) + .collect::>(); + SegmentPoint::build(state_per_tuple, Agg { func, arena }) + } } #[async_trait::async_trait] @@ -461,3 +510,42 @@ impl Processor for WindowFuncTransform { ))) } } + +struct Agg { + func: AggregateFunctionRef, + arena: Arc, +} + +impl Operation for Agg { + fn combine(&self, a: &StateAddr, b: &StateAddr) -> StateAddr { + let place = self.arena.alloc_layout(self.func.state_layout()); + let state_addr = place.into(); + self.func.init_state(state_addr); + self.func + .merge(state_addr, *a) + .expect("Failed to merge states"); + self.func + .merge(state_addr, *b) + .expect("Failed to merge states"); + state_addr + } + + fn combine_mut(&self, a: &mut StateAddr, b: &StateAddr) { + self.func.merge(*a, *b).expect("Failed to merge states"); + } + + fn combine_mut2(&self, a: &StateAddr, b: &mut StateAddr) { + self.func.merge(*b, *a).expect("Failed to merge states"); + } +} + +impl Commutative for Agg {} + +impl Identity for Agg { + fn identity(&self) -> StateAddr { + let place = self.arena.alloc_layout(self.func.state_layout()); + let state_addr = place.into(); + self.func.init_state(state_addr); + state_addr + } +} diff --git a/query/src/sql/plan_parser.rs b/query/src/sql/plan_parser.rs index 67e3a10d6288..372f858dc52a 100644 --- a/query/src/sql/plan_parser.rs +++ b/query/src/sql/plan_parser.rs @@ -264,9 +264,7 @@ impl PlanParser { false => { let exprs = data.window_expressions.to_vec(); exprs.into_iter().try_fold(plan, |input, window_func| { - PlanBuilder::from(&input) - .window_func(window_func.clone())? - .build() + PlanBuilder::from(&input).window_func(window_func)?.build() }) } } From 706e19a039145ed495d787ecb17b62b969a94571 Mon Sep 17 00:00:00 2001 From: doki Date: Tue, 7 Jun 2022 15:19:25 +0800 Subject: [PATCH 14/18] empty partition --- .../src/columns/column_with_field.rs | 6 --- .../aggregates/aggregate_function_state.rs | 2 +- common/planners/src/plan_node_rewriter.rs | 38 +++++++++++++++- query/src/optimizers/optimizer_scatters.rs | 43 +++++++++++-------- .../transforms/transform_window_func.rs | 25 +++++++---- 5 files changed, 80 insertions(+), 34 deletions(-) diff --git a/common/datavalues/src/columns/column_with_field.rs b/common/datavalues/src/columns/column_with_field.rs index 4a4aaa5ebe28..762d65eecf2d 100644 --- a/common/datavalues/src/columns/column_with_field.rs +++ b/common/datavalues/src/columns/column_with_field.rs @@ -34,12 +34,6 @@ impl ColumnWithField { pub fn data_type(&self) -> &DataTypeImpl { self.field.data_type() } - pub fn slice(&self, offset: usize, length: usize) -> Self { - Self { - column: self.column.slice(offset, length), - field: self.field.to_owned(), - } - } } pub type ColumnsWithField = [ColumnWithField]; diff --git a/common/functions/src/aggregates/aggregate_function_state.rs b/common/functions/src/aggregates/aggregate_function_state.rs index d3176facb805..d8b13847bce7 100644 --- a/common/functions/src/aggregates/aggregate_function_state.rs +++ b/common/functions/src/aggregates/aggregate_function_state.rs @@ -20,7 +20,7 @@ use common_exception::Result; use crate::aggregates::AggregateFunctionRef; -#[derive(Clone, Copy, Default)] +#[derive(Clone, Copy)] pub struct StateAddr { addr: usize, } diff --git a/common/planners/src/plan_node_rewriter.rs b/common/planners/src/plan_node_rewriter.rs index 94cdcb4a51de..7ce49ed0a132 100644 --- a/common/planners/src/plan_node_rewriter.rs +++ b/common/planners/src/plan_node_rewriter.rs @@ -673,7 +673,43 @@ impl RewriteHelper { } } - Expression::WindowFunction { .. } => todo!("not figure out yet @doki"), + Expression::WindowFunction { + op, + params, + args, + partition_by, + order_by, + window_frame, + } => { + let new_args: Result> = args + .iter() + .map(|v| RewriteHelper::expr_rewrite_alias(v, data)) + .collect(); + + let new_partition_by: Result> = partition_by + .iter() + .map(|v| RewriteHelper::expr_rewrite_alias(v, data)) + .collect(); + + let new_order_by: Result> = order_by + .iter() + .map(|v| RewriteHelper::expr_rewrite_alias(v, data)) + .collect(); + + match (new_args, new_partition_by, new_order_by) { + (Ok(new_args), Ok(new_partition_by), Ok(new_order_by)) => { + Ok(Expression::WindowFunction { + op: op.clone(), + params: params.clone(), + args: new_args, + partition_by: new_partition_by, + order_by: new_order_by, + window_frame: window_frame.clone(), + }) + } + (Err(e), _, _) | (_, Err(e), _) | (_, _, Err(e)) => Err(e), + } + } Expression::Alias(alias, plan) => { if data.inside_aliases.contains(alias) { diff --git a/query/src/optimizers/optimizer_scatters.rs b/query/src/optimizers/optimizer_scatters.rs index 332b4d576940..d76b0fee96ab 100644 --- a/query/src/optimizers/optimizer_scatters.rs +++ b/query/src/optimizers/optimizer_scatters.rs @@ -112,8 +112,6 @@ impl ScattersOptimizerImpl { } fn cluster_window(&mut self, plan: &WindowFuncPlan) -> Result { - self.running_mode = RunningMode::Cluster; - match self.input.take() { None => Err(ErrorCode::LogicalError("Cluster window input is None")), Some(input) => { @@ -130,25 +128,34 @@ impl ScattersOptimizerImpl { panic!() ); - let mut concat_ws_args = vec![Expression::create_literal(DataValue::String( - "#".as_bytes().to_vec(), - ))]; - concat_ws_args.extend(partition_by.to_owned()); - let concat_partition_by = - Expression::create_scalar_function("concat_ws", concat_ws_args); - - let scatters_expr = - Expression::create_scalar_function("sipHash", vec![concat_partition_by]); - - let scatter_plan = PlanNode::Stage(StagePlan { - scatters_expr, - kind: StageKind::Normal, - input, - }); + let stage_input = if !partition_by.is_empty() { + let mut concat_ws_args = vec![Expression::create_literal(DataValue::String( + "#".as_bytes().to_vec(), + ))]; + concat_ws_args.extend(partition_by.to_owned()); + let concat_partition_by = + Expression::create_scalar_function("concat_ws", concat_ws_args); + + let scatters_expr = + Expression::create_scalar_function("sipHash", vec![concat_partition_by]); + + PlanNode::Stage(StagePlan { + scatters_expr, + kind: StageKind::Normal, + input, + }) + } else { + self.running_mode = RunningMode::Standalone; + PlanNode::Stage(StagePlan { + scatters_expr: Expression::create_literal(DataValue::UInt64(0)), + kind: StageKind::Convergent, + input, + }) + }; Ok(PlanNode::WindowFunc(WindowFuncPlan { window_func: plan.window_func.to_owned(), - input: Arc::new(scatter_plan), + input: Arc::new(stage_input), schema: plan.schema.to_owned(), })) } diff --git a/query/src/pipelines/transforms/transform_window_func.rs b/query/src/pipelines/transforms/transform_window_func.rs index 7a256d3162a7..37cec6d20170 100644 --- a/query/src/pipelines/transforms/transform_window_func.rs +++ b/query/src/pipelines/transforms/transform_window_func.rs @@ -89,9 +89,9 @@ impl WindowFuncTransform { ); // sort block by partition_by and order_by exprs - let mut sort_exprs: Vec = + let mut partition_sort_exprs: Vec = Vec::with_capacity(partition_by.len() + order_by.len()); - sort_exprs.extend( + partition_sort_exprs.extend( partition_by .iter() .map(|part_by_expr| Expression::Sort { @@ -102,11 +102,14 @@ impl WindowFuncTransform { }) .collect::>(), ); - sort_exprs.extend(order_by.to_owned()); - let sort_column_desc = get_sort_descriptions(block.schema(), &sort_exprs)?; + partition_sort_exprs.extend(order_by.to_owned()); - // todo Fixme: empty partition by - let block = DataBlock::sort_block(block, &sort_column_desc, None)?; + let block = if !partition_sort_exprs.is_empty() { + let sort_column_desc = get_sort_descriptions(block.schema(), &partition_sort_exprs)?; + DataBlock::sort_block(block, &sort_column_desc, None)? + } else { + block.to_owned() + }; // set default window frame let window_frame = match window_frame { @@ -210,8 +213,14 @@ impl WindowFuncTransform { }) .collect::>(); - let mut partition_boundaries = - lexicographical_partition_ranges(&partition_by_sort_column).unwrap(); + let partition_boundaries = if !partition_by_sort_column.is_empty() { + lexicographical_partition_ranges(&partition_by_sort_column) + .unwrap() + .collect::>() + } else { + vec![0..block.num_rows(); 1] + }; + let mut partition_boundaries = partition_boundaries.into_iter(); let mut partition = partition_boundaries.next().unwrap(); match (frame_units, start, end) { From dd1597416c0ce7ea4e766279b3561acf14cc7ee9 Mon Sep 17 00:00:00 2001 From: doki Date: Tue, 7 Jun 2022 15:35:50 +0800 Subject: [PATCH 15/18] clippy --- common/planners/src/plan_node_rewriter.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/common/planners/src/plan_node_rewriter.rs b/common/planners/src/plan_node_rewriter.rs index 7ce49ed0a132..577a6154e2eb 100644 --- a/common/planners/src/plan_node_rewriter.rs +++ b/common/planners/src/plan_node_rewriter.rs @@ -704,7 +704,7 @@ impl RewriteHelper { args: new_args, partition_by: new_partition_by, order_by: new_order_by, - window_frame: window_frame.clone(), + window_frame: *window_frame, }) } (Err(e), _, _) | (_, Err(e), _) | (_, _, Err(e)) => Err(e), From ce0e752c7b55a9995d08abb9092e5704e2232603 Mon Sep 17 00:00:00 2001 From: doki Date: Wed, 8 Jun 2022 11:13:59 +0800 Subject: [PATCH 16/18] add stateless tests and fix some bugs of range frame current row bound --- common/functions/src/window/window_frame.rs | 29 +- query/src/interpreters/interpreter_select.rs | 3 +- query/src/pipelines/new/pipeline_builder.rs | 7 + query/src/pipelines/transforms/mod.rs | 1 - .../transforms/transform_window_func.rs | 28 +- .../pipelines/transforms/window_func/mod.rs | 13 - .../03_0024_select_window_function.result | 304 ++++++++++++++++++ .../03_dml/03_0024_select_window_function.sql | 56 ++++ .../1_stateful/02_query/02_0000_kill_query.py | 4 +- 9 files changed, 408 insertions(+), 37 deletions(-) delete mode 100644 query/src/pipelines/transforms/window_func/mod.rs create mode 100644 tests/suites/0_stateless/03_dml/03_0024_select_window_function.result create mode 100644 tests/suites/0_stateless/03_dml/03_0024_select_window_function.sql diff --git a/common/functions/src/window/window_frame.rs b/common/functions/src/window/window_frame.rs index 57e56298799a..468a3a3020ab 100644 --- a/common/functions/src/window/window_frame.rs +++ b/common/functions/src/window/window_frame.rs @@ -1,16 +1,23 @@ -// Copyright 2022 Datafuse Labs. +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 // -// 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 // -// 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. +// 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. + +// Code in this file is mainly copied from apache/arrow-datafusion +// Original project code: https://github.com/apache/arrow-datafusion/blob/a6e93a10ab2659500eb4f838b7b53f138e545be3/datafusion/expr/src/window_frame.rs#L39 +// PR: https://github.com/datafuselabs/databend/pull/5401 use std::cmp::Ordering; use std::fmt; diff --git a/query/src/interpreters/interpreter_select.rs b/query/src/interpreters/interpreter_select.rs index 92fb06921580..92b5793b223e 100644 --- a/query/src/interpreters/interpreter_select.rs +++ b/query/src/interpreters/interpreter_select.rs @@ -78,13 +78,14 @@ impl Interpreter for SelectInterpreter { { let async_runtime = self.ctx.get_storage_runtime(); let new_pipeline = self.create_new_pipeline()?; + tracing::debug!("get new_pipeline:\n{:?}", new_pipeline); let executor = PipelinePullingExecutor::try_create(async_runtime, new_pipeline)?; let (handler, stream) = ProcessorExecutorStream::create(executor)?; self.ctx.add_source_abort_handle(handler); return Ok(Box::pin(stream)); } let optimized_plan = self.rewrite_plan()?; - tracing::debug!("\nget optimized plan:\n{:?}", optimized_plan); + tracing::debug!("get optimized plan:\n{:?}", optimized_plan); plan_schedulers::schedule_query(&self.ctx, &optimized_plan).await } diff --git a/query/src/pipelines/new/pipeline_builder.rs b/query/src/pipelines/new/pipeline_builder.rs index f4b911500824..e032539441b5 100644 --- a/query/src/pipelines/new/pipeline_builder.rs +++ b/query/src/pipelines/new/pipeline_builder.rs @@ -30,6 +30,7 @@ use common_planners::ReadDataSourcePlan; use common_planners::SelectPlan; use common_planners::SortPlan; use common_planners::SubQueriesSetPlan; +use common_planners::WindowFuncPlan; use super::processors::SortMergeCompactor; use crate::pipelines::new::pipeline::NewPipeline; @@ -86,6 +87,7 @@ impl PlanVisitor for QueryPipelineBuilder { PlanNode::Expression(n) => self.visit_expression(n), PlanNode::AggregatorPartial(n) => self.visit_aggregate_partial(n), PlanNode::AggregatorFinal(n) => self.visit_aggregate_final(n), + PlanNode::WindowFunc(n) => self.visit_window_func(n), PlanNode::Filter(n) => self.visit_filter(n), PlanNode::Having(n) => self.visit_having(n), PlanNode::Sort(n) => self.visit_sort(n), @@ -147,6 +149,11 @@ impl PlanVisitor for QueryPipelineBuilder { }) } + fn visit_window_func(&mut self, plan: &WindowFuncPlan) -> Result<()> { + self.visit_plan_node(&plan.input)?; + unimplemented!("window function cannot work in new scheduler framework now") + } + fn visit_projection(&mut self, plan: &ProjectionPlan) -> Result<()> { self.visit_plan_node(&plan.input)?; self.pipeline diff --git a/query/src/pipelines/transforms/mod.rs b/query/src/pipelines/transforms/mod.rs index 363bd04f206f..416ec3ab0847 100644 --- a/query/src/pipelines/transforms/mod.rs +++ b/query/src/pipelines/transforms/mod.rs @@ -32,7 +32,6 @@ mod transform_window_func; pub mod group_by; mod streams; mod transform_sink; -pub mod window_func; pub use streams::AddOnStream; pub use transform_aggregator_final::AggregatorFinalTransform; diff --git a/query/src/pipelines/transforms/transform_window_func.rs b/query/src/pipelines/transforms/transform_window_func.rs index 37cec6d20170..f31c1db7884b 100644 --- a/query/src/pipelines/transforms/transform_window_func.rs +++ b/query/src/pipelines/transforms/transform_window_func.rs @@ -377,49 +377,57 @@ impl WindowFuncTransform { let mut end = partition.end; match frame_start { WindowFrameBound::Preceding(Some(preceding)) => { - let idx = order_by_values.partition_point(|x| { + let offset = order_by_values[start..end].partition_point(|x| { let min = DataValue::from( current_value.as_f64().unwrap() - preceding as f64, ); x.cmp(&min) == Ordering::Less }); - start = std::cmp::max(start, idx); + start += offset; } WindowFrameBound::CurrentRow => { - start = i; + let offset = order_by_values[start..end].partition_point(|x| { + let peer = DataValue::from(current_value.as_f64().unwrap()); + x.cmp(&peer) == Ordering::Less + }); + start += offset; } WindowFrameBound::Following(Some(following)) => { - let idx = order_by_values.partition_point(|x| { + let offset = order_by_values[start..end].partition_point(|x| { let max = DataValue::from( current_value.as_f64().unwrap() + following as f64, ); x.cmp(&max) == Ordering::Less }); - start = std::cmp::min(end, idx); + start += offset; } _ => (), } match frame_end { WindowFrameBound::Preceding(Some(preceding)) => { - let idx = order_by_values.partition_point(|x| { + let offset = order_by_values[start..end].partition_point(|x| { let min = DataValue::from( current_value.as_f64().unwrap() - preceding as f64, ); x.cmp(&min) != Ordering::Greater }); - end = std::cmp::max(start, idx); + end = start + offset; } WindowFrameBound::CurrentRow => { - end = i + 1; + let offset = order_by_values[start..end].partition_point(|x| { + let peer = DataValue::from(current_value.as_f64().unwrap()); + x.cmp(&peer) != Ordering::Greater + }); + end = start + offset; } WindowFrameBound::Following(Some(following)) => { - let idx = order_by_values.partition_point(|x| { + let offset = order_by_values[start..end].partition_point(|x| { let max = DataValue::from( current_value.as_f64().unwrap() + following as f64, ); x.cmp(&max) != Ordering::Greater }); - end = std::cmp::min(end, idx); + end = start + offset; } _ => (), } diff --git a/query/src/pipelines/transforms/window_func/mod.rs b/query/src/pipelines/transforms/window_func/mod.rs deleted file mode 100644 index ea0ed57e60ed..000000000000 --- a/query/src/pipelines/transforms/window_func/mod.rs +++ /dev/null @@ -1,13 +0,0 @@ -// Copyright 2022 Datafuse Labs. -// -// 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. diff --git a/tests/suites/0_stateless/03_dml/03_0024_select_window_function.result b/tests/suites/0_stateless/03_dml/03_0024_select_window_function.result new file mode 100644 index 000000000000..42730501b478 --- /dev/null +++ b/tests/suites/0_stateless/03_dml/03_0024_select_window_function.result @@ -0,0 +1,304 @@ +================sep================ +China 2001 7845 +China 2001 7845 +Finland 2000 7845 +Finland 2000 7845 +Finland 2001 7845 +India 2000 7845 +India 2000 7845 +India 2000 7845 +USA 2000 7845 +USA 2000 7845 +USA 2001 7845 +USA 2001 7845 +USA 2001 7845 +USA 2001 7845 +USA 2001 7845 +================sep================ +China 2001 310 +China 2001 310 +Finland 2000 1610 +Finland 2000 1610 +Finland 2001 1610 +India 2000 1350 +India 2000 1350 +India 2000 1350 +USA 2000 4575 +USA 2000 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +================sep================ +China 2001 310 +China 2001 310 +Finland 2000 1920 +Finland 2000 1920 +Finland 2001 1920 +India 2000 3270 +India 2000 3270 +India 2000 3270 +USA 2000 7845 +USA 2000 7845 +USA 2001 7845 +USA 2001 7845 +USA 2001 7845 +USA 2001 7845 +USA 2001 7845 +================sep================ +China 2001 310 +China 2001 310 +Finland 2000 1600 +Finland 2000 1600 +Finland 2001 1610 +India 2000 1350 +India 2000 1350 +India 2000 1350 +USA 2000 1575 +USA 2000 1575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +================sep================ +China 2001 310 +China 2001 310 +Finland 2000 1600 +Finland 2000 1610 +Finland 2001 110 +India 2000 150 +India 2000 1350 +India 2000 1275 +USA 2000 1575 +USA 2000 1625 +USA 2001 3050 +USA 2001 2750 +USA 2001 2850 +USA 2001 1450 +USA 2001 250 +================sep================ +China 2001 310 +China 2001 310 +Finland 2000 1600 +Finland 2000 1610 +Finland 2001 1610 +India 2000 150 +India 2000 1350 +India 2000 1350 +USA 2000 1575 +USA 2000 1625 +USA 2001 3125 +USA 2001 4325 +USA 2001 4475 +USA 2001 4575 +USA 2001 4575 +================sep================ +China 2001 310 +China 2001 310 +Finland 2000 1610 +Finland 2000 1610 +Finland 2001 110 +India 2000 1350 +India 2000 1350 +India 2000 1275 +USA 2000 4575 +USA 2000 4575 +USA 2001 4500 +USA 2001 3000 +USA 2001 2950 +USA 2001 1450 +USA 2001 250 +================sep================ +China 2001 110 +China 2001 310 +Finland 2000 1500 +Finland 2000 1600 +Finland 2001 1610 +India 2000 75 +India 2000 150 +India 2000 1350 +USA 2000 75 +USA 2000 1575 +USA 2001 1625 +USA 2001 3125 +USA 2001 4325 +USA 2001 4475 +USA 2001 4575 +================sep================ +China 2001 310 +China 2001 200 +Finland 2000 1610 +Finland 2000 110 +Finland 2001 10 +India 2000 1350 +India 2000 1275 +India 2000 1200 +USA 2000 4575 +USA 2000 4500 +USA 2001 3000 +USA 2001 2950 +USA 2001 1450 +USA 2001 250 +USA 2001 100 +================sep================ +China 2001 310 +China 2001 310 +Finland 2000 1610 +Finland 2000 1610 +Finland 2001 1610 +India 2000 1350 +India 2000 1350 +India 2000 1350 +USA 2000 4575 +USA 2000 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +================sep================ +China 2001 310 +China 2001 310 +Finland 2001 110 +Finland 2000 110 +Finland 2000 1500 +India 2000 150 +India 2000 150 +India 2000 1200 +USA 2001 375 +USA 2000 375 +USA 2001 375 +USA 2001 375 +USA 2001 4200 +USA 2000 4200 +USA 2001 4200 +================sep================ +China 2001 310 +China 2001 310 +Finland 2001 110 +Finland 2000 110 +Finland 2000 1610 +India 2000 150 +India 2000 150 +India 2000 1350 +USA 2001 375 +USA 2000 375 +USA 2001 375 +USA 2001 375 +USA 2001 4575 +USA 2000 4575 +USA 2001 4575 +================sep================ +China 2001 310 +China 2001 310 +Finland 2001 1610 +Finland 2000 1610 +Finland 2000 1500 +India 2000 1350 +India 2000 1350 +India 2000 1200 +USA 2001 4575 +USA 2000 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4200 +USA 2000 4200 +USA 2001 4200 +================sep================ +China 2001 310 +China 2001 200 +Finland 2001 110 +Finland 2000 100 +Finland 2000 1500 +India 2000 150 +India 2000 150 +India 2000 1200 +USA 2001 375 +USA 2000 325 +USA 2001 250 +USA 2001 150 +USA 2001 4200 +USA 2000 3000 +USA 2001 3000 +================sep================ +China 2001 110 +China 2001 310 +Finland 2001 10 +Finland 2000 110 +Finland 2000 1500 +India 2000 150 +India 2000 150 +India 2000 1200 +USA 2001 50 +USA 2000 125 +USA 2001 225 +USA 2001 375 +USA 2001 1200 +USA 2000 4200 +USA 2001 4200 +================sep================ +China 2001 110 +China 2001 310 +Finland 2001 10 +Finland 2000 110 +Finland 2000 1610 +India 2000 150 +India 2000 150 +India 2000 1350 +USA 2001 50 +USA 2000 125 +USA 2001 225 +USA 2001 375 +USA 2001 1575 +USA 2000 4575 +USA 2001 4575 +================sep================ +China 2001 310 +China 2001 200 +Finland 2001 1610 +Finland 2000 1600 +Finland 2000 1500 +India 2000 1350 +India 2000 1350 +India 2000 1200 +USA 2001 4575 +USA 2000 4525 +USA 2001 4450 +USA 2001 4350 +USA 2001 4200 +USA 2000 3000 +USA 2001 3000 +================sep================ +China 2001 310 +China 2001 310 +Finland 2001 1610 +Finland 2000 1610 +Finland 2000 1610 +India 2000 1350 +India 2000 1350 +India 2000 1350 +USA 2001 4575 +USA 2000 4575 +USA 2001 4575 +USA 2001 4575 +USA 2001 4575 +USA 2000 4575 +USA 2001 4575 +================sep================ +China 2001 310 155 +China 2001 310 155 +Finland 2001 110 55 +Finland 2000 110 55 +Finland 2000 1500 1500 +India 2000 150 75 +India 2000 150 75 +India 2000 1200 1200 +USA 2001 375 93.75 +USA 2000 375 93.75 +USA 2001 375 93.75 +USA 2001 375 93.75 +USA 2001 4200 1400 +USA 2000 4200 1400 +USA 2001 4200 1400 diff --git a/tests/suites/0_stateless/03_dml/03_0024_select_window_function.sql b/tests/suites/0_stateless/03_dml/03_0024_select_window_function.sql new file mode 100644 index 000000000000..d0174dd1d79a --- /dev/null +++ b/tests/suites/0_stateless/03_dml/03_0024_select_window_function.sql @@ -0,0 +1,56 @@ +DROP DATABASE IF EXISTS db1; +CREATE DATABASE db1; +USE db1; + +DROP TABLE IF EXISTS sales; +CREATE TABLE `sales` ( + `year` varchar(64) DEFAULT NULL, + `country` varchar(64) DEFAULT NULL, + `product` varchar(64) DEFAULT NULL, + `profit` int DEFAULT NULL +) Engine = Fuse; + +SET enable_new_processor_framework=0; + +INSERT INTO `sales` VALUES ('2000','Finland','Computer',1500),('2000','Finland','Phone',100),('2001','Finland','Phone',10),('2000','India','Calculator',75),('2000','India','Calculator',75),('2000','India','Computer',1200),('2000','USA','Calculator',75),('2000','USA','Computer',1500),('2001','USA','Calculator',50),('2001','USA','Computer',1500),('2001','USA','Computer',1200),('2001','USA','TV',150),('2001','USA','TV',100),('2001','China','TV',110),('2001','China','Computer',200); + +select '================sep================'; +select country, year, sum(profit) over() from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(order by country) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by year) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by year rows between 1 preceding and 1 following) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by year rows between unbounded preceding and 1 following) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by year rows between 1 preceding and unbounded following) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by year rows between unbounded preceding and current row) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by year rows between current row and unbounded following) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by year rows between unbounded preceding and unbounded following) from sales order by country, year; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between 500 preceding and 500 following) from sales order by country, profit; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between unbounded preceding and 500 following) from sales order by country, profit; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between 500 preceding and unbounded following) from sales order by country, profit; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between current row and 500 following) from sales order by country, profit; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between 500 preceding and current row) from sales order by country, profit; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between unbounded preceding and current row) from sales order by country, profit; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between current row and unbounded following) from sales order by country, profit; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between unbounded preceding and unbounded following) from sales order by country, profit; +select '================sep================'; +select country, year, sum(profit) over(partition by country order by profit range between 500 preceding and 500 following) as sum, avg(profit) over(partition by country order by profit range between 500 preceding and 500 following) as avg from sales order by country, profit; + +DROP DATABASE db1; \ No newline at end of file diff --git a/tests/suites/1_stateful/02_query/02_0000_kill_query.py b/tests/suites/1_stateful/02_query/02_0000_kill_query.py index 9b9c4bddc107..fc275e6e89af 100755 --- a/tests/suites/1_stateful/02_query/02_0000_kill_query.py +++ b/tests/suites/1_stateful/02_query/02_0000_kill_query.py @@ -26,7 +26,9 @@ client1.expect(prompt) client1.expect('') - client1.send('SELECT max(number), sum(number) FROM numbers_mt(100000000000) GROUP BY number % 3, number % 4, number % 5 LIMIT 10;') + client1.send( + 'SELECT max(number), sum(number) FROM numbers_mt(100000000000) GROUP BY number % 3, number % 4, number % 5 LIMIT 10;' + ) time.sleep(0.5) mycursor = mydb.cursor() From f358099dc24075c27efaf160bfa7c99f98d8ab16 Mon Sep 17 00:00:00 2001 From: doki Date: Wed, 8 Jun 2022 11:26:09 +0800 Subject: [PATCH 17/18] licenses --- common/functions/src/window/window_frame.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/common/functions/src/window/window_frame.rs b/common/functions/src/window/window_frame.rs index 468a3a3020ab..e4a0e6e1bf09 100644 --- a/common/functions/src/window/window_frame.rs +++ b/common/functions/src/window/window_frame.rs @@ -1,3 +1,17 @@ +// Copyright 2022 Datafuse Labs. +// +// 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. + // Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information From 5532bc823270cfb86c0c18623be9cfbea3e04739 Mon Sep 17 00:00:00 2001 From: doki Date: Wed, 8 Jun 2022 18:20:09 +0800 Subject: [PATCH 18/18] remove error case --- common/ast/tests/it/parser.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/common/ast/tests/it/parser.rs b/common/ast/tests/it/parser.rs index d19996525409..7fb95417f383 100644 --- a/common/ast/tests/it/parser.rs +++ b/common/ast/tests/it/parser.rs @@ -139,7 +139,7 @@ fn test_statements_in_legacy_suites() { // TODO(andylokandy): support all cases eventually // Remove currently unimplemented cases let file_str = regex::Regex::new( - "(?i).*(SLAVE|MASTER|COMMIT|START|ROLLBACK|FIELDS|GRANT|COPY|ROLE|STAGE|ENGINES|UNDROP).*\n", + "(?i).*(SLAVE|MASTER|COMMIT|START|ROLLBACK|FIELDS|GRANT|COPY|ROLE|STAGE|ENGINES|UNDROP|OVER).*\n", ) .unwrap() .replace_all(&file_str, "")