From 649202b0816bcebd05d17f43ec435bd69e7a7ac9 Mon Sep 17 00:00:00 2001
From: Cody Soyland <cody@soyland.com>
Date: Mon, 21 Dec 2020 13:25:15 -0600
Subject: [PATCH 1/2] Fix cluster size setter: expose failures

---
 executor_test.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/executor_test.go b/executor_test.go
index c0532b4f9..6a202ca3a 100644
--- a/executor_test.go
+++ b/executor_test.go
@@ -5517,7 +5517,7 @@ func TestExecutor_Execute_DistinctFailure(t *testing.T) {
 
 func TestExecutor_Execute_GroupBy(t *testing.T) {
 	groupByTest := func(t *testing.T, clusterSize int) {
-		c := test.MustRunCluster(t, 1)
+		c := test.MustRunCluster(t, clusterSize)
 		defer c.Close()
 		c.CreateField(t, "i", pilosa.IndexOptions{}, "general")
 		c.CreateField(t, "i", pilosa.IndexOptions{}, "sub")
@@ -5924,7 +5924,7 @@ func TestExecutor_Execute_GroupBy(t *testing.T) {
 		})
 
 	}
-	for size := range []int{1, 3} {
+	for _, size := range []int{1, 3} {
 		t.Run(fmt.Sprintf("%d_nodes", size), func(t *testing.T) {
 			groupByTest(t, size)
 		})

From dde318ac8cb929a2c2fb7b12c431d99e8fd648b6 Mon Sep 17 00:00:00 2001
From: Nia Weiss <nia.weiss@molecula.com>
Date: Mon, 25 Jan 2021 10:15:45 -0500
Subject: [PATCH 2/2] Move globally computed GroupBy rows calls into
 EmbeddedData

This fixes a bug where a globally computed Rows call would be computed with a subset of the shards.
---
 executor.go | 17 +++++++++++++++++
 pql/ast.go  |  1 +
 row.go      |  4 ++++
 3 files changed, 22 insertions(+)

diff --git a/executor.go b/executor.go
index 0c02fd6c8..b1a0ae1eb 100644
--- a/executor.go
+++ b/executor.go
@@ -2797,6 +2797,12 @@ func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c
 		}
 
 		if hasLimit || hasCol { // we need to perform this query cluster-wide ahead of executeGroupByShard
+			if idx, ok := child.Args["valueidx"].(int64); ok {
+				// The rows query was already completed on the initiating node.
+				childRows[i] = opt.EmbeddedData[idx].Columns()
+				continue
+			}
+
 			childRows[i], err = e.executeRows(ctx, qcx, index, child, shards, opt)
 			if err != nil {
 				return nil, errors.Wrap(err, "getting rows for ")
@@ -2804,6 +2810,13 @@ func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c
 			if len(childRows[i]) == 0 { // there are no results because this field has no values.
 				return &GroupCounts{}, nil
 			}
+
+			// Stuff the result into opt.EmbeddedData so that it gets sent to other nodes in the map-reduce.
+			// This is flagged as "NoSplit" to ensure that the entire row gets sent out.
+			rowsRow := NewRow(childRows[i]...)
+			rowsRow.NoSplit = true
+			child.Args["valueidx"] = int64(len(opt.EmbeddedData))
+			opt.EmbeddedData = append(opt.EmbeddedData, rowsRow)
 		}
 	}
 
@@ -5527,6 +5540,10 @@ func makeEmbeddedDataForShards(allRows []*Row, shards []uint64) []*Row {
 		if row == nil || len(row.segments) == 0 {
 			continue
 		}
+		if row.NoSplit {
+			newRows[i] = row
+			continue
+		}
 		segments := row.segments
 		segmentIndex := 0
 		newRows[i] = &Row{
diff --git a/pql/ast.go b/pql/ast.go
index 40e2acaba..54ae2dee7 100644
--- a/pql/ast.go
+++ b/pql/ast.go
@@ -388,6 +388,7 @@ var callInfoByFunc = map[string]callInfo{
 			"from":     nil,
 			"to":       nil,
 			"like":     "",
+			"valueidx": int64(0),
 		},
 	},
 	"Shift": {allowUnknown: false,
diff --git a/row.go b/row.go
index 46a0ec478..3ff47cba5 100644
--- a/row.go
+++ b/row.go
@@ -42,6 +42,10 @@ type Row struct {
 	// query. Knowing the index and field, we can figure out how to
 	// interpret the row data.
 	Field string
+
+	// NoSplit indicates that this row may not be split.
+	// This is used for `Rows` calls in a GroupBy.
+	NoSplit bool
 }
 
 // NewRow returns a new instance of Row.