combine match and if condition, add more test cases

apache · Jun 7, 2021 · c86c061 · c86c061
1 parent 6bcd83b
commit c86c061
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 10 deletions.
diff --git a/datafusion/src/sql/planner.rs b/datafusion/src/sql/planner.rs
@@ -2328,6 +2328,29 @@ mod tests {
              \n  Aggregate: groupBy=[[#state, #age]], aggr=[[COUNT(UInt8(1))]]\
              \n    TableScan: person projection=None",
         );
+        quick_test(
+            "SELECT state, age AS b, COUNT(1) FROM person GROUP BY 2, 1",
+            "Projection: #state, #age AS b, #COUNT(UInt8(1))\
+             \n  Aggregate: groupBy=[[#age, #state]], aggr=[[COUNT(UInt8(1))]]\
+             \n    TableScan: person projection=None",
+        );
+    }
+
+    #[test]
+    fn select_simple_aggregate_with_groupby_position_out_of_range() {
+        let sql = "SELECT state, MIN(age) FROM person GROUP BY 0";
+        let err = logical_plan(sql).expect_err("query should have failed");
+        assert_eq!(
+            "Plan(\"Projection references non-aggregate values\")",
+            format!("{:?}", err)
+        );
+
+        let sql2 = "SELECT state, MIN(age) FROM person GROUP BY 5";
+        let err2 = logical_plan(sql2).expect_err("query should have failed");
+        assert_eq!(
+            "Plan(\"Projection references non-aggregate values\")",
+            format!("{:?}", err2)
+        );
     }
 
     #[test]

diff --git a/datafusion/src/sql/utils.rs b/datafusion/src/sql/utils.rs
@@ -398,16 +398,14 @@ pub(crate) fn resolve_positions_to_exprs(
     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
-        Expr::Literal(ScalarValue::Int64(Some(pos))) => {
-            let index = (pos - 1) as usize;
-            if index < select_exprs.len() {
-                let select_expr = &select_exprs[index];
-                match select_expr {
-                    Expr::Alias(nested_expr, _alias_name) => Ok(*nested_expr.clone()),
-                    _ => Ok(select_expr.clone()),
-                }
-            } else {
-                Ok(expr.clone())
+        Expr::Literal(ScalarValue::Int64(Some(position)))
+            if position > &(0 as i64) && position <= &(select_exprs.len() as i64) =>
+        {
+            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()),
             }
         }
         _ => Ok(expr.clone()),