109 lines
2.5 KiB
Go
109 lines
2.5 KiB
Go
package dao
|
|
|
|
import (
|
|
"go-common/app/service/main/riot-search/model"
|
|
|
|
"github.com/go-ego/riot/types"
|
|
)
|
|
|
|
// SearchIDOnly only return aids
|
|
func (d *Dao) SearchIDOnly(arg *model.RiotSearchReq) *model.IDsResp {
|
|
if arg.Keyword == "" {
|
|
return nil
|
|
}
|
|
var docIDs map[uint64]bool
|
|
if len(arg.IDs) != 0 {
|
|
docIDs = make(map[uint64]bool, len(arg.IDs))
|
|
for _, id := range arg.IDs {
|
|
docIDs[id] = true
|
|
}
|
|
}
|
|
output := d.searcher.Search(types.SearchReq{
|
|
Text: arg.Keyword,
|
|
DocIds: docIDs,
|
|
Timeout: d.c.Riot.Timeout,
|
|
RankOpts: &types.RankOpts{
|
|
// 从第几条结果开始输出
|
|
OutputOffset: (arg.Pn - 1) * arg.Ps,
|
|
// 最大输出的搜索结果数，为 0 时无限制
|
|
MaxOutputs: arg.Ps,
|
|
},
|
|
})
|
|
docLength := len(output.Docs.(types.ScoredDocs))
|
|
tokenLength := len(output.Tokens)
|
|
res := &model.IDsResp{
|
|
IDs: make([]uint64, docLength),
|
|
Tokens: make([]string, tokenLength),
|
|
Page: &model.Page{
|
|
PageNum: arg.Pn,
|
|
PageSize: arg.Ps,
|
|
Total: docLength,
|
|
},
|
|
}
|
|
for i, doc := range output.Docs.(types.ScoredDocs) {
|
|
res.IDs[i] = doc.DocId
|
|
}
|
|
copy(res.Tokens, output.Tokens)
|
|
res.IDs, res.Page.Total = uniqueIDs(res.IDs)
|
|
return res
|
|
}
|
|
|
|
func uniqueIDs(IDs []uint64) (uIDs []uint64, length int) {
|
|
m := make(map[uint64]struct{})
|
|
for _, ID := range IDs {
|
|
if _, ok := m[ID]; !ok {
|
|
m[ID] = struct{}{}
|
|
uIDs = append(uIDs, ID)
|
|
length++
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
// Search return archives info
|
|
func (d *Dao) Search(arg *model.RiotSearchReq) *model.DocumentsResp {
|
|
if arg.Keyword == "" {
|
|
return nil
|
|
}
|
|
var docIDs map[uint64]bool
|
|
if len(arg.IDs) != 0 {
|
|
docIDs = make(map[uint64]bool, len(arg.IDs))
|
|
for _, id := range arg.IDs {
|
|
docIDs[id] = true
|
|
}
|
|
}
|
|
output := d.searcher.Search(types.SearchReq{
|
|
Text: arg.Keyword,
|
|
DocIds: docIDs,
|
|
Timeout: d.c.Riot.Timeout,
|
|
RankOpts: &types.RankOpts{
|
|
// 从第几条结果开始输出
|
|
OutputOffset: (arg.Pn - 1) * arg.Ps,
|
|
// 最大输出的搜索结果数，为 0 时无限制
|
|
MaxOutputs: arg.Ps,
|
|
},
|
|
})
|
|
docLength := len(output.Docs.(types.ScoredDocs))
|
|
tokenLength := len(output.Tokens)
|
|
res := &model.DocumentsResp{
|
|
Documents: make([]model.Document, docLength),
|
|
Tokens: make([]string, tokenLength),
|
|
Page: &model.Page{
|
|
PageNum: arg.Pn,
|
|
PageSize: arg.Ps,
|
|
Total: docLength,
|
|
},
|
|
}
|
|
for i, doc := range output.Docs.(types.ScoredDocs) {
|
|
res.Documents[i].ID = doc.DocId
|
|
res.Documents[i].Content = doc.Content
|
|
}
|
|
copy(res.Tokens, output.Tokens)
|
|
return res
|
|
}
|
|
|
|
// Has return DocId exists
|
|
func (d *Dao) Has(id uint64) bool {
|
|
return d.searcher.HasDoc(id)
|
|
}
|