package blevex
|
|
import (
|
"fmt"
|
"github.com/blevesearch/bleve/v2"
|
"github.com/blevesearch/bleve/v2/mapping"
|
"sync"
|
"time"
|
)
|
|
// InitAnalyzer 加载自定义分词器（sego）
|
|
var defaultAnalyzer *mapping.IndexMappingImpl
|
|
func InitAnalyzer() {
|
indexMapping := bleve.NewIndexMapping()
|
err := indexMapping.AddCustomTokenizer("sego",
|
map[string]interface{}{
|
"dictpath": "conf/dictionary.txt", // 替换为实际的字典路径
|
"type": "sego",
|
},
|
)
|
if err != nil {
|
panic(err)
|
}
|
err = indexMapping.AddCustomAnalyzer("sego",
|
map[string]interface{}{
|
"type": "sego",
|
"tokenizer": "sego",
|
},
|
)
|
if err != nil {
|
panic(err)
|
}
|
indexMapping.DefaultAnalyzer = "sego"
|
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, from, size int) (ids []string, total uint64, err error) {
|
index, err := LoadIndex(indexName)
|
if err != nil {
|
return
|
}
|
req := bleve.NewSearchRequest(bleve.NewMatchQuery(keyword))
|
req.From = from
|
req.Size = size
|
res, err := index.Search(req)
|
if err != nil {
|
return nil, 0, err
|
}
|
if res == nil {
|
return
|
}
|
for _, ret := range dealResult(res) {
|
ids = append(ids, ret.Id)
|
}
|
return ids, res.Total, nil
|
}
|
|
func ComplexSearch(indexName string, keyword string, conditions map[string]interface{}, from, size int) (ids []string, total uint64, err error) {
|
index, err := LoadIndex(indexName)
|
if err != nil {
|
return nil, 0, err
|
}
|
|
// Create a boolean query with a should clause for fuzzy search
|
boolQuery := bleve.NewBooleanQuery()
|
fuzzyQuery := bleve.NewMatchQuery(keyword)
|
boolQuery.AddMust(fuzzyQuery)
|
|
// Add a must clause for category filtering
|
for key, val := range conditions {
|
query := bleve.NewQueryStringQuery(fmt.Sprintf("%v:%v", key, val))
|
boolQuery.AddMust(query)
|
}
|
req := bleve.NewSearchRequest(boolQuery)
|
req.From = from
|
req.Size = size
|
res, err := index.Search(req)
|
if err != nil {
|
return nil, 0, err
|
}
|
if res == nil {
|
return
|
}
|
for _, ret := range dealResult(res) {
|
ids = append(ids, ret.Id)
|
}
|
return ids, res.Total, 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
|
}
|
|
func TimeSearch(indexName string, t time.Time, conditions map[string]interface{}, from, size int) (ids []string, total uint64, err error) {
|
index, err := LoadIndex(indexName)
|
if err != nil {
|
return nil, 0, err
|
}
|
startDate := t
|
endDate := t.Add(time.Hour * 24).Add(-time.Second * 1)
|
timeRangeQuery := bleve.NewDateRangeQuery(startDate, endDate)
|
boolQuery := bleve.NewBooleanQuery()
|
boolQuery.AddMust(timeRangeQuery)
|
for key, val := range conditions {
|
query := bleve.NewQueryStringQuery(fmt.Sprintf("%v:%v", key, val))
|
boolQuery.AddMust(query)
|
}
|
|
req := bleve.NewSearchRequest(boolQuery)
|
req.From = from
|
req.Size = size
|
res, err := index.Search(req)
|
if err != nil {
|
return nil, 0, err
|
}
|
if res == nil {
|
return
|
}
|
for _, ret := range dealResult(res) {
|
ids = append(ids, ret.Id)
|
}
|
return ids, res.Total, nil
|
}
|
