Merge pull request #584 from mschoch/more-collector-benchmarks

topn collector switch approach based on size+skip

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)
 }
 

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()
 }

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)
 }

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 {