diff --git a/cache.go b/cache.go index b102642be..ccd5d3fb4 100644 --- a/cache.go +++ b/cache.go @@ -255,6 +255,15 @@ func (p Pairs) Add(other []Pair) []Pair { return a } +// Keys returns a slice of all keys in p. +func (p Pairs) Keys() []uint64 { + a := make([]uint64, len(p)) + for i := range p { + a[i] = p[i].Key + } + return a +} + func encodePairs(a Pairs) []*internal.Pair { other := make([]*internal.Pair, len(a)) for i := range a { diff --git a/executor.go b/executor.go index d208c813a..39f305dbd 100644 --- a/executor.go +++ b/executor.go @@ -153,7 +153,31 @@ func (e *Executor) executeBitmapCallSlice(db string, c pql.BitmapCall, slice uin } // executeTopN executes a TopN() call. +// This first performs the TopN() to determine the top results and then +// requeries to retrieve the full counts for each of the top results. func (e *Executor) executeTopN(db string, c *pql.TopN, slices []uint64, opt *ExecOptions) ([]Pair, error) { + // Execute original query. + pairs, err := e.executeTopNSlices(db, c, slices, opt) + if err != nil { + return nil, err + } + + // If this call is against specific ids, or we didn't get results, + // or we are part of a larger distributed query then don't refetch. + if len(pairs) == 0 || len(c.BitmapIDs) > 0 || opt.Remote { + return pairs, nil + } + + // Only the original caller should refetch the full counts. + other := *c + other.N = 0 + other.BitmapIDs = Pairs(pairs).Keys() + sort.Sort(uint64Slice(other.BitmapIDs)) + + return e.executeTopNSlices(db, &other, slices, opt) +} + +func (e *Executor) executeTopNSlices(db string, c *pql.TopN, slices []uint64, opt *ExecOptions) ([]Pair, error) { var results []Pair for node, nodeSlices := range e.slicesByNode(slices) { // Execute locally if the hostname matches. @@ -180,7 +204,7 @@ func (e *Executor) executeTopN(db string, c *pql.TopN, slices []uint64, opt *Exe sort.Sort(Pairs(results)) // Only keep the top n after sorting. - if len(results) > c.N { + if c.N > 0 && len(results) > c.N { results = results[0:c.N] } @@ -210,7 +234,13 @@ func (e *Executor) executeTopNSlice(db string, c *pql.TopN, slice uint64) ([]Pai return nil, nil } - return f.TopN(c.N, src, c.Field, c.Filters) + return f.Top(TopOptions{ + N: c.N, + Src: src, + BitmapIDs: c.BitmapIDs, + FilterField: c.Field, + FilterValues: c.Filters, + }) } // executeDifferenceSlice executes a difference() call for a local slice. @@ -404,6 +434,7 @@ func (e *Executor) exec(node *Node, db string, q *pql.Query, slices []uint64, op Query: proto.String(q.String()), Slices: slices, Quantum: proto.Uint32(uint32(opt.Quantum)), + Remote: proto.Bool(true), } if opt.Timestamp != nil { pbreq.Timestamp = proto.Int64(opt.Timestamp.UnixNano()) @@ -489,6 +520,7 @@ func (e *Executor) slicesByNode(slices []uint64) map[*Node][]uint64 { type ExecOptions struct { Timestamp *time.Time Quantum TimeQuantum + Remote bool } // decodeError returns an error representation of s if s is non-blank. diff --git a/executor_test.go b/executor_test.go index 2c1be9358..49b70e057 100644 --- a/executor_test.go +++ b/executor_test.go @@ -402,15 +402,30 @@ func TestExecutor_Execute_Remote_TopN(t *testing.T) { c.Nodes[1].Host = s.Host() // Mock secondary server's executor to verify arguments and return a bitmap. + var remoteExecN int s.Handler.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) (interface{}, error) { if db != `d` { t.Fatalf("unexpected db: %s", db) - } else if query.String() != `TopN(frame=f, n=3)` { - t.Fatalf("unexpected query: %s", query.String()) } else if !reflect.DeepEqual(slices, []uint64{0, 2, 4, 6}) { t.Fatalf("unexpected slices: %+v", slices) } + // Query should be executed twice. Once to get the top bitmaps for the + // slices and a second time to get the counts for a set of bitmaps. + switch remoteExecN { + case 0: + if query.String() != `TopN(frame=f, n=3)` { + t.Fatalf("unexpected query(0): %s", query.String()) + } + case 1: + if query.String() != `TopN(frame=f, ids=[0,10,30])` { + t.Fatalf("unexpected query(1): %s", query.String()) + } + default: + t.Fatalf("too many remote exec calls") + } + remoteExecN++ + // Return pair counts. return []pilosa.Pair{ {Key: 0, Count: 5}, diff --git a/fragment.go b/fragment.go index c3aeabacc..66f816d50 100644 --- a/fragment.go +++ b/fragment.go @@ -397,27 +397,24 @@ func (f *Fragment) pos(bitmapID, profileID uint64) (uint64, error) { return (bitmapID * SliceWidth) + (profileID % SliceWidth), nil } -// TopN returns the top n bitmaps from the fragment. -// If src is specified then only bitmaps which intersect src are returned. -// If fieldValues exist then the bitmap attribute specified by field is matched. -func (f *Fragment) TopN(n int, src *Bitmap, field string, fieldValues []interface{}) ([]Pair, error) { - // Resort cache, if needed, and retrieve the top bitmaps. - f.mu.Lock() - f.cache.Invalidate() - pairs := f.cache.Top() - f.mu.Unlock() +// Top returns the top bitmaps from the fragment. +// If opt.Src is specified then only bitmaps which intersect src are returned. +// If opt.FilterValues exist then the bitmap attribute specified by field is matched. +func (f *Fragment) Top(opt TopOptions) ([]Pair, error) { + // Retrieve pairs. If no bitmap ids specified then return from cache. + pairs := f.topBitmapPairs(opt.BitmapIDs) // Create a fast lookup of filter values. var filters map[interface{}]struct{} - if len(fieldValues) > 0 { + if opt.FilterField != "" && len(opt.FilterValues) > 0 { filters = make(map[interface{}]struct{}) - for _, v := range fieldValues { + for _, v := range opt.FilterValues { filters[v] = struct{}{} } } // Iterate over rankings and add to results until we have enough. - results := make([]Pair, 0, n) + results := make([]Pair, 0, opt.N) for _, pair := range pairs { bitmapID, bm := pair.ID, pair.Bitmap @@ -433,7 +430,7 @@ func (f *Fragment) TopN(n int, src *Bitmap, field string, fieldValues []interfac return nil, err } else if attr == nil { continue - } else if attrValue := attr[field]; attrValue == nil { + } else if attrValue := attr[opt.FilterField]; attrValue == nil { continue } else if _, ok := filters[attrValue]; !ok { continue @@ -441,11 +438,11 @@ func (f *Fragment) TopN(n int, src *Bitmap, field string, fieldValues []interfac } // The initial n pairs should simply be added to the results. - if len(results) < n { + if opt.N == 0 || len(results) < opt.N { // Calculate count and append. count := bm.Count() - if src != nil { - count = src.IntersectionCount(bm) + if opt.Src != nil { + count = opt.Src.IntersectionCount(bm) } if count == 0 { continue @@ -455,8 +452,8 @@ func (f *Fragment) TopN(n int, src *Bitmap, field string, fieldValues []interfac // If we reach the requested number of pairs and we are not computing // intersections then simply exit. If we are intersecting then sort // and then only keep pairs that are higher than the lowest count. - if len(results) == n { - if src == nil { + if opt.N > 0 && len(results) == opt.N { + if opt.Src == nil { break } sort.Sort(Pairs(results)) @@ -479,7 +476,7 @@ func (f *Fragment) TopN(n int, src *Bitmap, field string, fieldValues []interfac // Calculate the intersecting bit count and skip if it's below our // last bitmap in our current result set. - count := src.IntersectionCount(bm) + count := opt.Src.IntersectionCount(bm) if count < threshold { continue } @@ -497,6 +494,42 @@ func (f *Fragment) TopN(n int, src *Bitmap, field string, fieldValues []interfac return results, nil } +func (f *Fragment) topBitmapPairs(bitmapIDs []uint64) []BitmapPair { + // If no specific bitmaps are requested, retrieve top bitmaps. + if len(bitmapIDs) == 0 { + f.mu.Lock() + defer f.mu.Unlock() + f.cache.Invalidate() + return f.cache.Top() + } + + // Otherwise retrieve specific bitmaps. + pairs := make([]BitmapPair, len(bitmapIDs)) + for i, bitmapID := range bitmapIDs { + pairs[i] = BitmapPair{ + ID: bitmapID, + Bitmap: f.Bitmap(bitmapID), + } + } + return pairs +} + +// TopOptions represents options passed into the Top() function. +type TopOptions struct { + // Number of bitmaps to return. + N int + + // Bitmap to intersect with. + Src *Bitmap + + // Specific bitmaps to filter against. + BitmapIDs []uint64 + + // Filter field name & values. + FilterField string + FilterValues []interface{} +} + func (f *Fragment) Range(bitmapID uint64, start, end time.Time) *Bitmap { f.mu.Lock() defer f.mu.Unlock() diff --git a/fragment_test.go b/fragment_test.go index 8dc13ef05..7e1465478 100644 --- a/fragment_test.go +++ b/fragment_test.go @@ -102,7 +102,7 @@ func TestFragment_Snapshot(t *testing.T) { } // Ensure a fragment can return the top n results. -func TestFragment_TopN(t *testing.T) { +func TestFragment_Top(t *testing.T) { f := MustOpenFragment("d", "f", 0) defer f.Close() @@ -112,7 +112,7 @@ func TestFragment_TopN(t *testing.T) { f.MustSetBits(102, 1, 2) // Retrieve top bitmaps. - if pairs, err := f.TopN(2, nil, "", nil); err != nil { + if pairs, err := f.Top(pilosa.TopOptions{N: 2}); err != nil { t.Fatal(err) } else if len(pairs) != 2 { t.Fatalf("unexpected count: %d", len(pairs)) @@ -124,7 +124,7 @@ func TestFragment_TopN(t *testing.T) { } // Ensure a fragment can filter bitmaps when retrieving the top n bitmaps. -func TestFragment_TopN_Filter(t *testing.T) { +func TestFragment_Top_Filter(t *testing.T) { f := MustOpenFragment("d", "f", 0) defer f.Close() @@ -138,7 +138,11 @@ func TestFragment_TopN_Filter(t *testing.T) { f.BitmapAttrStore.SetAttrs(102, map[string]interface{}{"x": uint64(20)}) // Retrieve top bitmaps. - if pairs, err := f.TopN(2, nil, "x", []interface{}{uint64(10), uint64(15), uint64(20)}); err != nil { + if pairs, err := f.Top(pilosa.TopOptions{ + N: 2, + FilterField: "x", + FilterValues: []interface{}{uint64(10), uint64(15), uint64(20)}, + }); err != nil { t.Fatal(err) } else if len(pairs) != 2 { t.Fatalf("unexpected count: %d", len(pairs)) @@ -164,7 +168,7 @@ func TestFragment_TopN_Intersect(t *testing.T) { f.MustSetBits(103, 1000, 1001, 1002) // no intersection // Retrieve top bitmaps. - if pairs, err := f.TopN(3, src, "", nil); err != nil { + if pairs, err := f.Top(pilosa.TopOptions{N: 3, Src: src}); err != nil { t.Fatal(err) } else if !reflect.DeepEqual(pairs, []pilosa.Pair{ {Key: 101, Count: 3}, @@ -198,7 +202,7 @@ func TestFragment_TopN_Intersect_Large(t *testing.T) { } // Retrieve top bitmaps. - if pairs, err := f.TopN(10, src, "", nil); err != nil { + if pairs, err := f.Top(pilosa.TopOptions{N: 10, Src: src}); err != nil { t.Fatal(err) } else if !reflect.DeepEqual(pairs, []pilosa.Pair{ {Key: 999, Count: 19}, @@ -216,6 +220,27 @@ func TestFragment_TopN_Intersect_Large(t *testing.T) { } } +// Ensure a fragment can return top bitmaps when specified by ID. +func TestFragment_TopN_BitmapIDs(t *testing.T) { + f := MustOpenFragment("d", "f", 0) + defer f.Close() + + // Set bits on various bitmaps. + f.MustSetBits(100, 1, 2, 3) + f.MustSetBits(101, 4, 5, 6, 7) + f.MustSetBits(102, 8, 9, 10, 11, 12) + + // Retrieve top bitmaps. + if pairs, err := f.Top(pilosa.TopOptions{BitmapIDs: []uint64{100, 101, 200}}); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(pairs, []pilosa.Pair{ + {Key: 101, Count: 4}, + {Key: 100, Count: 3}, + }) { + t.Fatalf("unexpected pairs: %s", spew.Sdump(pairs)) + } +} + // Ensure a fragment's cache can be persisted between restarts. func TestFragment_LRUCache_Persistence(t *testing.T) { f := MustOpenFragment("d", "f", 0) diff --git a/handler.go b/handler.go index bb59d40d1..6846766e6 100644 --- a/handler.go +++ b/handler.go @@ -115,6 +115,7 @@ func (h *Handler) handlePostQuery(w http.ResponseWriter, r *http.Request) { opt := &ExecOptions{ Timestamp: req.Timestamp, Quantum: req.Quantum, + Remote: req.Remote, } // Parse query string. @@ -475,6 +476,10 @@ type QueryRequest struct { // Time granularity to use with the timestamp. Quantum TimeQuantum + + // If true, indicates that query is part of a larger distributed query. + // If false, this request is on the originating node. + Remote bool } func decodeQueryRequest(pb *internal.QueryRequest) *QueryRequest { @@ -484,6 +489,7 @@ func decodeQueryRequest(pb *internal.QueryRequest) *QueryRequest { Slices: pb.GetSlices(), Profiles: pb.GetProfiles(), Quantum: TimeQuantum(pb.GetQuantum()), + Remote: pb.GetRemote(), } if pb.Timestamp != nil { diff --git a/internal/internal.pb.go b/internal/internal.pb.go index b28f58f3d..a1ab93a9c 100644 --- a/internal/internal.pb.go +++ b/internal/internal.pb.go @@ -26,17 +26,15 @@ It has these top-level messages: package internal import proto "github.com/gogo/protobuf/proto" -import fmt "fmt" import math "math" // Reference imports to suppress errors if they are not otherwise used. var _ = proto.Marshal -var _ = fmt.Errorf var _ = math.Inf type Bitmap struct { - Chunks []*Chunk `protobuf:"bytes,1,rep,name=Chunks" json:"Chunks,omitempty"` - Attrs []*Attr `protobuf:"bytes,2,rep,name=Attrs" json:"Attrs,omitempty"` + Chunks []*Chunk `protobuf:"bytes,1,rep" json:"Chunks,omitempty"` + Attrs []*Attr `protobuf:"bytes,2,rep" json:"Attrs,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -59,8 +57,8 @@ func (m *Bitmap) GetAttrs() []*Attr { } type Chunk struct { - Key *uint64 `protobuf:"varint,1,req,name=Key" json:"Key,omitempty"` - Value []uint64 `protobuf:"varint,2,rep,name=Value" json:"Value,omitempty"` + Key *uint64 `protobuf:"varint,1,req" json:"Key,omitempty"` + Value []uint64 `protobuf:"varint,2,rep" json:"Value,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -83,8 +81,8 @@ func (m *Chunk) GetValue() []uint64 { } type Pair struct { - Key *uint64 `protobuf:"varint,1,req,name=Key" json:"Key,omitempty"` - Count *uint64 `protobuf:"varint,2,req,name=Count" json:"Count,omitempty"` + Key *uint64 `protobuf:"varint,1,req" json:"Key,omitempty"` + Count *uint64 `protobuf:"varint,2,req" json:"Count,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -107,8 +105,8 @@ func (m *Pair) GetCount() uint64 { } type Bit struct { - BitmapID *uint64 `protobuf:"varint,1,req,name=BitmapID" json:"BitmapID,omitempty"` - ProfileID *uint64 `protobuf:"varint,2,req,name=ProfileID" json:"ProfileID,omitempty"` + BitmapID *uint64 `protobuf:"varint,1,req" json:"BitmapID,omitempty"` + ProfileID *uint64 `protobuf:"varint,2,req" json:"ProfileID,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -131,8 +129,8 @@ func (m *Bit) GetProfileID() uint64 { } type Profile struct { - ID *uint64 `protobuf:"varint,1,req,name=ID" json:"ID,omitempty"` - Attrs []*Attr `protobuf:"bytes,2,rep,name=Attrs" json:"Attrs,omitempty"` + ID *uint64 `protobuf:"varint,1,req" json:"ID,omitempty"` + Attrs []*Attr `protobuf:"bytes,2,rep" json:"Attrs,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -155,10 +153,10 @@ func (m *Profile) GetAttrs() []*Attr { } type Attr struct { - Key *string `protobuf:"bytes,1,req,name=Key" json:"Key,omitempty"` - StringValue *string `protobuf:"bytes,2,opt,name=StringValue" json:"StringValue,omitempty"` - UintValue *uint64 `protobuf:"varint,3,opt,name=UintValue" json:"UintValue,omitempty"` - BoolValue *bool `protobuf:"varint,4,opt,name=BoolValue" json:"BoolValue,omitempty"` + Key *string `protobuf:"bytes,1,req" json:"Key,omitempty"` + StringValue *string `protobuf:"bytes,2,opt" json:"StringValue,omitempty"` + UintValue *uint64 `protobuf:"varint,3,opt" json:"UintValue,omitempty"` + BoolValue *bool `protobuf:"varint,4,opt" json:"BoolValue,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -195,7 +193,7 @@ func (m *Attr) GetBoolValue() bool { } type AttrMap struct { - Attrs []*Attr `protobuf:"bytes,1,rep,name=Attrs" json:"Attrs,omitempty"` + Attrs []*Attr `protobuf:"bytes,1,rep" json:"Attrs,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -211,12 +209,13 @@ func (m *AttrMap) GetAttrs() []*Attr { } type QueryRequest struct { - DB *string `protobuf:"bytes,1,req,name=DB" json:"DB,omitempty"` - Query *string `protobuf:"bytes,2,req,name=Query" json:"Query,omitempty"` - Slices []uint64 `protobuf:"varint,3,rep,name=Slices" json:"Slices,omitempty"` - Profiles *bool `protobuf:"varint,4,opt,name=Profiles" json:"Profiles,omitempty"` - Timestamp *int64 `protobuf:"varint,5,opt,name=Timestamp" json:"Timestamp,omitempty"` - Quantum *uint32 `protobuf:"varint,6,opt,name=Quantum" json:"Quantum,omitempty"` + DB *string `protobuf:"bytes,1,req" json:"DB,omitempty"` + Query *string `protobuf:"bytes,2,req" json:"Query,omitempty"` + Slices []uint64 `protobuf:"varint,3,rep" json:"Slices,omitempty"` + Profiles *bool `protobuf:"varint,4,opt" json:"Profiles,omitempty"` + Timestamp *int64 `protobuf:"varint,5,opt" json:"Timestamp,omitempty"` + Quantum *uint32 `protobuf:"varint,6,opt" json:"Quantum,omitempty"` + Remote *bool `protobuf:"varint,7,opt" json:"Remote,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -266,13 +265,20 @@ func (m *QueryRequest) GetQuantum() uint32 { return 0 } +func (m *QueryRequest) GetRemote() bool { + if m != nil && m.Remote != nil { + return *m.Remote + } + return false +} + type QueryResponse struct { - Err *string `protobuf:"bytes,1,opt,name=Err" json:"Err,omitempty"` - Bitmap *Bitmap `protobuf:"bytes,2,opt,name=Bitmap" json:"Bitmap,omitempty"` - N *uint64 `protobuf:"varint,3,opt,name=N" json:"N,omitempty"` - Pairs []*Pair `protobuf:"bytes,4,rep,name=Pairs" json:"Pairs,omitempty"` - Profiles []*Profile `protobuf:"bytes,5,rep,name=Profiles" json:"Profiles,omitempty"` - Changed *bool `protobuf:"varint,6,opt,name=Changed" json:"Changed,omitempty"` + Err *string `protobuf:"bytes,1,opt" json:"Err,omitempty"` + Bitmap *Bitmap `protobuf:"bytes,2,opt" json:"Bitmap,omitempty"` + N *uint64 `protobuf:"varint,3,opt" json:"N,omitempty"` + Pairs []*Pair `protobuf:"bytes,4,rep" json:"Pairs,omitempty"` + Profiles []*Profile `protobuf:"bytes,5,rep" json:"Profiles,omitempty"` + Changed *bool `protobuf:"varint,6,opt" json:"Changed,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -323,11 +329,11 @@ func (m *QueryResponse) GetChanged() bool { } type ImportRequest struct { - DB *string `protobuf:"bytes,1,req,name=DB" json:"DB,omitempty"` - Frame *string `protobuf:"bytes,2,req,name=Frame" json:"Frame,omitempty"` - Slice *uint64 `protobuf:"varint,3,req,name=Slice" json:"Slice,omitempty"` - BitmapIDs []uint64 `protobuf:"varint,4,rep,name=BitmapIDs" json:"BitmapIDs,omitempty"` - ProfileIDs []uint64 `protobuf:"varint,5,rep,name=ProfileIDs" json:"ProfileIDs,omitempty"` + DB *string `protobuf:"bytes,1,req" json:"DB,omitempty"` + Frame *string `protobuf:"bytes,2,req" json:"Frame,omitempty"` + Slice *uint64 `protobuf:"varint,3,req" json:"Slice,omitempty"` + BitmapIDs []uint64 `protobuf:"varint,4,rep" json:"BitmapIDs,omitempty"` + ProfileIDs []uint64 `protobuf:"varint,5,rep" json:"ProfileIDs,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -371,7 +377,7 @@ func (m *ImportRequest) GetProfileIDs() []uint64 { } type ImportResponse struct { - Err *string `protobuf:"bytes,1,opt,name=Err" json:"Err,omitempty"` + Err *string `protobuf:"bytes,1,opt" json:"Err,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -387,7 +393,7 @@ func (m *ImportResponse) GetErr() string { } type Cache struct { - BitmapIDs []uint64 `protobuf:"varint,1,rep,name=BitmapIDs" json:"BitmapIDs,omitempty"` + BitmapIDs []uint64 `protobuf:"varint,1,rep" json:"BitmapIDs,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -403,7 +409,7 @@ func (m *Cache) GetBitmapIDs() []uint64 { } type SliceMaxResponse struct { - SliceMax *uint64 `protobuf:"varint,1,req,name=SliceMax" json:"SliceMax,omitempty"` + SliceMax *uint64 `protobuf:"varint,1,req" json:"SliceMax,omitempty"` XXX_unrecognized []byte `json:"-"` } @@ -419,17 +425,4 @@ func (m *SliceMaxResponse) GetSliceMax() uint64 { } func init() { - proto.RegisterType((*Bitmap)(nil), "internal.Bitmap") - proto.RegisterType((*Chunk)(nil), "internal.Chunk") - proto.RegisterType((*Pair)(nil), "internal.Pair") - proto.RegisterType((*Bit)(nil), "internal.Bit") - proto.RegisterType((*Profile)(nil), "internal.Profile") - proto.RegisterType((*Attr)(nil), "internal.Attr") - proto.RegisterType((*AttrMap)(nil), "internal.AttrMap") - proto.RegisterType((*QueryRequest)(nil), "internal.QueryRequest") - proto.RegisterType((*QueryResponse)(nil), "internal.QueryResponse") - proto.RegisterType((*ImportRequest)(nil), "internal.ImportRequest") - proto.RegisterType((*ImportResponse)(nil), "internal.ImportResponse") - proto.RegisterType((*Cache)(nil), "internal.Cache") - proto.RegisterType((*SliceMaxResponse)(nil), "internal.SliceMaxResponse") } diff --git a/internal/internal.proto b/internal/internal.proto index aaa34a6a2..de85efa30 100644 --- a/internal/internal.proto +++ b/internal/internal.proto @@ -43,6 +43,7 @@ message QueryRequest { optional bool Profiles = 4; optional int64 Timestamp = 5; optional uint32 Quantum = 6; + optional bool Remote = 7; } message QueryResponse { diff --git a/pql/ast.go b/pql/ast.go index 07f26e12d..aa27f06d5 100644 --- a/pql/ast.go +++ b/pql/ast.go @@ -3,6 +3,7 @@ package pql import ( "fmt" "sort" + "strconv" "strings" "time" ) @@ -295,12 +296,15 @@ func (c *SetProfileAttrs) String() string { type TopN struct { Frame string + // Maximum number of results to return. + N int + // Bitmap to use for intersection while computing top results. // Original bitmap counts are used if no Src is provided. Src BitmapCall - // Maximum number of results to return. - N int + // Specific bitmaps to retrieve. + BitmapIDs []uint64 // Field name and values to filter on. Field string @@ -319,6 +323,13 @@ func (c *TopN) String() string { if c.N > 0 { args = append(args, fmt.Sprintf("n=%d", c.N)) } + if len(c.BitmapIDs) > 0 { + strs := make([]string, len(c.BitmapIDs)) + for i := range c.BitmapIDs { + strs[i] = strconv.FormatUint(c.BitmapIDs[i], 10) + } + args = append(args, fmt.Sprintf("ids=[%s]", strings.Join(strs, ","))) + } if c.Field != "" { args = append(args, fmt.Sprintf("field=%q", c.Field)) } diff --git a/pql/parser.go b/pql/parser.go index 4761aeb88..5e44c07f4 100644 --- a/pql/parser.go +++ b/pql/parser.go @@ -473,9 +473,13 @@ func (p *Parser) parseTopNCall() (*TopN, error) { c.Src = v continue } + + // Assign filter values if there's only a value and no named key. if v, ok := arg.value.([]interface{}); ok { - c.Filters = v - continue + if _, ok := arg.key.(string); !ok { + c.Filters = v + continue + } } switch arg.key { @@ -491,6 +495,10 @@ func (p *Parser) parseTopNCall() (*TopN, error) { if err := decodeString(arg.value, &c.Field); err != nil { return nil, parseErrorf(pos, "n: %s", err) } + case "ids": + if err := decodeUint64Slice(arg.value, &c.BitmapIDs); err != nil { + return nil, parseErrorf(pos, "n: %s", err) + } default: return nil, parseErrorf(pos, "invalid TopN() arg: %v", arg.key) } @@ -745,6 +753,26 @@ func decodeUint64(v interface{}, target *uint64) error { return fmt.Errorf("invalid int value: %v", v) } +// decodeUint64Slice type converts v to target. +func decodeUint64Slice(v interface{}, target *[]uint64) error { + input, ok := v.([]interface{}) + if !ok { + return fmt.Errorf("invalid array value: %v", v) + } + + a := make([]uint64, len(input)) + for i := range input { + elem, ok := input[i].(uint64) + if !ok { + return fmt.Errorf("invalid int element: %v", input[i]) + } + a[i] = elem + } + + *target = a + return nil +} + // decodeString type converts v to target. func decodeString(v interface{}, target *string) error { if v, ok := v.(string); ok { diff --git a/pql/parser_test.go b/pql/parser_test.go index df4bcf810..59628af49 100644 --- a/pql/parser_test.go +++ b/pql/parser_test.go @@ -254,22 +254,23 @@ func TestParser_Parse_SetBitmapAttrs_Array(t *testing.T) { // Ensure the parser can parse a "TopN()" function with keyed args. func TestParser_Parse_TopN_Key(t *testing.T) { - q, err := pql.ParseString(`TopN(Bitmap(100), frame="b.n", n=2, field="XXX", [5,10,15])`) + q, err := pql.ParseString(`TopN(Bitmap(100), frame="b.n", n=2, ids=[1,2,3], field="XXX", [5,10,15])`) if err != nil { t.Fatal(err) } else if !reflect.DeepEqual(q, &pql.Query{ Root: &pql.TopN{ - Src: &pql.Bitmap{ID: 100}, - Frame: "b.n", - N: 2, - Field: "XXX", - Filters: []interface{}{uint64(5), uint64(10), uint64(15)}, + Src: &pql.Bitmap{ID: 100}, + Frame: "b.n", + N: 2, + BitmapIDs: []uint64{1, 2, 3}, + Field: "XXX", + Filters: []interface{}{uint64(5), uint64(10), uint64(15)}, }, }) { t.Fatalf("unexpected query: %s", spew.Sdump(q)) } - if s := q.String(); s != `TopN(Bitmap(id=100), frame=b.n, n=2, field="XXX", [5,10,15])` { + if s := q.String(); s != `TopN(Bitmap(id=100), frame=b.n, n=2, ids=[1,2,3], field="XXX", [5,10,15])` { t.Fatalf("unexpected string encoding: %s", s) } }