Merge pull request #1369 from niaow/count-global

CORE-29 Invoke pre-calls directly in count operations (fixes Count(Distinct()) on negative integers)

executor.go

@@ -555,11 +555,22 @@ func (e *executor) execute(ctx context.Context, qcx *Qcx, index string, q *pql.Q
 		// about the positive values, because only positive values
 		// are valid column IDs. So we don't actually eat top-level
 		// pre calls.
-		err := e.handlePreCallChildren(ctx, qcx, index, call, shards, opt)
-		if err != nil {
-			return nil, err
+		if call.Name == "Count" {
+			// Handle count specially, skipping the level directly underneath it.
+			for _, child := range call.Children {
+				err := e.handlePreCallChildren(ctx, qcx, index, child, shards, opt)
+				if err != nil {
+					return nil, err
+				}
+			}
+		} else {
+			err := e.handlePreCallChildren(ctx, qcx, index, call, shards, opt)
+			if err != nil {
+				return nil, err
+			}
 		}
 		var v interface{}
+		var err error
 		// Top-level calls don't need to precompute cross-index things,
 		// because we can just pick whatever index we want, but we
 		// still need to handle them. Since everything else was
@@ -4628,28 +4639,21 @@ func (e *executor) executeCount(ctx context.Context, qcx *Qcx, index string, c *
 
 	child := c.Children[0]
 
-	// If the child is precomputed, we'll bypass mapreduce, ignore
-	// shards, and just count the number of bits.
-	if child.Name == "Precomputed" {
-		count := uint64(0)
-		for _, irow := range child.Precomputed {
-			switch row := irow.(type) {
-			case *Row:
-				for _, seg := range row.segments {
-					count += seg.n
-				}
-			case SignedRow:
-				for _, seg := range row.Pos.segments {
-					count += seg.n
-				}
-				for _, seg := range row.Neg.segments {
-					count += seg.n
-				}
-			default:
-				return 0, errors.Errorf("unexpected precomputed value type inside count: %+v", row)
-			}
+	// If the child is distinct/similar, execute it directly here and count the result.
+	if child.Type == pql.PrecallGlobal {
+		result, err := e.executeCall(ctx, qcx, index, child, shards, opt)
+		if err != nil {
+			return 0, err
+		}
+
+		switch row := result.(type) {
+		case *Row:
+			return row.Count(), nil
+		case SignedRow:
+			return row.Pos.Count() + row.Neg.Count(), nil
+		default:
+			return 0, errors.Errorf("cannot count result of type %T from call %q", row, child.String())
 		}
-		return count, nil
 	}
 
 	// Execute calls in bulk on each remote node and merge.

executor_test.go

@@ -6837,6 +6837,8 @@ func TestMissingKeyRegression(t *testing.T) {
 func TestVariousQueries(t *testing.T) {
 	for _, clusterSize := range []int{1, 3, 4, 7} {
 		t.Run(fmt.Sprintf("%d-node", clusterSize), func(t *testing.T) {
+			t.Parallel()
+
 			variousQueries(t, clusterSize)
 		})
 	}
@@ -6916,8 +6918,7 @@ func variousQueries(t *testing.T, clusterSize int) {
 		{Val: 0, Key: "userE"},
 	})
 
-	// Create and populate "affinity" int field with negative, positive, zero and null values.
-
+	// Create and populate "net_worth" int field with positive values.
 	c.CreateField(t, "users", pilosa.IndexOptions{Keys: true, TrackExistence: true}, "net_worth", pilosa.OptFieldTypeInt(-100000000, 100000000))
 	c.ImportIntKey(t, "users", "net_worth", []test.IntKey{
 		{Val: 1, Key: "userA"},
@@ -7051,6 +7052,15 @@ toronto,2,11
 			},
 			csvVerifier: "-10\n-5\n0\n5\n10\n",
 		},
+		{
+			query: "Count(Distinct(field=affinity))",
+			qrVerifier: func(t *testing.T, resp pilosa.QueryResponse) {
+				if resp.Results[0].(uint64) != 5 {
+					t.Errorf("wrong number of values: %+v", resp.Results[0])
+				}
+			},
+			csvVerifier: "5\n",
+		},
 		{
 			query: "Distinct(Row(affinity>=0),field=affinity)",
 			qrVerifier: func(t *testing.T, resp pilosa.QueryResponse) {

pql/ast.go

@@ -274,14 +274,17 @@ type callStackElem struct {
 type CallType byte
 
 const (
-	// Normal calls can be executed per shard.
+	// PrecallNone calls can be executed per shard.
 	PrecallNone = CallType(iota)
-	// PreCallGlobal indicates a call which must be run globally *before*
+
+	// PrecallGlobal indicates a call which must be run globally *before*
 	// distributing the call to other shards. Example: A Distinct query,
 	// where every shard could potentially produce results for any shard,
 	// so you have to produce the results up front.
+	// These are processed directly when inside of a count operation.
 	PrecallGlobal
-	// PreCallPerNode indicates a call which needs to be run per-shard
+
+	// PrecallPerNode indicates a call which needs to be run per-shard
 	// in a way that lets it be done on each shard, but where it should
 	// be done prior to spawning per-shard goroutines. Example:
 	// A cross-index query, where each local shard may or may not need