From 586a13e9423a3d79ef03ff597d0bc4e910e75352 Mon Sep 17 00:00:00 2001
From: Travis <travis@pilosa.com>
Date: Fri, 20 Dec 2019 22:43:50 -0600
Subject: [PATCH] Allow All() to be called at the shard level

---
 executor.go      | 10 ++++++----
 executor_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/executor.go b/executor.go
index f9586f958..d49247adf 100644
--- a/executor.go
+++ b/executor.go
@@ -720,7 +720,7 @@ func (e *executor) executeAllCall(ctx context.Context, index string, c *pql.Call
 func (e *executor) executeAllCallMapReduce(ctx context.Context, index string, c *pql.Call, shard uint64, opt *execOptions) (*Row, error) {
 	// Execute calls in bulk on each remote node and merge.
 	mapFn := func(shard uint64) (interface{}, error) {
-		return e.executeAllShard(ctx, index, c, shard)
+		return e.executeAllCallShard(ctx, index, c, shard)
 	}
 
 	// Merge returned results at coordinating node.
@@ -1113,6 +1113,8 @@ func (e *executor) executeBitmapCallShard(ctx context.Context, index string, c *
 		return e.executeNotShard(ctx, index, c, shard)
 	case "Shift":
 		return e.executeShiftShard(ctx, index, c, shard)
+	case "All": // Allow a shard computation to use All() (note, limit/offset not applied)
+		return e.executeAllCallShard(ctx, index, c, shard)
 	case "Precomputed":
 		return e.executePrecomputedCallShard(ctx, index, c, shard)
 	default:
@@ -2524,9 +2526,9 @@ func (e *executor) executeNotShard(ctx context.Context, index string, c *pql.Cal
 	return existenceRow.Difference(row), nil
 }
 
-// executeAllShard executes an All() call for a local shard.
-func (e *executor) executeAllShard(ctx context.Context, index string, c *pql.Call, shard uint64) (*Row, error) {
-	span, _ := tracing.StartSpanFromContext(ctx, "Executor.executeAllShard")
+// executeAllCallShard executes an All() call for a local shard.
+func (e *executor) executeAllCallShard(ctx context.Context, index string, c *pql.Call, shard uint64) (*Row, error) {
+	span, _ := tracing.StartSpanFromContext(ctx, "Executor.executeAllCallShard")
 	defer span.Finish()
 
 	if len(c.Children) > 0 {
diff --git a/executor_test.go b/executor_test.go
index 3750a04ad..ec8865c8a 100644
--- a/executor_test.go
+++ b/executor_test.go
@@ -3039,6 +3039,32 @@ func TestExecutor_Execute_All(t *testing.T) {
 			}
 		}
 	})
+
+	// Ensure that a query which uses All() at the shard level can call it.
+	t.Run("AllShard", func(t *testing.T) {
+		c := test.MustRunCluster(t, 1)
+		defer c.Close()
+		hldr := test.Holder{Holder: c[0].Server.Holder()}
+		index := hldr.MustCreateIndexIfNotExists("i", pilosa.IndexOptions{TrackExistence: true})
+		_, err := index.CreateField("f", pilosa.OptFieldTypeDefault())
+		if err != nil {
+			t.Fatal(err)
+		}
+		if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `
+Set(3001, f=3)
+Set(5001, f=5)
+Set(5002, f=5)
+`}); err != nil {
+			t.Fatalf("querying remote: %v", err)
+		}
+
+		expCols := []uint64{5001, 5002}
+		if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: "Intersect(All(), Row(f=5))"}); err != nil {
+			t.Fatal(err)
+		} else if cols := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(cols, expCols) {
+			t.Fatalf("unexpected columns, got: %v, but expected: %v", cols, expCols)
+		}
+	})
 }
 
 // Ensure a row can be cleared.