diff --git a/search/collector/heap.go b/search/collector/heap.go index 19f3a059..e7fd21c6 100644 --- a/search/collector/heap.go +++ b/search/collector/heap.go @@ -34,11 +34,20 @@ func newStoreHeap(cap int, compare collectorCompare) *collectStoreHeap { return rv } -func (c *collectStoreHeap) Add(doc *search.DocumentMatch) { +func (c *collectStoreHeap) AddNotExceedingSize(doc *search.DocumentMatch, + size int) *search.DocumentMatch { + c.add(doc) + if c.Len() > size { + return c.removeLast() + } + return nil +} + +func (c *collectStoreHeap) add(doc *search.DocumentMatch) { heap.Push(c, doc) } -func (c *collectStoreHeap) RemoveLast() *search.DocumentMatch { +func (c *collectStoreHeap) removeLast() *search.DocumentMatch { return heap.Pop(c).(*search.DocumentMatch) } diff --git a/search/collector/list.go b/search/collector/list.go index 51b47b12..ec2f69cb 100644 --- a/search/collector/list.go +++ b/search/collector/list.go @@ -34,7 +34,16 @@ func newStoreList(cap int, compare collectorCompare) *collectStoreList { return rv } -func (c *collectStoreList) Add(doc *search.DocumentMatch) { +func (c *collectStoreList) AddNotExceedingSize(doc *search.DocumentMatch, + size int) *search.DocumentMatch { + c.add(doc) + if c.len() > size { + return c.removeLast() + } + return nil +} + +func (c *collectStoreList) add(doc *search.DocumentMatch) { for e := c.results.Front(); e != nil; e = e.Next() { curr := e.Value.(*search.DocumentMatch) if c.compare(doc, curr) >= 0 { @@ -46,7 +55,7 @@ func (c *collectStoreList) Add(doc *search.DocumentMatch) { c.results.PushBack(doc) } -func (c *collectStoreList) RemoveLast() *search.DocumentMatch { +func (c *collectStoreList) removeLast() *search.DocumentMatch { return c.results.Remove(c.results.Front()).(*search.DocumentMatch) } @@ -73,6 +82,6 @@ func (c *collectStoreList) Final(skip int, fixup collectorFixup) (search.Documen return search.DocumentMatchCollection{}, nil } -func (c *collectStoreList) Len() int { +func (c *collectStoreList) len() int { return c.results.Len() } diff --git a/search/collector/slice.go b/search/collector/slice.go index b061643e..32cb8624 100644 --- a/search/collector/slice.go +++ b/search/collector/slice.go @@ -29,7 +29,16 @@ func newStoreSlice(cap int, compare collectorCompare) *collectStoreSlice { return rv } -func (c *collectStoreSlice) Add(doc *search.DocumentMatch) { +func (c *collectStoreSlice) AddNotExceedingSize(doc *search.DocumentMatch, + size int) *search.DocumentMatch { + c.add(doc) + if c.len() > size { + return c.removeLast() + } + return nil +} + +func (c *collectStoreSlice) add(doc *search.DocumentMatch) { // find where to insert, starting at end (lowest) i := len(c.slice) for ; i > 0; i-- { @@ -44,7 +53,7 @@ func (c *collectStoreSlice) Add(doc *search.DocumentMatch) { c.slice[i] = doc } -func (c *collectStoreSlice) RemoveLast() *search.DocumentMatch { +func (c *collectStoreSlice) removeLast() *search.DocumentMatch { var rv *search.DocumentMatch rv, c.slice = c.slice[len(c.slice)-1], c.slice[:len(c.slice)-1] return rv @@ -63,6 +72,6 @@ func (c *collectStoreSlice) Final(skip int, fixup collectorFixup) (search.Docume return search.DocumentMatchCollection{}, nil } -func (c *collectStoreSlice) Len() int { +func (c *collectStoreSlice) len() int { return len(c.slice) } diff --git a/search/collector/topn.go b/search/collector/topn.go index 946ca3ec..2c7c6752 100644 --- a/search/collector/topn.go +++ b/search/collector/topn.go @@ -22,6 +22,15 @@ import ( "golang.org/x/net/context" ) +type collectorStore interface { + // Add the document, and if the new store size exceeds the provided size + // the last element is removed and returned. If the size has not been + // exceeded, nil is returned. + AddNotExceedingSize(doc *search.DocumentMatch, size int) *search.DocumentMatch + + Final(skip int, fixup collectorFixup) (search.DocumentMatchCollection, error) +} + // PreAllocSizeSkipCap will cap preallocation to this amount when // size+skip exceeds this value var PreAllocSizeSkipCap = 1000 @@ -41,7 +50,7 @@ type TopNCollector struct { results search.DocumentMatchCollection facetsBuilder *search.FacetsBuilder - store *collectStoreHeap + store collectorStore needDocIds bool neededFields []string @@ -68,9 +77,15 @@ func NewTopNCollector(size int, skip int, sort search.SortOrder) *TopNCollector backingSize = PreAllocSizeSkipCap + 1 } - hc.store = newStoreHeap(backingSize, func(i, j *search.DocumentMatch) int { - return hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, i, j) - }) + if size+skip > 10 { + hc.store = newStoreHeap(backingSize, func(i, j *search.DocumentMatch) int { + return hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, i, j) + }) + } else { + hc.store = newStoreSlice(backingSize, func(i, j *search.DocumentMatch) int { + return hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, i, j) + }) + } // these lookups traverse an interface, so do once up-front if sort.RequiresDocID() { @@ -184,9 +199,8 @@ func (hc *TopNCollector) collectSingle(ctx *search.SearchContext, reader index.I } } - hc.store.Add(d) - if hc.store.Len() > hc.size+hc.skip { - removed := hc.store.RemoveLast() + removed := hc.store.AddNotExceedingSize(d, hc.size+hc.skip) + if removed != nil { if hc.lowestMatchOutsideResults == nil { hc.lowestMatchOutsideResults = removed } else {