feat: Window func

common/ast/src/udfs/udf_expr_visitor.rs

@@ -163,6 +163,15 @@ pub trait UDFExprVisitor: Sized + Send {
             };
         }
 
+        if let Some(over) = &function.over {
+            for partition_by in &over.partition_by {
+                UDFExprTraverser::accept(partition_by, self).await?;
+            }
+            for order_by in &over.order_by {
+                UDFExprTraverser::accept(&order_by.expr, self).await?;
+            }
+        }
+
         Ok(())
     }
 

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;

common/functions/src/window/mod.rs

@@ -0,0 +1,5 @@
+mod window_frame;
+
+pub use window_frame::WindowFrame;
+pub use window_frame::WindowFrameBound;
+pub use window_frame::WindowFrameUnits;

common/functions/src/window/window_frame.rs

@@ -0,0 +1,223 @@
+// 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<ast::WindowFrame> for WindowFrame {
+    type Error = ErrorCode;
+
+    fn try_from(value: ast::WindowFrame) -> Result<Self, Self::Error> {
+        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 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,
+}
+
+impl fmt::Display for WindowFrameUnits {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            WindowFrameUnits::Range => "RANGE",
+            WindowFrameUnits::Rows => "ROWS",
+        })
+    }
+}
+
+impl From<ast::WindowFrameUnits> for WindowFrameUnits {
+    fn from(value: ast::WindowFrameUnits) -> Self {
+        match value {
+            ast::WindowFrameUnits::Range => Self::Range,
+            ast::WindowFrameUnits::Rows => Self::Rows,
+            _ => unimplemented!(),
+        }
+    }
+}
+
+#[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. <expr> PRECEDING
+    /// <expr> must be a non-negative constant numeric expression. The boundary is a row that
+    /// is <expr> "units" prior to the current row.
+    Preceding(Option<u64>),
+    /// 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 "<expr> PRECEDING" except that the boundary is <expr> 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<u64>),
+}
+
+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<ast::WindowFrameBound> 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<Ordering> {
+        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<H: Hasher>(&self, state: &mut H) {
+        self.get_rank().hash(state)
+    }
+}

common/planners/src/lib.rs

@@ -99,6 +99,7 @@ mod plan_user_udf_drop;
 mod plan_view_alter;
 mod plan_view_create;
 mod plan_view_drop;
+mod plan_window_func;
 
 pub use plan_aggregator_final::AggregatorFinalPlan;
 pub use plan_aggregator_partial::AggregatorPartialPlan;
@@ -127,6 +128,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;
@@ -224,3 +227,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;

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,22 @@ pub enum Expression {
         args: Vec<Expression>,
     },
 
+    /// WindowFunction
+    WindowFunction {
+        /// operation performed
+        op: String,
+        /// params
+        params: Vec<DataValue>,
+        /// arguments
+        args: Vec<Expression>,
+        /// partition by
+        partition_by: Vec<Expression>,
+        /// order by
+        order_by: Vec<Expression>,
+        /// window frame
+        window_frame: Option<WindowFrame>,
+    },
+
     /// A sort expression, that can be used to sort values.
     Sort {
         /// The expression to sort on
@@ -421,6 +438,53 @@ impl fmt::Debug for Expression {
                 Ok(())
             }
 
+            Expression::WindowFunction {
+                op,
+                params,
+                args,
+                partition_by,
+                order_by,
+                window_frame,
+            } => {
+                let args_column_name = args.iter().map(Expression::column_name).collect::<Vec<_>>();
+                let params_name = params
+                    .iter()
+                    .map(|v| DataValue::custom_display(v, true))
+                    .collect::<Vec<_>>();
+
+                if params.is_empty() {
+                    write!(f, "{}", op)?;
+                } else {
+                    write!(f, "{}({})", op, params_name.join(", "))?;
+                }
+
+                write!(f, "({})", args_column_name.join(","))?;
+
+                write!(f, " OVER(")?;
+                if !partition_by.is_empty() {
+                    write!(f, "PARTITION BY {:?}", partition_by)?;
+                }
+                if !order_by.is_empty() {
+                    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 {}",
+                        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 {

common/planners/src/plan_expression_chain.rs

@@ -192,7 +192,11 @@ impl ExpressionChain {
                     "Action must be a non-aggregated function.",
                 ));
             }
+
+            Expression::WindowFunction { .. } => {}
+
             Expression::Wildcard | Expression::Sort { .. } => {}
+
             Expression::Cast {
                 expr: sub_expr,
                 data_type,