feat: support COUNT()

datafusion/core/src/execution/session_state.rs

@@ -242,6 +242,9 @@ impl SessionState {
             Arc::new(functions::datetime::planner::ExtractPlanner),
             #[cfg(feature = "unicode_expressions")]
             Arc::new(functions::unicode::planner::PositionPlanner),
+            Arc::new(
+                functions_aggregate::aggregate_function_planner::AggregateFunctionPlanner,
+            ),
         ];
 
         let mut new_self = SessionState {

datafusion/expr/src/planner.rs

@@ -24,6 +24,7 @@ use datafusion_common::{
     config::ConfigOptions, file_options::file_type::FileType, not_impl_err, DFSchema,
     Result, TableReference,
 };
+use sqlparser::ast::NullTreatment;
 
 use crate::{AggregateUDF, Expr, GetFieldAccess, ScalarUDF, TableSource, WindowUDF};
 
@@ -107,7 +108,7 @@ pub trait UserDefinedSQLPlanner: Send + Sync {
 
     /// Plan the array literal, returns OriginalArray if not possible
     ///
-    /// Returns origin expression arguments if not possible
+    /// Returns original expression arguments if not possible
     fn plan_array_literal(
         &self,
         exprs: Vec<Expr>,
@@ -124,7 +125,7 @@ pub trait UserDefinedSQLPlanner: Send + Sync {
 
     /// Plan the dictionary literal `{ key: value, ...}`
     ///
-    /// Returns origin expression arguments if not possible
+    /// Returns original expression arguments if not possible
     fn plan_dictionary_literal(
         &self,
         expr: RawDictionaryExpr,
@@ -135,10 +136,20 @@ pub trait UserDefinedSQLPlanner: Send + Sync {
 
     /// Plan an extract expression, e.g., `EXTRACT(month FROM foo)`
     ///
-    /// Returns origin expression arguments if not possible
+    /// Returns original expression arguments if not possible
     fn plan_extract(&self, args: Vec<Expr>) -> Result<PlannerResult<Vec<Expr>>> {
         Ok(PlannerResult::Original(args))
     }
+
+    /// Plan an aggregate function, e.g., `SUM(foo)`
+    ///
+    /// Returns original expression arguments if not possible
+    fn plan_aggregate_function(
+        &self,
+        aggregate_function: RawAggregateFunction,
+    ) -> Result<PlannerResult<RawAggregateFunction>> {
+        Ok(PlannerResult::Original(aggregate_function))
+    }
 }
 
 /// An operator with two arguments to plan
@@ -183,3 +194,14 @@ pub enum PlannerResult<T> {
     /// The raw expression could not be planned, and is returned unmodified
     Original(T),
 }
+
+// An aggregate function to plan.
+#[derive(Debug, Clone)]
+pub struct RawAggregateFunction {
+    pub udf: Arc<crate::AggregateUDF>,
+    pub args: Vec<Expr>,
+    pub distinct: bool,
+    pub filter: Option<Box<Expr>>,
+    pub order_by: Option<Vec<Expr>>,
+    pub null_treatment: Option<NullTreatment>,
+}

datafusion/functions-aggregate/src/aggregate_function_planner.rs

@@ -0,0 +1,80 @@
+// 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
+//
+//   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 datafusion_expr::{
+    expr, lit,
+    planner::{PlannerResult, RawAggregateFunction, UserDefinedSQLPlanner},
+    utils::COUNT_STAR_EXPANSION,
+    Expr,
+};
+
+fn is_wildcard(expr: &Expr) -> bool {
+    matches!(expr, Expr::Wildcard { qualifier: None })
+}
+
+pub struct AggregateFunctionPlanner;
+
+impl AggregateFunctionPlanner {
+    fn plan_count(
+        &self,
+        aggregate_function: RawAggregateFunction,
+    ) -> datafusion_common::Result<PlannerResult<RawAggregateFunction>> {
+        if aggregate_function.args.is_empty() {
+            return Ok(PlannerResult::Planned(Expr::AggregateFunction(
+                expr::AggregateFunction::new_udf(
+                    aggregate_function.udf,
+                    vec![lit(COUNT_STAR_EXPANSION).alias("count()")],
+                    aggregate_function.distinct,
+                    aggregate_function.filter,
+                    aggregate_function.order_by,
+                    aggregate_function.null_treatment,
+                ),
+            )));
+        }
+
+        if aggregate_function.udf.name() == "count"
+            && aggregate_function.args.len() == 1
+            && is_wildcard(&aggregate_function.args[0])
+        {
+            return Ok(PlannerResult::Planned(Expr::AggregateFunction(
+                expr::AggregateFunction::new_udf(
+                    aggregate_function.udf,
+                    vec![lit(COUNT_STAR_EXPANSION).alias("*")],
+                    aggregate_function.distinct,
+                    aggregate_function.filter,
+                    aggregate_function.order_by,
+                    aggregate_function.null_treatment,
+                ),
+            )));
+        }
+
+        Ok(PlannerResult::Original(aggregate_function))
+    }
+}
+
+impl UserDefinedSQLPlanner for AggregateFunctionPlanner {
+    fn plan_aggregate_function(
+        &self,
+        aggregate_function: RawAggregateFunction,
+    ) -> datafusion_common::Result<PlannerResult<RawAggregateFunction>> {
+        if aggregate_function.udf.name() == "count" {
+            return self.plan_count(aggregate_function);
+        }
+
+        Ok(PlannerResult::Original(aggregate_function))
+    }
+}

datafusion/functions-aggregate/src/lib.rs

@@ -84,6 +84,8 @@ use datafusion_expr::AggregateUDF;
 use log::debug;
 use std::sync::Arc;
 
+pub mod aggregate_function_planner;
+
 /// Fluent-style API for creating `Expr`s
 pub mod expr_fn {
     pub use super::approx_distinct;