Implement group by index

Signed-off-by: kuba-- <kuba@sourced.tech>

Author: kuba--

engine_test.go

@@ -159,6 +159,18 @@ var queries = []struct {
 			{int32(1), "third row"},
 		},
 	},
+	{
+		`SELECT COUNT(*) as cnt, fi FROM (
+			SELECT tbl.s AS fi
+			FROM mytable tbl
+		) t
+		GROUP BY 2`,
+		[]sql.Row{
+			{int32(1), "first row"},
+			{int32(1), "second row"},
+			{int32(1), "third row"},
+		},
+	},
 	{
 		`SELECT COUNT(*) as cnt, s as fi FROM mytable GROUP BY fi`,
 		[]sql.Row{
@@ -167,6 +179,14 @@ var queries = []struct {
 			{int32(1), "third row"},
 		},
 	},
+	{
+		`SELECT COUNT(*) as cnt, s as fi FROM mytable GROUP BY 2`,
+		[]sql.Row{
+			{int32(1), "first row"},
+			{int32(1), "second row"},
+			{int32(1), "third row"},
+		},
+	},
 	{
 		"SELECT CAST(-3 AS UNSIGNED) FROM mytable",
 		[]sql.Row{
@@ -296,6 +316,14 @@ var queries = []struct {
 			{int32(1), int64(1)},
 		},
 	},
+	{
+		`SELECT COUNT(*) c, i as foo FROM mytable GROUP BY 2 ORDER BY 2 DESC`,
+		[]sql.Row{
+			{int32(1), int64(3)},
+			{int32(1), int64(2)},
+			{int32(1), int64(1)},
+		},
+	},
 	{
 		`SELECT COUNT(*) c, i as foo FROM mytable GROUP BY i ORDER BY foo, i DESC`,
 		[]sql.Row{
@@ -304,6 +332,30 @@ var queries = []struct {
 			{int32(1), int64(1)},
 		},
 	},
+	{
+		`SELECT COUNT(*) c, i as foo FROM mytable GROUP BY 2 ORDER BY foo, i DESC`,
+		[]sql.Row{
+			{int32(1), int64(3)},
+			{int32(1), int64(2)},
+			{int32(1), int64(1)},
+		},
+	},
+	{
+		`SELECT COUNT(*) c, i as i FROM mytable GROUP BY 2`,
+		[]sql.Row{
+			{int32(1), int64(3)},
+			{int32(1), int64(2)},
+			{int32(1), int64(1)},
+		},
+	},
+	{
+		`SELECT i as i FROM mytable GROUP BY 1`,
+		[]sql.Row{
+			{int64(3)},
+			{int64(2)},
+			{int64(1)},
+		},
+	},
 	{
 		`SELECT CONCAT("a", "b", "c")`,
 		[]sql.Row{
@@ -370,6 +422,12 @@ var queries = []struct {
 			{"third row"},
 		},
 	},
+	{
+		`SELECT SUBSTRING(s, -3, 3) as s FROM mytable WHERE s LIKE '%d row' GROUP BY 1`,
+		[]sql.Row{
+			{"row"},
+		},
+	},
 	{
 		`SELECT s FROM mytable WHERE s NOT LIKE '%d row'`,
 		[]sql.Row{
@@ -560,6 +618,32 @@ var queries = []struct {
 			{"i2"},
 		},
 	},
+	{
+		`
+		SELECT COLUMN_NAME FROM information_schema.COLUMNS
+		WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME LIKE '%table'
+		GROUP BY 1
+		`,
+		[]sql.Row{
+			{"s"},
+			{"i"},
+			{"s2"},
+			{"i2"},
+		},
+	},
+	{
+		`
+		SELECT COLUMN_NAME as COLUMN_NAME FROM information_schema.COLUMNS
+		WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME LIKE '%table'
+		GROUP BY 1
+		`,
+		[]sql.Row{
+			{"s"},
+			{"i"},
+			{"s2"},
+			{"i2"},
+		},
+	},
 	{
 		`SHOW CREATE DATABASE mydb`,
 		[]sql.Row{{
@@ -1491,6 +1575,17 @@ func TestOrderByGroupBy(t *testing.T) {
 	}
 
 	require.Equal(expected, rows)
+
+	_, iter, err = e.Query(
+		newCtx(),
+		"SELECT team, COUNT(*) FROM members GROUP BY 1 ORDER BY 2",
+	)
+	require.NoError(err)
+
+	rows, err = sql.RowIterToRows(iter)
+	require.NoError(err)
+
+	require.Equal(expected, rows)
 }
 
 func TestTracing(t *testing.T) {

sql/analyzer/resolve_columns.go

@@ -303,7 +303,8 @@ func resolveColumns(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error)
 			var col *sql.Column
 			var found bool
 			for _, c := range columns {
-				if strings.ToLower(c.Source) == table {
+				_, ok := n.(*plan.GroupBy)
+				if ok || (strings.ToLower(c.Source) == table) {
 					col = c
 					found = true
 					break

sql/analyzer/validation_rules.go

@@ -77,7 +77,7 @@ func validateOrderBy(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error
 }
 
 func validateGroupBy(ctx *sql.Context, a *Analyzer, n sql.Node) (sql.Node, error) {
-	span, ctx := ctx.Span("validate_order_by")
+	span, ctx := ctx.Span("validate_group_by")
 	defer span.Finish()
 
 	switch n := n.(type) {

sql/parse/parse.go

@@ -615,6 +615,20 @@ func selectToProjectOrGroupBy(se sqlparser.SelectExprs, g sqlparser.GroupBy, chi
 			return nil, err
 		}
 
+		agglen := int64(len(selectExprs))
+		for i, ge := range groupingExprs {
+			// if GROUP BY index
+			if l, ok := ge.(*expression.Literal); ok && sql.IsNumber(l.Type()) {
+				if idx, ok := l.Value().(int64); ok && idx > 0 && idx <= agglen {
+					aggexpr := selectExprs[idx-1]
+					if alias, ok := aggexpr.(*expression.Alias); ok {
+						aggexpr = expression.NewUnresolvedColumn(alias.Name())
+					}
+					groupingExprs[i] = aggexpr
+				}
+			}
+		}
+
 		return plan.NewGroupBy(selectExprs, groupingExprs, child), nil
 	}
 

sql/parse/parse_test.go

@@ -193,6 +193,17 @@ var fixtures = map[string]sql.Node{
 		},
 		plan.NewUnresolvedTable("t1", ""),
 	),
+	`SELECT foo, bar FROM t1 GROUP BY 1, 2;`: plan.NewGroupBy(
+		[]sql.Expression{
+			expression.NewUnresolvedColumn("foo"),
+			expression.NewUnresolvedColumn("bar"),
+		},
+		[]sql.Expression{
+			expression.NewUnresolvedColumn("foo"),
+			expression.NewUnresolvedColumn("bar"),
+		},
+		plan.NewUnresolvedTable("t1", ""),
+	),
 	`SELECT COUNT(*) FROM t1;`: plan.NewGroupBy(
 		[]sql.Expression{
 			expression.NewUnresolvedFunction("count", true,