package blevex
|
|
import (
|
"github.com/blevesearch/bleve/v2"
|
"github.com/blevesearch/bleve/v2/mapping"
|
"github.com/yanyiwu/gojieba"
|
"sync"
|
)
|
|
// InitAnalyzer 加载自定义分词器（结巴分词）
|
|
var defaultAnalyzer *mapping.IndexMappingImpl
|
|
func InitAnalyzer() {
|
indexMapping := bleve.NewIndexMapping()
|
//os.RemoveAll(IndexDir)
|
//// clean index when example finished
|
//defer os.RemoveAll(IndexDir)
|
|
err := indexMapping.AddCustomTokenizer("gojieba",
|
map[string]interface{}{
|
"dictpath": gojieba.DICT_PATH,
|
"hmmpath": gojieba.HMM_PATH,
|
"userdictpath": gojieba.USER_DICT_PATH,
|
"idf": gojieba.IDF_PATH,
|
"stop_words": gojieba.STOP_WORDS_PATH,
|
"type": "gojieba",
|
},
|
)
|
if err != nil {
|
panic(err)
|
}
|
err = indexMapping.AddCustomAnalyzer("gojieba",
|
map[string]interface{}{
|
"type": "gojieba",
|
"tokenizer": "gojieba",
|
},
|
)
|
if err != nil {
|
panic(err)
|
}
|
indexMapping.DefaultAnalyzer = "gojieba"
|
defaultAnalyzer = indexMapping
|
}
|
|
var indexMap sync.Map
|
|
func NewIndex(indexName string) (bleve.Index, error) {
|
if defaultAnalyzer == nil {
|
InitAnalyzer()
|
}
|
index, err := bleve.New(indexName, defaultAnalyzer)
|
if err == bleve.ErrorIndexPathExists {
|
index, err = bleve.Open(indexName)
|
if err != nil {
|
return nil, err
|
}
|
return index, nil
|
} else if err != nil {
|
return nil, err
|
}
|
return index, nil
|
}
|
|
func LoadIndex(indexName string) (bleve.Index, error) {
|
if v, ok := indexMap.Load(indexName); ok {
|
return v.(bleve.Index), nil
|
}
|
index, err := NewIndex(indexName)
|
if err == nil {
|
indexMap.Store(indexName, index)
|
}
|
return index, err
|
}
|
func Search(indexName string, keyword string) (ids []string, err error) {
|
index, err := LoadIndex(indexName)
|
if err != nil {
|
return nil, err
|
}
|
req := bleve.NewSearchRequest(bleve.NewQueryStringQuery(keyword))
|
res, err := index.Search(req)
|
if err != nil {
|
panic(err)
|
}
|
if res == nil || res.Total == 0 {
|
return ids, nil
|
}
|
for _, ret := range dealResult(res) {
|
ids = append(ids, ret.Id)
|
}
|
return ids, nil
|
}
|
|
type Result struct {
|
Id string `json:"id"`
|
Score float64 `json:"score"`
|
}
|
|
func dealResult(res *bleve.SearchResult) []Result {
|
var results []Result
|
for _, item := range res.Hits {
|
results = append(results, Result{item.ID, item.Score})
|
}
|
return results
|
}
|
