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

Author: jiangzhx

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,58 @@ 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 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();

datafusion/expr/src/expr_schema.rs

@@ -136,9 +136,10 @@ impl ExprSchemable for Expr {
             Expr::Placeholder { data_type, .. } => data_type.clone().ok_or_else(|| {
                 DataFusionError::Plan("Placeholder type could not be resolved".to_owned())
             }),
-            Expr::Wildcard => Err(DataFusionError::Internal(
-                "Wildcard expressions are not valid in a logical query plan".to_owned(),
-            )),
+            Expr::Wildcard => {
+                // Wildcard do not really have a type and do not appear in projections
+                Ok(DataType::Null)
+            }
             Expr::QualifiedWildcard { .. } => Err(DataFusionError::Internal(
                 "QualifiedWildcard expressions are not valid in a logical query plan"
                     .to_owned(),

datafusion/expr/src/logical_plan/builder.rs

@@ -1371,7 +1371,7 @@ 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::{OwnedTableReference, SchemaError, TableReference};
 

datafusion/optimizer/src/analyzer.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::count_wildcard_rule::CountWildcardRule;
 use crate::rewrite::TreeNodeRewritable;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::{DataFusionError, Result};
@@ -49,7 +50,8 @@ impl Default for Analyzer {
 impl Analyzer {
     /// Create a new analyzer using the recommended list of rules
     pub fn new() -> Self {
-        let rules = vec![];
+        let rules: Vec<Arc<dyn AnalyzerRule + Send + Sync>> =
+            vec![Arc::new(CountWildcardRule::new())];
         Self::with_rules(rules)
     }
 

datafusion/optimizer/src/count_wildcard_rule.rs

@@ -0,0 +1,72 @@
+use crate::analyzer::AnalyzerRule;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::Result;
+use datafusion_expr::expr::AggregateFunction;
+use datafusion_expr::utils::COUNT_STAR_EXPANSION;
+use datafusion_expr::{aggregate_function, lit, Aggregate, Expr, LogicalPlan};
+use std::sync::Arc;
+
+pub struct CountWildcardRule {}
+
+impl CountWildcardRule {
+    pub fn new() -> Self {
+        CountWildcardRule {}
+    }
+}
+
+impl AnalyzerRule for CountWildcardRule {
+    fn analyze(
+        &self,
+        plan: &LogicalPlan,
+        _: &ConfigOptions,
+    ) -> datafusion_common::Result<LogicalPlan> {
+        let new_plan = match plan {
+            LogicalPlan::Window(_window) => plan.clone(),
+            LogicalPlan::Aggregate(aggregate) => {
+                let aggr_expr = aggregate.clone().aggr_expr;
+                let aggr_expr = handle_wildcard(aggr_expr).unwrap();
+
+                LogicalPlan::Aggregate(
+                    Aggregate::try_new_with_schema(
+                        Arc::new(plan.inputs().get(0).unwrap().clone().clone()),
+                        aggregate.clone().group_expr,
+                        aggr_expr,
+                        plan.schema().clone(),
+                    )
+                    .unwrap(),
+                )
+            }
+            _ => plan.clone(),
+        };
+        Ok(new_plan)
+    }
+
+    fn name(&self) -> &str {
+        "count_wildcard_rule"
+    }
+}
+
+//handle Count(Expr:Wildcard) with DataFrame API
+pub fn handle_wildcard(exprs: Vec<Expr>) -> Result<Vec<Expr>> {
+    let exprs: Vec<Expr> = exprs
+        .iter()
+        .map(|expr| match expr {
+            Expr::AggregateFunction(AggregateFunction {
+                fun: aggregate_function::AggregateFunction::Count,
+                args,
+                distinct,
+                filter,
+            }) if args.len() == 1 => match args[0] {
+                Expr::Wildcard => Expr::AggregateFunction(AggregateFunction {
+                    fun: aggregate_function::AggregateFunction::Count,
+                    args: vec![lit(COUNT_STAR_EXPANSION)],
+                    distinct: *distinct,
+                    filter: filter.clone(),
+                }),
+                _ => expr.clone(),
+            },
+            _ => expr.clone(),
+        })
+        .collect();
+    Ok(exprs)
+}

datafusion/optimizer/src/lib.rs

@@ -45,6 +45,7 @@ pub mod type_coercion;
 pub mod unwrap_cast_in_comparison;
 pub mod utils;
 
+pub mod count_wildcard_rule;
 #[cfg(test)]
 pub mod test;