indexer_worker.go

// Copyright 2013 Hui Chen
// Copyright 2016 ego authors
//
// 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 riot

import (
	"sync/atomic"

	"github.com/go-ego/riot/types"
)

type indexerAddDocRequest struct {
	document    *types.DocumentIndex
	forceUpdate bool
}

type indexerLookupRequest struct {
	countDocsOnly       bool
	tokens              []string
	labels              []string
	docIds              map[uint64]bool
	options             types.RankOpts
	rankerReturnChannel chan rankerReturnRequest
	orderless           bool
	logic               types.Logic
}

type indexerRemoveDocRequest struct {
	docId       uint64
	forceUpdate bool
}

func (engine *Engine) indexerAddDocWorker(shard int) {
	for {
		request := <-engine.indexerAddDocChannels[shard]
		engine.indexers[shard].AddDocToCache(request.document, request.forceUpdate)
		if request.document != nil {
			atomic.AddUint64(&engine.numTokenIndexAdded,
				uint64(len(request.document.Keywords)))
			atomic.AddUint64(&engine.numDocumentsIndexed, 1)
		}
		if request.forceUpdate {
			atomic.AddUint64(&engine.numDocumentsForceUpdated, 1)
		}
	}
}

func (engine *Engine) indexerRemoveDocWorker(shard int) {
	for {
		request := <-engine.indexerRemoveDocChannels[shard]
		engine.indexers[shard].RemoveDocToCache(request.docId, request.forceUpdate)
		if request.docId != 0 {
			atomic.AddUint64(&engine.numDocumentsRemoved, 1)
		}
		if request.forceUpdate {
			atomic.AddUint64(&engine.numDocumentsForceUpdated, 1)
		}
	}
}

func (engine *Engine) indexerLookupWorker(shard int) {
	for {
		request := <-engine.indexerLookupChannels[shard]

		var (
			docs    []types.IndexedDocument
			numDocs int
		)
		if request.docIds == nil {
			docs, numDocs = engine.indexers[shard].Lookup(request.tokens, request.labels, nil, request.countDocsOnly, request.logic)
			// docs, numDocs = engine.indexers[shard].Lookup(request.tokens, request.labels, nil, request.countDocsOnly)
		} else {
			docs, numDocs = engine.indexers[shard].Lookup(request.tokens, request.labels, request.docIds, request.countDocsOnly, request.logic)
			// docs, numDocs = engine.indexers[shard].Lookup(request.tokens, request.labels, request.docIds, request.countDocsOnly)
		}

		if request.countDocsOnly {
			request.rankerReturnChannel <- rankerReturnRequest{numDocs: numDocs}
			continue
		}

		if len(docs) == 0 {
			request.rankerReturnChannel <- rankerReturnRequest{}
			continue
		}

		if request.orderless {
			var outputDocs []types.ScoredDoc
			for _, d := range docs {
				outputDocs = append(outputDocs, types.ScoredDoc{
					DocId: d.DocId,
					TokenSnippetLocations: d.TokenSnippetLocations,
					TokenLocations:        d.TokenLocations})
			}
			request.rankerReturnChannel <- rankerReturnRequest{
				docs:    outputDocs,
				numDocs: len(outputDocs),
			}
			continue
		}

		rankerRequest := rankerRankRequest{
			countDocsOnly:       request.countDocsOnly,
			docs:                docs,
			options:             request.options,
			rankerReturnChannel: request.rankerReturnChannel,
		}
		engine.rankerRankChannels[shard] <- rankerRequest
	}
}