refactor: type aggregate and window binding

Author: KKould

src/binder/expr.rs

@@ -476,76 +476,60 @@ impl<'a, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'a, '_, T
         })
     }
 
-    pub(crate) fn bind_function_call(
+    pub(crate) fn bind_aggregate_function(
         &mut self,
-        function_name: String,
-        mut args: Vec<ScalarExpression>,
+        kind: AggKind,
+        args: Vec<ScalarExpression>,
         is_distinct: bool,
         arena: &mut PlanArena,
     ) -> Result<ScalarExpression, DatabaseError> {
-        match function_name.as_str() {
-            "count" => {
+        let ty = match kind {
+            AggKind::Count => {
                 if args.len() != 1 {
                     return Err(DatabaseError::MisMatch("number of count() parameters", "1"));
                 }
-                return Ok(ScalarExpression::AggCall {
-                    distinct: is_distinct,
-                    kind: AggKind::Count,
-                    args,
-                    ty: LogicalType::Integer,
-                });
+                LogicalType::Integer
             }
-            "sum" => {
+            AggKind::Sum => {
                 if args.len() != 1 {
                     return Err(DatabaseError::MisMatch("number of sum() parameters", "1"));
                 }
-                let ty = args[0].return_type(arena).into_owned();
-
-                return Ok(ScalarExpression::AggCall {
-                    distinct: is_distinct,
-                    kind: AggKind::Sum,
-                    args,
-                    ty,
-                });
+                args[0].return_type(arena).into_owned()
             }
-            "min" => {
+            AggKind::Min => {
                 if args.len() != 1 {
                     return Err(DatabaseError::MisMatch("number of min() parameters", "1"));
                 }
-                let ty = args[0].return_type(arena).into_owned();
-
-                return Ok(ScalarExpression::AggCall {
-                    distinct: is_distinct,
-                    kind: AggKind::Min,
-                    args,
-                    ty,
-                });
+                args[0].return_type(arena).into_owned()
             }
-            "max" => {
+            AggKind::Max => {
                 if args.len() != 1 {
                     return Err(DatabaseError::MisMatch("number of max() parameters", "1"));
                 }
-                let ty = args[0].return_type(arena).into_owned();
-
-                return Ok(ScalarExpression::AggCall {
-                    distinct: is_distinct,
-                    kind: AggKind::Max,
-                    args,
-                    ty,
-                });
+                args[0].return_type(arena).into_owned()
             }
-            "avg" => {
+            AggKind::Avg => {
                 if args.len() != 1 {
                     return Err(DatabaseError::MisMatch("number of avg() parameters", "1"));
                 }
-
-                return Ok(ScalarExpression::AggCall {
-                    distinct: is_distinct,
-                    kind: AggKind::Avg,
-                    args,
-                    ty: LogicalType::Double,
-                });
+                LogicalType::Double
             }
+        };
+        Ok(ScalarExpression::AggCall {
+            distinct: is_distinct,
+            kind,
+            args,
+            ty,
+        })
+    }
+
+    pub(crate) fn bind_function_call(
+        &mut self,
+        function_name: String,
+        mut args: Vec<ScalarExpression>,
+        arena: &mut PlanArena,
+    ) -> Result<ScalarExpression, DatabaseError> {
+        match function_name.as_str() {
             "if" => {
                 if args.len() != 3 {
                     return Err(DatabaseError::MisMatch("number of if() parameters", "3"));

src/binder/parser.rs

@@ -23,8 +23,10 @@ use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef, TableName};
 use crate::db::{BindSource, DBTransaction, Database, DatabaseIter, TransactionIter};
 use crate::errors::{DatabaseError, SqlErrorSpan};
 use crate::expression;
+use crate::expression::agg::AggKind;
 use crate::expression::simplify::ConstantCalculator;
 use crate::expression::visitor_mut::ExprVisitorMut;
+use crate::expression::window::WindowFunctionKind;
 use crate::expression::{AliasType, ScalarExpression};
 use crate::iter_ext::Itertools;
 use crate::parser::parse_sql;
@@ -2440,18 +2442,30 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
                         .to_string(),
                 ));
             }
+            let Some(kind) = WindowFunctionKind::from_name(&function_name) else {
+                return Err(attach_span_if_absent(
+                    DatabaseError::UnsupportedStmt(format!(
+                        "window function `{function_name}` is not supported"
+                    )),
+                    func_span,
+                ));
+            };
             return self
-                .bind_window_call(function_name, args, is_distinct, over, arena)
+                .bind_window_call(kind, args, is_distinct, over, arena)
                 .map_err(|err| attach_span_if_absent(err, func_span));
         }
 
-        self.bind_function_call(function_name, args, is_distinct, arena)
-            .map_err(|err| attach_span_if_absent(err, func_span))
+        let result = if let Some(kind) = AggKind::from_name(&function_name) {
+            self.bind_aggregate_function(kind, args, is_distinct, arena)
+        } else {
+            self.bind_function_call(function_name, args, arena)
+        };
+        result.map_err(|err| attach_span_if_absent(err, func_span))
     }
 
     fn bind_window_call(
         &mut self,
-        function_name: String,
+        kind: WindowFunctionKind,
         args: Vec<ScalarExpression>,
         is_distinct: bool,
         over: &WindowType,
@@ -2502,7 +2516,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
                 ))
             })
             .collect::<Result<Vec<_>, DatabaseError>>()?;
-        self.bind_window_function(function_name, args, partition_by, order_by, arena)
+        self.bind_window_function(kind, args, partition_by, order_by, arena)
     }
 
     pub fn bind_set_expr(

src/binder/window.rs

@@ -94,7 +94,7 @@ struct WindowGroup {
 impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A> {
     pub(crate) fn bind_window_function(
         &mut self,
-        function_name: String,
+        kind: WindowFunctionKind,
         args: Vec<ScalarExpression>,
         partition_by: Vec<ScalarExpression>,
         order_by: Vec<SortField>,
@@ -120,34 +120,25 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
             }
         }
 
-        let (kind, args, ty) = match function_name.as_str() {
-            "row_number" | "rank" | "dense_rank" => {
+        let (args, ty) = match kind {
+            WindowFunctionKind::RowNumber
+            | WindowFunctionKind::Rank
+            | WindowFunctionKind::DenseRank => {
                 if !args.is_empty() {
                     return Err(DatabaseError::MisMatch(
                         "number of ranking function parameters",
                         "0",
                     ));
                 }
-                let kind = match function_name.as_str() {
-                    "row_number" => WindowFunctionKind::RowNumber,
-                    "rank" => WindowFunctionKind::Rank,
-                    "dense_rank" => WindowFunctionKind::DenseRank,
-                    _ => unreachable!(),
-                };
-                (kind, args, LogicalType::Bigint)
+                (args, LogicalType::Bigint)
             }
-            "count" | "sum" | "avg" | "min" | "max" => {
-                let ScalarExpression::AggCall { kind, args, ty, .. } =
-                    self.bind_function_call(function_name, args, false, arena)?
+            WindowFunctionKind::Aggregate(agg_kind) => {
+                let ScalarExpression::AggCall { args, ty, .. } =
+                    self.bind_aggregate_function(agg_kind, args, false, arena)?
                 else {
                     unreachable!()
                 };
-                (WindowFunctionKind::Aggregate(kind), args, ty)
-            }
-            _ => {
-                return Err(DatabaseError::UnsupportedStmt(format!(
-                    "window function `{function_name}` is not supported"
-                )))
+                (args, ty)
             }
         };
 

src/expression/agg.rs

@@ -24,6 +24,27 @@ pub enum AggKind {
 }
 
 impl AggKind {
+    pub(crate) fn from_name(name: &str) -> Option<Self> {
+        match name {
+            "avg" => Some(Self::Avg),
+            "max" => Some(Self::Max),
+            "min" => Some(Self::Min),
+            "sum" => Some(Self::Sum),
+            "count" => Some(Self::Count),
+            _ => None,
+        }
+    }
+
+    pub(crate) fn name(self) -> &'static str {
+        match self {
+            Self::Avg => "avg",
+            Self::Max => "max",
+            Self::Min => "min",
+            Self::Sum => "sum",
+            Self::Count => "count",
+        }
+    }
+
     pub fn allow_distinct(&self) -> bool {
         match self {
             AggKind::Avg => false,

src/expression/mod.rs

@@ -784,14 +784,7 @@ impl ScalarExpression {
                     .iter()
                     .map(|expr| expr.output_name_by(fn_display))
                     .join(", ");
-                let function = match window.function.kind {
-                    window::WindowFunctionKind::RowNumber => "row_number".to_string(),
-                    window::WindowFunctionKind::Rank => "rank".to_string(),
-                    window::WindowFunctionKind::DenseRank => "dense_rank".to_string(),
-                    window::WindowFunctionKind::Aggregate(kind) => {
-                        format!("{kind:?}").to_lowercase()
-                    }
-                };
+                let function = window.function.kind.name();
                 let mut spec = Vec::new();
                 if !window.spec.partition_by.is_empty() {
                     spec.push(format!(

src/expression/window.rs

@@ -26,6 +26,26 @@ pub enum WindowFunctionKind {
     Aggregate(AggKind),
 }
 
+impl WindowFunctionKind {
+    pub(crate) fn from_name(name: &str) -> Option<Self> {
+        match name {
+            "row_number" => Some(Self::RowNumber),
+            "rank" => Some(Self::Rank),
+            "dense_rank" => Some(Self::DenseRank),
+            name => AggKind::from_name(name).map(Self::Aggregate),
+        }
+    }
+
+    pub(crate) fn name(self) -> &'static str {
+        match self {
+            Self::RowNumber => "row_number",
+            Self::Rank => "rank",
+            Self::DenseRank => "dense_rank",
+            Self::Aggregate(kind) => kind.name(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash, ReferenceSerialization)]
 pub struct WindowFunction {
     pub kind: WindowFunctionKind,