fix erro on Count(Expr:Wildcard) with DataFrame API

Author: jiangzhx

Diff for: datafusion/core/tests/dataframe.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use arrow::datatypes::{DataType, Field, Schema};
+use arrow::util::pretty::pretty_format_batches;
 use arrow::{
     array::{
         ArrayRef, Int32Array, Int32Builder, ListBuilder, StringArray, StringBuilder,
@@ -35,6 +36,70 @@ use datafusion::{assert_batches_eq, assert_batches_sorted_eq};
 use datafusion_expr::expr::{GroupingSet, Sort};
 use datafusion_expr::{avg, col, count, lit, max, sum, Expr, ExprSchemable};
 
+#[tokio::test]
+async fn count_wildcard() -> Result<()> {
+    let ctx = SessionContext::new();
+    let testdata = datafusion::test_util::parquet_test_data();
+
+    ctx.register_parquet(
+        "alltypes_tiny_pages",
+        &format!("{testdata}/alltypes_tiny_pages.parquet"),
+        ParquetReadOptions::default(),
+    )
+    .await?;
+
+    let expected = vec![
+        "+---------------+---------------------------------------------------+",
+        "| plan_type     | plan                                              |",
+        "+---------------+---------------------------------------------------+",
+        "|               |                                                   |",
+        "|               |   EmptyExec: produce_one_row=true                 |",
+        "|               |   TableScan: alltypes_tiny_pages projection=[id]  |",
+        "| logical_plan  | Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1))]] |",
+        "| physical_plan | ProjectionExec: expr=[7300 as COUNT(UInt8(1))]    |",
+        "+---------------+---------------------------------------------------+",
+    ];
+    let sql_results = ctx
+        .sql("select count(*) from alltypes_tiny_pages")
+        .await?
+        .explain(false, false)?
+        .collect()
+        .await?;
+
+    let df_results = ctx
+        .table("alltypes_tiny_pages")
+        .await?
+        .aggregate(vec![], vec![count(Expr::Wildcard)])?
+        .explain(false, false)
+        .unwrap()
+        .collect()
+        .await?;
+
+    //make sure sql plan same with df plan
+    assert_eq!(
+        pretty_format_batches(&sql_results)?.to_string(),
+        pretty_format_batches(&df_results)?.to_string()
+    );
+
+    let results = ctx
+        .table("alltypes_tiny_pages")
+        .await?
+        .aggregate(vec![], vec![count(Expr::Wildcard)])?
+        .collect()
+        .await?;
+
+    let expected = vec![
+        "+-----------------+",
+        "| COUNT(UInt8(1)) |",
+        "+-----------------+",
+        "| 7300            |",
+        "+-----------------+",
+    ];
+    assert_batches_sorted_eq!(expected, &results);
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn describe() -> Result<()> {
     let ctx = SessionContext::new();

Diff for: datafusion/expr/src/logical_plan/builder.rs

@@ -17,14 +17,15 @@
 
 //! This module provides a builder for creating LogicalPlans
 
+use crate::expr::AggregateFunction;
 use crate::expr_rewriter::{
     coerce_plan_expr_for_schema, normalize_col,
     normalize_col_with_schemas_and_ambiguity_check, normalize_cols,
     rewrite_sort_cols_by_aggs,
 };
 use crate::type_coercion::binary::comparison_coercion;
 use crate::utils::{columnize_expr, compare_sort_expr, exprlist_to_fields, from_plan};
-use crate::{and, binary_expr, Operator};
+use crate::{aggregate_function, and, binary_expr, lit, Operator};
 use crate::{
     logical_plan::{
         Aggregate, Analyze, CrossJoin, Distinct, EmptyRelation, Explain, Filter, Join,
@@ -801,6 +802,43 @@ impl LogicalPlanBuilder {
     ) -> Result<Self> {
         let group_expr = normalize_cols(group_expr, &self.plan)?;
         let aggr_expr = normalize_cols(aggr_expr, &self.plan)?;
+
+        //handle Count(Expr:Wildcard) with DataFrame API
+        let aggr_expr = aggr_expr
+            .iter()
+            .map(|expr| {
+                if let Expr::AggregateFunction(AggregateFunction {
+                    fun,
+                    args,
+                    distinct,
+                    filter,
+                }) = expr
+                {
+                    if let aggregate_function::AggregateFunction::Count = fun {
+                        if args.len() == 1 {
+                            let arg = args.get(0).unwrap().clone();
+                            match arg {
+                                Expr::Wildcard => {
+                                    Expr::AggregateFunction(AggregateFunction {
+                                        fun: fun.clone(),
+                                        args: vec![lit(ScalarValue::UInt8(Some(1)))],
+                                        distinct: *distinct,
+                                        filter: filter.clone(),
+                                    })
+                                }
+                                _ => expr.clone(),
+                            }
+                        } else {
+                            expr.clone()
+                        }
+                    } else {
+                        expr.clone()
+                    }
+                } else {
+                    expr.clone()
+                }
+            })
+            .collect();
         Ok(Self::from(LogicalPlan::Aggregate(Aggregate::try_new(
             Arc::new(self.plan),
             group_expr,
@@ -1315,15 +1353,31 @@ pub fn unnest(input: LogicalPlan, column: Column) -> Result<LogicalPlan> {
 
 #[cfg(test)]
 mod tests {
-    use crate::{expr, expr_fn::exists};
+    use crate::{count, expr, expr_fn::exists};
     use arrow::datatypes::{DataType, Field};
     use datafusion_common::SchemaError;
 
     use crate::logical_plan::StringifiedPlan;
 
     use super::*;
+    use crate::BuiltinScalarFunction::Exp;
     use crate::{col, in_subquery, lit, scalar_subquery, sum};
 
+    #[test]
+    fn count_wildcard() -> Result<()> {
+        let mut group_expr: Vec<Expr> = Vec::new();
+        let plan = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![]))?
+            .aggregate(group_expr, vec![count(Expr::Wildcard)])?
+            .build()?;
+
+        let expected = "Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1))]]\
+                \n  TableScan: employee_csv projection=[]";
+
+        assert_eq!(expected, format!("{plan:?}"));
+
+        Ok(())
+    }
+
     #[test]
     fn plan_builder_simple() -> Result<()> {
         let plan =