From c45fc418dbc156c0067b5ae219fb44e8c3083d11 Mon Sep 17 00:00:00 2001
From: HuSen <husen.xjtu@gmail.com>
Date: Thu, 3 Oct 2024 00:55:47 +0800
Subject: [PATCH] Fix: support Qualified Wildcard in count aggregate function
 (#12673)

---
 .../optimizer/src/analyzer/count_wildcard_rule.rs    |  8 +-------
 datafusion/sql/src/expr/function.rs                  | 12 ++++++++++++
 datafusion/sqllogictest/test_files/aggregate.slt     |  8 ++++++++
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs b/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs
index 86520b3587cd..b3b24724552a 100644
--- a/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs
+++ b/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs
@@ -48,13 +48,7 @@ impl AnalyzerRule for CountWildcardRule {
 }
 
 fn is_wildcard(expr: &Expr) -> bool {
-    matches!(
-        expr,
-        Expr::Wildcard {
-            qualifier: None,
-            ..
-        }
-    )
+    matches!(expr, Expr::Wildcard { .. })
 }
 
 fn is_count_star_aggregate(aggregate_function: &AggregateFunction) -> bool {
diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs
index ddafc4e3a03a..20a772cdd088 100644
--- a/datafusion/sql/src/expr/function.rs
+++ b/datafusion/sql/src/expr/function.rs
@@ -432,6 +432,18 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 qualifier: None,
                 options: WildcardOptions::default(),
             }),
+            FunctionArg::Unnamed(FunctionArgExpr::QualifiedWildcard(object_name)) => {
+                let qualifier = self.object_name_to_table_reference(object_name)?;
+                // sanity check on qualifier with schema
+                let qualified_indices = schema.fields_indices_with_qualified(&qualifier);
+                if qualified_indices.is_empty() {
+                    return plan_err!("Invalid qualifier {qualifier}");
+                }
+                Ok(Expr::Wildcard {
+                    qualifier: Some(qualifier),
+                    options: WildcardOptions::default(),
+                })
+            }
             _ => not_impl_err!("Unsupported qualified wildcard argument: {sql:?}"),
         }
     }
diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt
index 46327534e7de..a78ade81eeba 100644
--- a/datafusion/sqllogictest/test_files/aggregate.slt
+++ b/datafusion/sqllogictest/test_files/aggregate.slt
@@ -1124,6 +1124,14 @@ SELECT COUNT(*) FROM aggregate_test_100
 ----
 100
 
+query I
+SELECT COUNT(aggregate_test_100.*) FROM aggregate_test_100
+----
+100
+
+query error Error during planning: Invalid qualifier foo
+SELECT COUNT(foo.*) FROM aggregate_test_100
+
 # csv_query_count_literal
 query I
 SELECT COUNT(2) FROM aggregate_test_100