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

Author: jiangzhx

datafusion/core/tests/dataframe.rs

@@ -36,7 +36,7 @@ use datafusion_expr::expr::{GroupingSet, Sort};
 use datafusion_expr::{avg, col, count, lit, max, sum, Expr, ExprSchemable};
 
 #[tokio::test]
-async fn count_star() -> Result<()> {
+async fn count_wildcard() -> Result<()> {
     let ctx = SessionContext::new();
 
     let testdata = datafusion::test_util::parquet_test_data();
@@ -48,60 +48,38 @@ async fn count_star() -> Result<()> {
         .await?;
 
     let results = df
+        .clone()
         .aggregate(vec![], vec![count(Expr::Wildcard)])?
+        .explain(false, false)
+        .unwrap()
         .collect()
         .await?;
 
     let expected = vec![
-        "+-----------------+",
-        "| COUNT(UInt8(1)) |",
-        "+-----------------+",
-        "| 7300            |",
-        "+-----------------+",
+        "+---------------+---------------------------------------------------+",
+        "| plan_type     | plan                                              |",
+        "+---------------+---------------------------------------------------+",
+        "|               |                                                   |",
+        "|               |   EmptyExec: produce_one_row=true                 |",
+        "|               |   TableScan: ?table? projection=[id]              |",
+        "| logical_plan  | Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1))]] |",
+        "| physical_plan | ProjectionExec: expr=[7300 as COUNT(UInt8(1))]    |",
+        "+---------------+---------------------------------------------------+",
     ];
     assert_batches_sorted_eq!(expected, &results);
 
-    Ok(())
-}
-
-#[tokio::test]
-async fn count_test() -> Result<()> {
-    let mut ctx = SessionContext::new();
-    let disable_aggregate_statistics = true;
-
-    //disable or not,this test case should be passed
-    if disable_aggregate_statistics {
-        let with_out_aggregate_statistics = ctx
-            .state()
-            .physical_optimizers()
-            .iter()
-            .filter(|optimizer| optimizer.as_ref().name().ne("aggregate_statistics"))
-            .map(|optimizer| optimizer.clone())
-            .collect();
-        let state = ctx
-            .state()
-            .with_physical_optimizer_rules(with_out_aggregate_statistics);
-        ctx = SessionContext::with_state(state);
-    }
-
-    let testdata = datafusion::test_util::parquet_test_data();
-    ctx.register_parquet(
-        "alltypes_tiny_pages",
-        &format!("{testdata}/alltypes_tiny_pages.parquet"),
-        ParquetReadOptions::default(),
-    )
-    .await?;
-    let results = ctx
-        .sql("SELECT count(id), max(id), min(id) FROM alltypes_tiny_pages")
-        .await?
+    let results = df
+        .clone()
+        .aggregate(vec![], vec![count(Expr::Wildcard)])?
         .collect()
         .await?;
+
     let expected = vec![
-        "+-------------------------------+-----------------------------+-----------------------------+",
-        "| COUNT(alltypes_tiny_pages.id) | MAX(alltypes_tiny_pages.id) | MIN(alltypes_tiny_pages.id) |",
-        "+-------------------------------+-----------------------------+-----------------------------+",
-        "| 7300                          | 7299                        | 0                           |",
-        "+-------------------------------+-----------------------------+-----------------------------+",
+        "+-----------------+",
+        "| COUNT(UInt8(1)) |",
+        "+-----------------+",
+        "| 7300            |",
+        "+-----------------+",
     ];
     assert_batches_sorted_eq!(expected, &results);
 

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,38 @@ impl LogicalPlanBuilder {
     ) -> Result<Self> {
         let group_expr = normalize_cols(group_expr, &self.plan)?;
         let aggr_expr = normalize_cols(aggr_expr, &self.plan)?;
+
+        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 {
+                        let w = args.get(0).unwrap();
+                        match w {
+                            Expr::Wildcard => {
+                                Expr::AggregateFunction(AggregateFunction {
+                                    fun: fun.clone(),
+                                    args: vec![lit(ScalarValue::UInt8(Some(1)))],
+                                    distinct: distinct.clone(),
+                                    filter: filter.clone(),
+                                })
+                            }
+                            _ => expr.clone(),
+                        }
+                    } else {
+                        expr.clone()
+                    }
+                } else {
+                    expr.clone()
+                }
+            })
+            .collect();
         Ok(Self::from(LogicalPlan::Aggregate(Aggregate::try_new(
             Arc::new(self.plan),
             group_expr,