From 6505b4dd3f3e2072135d9de8338af996d8164c4f Mon Sep 17 00:00:00 2001
From: Andy Grove <andygrove73@gmail.com>
Date: Fri, 2 Sep 2022 12:38:03 -0600
Subject: [PATCH] Remove Result.unwrap from single_distinct_to_groupby

---
 .../src/single_distinct_to_groupby.rs         | 31 +++++++++----------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs
index 45d3e6ff13ca..dd220e16e6b2 100644
--- a/datafusion/optimizer/src/single_distinct_to_groupby.rs
+++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs
@@ -63,7 +63,7 @@ fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
             schema,
             group_expr,
         }) => {
-            if is_single_distinct_agg(plan) && !contains_grouping_set(group_expr) {
+            if is_single_distinct_agg(plan)? && !contains_grouping_set(group_expr) {
                 let mut group_fields_set = HashSet::new();
                 let base_group_expr = grouping_set_to_exprlist(group_expr)?;
                 let mut all_group_args: Vec<Expr> = group_expr.clone();
@@ -159,27 +159,24 @@ fn optimize_children(plan: &LogicalPlan) -> Result<LogicalPlan> {
     from_plan(plan, &expr, &new_inputs)
 }
 
-fn is_single_distinct_agg(plan: &LogicalPlan) -> bool {
+fn is_single_distinct_agg(plan: &LogicalPlan) -> Result<bool> {
     match plan {
         LogicalPlan::Aggregate(Aggregate { aggr_expr, .. }) => {
             let mut fields_set = HashSet::new();
-            aggr_expr
-                .iter()
-                .filter(|expr| {
-                    let mut is_distinct = false;
-                    if let Expr::AggregateFunction { distinct, args, .. } = expr {
-                        is_distinct = *distinct;
-                        args.iter().for_each(|expr| {
-                            fields_set.insert(expr.name().unwrap());
-                        })
+            let mut count = 0;
+            for expr in aggr_expr {
+                if let Expr::AggregateFunction { distinct, args, .. } = expr {
+                    if *distinct {
+                        count += 1;
                     }
-                    is_distinct
-                })
-                .count()
-                == aggr_expr.len()
-                && fields_set.len() == 1
+                    for expr in args {
+                        fields_set.insert(expr.name()?);
+                    }
+                }
+            }
+            Ok(count == aggr_expr.len() && fields_set.len() == 1)
         }
-        _ => false,
+        _ => Ok(false),
     }
 }