From 5e6fd661ea1525f192ba8789c52fe9b158fef784 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           | 17 ++++++++++++++++-
 sql/testdata/aggregate | 17 +++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/sql/group.go b/sql/group.go
index 0936e2f77a71..52ebe0617fb4 100644
--- a/sql/group.go
+++ b/sql/group.go
@@ -46,6 +46,13 @@ func (p *planner) groupBy(n *parser.Select, s *scanNode) (*groupNode, error) {
 	// that determination is made during validation, which will require matching
 	// expressions.
 	for i := range n.GroupBy {
+		// If a col offset is specified, replace it with that expression first.
+		if col, err := s.colOffset(n.GroupBy[i], false); err != nil {
+			return nil, err
+		} else if col >= 0 {
+			n.GroupBy[i] = s.render[col]
+		}
+
 		norm, err := s.resolveQNames(n.GroupBy[i])
 		if err != nil {
 			return nil, err
@@ -192,7 +199,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 +235,8 @@ func (n *groupNode) computeAggregates() {
 		n.buckets[""] = struct{}{}
 	}
 
+	n.intialized = true
+
 	// Render the results.
 	n.values.rows = make([]parser.DTuple, 0, len(n.buckets))
 	for k := range n.buckets {
@@ -562,6 +570,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 (ie bucket filling), grouped expressions (ie wrapped
+	// qvalues) are Eval()'ed to determine the bucket for a row, so pass these
+	// calls through to the underlying `arg` expr eval until the group is init'ed.
+	if !a.group.intialized {
+		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..939b46f1f6c0 100644
--- a/sql/testdata/aggregate
+++ b/sql/testdata/aggregate
@@ -28,6 +28,16 @@ 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 I rowsort
 SELECT 1 GROUP BY 'a';
 ----
@@ -56,6 +66,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
 ----