From 91391e884dadda504f5275962b7d1c9b0d6630c6 Mon Sep 17 00:00:00 2001
From: David Taylor <tinystatemachine@gmail.com>
Date: Mon, 4 Jan 2016 21:23:06 +0000
Subject: [PATCH] support 'GROUP BY' column number

Fixes #3549.
---
 sql/group.go           | 26 +++++++++++++++++++++-----
 sql/testdata/aggregate | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/sql/group.go b/sql/group.go
index 0936e2f77a71..4464e318d386 100644
--- a/sql/group.go
+++ b/sql/group.go
@@ -54,7 +54,15 @@ func (p *planner) groupBy(n *parser.Select, s *scanNode) (*groupNode, error) {
 		if err != nil {
 			return nil, err
 		}
-		n.GroupBy[i] = norm
+		// If a col index is specified, replace it with that expression first.
+		if col, err := s.colIndex(norm); err != nil {
+			return nil, err
+		} else if col >= 0 {
+			n.GroupBy[i] = s.render[col]
+		} else {
+			n.GroupBy[i] = norm
+		}
+
 	}
 
 	if err := checkAggregateExprs(n.GroupBy, s.render); err != nil {
@@ -158,8 +166,8 @@ type groupNode struct {
 
 	addNullBucketIfEmpty bool
 
-	values     valuesNode
-	intialized bool
+	values      valuesNode
+	initialized bool
 
 	// During rendering, aggregateFuncs compute their result for group.currentBucket.
 	currentBucket string
@@ -182,7 +190,7 @@ func (n *groupNode) Values() parser.DTuple {
 }
 
 func (n *groupNode) Next() bool {
-	if !n.intialized && n.err == nil {
+	if !n.initialized && n.err == nil {
 		n.computeAggregates()
 	}
 	if n.err != nil {
@@ -192,7 +200,6 @@ func (n *groupNode) Next() bool {
 }
 
 func (n *groupNode) computeAggregates() {
-	n.intialized = true
 	var scratch []byte
 
 	// Loop over the rows passing the values into the corresponding aggregation
@@ -229,6 +236,8 @@ func (n *groupNode) computeAggregates() {
 		n.buckets[""] = struct{}{}
 	}
 
+	n.initialized = true
+
 	// Render the results.
 	n.values.rows = make([]parser.DTuple, 0, len(n.buckets))
 	for k := range n.buckets {
@@ -562,6 +571,13 @@ func (a *aggregateFunc) TypeCheck(args parser.MapArgs) (parser.Datum, error) {
 }
 
 func (a *aggregateFunc) Eval(ctx parser.EvalContext) (parser.Datum, error) {
+	// During init of the group buckets, grouped expressions (i.e. wrapped
+	// qvalues) are Eval()'ed to determine the bucket for a row, so pass these
+	// calls through to the underlying `arg` expr Eval until init is done.
+	if !a.group.initialized {
+		return a.arg.Eval(ctx)
+	}
+
 	found, ok := a.buckets[a.group.currentBucket]
 	if !ok {
 		found = a.create()
diff --git a/sql/testdata/aggregate b/sql/testdata/aggregate
index c1710d6a5d82..9c6b376b73b2 100644
--- a/sql/testdata/aggregate
+++ b/sql/testdata/aggregate
@@ -28,10 +28,35 @@ SELECT COUNT(*), k FROM kv GROUP BY k
 1 7
 1 8
 
+query II rowsort
+SELECT COUNT(*), k FROM kv GROUP BY 2
+----
+1 1
+1 3
+1 5
+1 6
+1 7
+1 8
+
+query error invalid column index: 5 not in range \[1, 2\]
+SELECT COUNT(*), k FROM kv GROUP BY 5
+
+query error invalid column index: 0 not in range \[1, 2\]
+SELECT COUNT(*), k FROM kv GROUP BY 0
+
+query error invalid column index: -4 not in range \[1, 2\]
+SELECT COUNT(*), k FROM kv GROUP BY -4
+
 query I rowsort
-SELECT 1 GROUP BY 'a';
+SELECT k FROM kv GROUP BY k
 ----
 1
+3
+5
+6
+7
+8
+
 
 query I rowsort
 SELECT MIN(1) FROM kv;
@@ -56,6 +81,13 @@ SELECT COUNT(*), UPPER(s) FROM kv GROUP BY UPPER(s)
 2 B
 3 A
 
+query IT rowsort
+SELECT COUNT(*), UPPER(s) FROM kv GROUP BY 2
+----
+1 NULL
+2 B
+3 A
+
 query IT rowsort
 SELECT COUNT(*), UPPER(s) FROM kv GROUP BY s
 ----