From a40f0b08a88df8265606cd35e07951d3f6e02cf1 Mon Sep 17 00:00:00 2001
From: Jiayu Liu <jiayu.liu@airbnb.com>
Date: Wed, 23 Jun 2021 16:09:52 +0800
Subject: [PATCH] reuse alias map in aggregate logical planning and refactor
 position resolving

---
 datafusion/src/sql/planner.rs | 15 ++++++---------
 datafusion/src/sql/utils.rs   | 15 +++++++++------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/datafusion/src/sql/planner.rs b/datafusion/src/sql/planner.rs
index 7912241329a3..961c3edb912c 100644
--- a/datafusion/src/sql/planner.rs
+++ b/datafusion/src/sql/planner.rs
@@ -535,6 +535,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         let mut combined_schema = (**projected_plan.schema()).clone();
         combined_schema.merge(plan.schema());
 
+        // this alias map is resolved and looked up in both having exprs and group by exprs
+        let alias_map = extract_aliases(&select_exprs);
+
         // Optionally the HAVING expression.
         let having_expr_opt = select
             .having
@@ -542,7 +545,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             .map::<Result<Expr>, _>(|having_expr| {
                 let having_expr =
                     self.sql_expr_to_logical_expr(having_expr, &combined_schema)?;
-
                 // This step "dereferences" any aliases in the HAVING clause.
                 //
                 // This is how we support queries with HAVING expressions that
@@ -558,12 +560,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 //   SELECT c1 AS m FROM t HAVING c1 > 10;
                 //   SELECT c1, MAX(c2) AS m FROM t GROUP BY c1 HAVING MAX(c2) > 10;
                 //
-                let having_expr = resolve_aliases_to_exprs(
-                    &having_expr,
-                    &extract_aliases(&select_exprs),
-                )?;
-
-                Ok(having_expr)
+                resolve_aliases_to_exprs(&having_expr, &alias_map)
             })
             .transpose()?;
 
@@ -578,7 +575,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         // All of the aggregate expressions (deduplicated).
         let aggr_exprs = find_aggregate_exprs(&aggr_expr_haystack);
 
-        let alias_map = extract_aliases(&select_exprs);
         let group_by_exprs = select
             .group_by
             .iter()
@@ -586,7 +582,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 let group_by_expr = self.sql_expr_to_logical_expr(e, &combined_schema)?;
                 let group_by_expr = resolve_aliases_to_exprs(&group_by_expr, &alias_map)?;
                 let group_by_expr =
-                    resolve_positions_to_exprs(&group_by_expr, &select_exprs)?;
+                    resolve_positions_to_exprs(&group_by_expr, &select_exprs)
+                        .unwrap_or(group_by_expr);
                 self.validate_schema_satisfies_exprs(
                     plan.schema(),
                     &[group_by_expr.clone()],
diff --git a/datafusion/src/sql/utils.rs b/datafusion/src/sql/utils.rs
index 7702748df44f..b3da98ee30da 100644
--- a/datafusion/src/sql/utils.rs
+++ b/datafusion/src/sql/utils.rs
@@ -398,10 +398,13 @@ pub(crate) fn extract_aliases(exprs: &[Expr]) -> HashMap<String, Expr> {
         .collect::<HashMap<String, Expr>>()
 }
 
+/// Given an expression that's literal int encoding position, lookup the corresponding expression
+/// in the select_exprs list, if the index is within the bounds and it is indeed a position literal;
+/// Otherwise, return None
 pub(crate) fn resolve_positions_to_exprs(
     expr: &Expr,
     select_exprs: &[Expr],
-) -> Result<Expr> {
+) -> Option<Expr> {
     match expr {
         // sql_expr_to_logical_expr maps number to i64
         // https://github.com/apache/arrow-datafusion/blob/8d175c759e17190980f270b5894348dc4cff9bbf/datafusion/src/sql/planner.rs#L882-L887
@@ -410,12 +413,12 @@ pub(crate) fn resolve_positions_to_exprs(
         {
             let index = (position - 1) as usize;
             let select_expr = &select_exprs[index];
-            match select_expr {
-                Expr::Alias(nested_expr, _alias_name) => Ok(*nested_expr.clone()),
-                _ => Ok(select_expr.clone()),
-            }
+            Some(match select_expr {
+                Expr::Alias(nested_expr, _alias_name) => *nested_expr.clone(),
+                _ => select_expr.clone(),
+            })
         }
-        _ => Ok(expr.clone()),
+        _ => None,
     }
 }