From 8f6b876af0af947c45ffacb6ce5ad75a45efe00a Mon Sep 17 00:00:00 2001
From: Travis <travis@pilosa.com>
Date: Fri, 10 Jul 2020 22:41:55 -0500
Subject: [PATCH] get ForeignIndex keys in GroupBy

---
 executor.go      | 25 +++++++++++++++++++------
 executor_test.go | 39 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/executor.go b/executor.go
index 42abe099c..3e50c3c95 100644
--- a/executor.go
+++ b/executor.go
@@ -4350,7 +4350,7 @@ func (e *executor) translateResults(ctx context.Context, index string, idx *Inde
 	}
 
 	for i := range results {
-		results[i], err = e.translateResult(index, idx, calls[i], results[i], idMap)
+		results[i], err = e.translateResult(ctx, index, idx, calls[i], results[i], idMap)
 		if err != nil {
 			return err
 		}
@@ -4374,7 +4374,7 @@ func (e *executor) collectResultIDs(index string, idx *Index, call *pql.Call, re
 	return nil
 }
 
-func (e *executor) translateResult(index string, idx *Index, call *pql.Call, result interface{}, idSet map[uint64]string) (interface{}, error) {
+func (e *executor) translateResult(ctx context.Context, index string, idx *Index, call *pql.Call, result interface{}, idSet map[uint64]string) (interface{}, error) {
 	switch result := result.(type) {
 	case *Row:
 		if idx.Keys() {
@@ -4481,10 +4481,23 @@ func (e *executor) translateResult(index string, idx *Index, call *pql.Call, res
 					return nil, ErrFieldNotFound
 				}
 				if field.Keys() {
-					// TODO: does this need to take field.ForeignIndex() into consideration?
-					key, err := field.TranslateStore().TranslateID(g.RowID)
-					if err != nil {
-						return nil, errors.Wrap(err, "translating row ID in Group")
+					var key string
+					var err error
+					if fi := field.ForeignIndex(); fi != "" && g.Value != nil {
+						val := uint64(*g.Value) // not worried about overflow here because it's a foreign key
+						keys, err := e.Cluster.translateIndexIDs(ctx, fi, []uint64{val})
+						if err != nil {
+							return nil, errors.Wrap(err, "translating foreign index in Group")
+						}
+						if len(keys) == 1 {
+							key = keys[0]
+							group[i].Value = nil // Remove value now that it has been translated.
+						}
+					} else {
+						key, err = field.TranslateStore().TranslateID(g.RowID)
+						if err != nil {
+							return nil, errors.Wrap(err, "translating row ID in Group")
+						}
 					}
 					group[i].RowKey = key
 				}
diff --git a/executor_test.go b/executor_test.go
index 4a39a7713..32fc01b27 100644
--- a/executor_test.go
+++ b/executor_test.go
@@ -4743,7 +4743,7 @@ func TestExecutor_Execute_Rows_Keys(t *testing.T) {
 			exp: []string{},
 		},
 		{
-			q: `Rows(f, like="__")`,
+			q:   `Rows(f, like="__")`,
 			exp: []string{"10", "11", "12", "13", "14", "15", "16", "17", "18"},
 		},
 	}
@@ -5189,6 +5189,43 @@ func TestExecutor_Execute_GroupBy(t *testing.T) {
 
 		})
 
+		// Foreign Index
+		c.CreateField(t, "fip", pilosa.IndexOptions{Keys: true}, "parent")
+		c.CreateField(t, "fic", pilosa.IndexOptions{}, "child",
+			pilosa.OptFieldTypeInt(0, math.MaxInt64),
+			pilosa.OptFieldForeignIndex("fip"),
+		)
+		// Set data on the parent so we have some index keys.
+		c.Query(t, "fip", `
+			Set("one", parent=1)
+			Set("two", parent=2)
+			Set("three", parent=3)
+			Set("four", parent=4)
+			Set("five", parent=5)
+		`)
+		// Set data on the child to align with the foreign index keys.
+		c.Query(t, "fic", `
+			Set(1, child="one")
+			Set(2, child="one")
+			Set(3, child="one")
+			Set(4, child="three")
+			Set(5, child="three")
+			Set(6, child="five")
+		`)
+
+		t.Run("test foreign index with keys", func(t *testing.T) {
+			// the execututor returns row IDs when the field has keys, so they should be included in the target.
+			// because the order is determined by the partitioned index key, they seem out of order.
+			expected := []pilosa.GroupCount{
+				{Group: []pilosa.FieldRow{{Field: "child", RowID: 0, RowKey: "one"}}, Count: 3},
+				{Group: []pilosa.FieldRow{{Field: "child", RowID: 1, RowKey: "five"}}, Count: 1},
+				{Group: []pilosa.FieldRow{{Field: "child", RowID: 2, RowKey: "three"}}, Count: 2},
+			}
+
+			results := c.Query(t, "fic", `GroupBy(Rows(child))`).Results[0].([]pilosa.GroupCount)
+			test.CheckGroupBy(t, expected, results)
+		})
+
 	}
 	for size := range []int{1, 3} {
 		t.Run(fmt.Sprintf("%d_nodes", size), func(t *testing.T) {