diff --git a/executor.go b/executor.go index ae52457c7..52b0a20b7 100644 --- a/executor.go +++ b/executor.go @@ -361,7 +361,7 @@ func (e *executor) readColumnAttrSets(index *Index, ids []uint64) ([]*ColumnAttr } // handlePreCalls traverses the call tree looking for calls that need -// precomputed values. Right now, that's just Distinct. +// precomputed values (e.g. Distinct, UnionRows, ConstRow...) func (e *executor) handlePreCalls(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) error { if c.Name == "Precomputed" { idx := c.Args["valueidx"].(int64) @@ -4235,9 +4235,27 @@ func (e *executor) executeCount(ctx context.Context, qcx *Qcx, index string, c * return 0, errors.New("Count() only accepts a single bitmap input") } + 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 { + if row, ok := irow.(*Row); !ok { + return 0, errors.Errorf("unexpected precomputed value type inside count: %+v", irow) + } else { + for _, seg := range row.segments { + count += seg.n + } + } + } + return count, nil + } + // Execute calls in bulk on each remote node and merge. mapFn := func(ctx context.Context, shard uint64) (_ interface{}, err error) { - row, err := e.executeBitmapCallShard(ctx, qcx, index, c.Children[0], shard) + row, err := e.executeBitmapCallShard(ctx, qcx, index, child, shard) if err != nil { return 0, err } diff --git a/executor_test.go b/executor_test.go index 11a774014..cec51220f 100644 --- a/executor_test.go +++ b/executor_test.go @@ -6742,3 +6742,96 @@ func TestMissingKeyRegression(t *testing.T) { }) } } + +func TestDistinctOnSetsKeyedIndex(t *testing.T) { + c := test.MustRunCluster(t, 3) + defer c.Close() + + // create and populate "likenums" similar to "likes", but no keys on the field + c.CreateField(t, "users", pilosa.IndexOptions{Keys: true, TrackExistence: true}, "likenums") + c.ImportIDKey(t, "users", "likenums", []test.KeyID{ + {ID: 1, Key: "userA"}, + {ID: 2, Key: "userB"}, + {ID: 3, Key: "userC"}, + {ID: 4, Key: "userD"}, + {ID: 5, Key: "userE"}, + {ID: 6, Key: "userF"}, + {ID: 7, Key: "userA"}, + {ID: 7, Key: "userB"}, + {ID: 7, Key: "userC"}, + {ID: 7, Key: "userD"}, + {ID: 7, Key: "userE"}, + {ID: 7, Key: "userF"}, + }) + + // create and populate "likes" field + c.CreateField(t, "users", pilosa.IndexOptions{Keys: true, TrackExistence: true}, "likes", pilosa.OptFieldKeys()) + c.ImportKeyKey(t, "users", "likes", [][2]string{ + {"molecula", "userA"}, + {"pilosa", "userB"}, + {"pangolin", "userC"}, + {"zebra", "userD"}, + {"toucan", "userE"}, + {"dog", "userF"}, + {"icecream", "userA"}, + {"icecream", "userB"}, + {"icecream", "userC"}, + {"icecream", "userD"}, + {"icecream", "userE"}, + {"icecream", "userF"}, + }) + + tests := []struct { + query string + verifier func(t *testing.T, resp pilosa.QueryResponse) + }{ + { + query: "Count(All())", + verifier: func(t *testing.T, resp pilosa.QueryResponse) { + if resp.Results[0].(uint64) != 6 { + t.Errorf("expected 6, got %+v", resp.Results[0]) + } + }, + }, + { + query: "Count(Distinct(field=likenums))", + verifier: func(t *testing.T, resp pilosa.QueryResponse) { + if resp.Results[0].(uint64) != 7 { + t.Errorf("wrong count: %+v", resp.Results[0]) + } + }, + }, + { + query: "Distinct(field=likenums)", + verifier: func(t *testing.T, resp pilosa.QueryResponse) { + if !reflect.DeepEqual(resp.Results[0].(pilosa.SignedRow).Pos.Columns(), []uint64{1, 2, 3, 4, 5, 6, 7}) { + t.Errorf("wrong values: %+v", resp.Results[0].(pilosa.SignedRow).Pos.Columns()) + } + }, + }, + { + query: "Count(Distinct(field=likes))", + verifier: func(t *testing.T, resp pilosa.QueryResponse) { + if resp.Results[0].(uint64) != 7 { + t.Errorf("wrong count: %+v", resp.Results[0]) + } + }, + }, + // { + // query: "Distinct(field=likes)", + // verifier: func(t *testing.T, resp pilosa.QueryResponse) { + // if !reflect.DeepEqual(resp.Results[0].(*pilosa.Row).Keys, []string{"molecula", "pilosa", "pangolin", "zebra", "toucan", "dog", "icecream"}) { + // t.Errorf("wrong values: %+v", resp.Results[0]) + // } + // }, + // }, + } + + for i, tst := range tests { + t.Run(fmt.Sprintf("%d-%s", i, tst.query), func(t *testing.T) { + resp := c.Query(t, "users", tst.query) + fmt.Println(resp.Results[0]) + tst.verifier(t, resp) + }) + } +} diff --git a/test/cluster.go b/test/cluster.go index 5dfb4edfb..0a5d2c451 100644 --- a/test/cluster.go +++ b/test/cluster.go @@ -108,6 +108,53 @@ func (c *Cluster) ImportBits(t testing.TB, index, field string, rowcols [][2]uin } } +// ImportKeyKey imports data into an index where both the index and +// the field are using string keys. +func (c *Cluster) ImportKeyKey(t testing.TB, index, field string, valAndRecKeys [][2]string) { + t.Helper() + importRequest := &pilosa.ImportRequest{ + Index: index, + Field: field, + RowKeys: make([]string, len(valAndRecKeys)), + ColumnKeys: make([]string, len(valAndRecKeys)), + } + for i, vk := range valAndRecKeys { + importRequest.RowKeys[i] = vk[0] + importRequest.ColumnKeys[i] = vk[1] + } + err := c.Nodes[0].API.Import(context.Background(), nil, importRequest) + if err != nil { + t.Fatalf("importing keykey data: %v", err) + } +} + +// KeyID represents a key and an ID for importing data into an index +// and field where one uses string keys and the other does not. +type KeyID struct { + Key string + ID uint64 +} + +// ImportIDKey imports data into an index where the index is using +// keys, but the field is not. +func (c *Cluster) ImportIDKey(t testing.TB, index, field string, pairs []KeyID) { + t.Helper() + importRequest := &pilosa.ImportRequest{ + Index: index, + Field: field, + RowIDs: make([]uint64, len(pairs)), + ColumnKeys: make([]string, len(pairs)), + } + for i, pair := range pairs { + importRequest.RowIDs[i] = pair.ID + importRequest.ColumnKeys[i] = pair.Key + } + err := c.Nodes[0].API.Import(context.Background(), nil, importRequest) + if err != nil { + t.Fatalf("importing IDKey data: %v", err) + } +} + // CreateField creates the index (if necessary) and field specified. func (c *Cluster) CreateField(t testing.TB, index string, iopts pilosa.IndexOptions, field string, fopts ...pilosa.FieldOption) *pilosa.Field { t.Helper()