feat: #110 support group by column positions

datafusion/src/sql/planner.rs

@@ -54,8 +54,9 @@ use super::{
     parser::DFParser,
     utils::{
         can_columns_satisfy_exprs, expand_wildcard, expr_as_column_expr, extract_aliases,
-        find_aggregate_exprs, find_column_exprs, find_window_exprs,
+        extract_positions, find_aggregate_exprs, find_column_exprs, find_window_exprs,
         group_window_expr_by_sort_keys, rebase_expr, resolve_aliases_to_exprs,
+        resolve_positions_to_exprs,
     },
 };
 
@@ -591,6 +592,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                     &group_by_expr,
                     &extract_aliases(&select_exprs),
                 )?;
+                let group_by_expr = resolve_positions_to_exprs(
+                    &group_by_expr,
+                    &extract_positions(&select_exprs),
+                )?;
                 self.validate_schema_satisfies_exprs(
                     plan.schema(),
                     &[group_by_expr.clone()],
@@ -2319,6 +2324,16 @@ mod tests {
         );
     }
 
+    #[test]
+    fn select_simple_aggregate_with_groupby_can_use_positions() {
+        quick_test(
+            "SELECT state, age AS b, COUNT(1) FROM person GROUP BY 1, 2",
+            "Projection: #state, #age AS b, #COUNT(UInt8(1))\
+             \n  Aggregate: groupBy=[[#state, #age]], aggr=[[COUNT(UInt8(1))]]\
+             \n    TableScan: person projection=None",
+        );
+    }
+
     #[test]
     fn select_simple_aggregate_with_groupby_can_use_alias() {
         quick_test(

datafusion/src/sql/utils.rs

@@ -390,6 +390,42 @@ pub(crate) fn extract_aliases(exprs: &[Expr]) -> HashMap<String, Expr> {
         .collect::<HashMap<String, Expr>>()
 }
 
+/// Returns mapping of each position (`String`) to the expression (`Expr`) it is
+/// aliasing.
+pub(crate) fn extract_positions(exprs: &[Expr]) -> HashMap<String, Expr> {
+    let mut position = 0;
+    exprs
+        .iter()
+        .filter_map(|expr| match expr {
+            // position starts with 1
+            Expr::Alias(nested_expr, _alias_name) => {
+                position += 1;
+                Some((position.clone().to_string(), *nested_expr.clone()))
+            }
+            _ => {
+                position += 1;
+                Some((position.clone().to_string(), expr.clone()))
+            }
+        })
+        .collect::<HashMap<String, Expr>>()
+}
+
+pub(crate) fn resolve_positions_to_exprs(
+    expr: &Expr,
+    positions: &HashMap<String, Expr>,
+) -> Result<Expr> {
+    match expr {
+        Expr::Literal(value) => {
+            if let Some(position_expr) = positions.get(&(value.to_string())) {
+                Ok(position_expr.clone())
+            } else {
+                Ok(Expr::Literal(value.clone()))
+            }
+        }
+        default => Ok(default.clone()),
+    }
+}
+
 /// Rebuilds an `Expr` with columns that refer to aliases replaced by the
 /// alias' underlying `Expr`.
 pub(crate) fn resolve_aliases_to_exprs(