reuse alias map in aggregate logical planning and refactor position resolution

datafusion/src/sql/planner.rs

@@ -535,14 +535,16 @@ 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
             .as_ref()
             .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,15 +575,15 @@ 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()
             .map(|e| {
                 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()],

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,
     }
 }