initial commit

vendor/github.com/blevesearch/bleve/search/collectors/collector_top_score.go

@@ -0,0 +1,149 @@
+//  Copyright (c) 2014 Couchbase, Inc.
+//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+//  except in compliance with the License. You may obtain a copy of the License at
+//    http://www.apache.org/licenses/LICENSE-2.0
+//  Unless required by applicable law or agreed to in writing, software distributed under the
+//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+//  either express or implied. See the License for the specific language governing permissions
+//  and limitations under the License.
+
+package collectors
+
+import (
+	"container/list"
+	"time"
+
+	"golang.org/x/net/context"
+
+	"github.com/blevesearch/bleve/search"
+)
+
+type TopScoreCollector struct {
+	k             int
+	skip          int
+	results       *list.List
+	took          time.Duration
+	maxScore      float64
+	total         uint64
+	facetsBuilder *search.FacetsBuilder
+}
+
+func NewTopScorerCollector(k int) *TopScoreCollector {
+	return &TopScoreCollector{
+		k:       k,
+		skip:    0,
+		results: list.New(),
+	}
+}
+
+func NewTopScorerSkipCollector(k, skip int) *TopScoreCollector {
+	return &TopScoreCollector{
+		k:       k,
+		skip:    skip,
+		results: list.New(),
+	}
+}
+
+func (tksc *TopScoreCollector) Total() uint64 {
+	return tksc.total
+}
+
+func (tksc *TopScoreCollector) MaxScore() float64 {
+	return tksc.maxScore
+}
+
+func (tksc *TopScoreCollector) Took() time.Duration {
+	return tksc.took
+}
+
+func (tksc *TopScoreCollector) Collect(ctx context.Context, searcher search.Searcher) error {
+	startTime := time.Now()
+	var err error
+	var next *search.DocumentMatch
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		next, err = searcher.Next()
+	}
+	for err == nil && next != nil {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			tksc.collectSingle(next)
+			if tksc.facetsBuilder != nil {
+				err = tksc.facetsBuilder.Update(next)
+				if err != nil {
+					break
+				}
+			}
+			next, err = searcher.Next()
+		}
+	}
+	// compute search duration
+	tksc.took = time.Since(startTime)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (tksc *TopScoreCollector) collectSingle(dm *search.DocumentMatch) {
+	// increment total hits
+	tksc.total++
+
+	// update max score
+	if dm.Score > tksc.maxScore {
+		tksc.maxScore = dm.Score
+	}
+
+	for e := tksc.results.Front(); e != nil; e = e.Next() {
+		curr := e.Value.(*search.DocumentMatch)
+		if dm.Score < curr.Score {
+
+			tksc.results.InsertBefore(dm, e)
+			// if we just made the list too long
+			if tksc.results.Len() > (tksc.k + tksc.skip) {
+				// remove the head
+				tksc.results.Remove(tksc.results.Front())
+			}
+			return
+		}
+	}
+	// if we got to the end, we still have to add it
+	tksc.results.PushBack(dm)
+	if tksc.results.Len() > (tksc.k + tksc.skip) {
+		// remove the head
+		tksc.results.Remove(tksc.results.Front())
+	}
+}
+
+func (tksc *TopScoreCollector) Results() search.DocumentMatchCollection {
+	if tksc.results.Len()-tksc.skip > 0 {
+		rv := make(search.DocumentMatchCollection, tksc.results.Len()-tksc.skip)
+		i := 0
+		skipped := 0
+		for e := tksc.results.Back(); e != nil; e = e.Prev() {
+			if skipped < tksc.skip {
+				skipped++
+				continue
+			}
+			rv[i] = e.Value.(*search.DocumentMatch)
+			i++
+		}
+		return rv
+	}
+	return search.DocumentMatchCollection{}
+}
+
+func (tksc *TopScoreCollector) SetFacetsBuilder(facetsBuilder *search.FacetsBuilder) {
+	tksc.facetsBuilder = facetsBuilder
+}
+
+func (tksc *TopScoreCollector) FacetResults() search.FacetResults {
+	if tksc.facetsBuilder != nil {
+		return tksc.facetsBuilder.Results()
+	}
+	return search.FacetResults{}
+}