package core
import (
"github.com/huichen/wukong/types"
"github.com/huichen/wukong/utils"
"log"
"sort"
"sync"
)
type Ranker struct { lock struct { sync.RWMutex
fields map[uint64]interface{} docs map[uint64]bool
}
initialized bool
}
func (ranker *Ranker) Init() { if ranker.initialized == true { log.Fatal("排序器不能初始化两次") }
ranker.initialized = true
ranker.lock.fields = make(map[uint64]interface{}) ranker.lock.docs = make(map[uint64]bool)
}
// 给某个文档添加评分字段
func (ranker *Ranker) AddDoc(docId uint64, fields interface{}) { if ranker.initialized == false { log.Fatal("排序器尚未初始化") }
ranker.lock.Lock()
ranker.lock.fields[docId] = fields
ranker.lock.docs[docId] = true
ranker.lock.Unlock()
}
// 删除某个文档的评分字段
func (ranker *Ranker) RemoveDoc(docId uint64) { if ranker.initialized == false { log.Fatal("排序器尚未初始化") }
ranker.lock.Lock()
delete(ranker.lock.fields, docId)
delete(ranker.lock.docs, docId)
ranker.lock.Unlock()
}
// 给文档评分并排序
func (ranker *Ranker) Rank(
docs []types.IndexedDocument, options types.RankOptions, countDocsOnly bool) (types.ScoredDocuments, int) { if ranker.initialized == false { log.Fatal("排序器尚未初始化") }
// 对每个文档评分
var outputDocs types.ScoredDocuments
numDocs := 0
for _, d := range docs { ranker.lock.RLock()
// 判断doc是否存在
if _, ok := ranker.lock.docs[d.DocId]; ok { fs := ranker.lock.fields[d.DocId]
ranker.lock.RUnlock()
// 计算评分并剔除没有分值的文档
scores := options.ScoringCriteria.Score(d, fs)
if len(scores) > 0 { if !countDocsOnly { outputDocs = append(outputDocs, types.ScoredDocument{ DocId: d.DocId,
Scores: scores,
TokenSnippetLocations: d.TokenSnippetLocations,
TokenLocations: d.TokenLocations})
}
numDocs++
}
} else { ranker.lock.RUnlock()
}
}
// 排序
if !countDocsOnly { if options.ReverseOrder { sort.Sort(sort.Reverse(outputDocs))
} else { sort.Sort(outputDocs)
}
// 当用户要求只返回部分结果时返回部分结果
var start, end int
if options.MaxOutputs != 0 { start = utils.MinInt(options.OutputOffset, len(outputDocs))
end = utils.MinInt(options.OutputOffset+options.MaxOutputs, len(outputDocs))
} else { start = utils.MinInt(options.OutputOffset, len(outputDocs))
end = len(outputDocs)
}
return outputDocs[start:end], numDocs
}
return outputDocs, numDocs
}