├── README.md ├── core ├── data.go ├── indexer.go ├── indexer_test.go ├── ranker.go ├── ranker_test.go └── test_utils.go ├── data ├── README.md ├── dictionary.txt └── stop_tokens.txt ├── docs ├── benchmarking.md ├── bm25.md ├── codelab.md ├── cpu.png ├── custom_scoring_criteria.md ├── distributed_indexing_and_search.md ├── docker.md ├── feedback.md ├── persistent_storage.md ├── realtime_indexing.md ├── token_proximity.md ├── why_wukong.md └── wukong.png ├── engine ├── counters.go ├── engine.go ├── engine_test.go ├── indexer_worker.go ├── persistent_storage_worker.go ├── ranker_worker.go ├── segmenter_worker.go └── stop_tokens.go ├── examples ├── benchmark.go ├── codelab │ ├── Dockerfile │ ├── build_docker_image.sh │ ├── search_server.go │ └── static │ │ ├── index.html │ │ └── jquery.min.js ├── custom_scoring_criteria.go └── simplest_example.go ├── license.txt ├── storage ├── bolt_storage.go ├── bolt_storage_test.go ├── kv_storage.go ├── kv_storage_test.go ├── ldb_storage.go └── storage.go ├── testdata ├── crawl_weibo_data.go ├── test_dict.txt ├── users.txt └── weibo_data.txt ├── types ├── doc_info.go ├── document_index_data.go ├── engine_init_options.go ├── index.go ├── indexer_init_options.go ├── inverted_index.go ├── scoring_criteria.go ├── search_request.go └── search_response.go ├── utils ├── test_utils.go └── utils.go └── wukong.go /README.md: -------------------------------------------------------------------------------- 1 | 悟空全文搜索引擎 2 | ====== 3 | 4 | * [高效索引和搜索](/docs/benchmarking.md)（1M条微博500M数据28秒索引完，1.65毫秒搜索响应时间，19K搜索QPS） 5 | * 支持中文分词（使用[sego分词包](https://github.com/huichen/sego)并发分词，速度27MB/秒） 6 | * 支持计算关键词在文本中的[紧邻距离](/docs/token_proximity.md)（token proximity） 7 | * 支持计算[BM25相关度](/docs/bm25.md) 8 | * 支持[自定义评分字段和评分规则](/docs/custom_scoring_criteria.md) 9 | * 支持[在线添加、删除索引](/docs/realtime_indexing.md) 10 | * 支持[持久存储](/docs/persistent_storage.md) 11 | * 可实现[分布式索引和搜索](/docs/distributed_indexing_and_search.md) 12 | * 采用对商业应用友好的[Apache License v2](/license.txt)发布 13 | 14 | [微博搜索demo](http://vhaa7.fmt.tifan.net:8080/) 15 | 16 | # 安装/更新 17 | 18 | ``` 19 | go get -u -v github.com/huichen/wukong 20 | ``` 21 | 22 | 需要Go版本至少1.1.1 23 | 24 | # 使用 25 | 26 | 先看一个例子（来自[examples/simplest_example.go](/examples/simplest_example.go)） 27 | ```go 28 | package main 29 | 30 | import ( 31 | "github.com/huichen/wukong/engine" 32 | "github.com/huichen/wukong/types" 33 | "log" 34 | ) 35 | 36 | var ( 37 | // searcher是协程安全的 38 | searcher = engine.Engine{} 39 | ) 40 | 41 | func main() { 42 | // 初始化 43 | searcher.Init(types.EngineInitOptions{ 44 | SegmenterDictionaries: "github.com/huichen/wukong/data/dictionary.txt"}) 45 | defer searcher.Close() 46 | 47 | // 将文档加入索引 48 | searcher.IndexDocument(0, types.DocumentIndexData{Content: "此次百度收购将成中国互联网最大并购"}) 49 | searcher.IndexDocument(1, types.DocumentIndexData{Content: "百度宣布拟全资收购91无线业务"}) 50 | searcher.IndexDocument(2, types.DocumentIndexData{Content: "百度是中国最大的搜索引擎"}) 51 | 52 | // 等待索引刷新完毕 53 | searcher.FlushIndex() 54 | 55 | // 搜索输出格式见types.SearchResponse结构体 56 | log.Print(searcher.Search(types.SearchRequest{Text:"百度中国"})) 57 | } 58 | ``` 59 | 60 | 是不是很简单！ 61 | 62 | 然后看看一个[入门教程](/docs/codelab.md)，教你用不到200行Go代码实现一个微博搜索网站。 63 | 64 | # 其它 65 | 66 | * [为什么要有悟空引擎](/docs/why_wukong.md) 67 | * [联系方式](/docs/feedback.md) 68 | -------------------------------------------------------------------------------- /core/data.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "github.com/huichen/wukong/types" 5 | "sync" 6 | ) 7 | 8 | // 文档信息[shard][id]info 9 | var DocInfoGroup = make(map[int]*types.DocInfosShard) 10 | var docInfosGroupRWMutex sync.RWMutex 11 | 12 | func AddDocInfosShard(shard int) { 13 | docInfosGroupRWMutex.Lock() 14 | defer docInfosGroupRWMutex.Unlock() 15 | if _, found := DocInfoGroup[shard]; !found { 16 | DocInfoGroup[shard] = &types.DocInfosShard{ 17 | DocInfos: make(map[uint64]*types.DocInfo), 18 | } 19 | } 20 | } 21 | 22 | func AddDocInfo(shard int, docId uint64, docinfo *types.DocInfo) { 23 | docInfosGroupRWMutex.Lock() 24 | defer docInfosGroupRWMutex.Unlock() 25 | if _, ok := DocInfoGroup[shard]; !ok { 26 | DocInfoGroup[shard] = &types.DocInfosShard{ 27 | DocInfos: make(map[uint64]*types.DocInfo), 28 | } 29 | } 30 | DocInfoGroup[shard].DocInfos[docId] = docinfo 31 | DocInfoGroup[shard].NumDocuments++ 32 | } 33 | 34 | // func IsDocExist(docId uint64) bool { 35 | // docInfosGroupRWMutex.RLock() 36 | // defer docInfosGroupRWMutex.RUnlock() 37 | // for _, docInfosShard := range DocInfoGroup { 38 | // _, found := docInfosShard.DocInfos[docId] 39 | // if found { 40 | // return true 41 | // } 42 | // } 43 | // return false 44 | // } 45 | 46 | // 反向索引表([shard][关键词]反向索引表) 47 | var InvertedIndexGroup = make(map[int]*types.InvertedIndexShard) 48 | var invertedIndexGroupRWMutex sync.RWMutex 49 | 50 | func AddInvertedIndexShard(shard int) { 51 | invertedIndexGroupRWMutex.Lock() 52 | defer invertedIndexGroupRWMutex.Unlock() 53 | if _, found := InvertedIndexGroup[shard]; !found { 54 | InvertedIndexGroup[shard] = &types.InvertedIndexShard{ 55 | InvertedIndex: make(map[string]*types.KeywordIndices), 56 | } 57 | } 58 | } 59 | 60 | func AddKeywordIndices(shard int, keyword string, keywordIndices *types.KeywordIndices) { 61 | invertedIndexGroupRWMutex.Lock() 62 | defer invertedIndexGroupRWMutex.Unlock() 63 | if _, ok := InvertedIndexGroup[shard]; !ok { 64 | InvertedIndexGroup[shard] = &types.InvertedIndexShard{ 65 | InvertedIndex: make(map[string]*types.KeywordIndices), 66 | } 67 | } 68 | InvertedIndexGroup[shard].InvertedIndex[keyword] = keywordIndices 69 | InvertedIndexGroup[shard].TotalTokenLength++ 70 | } 71 | -------------------------------------------------------------------------------- /core/indexer.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "github.com/huichen/wukong/types" 5 | "github.com/huichen/wukong/utils" 6 | "log" 7 | "math" 8 | ) 9 | 10 | // 索引器 11 | type Indexer struct { 12 | shard int 13 | initOptions types.IndexerInitOptions 14 | initialized bool 15 | // 文档信息 16 | *types.DocInfosShard 17 | // 反向索引 18 | *types.InvertedIndexShard 19 | } 20 | 21 | // 初始化索引器 22 | func (indexer *Indexer) Init(shard int, options types.IndexerInitOptions) { 23 | if indexer.initialized == true { 24 | log.Fatal("索引器不能初始化两次") 25 | } 26 | indexer.initialized = true 27 | 28 | indexer.shard = shard 29 | 30 | AddDocInfosShard(shard) 31 | indexer.DocInfosShard = DocInfoGroup[shard] 32 | 33 | AddInvertedIndexShard(shard) 34 | indexer.InvertedIndexShard = InvertedIndexGroup[shard] 35 | 36 | indexer.initOptions = options 37 | } 38 | 39 | // 向反向索引表中加入一个文档 40 | func (indexer *Indexer) AddDocument(document *types.DocumentIndex, dealDocInfoChan chan<- bool) (addInvertedIndex map[string]*types.KeywordIndices) { 41 | if indexer.initialized == false { 42 | log.Fatal("索引器尚未初始化") 43 | } 44 | 45 | indexer.InvertedIndexShard.Lock() 46 | defer indexer.InvertedIndexShard.Unlock() 47 | 48 | // 更新文档总数及关键词总长度 49 | indexer.DocInfosShard.Lock() 50 | if _, found := indexer.DocInfosShard.DocInfos[document.DocId]; !found { 51 | indexer.DocInfosShard.DocInfos[document.DocId] = new(types.DocInfo) 52 | indexer.DocInfosShard.NumDocuments++ 53 | } 54 | if document.TokenLength != 0 { 55 | originalLength := indexer.DocInfosShard.DocInfos[document.DocId].TokenLengths 56 | indexer.DocInfosShard.DocInfos[document.DocId].TokenLengths = float32(document.TokenLength) 57 | indexer.InvertedIndexShard.TotalTokenLength += document.TokenLength - originalLength 58 | } 59 | indexer.DocInfosShard.Unlock() 60 | close(dealDocInfoChan) 61 | 62 | // docIdIsNew := true 63 | foundKeyword := false 64 | addInvertedIndex = make(map[string]*types.KeywordIndices) 65 | for _, keyword := range document.Keywords { 66 | addInvertedIndex[keyword.Text], foundKeyword = indexer.InvertedIndexShard.InvertedIndex[keyword.Text] 67 | if !foundKeyword { 68 | addInvertedIndex[keyword.Text] = new(types.KeywordIndices) 69 | } 70 | indices := addInvertedIndex[keyword.Text] 71 | 72 | if !foundKeyword { 73 | // 如果没找到该搜索键则加入 74 | switch indexer.initOptions.IndexType { 75 | case types.LocationsIndex: 76 | indices.Locations = [][]int{keyword.Starts} 77 | case types.FrequenciesIndex: 78 | indices.Frequencies = []float32{keyword.Frequency} 79 | } 80 | indices.DocIds = []uint64{document.DocId} 81 | indexer.InvertedIndexShard.InvertedIndex[keyword.Text] = indices 82 | continue 83 | } 84 | 85 | // 查找应该插入的位置 86 | position, found := indexer.searchIndex( 87 | indices, 0, indexer.getIndexLength(indices)-1, document.DocId) 88 | if found { 89 | // docIdIsNew = false 90 | 91 | // 覆盖已有的索引项 92 | switch indexer.initOptions.IndexType { 93 | case types.LocationsIndex: 94 | indices.Locations[position] = keyword.Starts 95 | case types.FrequenciesIndex: 96 | indices.Frequencies[position] = keyword.Frequency 97 | } 98 | continue 99 | } 100 | 101 | // 当索引不存在时，插入新索引项 102 | switch indexer.initOptions.IndexType { 103 | case types.LocationsIndex: 104 | indices.Locations = append(indices.Locations, []int{}) 105 | copy(indices.Locations[position+1:], indices.Locations[position:]) 106 | indices.Locations[position] = keyword.Starts 107 | case types.FrequenciesIndex: 108 | indices.Frequencies = append(indices.Frequencies, float32(0)) 109 | copy(indices.Frequencies[position+1:], indices.Frequencies[position:]) 110 | indices.Frequencies[position] = keyword.Frequency 111 | } 112 | indices.DocIds = append(indices.DocIds, 0) 113 | copy(indices.DocIds[position+1:], indices.DocIds[position:]) 114 | indices.DocIds[position] = document.DocId 115 | } 116 | return 117 | } 118 | 119 | // 查找包含全部搜索键(AND操作)的文档 120 | // 当docIds不为nil时仅从docIds指定的文档中查找 121 | func (indexer *Indexer) Lookup( 122 | tokens []string, labels []string, docIds map[uint64]bool, countDocsOnly bool) (docs []types.IndexedDocument, numDocs int) { 123 | if indexer.initialized == false { 124 | log.Fatal("索引器尚未初始化") 125 | } 126 | 127 | indexer.DocInfosShard.RLock() 128 | defer indexer.DocInfosShard.RUnlock() 129 | 130 | if indexer.DocInfosShard.NumDocuments == 0 { 131 | return 132 | } 133 | numDocs = 0 134 | 135 | // 合并关键词和标签为搜索键 136 | keywords := make([]string, len(tokens)+len(labels)) 137 | copy(keywords, tokens) 138 | copy(keywords[len(tokens):], labels) 139 | 140 | indexer.InvertedIndexShard.RLock() 141 | 142 | table := make([]*types.KeywordIndices, len(keywords)) 143 | for i, keyword := range keywords { 144 | indices, found := indexer.InvertedIndexShard.InvertedIndex[keyword] 145 | if !found { 146 | // 当反向索引表中无此搜索键时直接返回 147 | indexer.InvertedIndexShard.RUnlock() 148 | return 149 | } else { 150 | // 否则加入反向表中 151 | table[i] = indices 152 | } 153 | } 154 | 155 | // 当没有找到时直接返回 156 | if len(table) == 0 { 157 | indexer.InvertedIndexShard.RUnlock() 158 | return 159 | } 160 | 161 | // 归并查找各个搜索键出现文档的交集 162 | // 从后向前查保证先输出DocId较大文档 163 | indexPointers := make([]int, len(table)) 164 | for iTable := 0; iTable < len(table); iTable++ { 165 | indexPointers[iTable] = indexer.getIndexLength(table[iTable]) - 1 166 | } 167 | // 平均文本关键词长度，用于计算BM25 168 | avgDocLength := indexer.InvertedIndexShard.TotalTokenLength / float32(indexer.DocInfosShard.NumDocuments) 169 | indexer.InvertedIndexShard.RUnlock() 170 | 171 | for ; indexPointers[0] >= 0; indexPointers[0]-- { 172 | // 以第一个搜索键出现的文档作为基准，并遍历其他搜索键搜索同一文档 173 | baseDocId := indexer.getDocId(table[0], indexPointers[0]) 174 | 175 | // 全局范围查找目标文档是否存在 176 | if _, ok := indexer.DocInfosShard.DocInfos[baseDocId]; !ok { 177 | // if !IsDocExist(baseDocId) { 178 | // 文档信息中不存在反向索引文档时，跳过 179 | // 该情况由不对称删除操作所造成 180 | continue 181 | } 182 | 183 | if docIds != nil { 184 | _, found := docIds[baseDocId] 185 | if !found { 186 | continue 187 | } 188 | } 189 | iTable := 1 190 | found := true 191 | for ; iTable < len(table); iTable++ { 192 | // 二分法比简单的顺序归并效率高，也有更高效率的算法， 193 | // 但顺序归并也许是更好的选择，考虑到将来需要用链表重新实现 194 | // 以避免反向表添加新文档时的写锁。 195 | // TODO: 进一步研究不同求交集算法的速度和可扩展性。 196 | position, foundBaseDocId := indexer.searchIndex(table[iTable], 197 | 0, indexPointers[iTable], baseDocId) 198 | if foundBaseDocId { 199 | indexPointers[iTable] = position 200 | } else { 201 | if position == 0 { 202 | // 该搜索键中所有的文档ID都比baseDocId大，因此已经没有 203 | // 继续查找的必要。 204 | return 205 | } else { 206 | // 继续下一indexPointers[0]的查找 207 | indexPointers[iTable] = position - 1 208 | found = false 209 | break 210 | } 211 | } 212 | } 213 | 214 | if found { 215 | indexedDoc := types.IndexedDocument{} 216 | 217 | // 当为LocationsIndex时计算关键词紧邻距离 218 | if indexer.initOptions.IndexType == types.LocationsIndex { 219 | // 计算有多少关键词是带有距离信息的 220 | numTokensWithLocations := 0 221 | for i, t := range table[:len(tokens)] { 222 | if len(t.Locations[indexPointers[i]]) > 0 { 223 | numTokensWithLocations++ 224 | } 225 | } 226 | if numTokensWithLocations != len(tokens) { 227 | if !countDocsOnly { 228 | docs = append(docs, types.IndexedDocument{ 229 | DocId: baseDocId, 230 | }) 231 | } 232 | numDocs++ 233 | break 234 | } 235 | 236 | // 计算搜索键在文档中的紧邻距离 237 | tokenProximity, tokenLocations := computeTokenProximity(table[:len(tokens)], indexPointers, tokens) 238 | indexedDoc.TokenProximity = int32(tokenProximity) 239 | indexedDoc.TokenSnippetLocations = tokenLocations 240 | 241 | // 添加TokenLocations 242 | indexedDoc.TokenLocations = make([][]int, len(tokens)) 243 | for i, t := range table[:len(tokens)] { 244 | indexedDoc.TokenLocations[i] = t.Locations[indexPointers[i]] 245 | } 246 | } 247 | 248 | // 当为LocationsIndex或者FrequenciesIndex时计算BM25 249 | if indexer.initOptions.IndexType == types.LocationsIndex || 250 | indexer.initOptions.IndexType == types.FrequenciesIndex { 251 | bm25 := float32(0) 252 | d := indexer.DocInfosShard.DocInfos[baseDocId].TokenLengths 253 | for i, t := range table[:len(tokens)] { 254 | var frequency float32 255 | if indexer.initOptions.IndexType == types.LocationsIndex { 256 | frequency = float32(len(t.Locations[indexPointers[i]])) 257 | } else { 258 | frequency = t.Frequencies[indexPointers[i]] 259 | } 260 | 261 | // 计算BM25 262 | if len(t.DocIds) > 0 && frequency > 0 && indexer.initOptions.BM25Parameters != nil && avgDocLength != 0 { 263 | // 带平滑的idf 264 | idf := float32(math.Log2(float64(indexer.DocInfosShard.NumDocuments)/float64(len(t.DocIds)) + 1)) 265 | k1 := indexer.initOptions.BM25Parameters.K1 266 | b := indexer.initOptions.BM25Parameters.B 267 | bm25 += idf * frequency * (k1 + 1) / (frequency + k1*(1-b+b*d/avgDocLength)) 268 | } 269 | } 270 | indexedDoc.BM25 = float32(bm25) 271 | } 272 | 273 | indexedDoc.DocId = baseDocId 274 | if !countDocsOnly { 275 | docs = append(docs, indexedDoc) 276 | } 277 | numDocs++ 278 | } 279 | } 280 | return 281 | } 282 | 283 | // 二分法查找indices中某文档的索引项 284 | // 第一个返回参数为找到的位置或需要插入的位置 285 | // 第二个返回参数标明是否找到 286 | func (indexer *Indexer) searchIndex( 287 | indices *types.KeywordIndices, start int, end int, docId uint64) (int, bool) { 288 | // 特殊情况 289 | if indexer.getIndexLength(indices) == start { 290 | return start, false 291 | } 292 | if docId < indexer.getDocId(indices, start) { 293 | return start, false 294 | } else if docId == indexer.getDocId(indices, start) { 295 | return start, true 296 | } 297 | if docId > indexer.getDocId(indices, end) { 298 | return end + 1, false 299 | } else if docId == indexer.getDocId(indices, end) { 300 | return end, true 301 | } 302 | 303 | // 二分 304 | var middle int 305 | for end-start > 1 { 306 | middle = (start + end) / 2 307 | if docId == indexer.getDocId(indices, middle) { 308 | return middle, true 309 | } else if docId > indexer.getDocId(indices, middle) { 310 | start = middle 311 | } else { 312 | end = middle 313 | } 314 | } 315 | return end, false 316 | } 317 | 318 | // 计算搜索键在文本中的紧邻距离 319 | // 320 | // 假定第 i 个搜索键首字节出现在文本中的位置为 P_i，长度 L_i 321 | // 紧邻距离计算公式为 322 | // 323 | // ArgMin(Sum(Abs(P_(i+1) - P_i - L_i))) 324 | // 325 | // 具体由动态规划实现，依次计算前 i 个 token 在每个出现位置的最优值。 326 | // 选定的 P_i 通过 tokenLocations 参数传回。 327 | func computeTokenProximity(table []*types.KeywordIndices, indexPointers []int, tokens []string) ( 328 | minTokenProximity int, tokenLocations []int) { 329 | minTokenProximity = -1 330 | tokenLocations = make([]int, len(tokens)) 331 | 332 | var ( 333 | currentLocations, nextLocations []int 334 | currentMinValues, nextMinValues []int 335 | path [][]int 336 | ) 337 | 338 | // 初始化路径数组 339 | path = make([][]int, len(tokens)) 340 | for i := 1; i < len(path); i++ { 341 | path[i] = make([]int, len(table[i].Locations[indexPointers[i]])) 342 | } 343 | 344 | // 动态规划 345 | currentLocations = table[0].Locations[indexPointers[0]] 346 | currentMinValues = make([]int, len(currentLocations)) 347 | for i := 1; i < len(tokens); i++ { 348 | nextLocations = table[i].Locations[indexPointers[i]] 349 | nextMinValues = make([]int, len(nextLocations)) 350 | for j, _ := range nextMinValues { 351 | nextMinValues[j] = -1 352 | } 353 | 354 | var iNext int 355 | for iCurrent, currentLocation := range currentLocations { 356 | if currentMinValues[iCurrent] == -1 { 357 | continue 358 | } 359 | for iNext+1 < len(nextLocations) && nextLocations[iNext+1] < currentLocation { 360 | iNext++ 361 | } 362 | 363 | update := func(from int, to int) { 364 | if to >= len(nextLocations) { 365 | return 366 | } 367 | value := currentMinValues[from] + utils.AbsInt(nextLocations[to]-currentLocations[from]-len(tokens[i-1])) 368 | if nextMinValues[to] == -1 || value < nextMinValues[to] { 369 | nextMinValues[to] = value 370 | path[i][to] = from 371 | } 372 | } 373 | 374 | // 最优解的状态转移只发生在左右最接近的位置 375 | update(iCurrent, iNext) 376 | update(iCurrent, iNext+1) 377 | } 378 | 379 | currentLocations = nextLocations 380 | currentMinValues = nextMinValues 381 | } 382 | 383 | // 找出最优解 384 | var cursor int 385 | for i, value := range currentMinValues { 386 | if value == -1 { 387 | continue 388 | } 389 | if minTokenProximity == -1 || value < minTokenProximity { 390 | minTokenProximity = value 391 | cursor = i 392 | } 393 | } 394 | 395 | // 从路径倒推出最优解的位置 396 | for i := len(tokens) - 1; i >= 0; i-- { 397 | if i != len(tokens)-1 { 398 | cursor = path[i+1][cursor] 399 | } 400 | tokenLocations[i] = table[i].Locations[indexPointers[i]][cursor] 401 | } 402 | return 403 | } 404 | 405 | // 从KeywordIndices中得到第i个文档的DocId 406 | func (indexer *Indexer) getDocId(ti *types.KeywordIndices, i int) uint64 { 407 | return ti.DocIds[i] 408 | } 409 | 410 | // 得到KeywordIndices中文档总数 411 | func (indexer *Indexer) getIndexLength(ti *types.KeywordIndices) int { 412 | return len(ti.DocIds) 413 | } 414 | 415 | // 删除某个文档（反向索引的删除太复杂故而不做，只在排序器中删除文档即可） 416 | func (indexer *Indexer) RemoveDoc(docId uint64) { 417 | if indexer.initialized == false { 418 | log.Fatal("排序器尚未初始化") 419 | } 420 | } 421 | -------------------------------------------------------------------------------- /core/indexer_test.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "github.com/huichen/wukong/types" 5 | "github.com/huichen/wukong/utils" 6 | "testing" 7 | ) 8 | 9 | func TestAddKeywords(t *testing.T) { 10 | var indexer Indexer 11 | indexer.Init(1, types.IndexerInitOptions{IndexType: types.LocationsIndex}) 12 | indexer.AddDocument(&types.DocumentIndex{ 13 | DocId: 1, 14 | Keywords: []types.KeywordIndex{{"token1", 0, []int{}}}, 15 | }, 16 | make(chan<- bool), 17 | ) 18 | indexer.AddDocument(&types.DocumentIndex{ 19 | DocId: 7, 20 | Keywords: []types.KeywordIndex{{"token1", 0, []int{}}}, 21 | }, 22 | make(chan<- bool), 23 | ) 24 | indexer.AddDocument(&types.DocumentIndex{ 25 | DocId: 2, 26 | Keywords: []types.KeywordIndex{{"token1", 0, []int{}}}, 27 | }, 28 | make(chan<- bool), 29 | ) 30 | indexer.AddDocument(&types.DocumentIndex{ 31 | DocId: 3, 32 | Keywords: []types.KeywordIndex{{"token2", 0, []int{}}}, 33 | }, 34 | make(chan<- bool), 35 | ) 36 | indexer.AddDocument(&types.DocumentIndex{ 37 | DocId: 1, 38 | Keywords: []types.KeywordIndex{{"token1", 0, []int{}}}, 39 | }, 40 | make(chan<- bool), 41 | ) 42 | indexer.AddDocument(&types.DocumentIndex{ 43 | DocId: 1, 44 | Keywords: []types.KeywordIndex{{"token2", 0, []int{}}}, 45 | }, 46 | make(chan<- bool), 47 | ) 48 | indexer.AddDocument(&types.DocumentIndex{ 49 | DocId: 2, 50 | Keywords: []types.KeywordIndex{{"token2", 0, []int{}}}, 51 | }, 52 | make(chan<- bool), 53 | ) 54 | indexer.AddDocument(&types.DocumentIndex{ 55 | DocId: 0, 56 | Keywords: []types.KeywordIndex{{"token2", 0, []int{}}}, 57 | }, 58 | make(chan<- bool), 59 | ) 60 | 61 | utils.Expect(t, "1 2 7 ", indicesToString(&indexer, "token1")) 62 | utils.Expect(t, "0 1 2 3 ", indicesToString(&indexer, "token2")) 63 | } 64 | 65 | func TestLookup(t *testing.T) { 66 | var indexer Indexer 67 | indexer.Init(2, types.IndexerInitOptions{IndexType: types.LocationsIndex}) 68 | // doc0 = "token2 token3" 69 | indexer.AddDocument(&types.DocumentIndex{ 70 | DocId: 0, 71 | Keywords: []types.KeywordIndex{ 72 | {"token2", 0, []int{0}}, 73 | {"token3", 0, []int{7}}, 74 | }, 75 | }, 76 | make(chan<- bool), 77 | ) 78 | // doc1 = "token1 token2 token3" 79 | indexer.AddDocument(&types.DocumentIndex{ 80 | DocId: 1, 81 | Keywords: []types.KeywordIndex{ 82 | {"token1", 0, []int{0}}, 83 | {"token2", 0, []int{7}}, 84 | {"token3", 0, []int{14}}, 85 | }, 86 | }, 87 | make(chan<- bool), 88 | ) 89 | // doc2 = "token1 token2" 90 | indexer.AddDocument(&types.DocumentIndex{ 91 | DocId: 2, 92 | Keywords: []types.KeywordIndex{ 93 | {"token1", 0, []int{0}}, 94 | {"token2", 0, []int{7}}, 95 | }, 96 | }, 97 | make(chan<- bool), 98 | ) 99 | // doc3 = "token2" 100 | indexer.AddDocument(&types.DocumentIndex{ 101 | DocId: 3, 102 | Keywords: []types.KeywordIndex{ 103 | {"token2", 0, []int{0}}, 104 | }, 105 | }, 106 | make(chan<- bool), 107 | ) 108 | // doc7 = "token1 token3" 109 | indexer.AddDocument(&types.DocumentIndex{ 110 | DocId: 7, 111 | Keywords: []types.KeywordIndex{ 112 | {"token1", 0, []int{0}}, 113 | {"token3", 0, []int{7}}, 114 | }, 115 | }, 116 | make(chan<- bool), 117 | ) 118 | // doc9 = "token3" 119 | indexer.AddDocument(&types.DocumentIndex{ 120 | DocId: 9, 121 | Keywords: []types.KeywordIndex{ 122 | {"token3", 0, []int{0}}, 123 | }, 124 | }, 125 | make(chan<- bool), 126 | ) 127 | 128 | utils.Expect(t, "1 2 7 ", indicesToString(&indexer, "token1")) 129 | utils.Expect(t, "0 1 2 3 ", indicesToString(&indexer, "token2")) 130 | utils.Expect(t, "0 1 7 9 ", indicesToString(&indexer, "token3")) 131 | 132 | utils.Expect(t, "", indexedDocsToString(indexer.Lookup([]string{"token4"}, []string{}, nil, false))) 133 | 134 | utils.Expect(t, "[7 0 [0]] [2 0 [0]] [1 0 [0]] ", 135 | indexedDocsToString(indexer.Lookup([]string{"token1"}, []string{}, nil, false))) 136 | utils.Expect(t, "", indexedDocsToString(indexer.Lookup([]string{"token1", "token4"}, []string{}, nil, false))) 137 | 138 | utils.Expect(t, "[2 1 [0 7]] [1 1 [0 7]] ", 139 | indexedDocsToString(indexer.Lookup([]string{"token1", "token2"}, []string{}, nil, false))) 140 | utils.Expect(t, "[2 13 [7 0]] [1 13 [7 0]] ", 141 | indexedDocsToString(indexer.Lookup([]string{"token2", "token1"}, []string{}, nil, false))) 142 | utils.Expect(t, "[7 1 [0 7]] [1 8 [0 14]] ", 143 | indexedDocsToString(indexer.Lookup([]string{"token1", "token3"}, []string{}, nil, false))) 144 | utils.Expect(t, "[7 13 [7 0]] [1 20 [14 0]] ", 145 | indexedDocsToString(indexer.Lookup([]string{"token3", "token1"}, []string{}, nil, false))) 146 | utils.Expect(t, "[1 1 [7 14]] [0 1 [0 7]] ", 147 | indexedDocsToString(indexer.Lookup([]string{"token2", "token3"}, []string{}, nil, false))) 148 | utils.Expect(t, "[1 13 [14 7]] [0 13 [7 0]] ", 149 | indexedDocsToString(indexer.Lookup([]string{"token3", "token2"}, []string{}, nil, false))) 150 | 151 | utils.Expect(t, "[1 2 [0 7 14]] ", 152 | indexedDocsToString(indexer.Lookup([]string{"token1", "token2", "token3"}, []string{}, nil, false))) 153 | utils.Expect(t, "[1 26 [14 7 0]] ", 154 | indexedDocsToString(indexer.Lookup([]string{"token3", "token2", "token1"}, []string{}, nil, false))) 155 | } 156 | 157 | func TestDocIdsIndex(t *testing.T) { 158 | var indexer Indexer 159 | indexer.Init(1, types.IndexerInitOptions{IndexType: types.DocIdsIndex}) 160 | // doc0 = "token2 token3" 161 | indexer.AddDocument(&types.DocumentIndex{ 162 | DocId: 0, 163 | Keywords: []types.KeywordIndex{ 164 | {"token2", 0, []int{0}}, 165 | {"token3", 0, []int{7}}, 166 | }, 167 | }, 168 | make(chan<- bool), 169 | ) 170 | // doc1 = "token1 token2 token3" 171 | indexer.AddDocument(&types.DocumentIndex{ 172 | DocId: 1, 173 | Keywords: []types.KeywordIndex{ 174 | {"token1", 0, []int{0}}, 175 | {"token2", 0, []int{7}}, 176 | {"token3", 0, []int{14}}, 177 | }, 178 | }, 179 | make(chan<- bool), 180 | ) 181 | // doc2 = "token1 token2" 182 | indexer.AddDocument(&types.DocumentIndex{ 183 | DocId: 2, 184 | Keywords: []types.KeywordIndex{ 185 | {"token1", 0, []int{0}}, 186 | {"token2", 0, []int{7}}, 187 | }, 188 | }, 189 | make(chan<- bool), 190 | ) 191 | // doc3 = "token2" 192 | indexer.AddDocument(&types.DocumentIndex{ 193 | DocId: 3, 194 | Keywords: []types.KeywordIndex{ 195 | {"token2", 0, []int{0}}, 196 | }, 197 | }, 198 | make(chan<- bool), 199 | ) 200 | // doc7 = "token1 token3" 201 | indexer.AddDocument(&types.DocumentIndex{ 202 | DocId: 7, 203 | Keywords: []types.KeywordIndex{ 204 | {"token1", 0, []int{0}}, 205 | {"token3", 0, []int{7}}, 206 | }, 207 | }, 208 | make(chan<- bool), 209 | ) 210 | // doc9 = "token3" 211 | indexer.AddDocument(&types.DocumentIndex{ 212 | DocId: 9, 213 | Keywords: []types.KeywordIndex{ 214 | {"token3", 0, []int{0}}, 215 | }, 216 | }, 217 | make(chan<- bool), 218 | ) 219 | 220 | utils.Expect(t, "1 2 7 ", indicesToString(&indexer, "token1")) 221 | utils.Expect(t, "0 1 2 3 ", indicesToString(&indexer, "token2")) 222 | utils.Expect(t, "0 1 7 9 ", indicesToString(&indexer, "token3")) 223 | 224 | utils.Expect(t, "", indexedDocsToString(indexer.Lookup([]string{"token4"}, []string{}, nil, false))) 225 | 226 | utils.Expect(t, "[7 0 []] [2 0 []] [1 0 []] ", 227 | indexedDocsToString(indexer.Lookup([]string{"token1"}, []string{}, nil, false))) 228 | utils.Expect(t, "", indexedDocsToString(indexer.Lookup([]string{"token1", "token4"}, []string{}, nil, false))) 229 | 230 | utils.Expect(t, "[2 0 []] [1 0 []] ", 231 | indexedDocsToString(indexer.Lookup([]string{"token1", "token2"}, []string{}, nil, false))) 232 | utils.Expect(t, "[2 0 []] [1 0 []] ", 233 | indexedDocsToString(indexer.Lookup([]string{"token2", "token1"}, []string{}, nil, false))) 234 | utils.Expect(t, "[7 0 []] [1 0 []] ", 235 | indexedDocsToString(indexer.Lookup([]string{"token1", "token3"}, []string{}, nil, false))) 236 | utils.Expect(t, "[7 0 []] [1 0 []] ", 237 | indexedDocsToString(indexer.Lookup([]string{"token3", "token1"}, []string{}, nil, false))) 238 | utils.Expect(t, "[1 0 []] [0 0 []] ", 239 | indexedDocsToString(indexer.Lookup([]string{"token2", "token3"}, []string{}, nil, false))) 240 | utils.Expect(t, "[1 0 []] [0 0 []] ", 241 | indexedDocsToString(indexer.Lookup([]string{"token3", "token2"}, []string{}, nil, false))) 242 | 243 | utils.Expect(t, "[1 0 []] ", 244 | indexedDocsToString(indexer.Lookup([]string{"token1", "token2", "token3"}, []string{}, nil, false))) 245 | utils.Expect(t, "[1 0 []] ", 246 | indexedDocsToString(indexer.Lookup([]string{"token3", "token2", "token1"}, []string{}, nil, false))) 247 | } 248 | 249 | func TestLookupWithProximity(t *testing.T) { 250 | var indexer Indexer 251 | indexer.Init(3, types.IndexerInitOptions{IndexType: types.LocationsIndex}) 252 | 253 | // doc0 = "token2 token4 token4 token2 token3 token4" 254 | indexer.AddDocument(&types.DocumentIndex{ 255 | DocId: 0, 256 | Keywords: []types.KeywordIndex{ 257 | {"token2", 0, []int{0, 21}}, 258 | {"token3", 0, []int{28}}, 259 | {"token4", 0, []int{7, 14, 35}}, 260 | }, 261 | }, 262 | make(chan<- bool), 263 | ) 264 | utils.Expect(t, "[0 1 [21 28]] ", 265 | indexedDocsToString(indexer.Lookup([]string{"token2", "token3"}, []string{}, nil, false))) 266 | 267 | // doc0 = "t2 t1 . . . t2 t3" 268 | indexer.AddDocument(&types.DocumentIndex{ 269 | DocId: 0, 270 | Keywords: []types.KeywordIndex{ 271 | {"t1", 0, []int{3}}, 272 | {"t2", 0, []int{0, 12}}, 273 | {"t3", 0, []int{15}}, 274 | }, 275 | }, 276 | make(chan<- bool), 277 | ) 278 | utils.Expect(t, "[0 8 [3 12 15]] ", 279 | indexedDocsToString(indexer.Lookup([]string{"t1", "t2", "t3"}, []string{}, nil, false))) 280 | 281 | // doc0 = "t3 t2 t1 . . . . . t2 t3" 282 | indexer.AddDocument(&types.DocumentIndex{ 283 | DocId: 0, 284 | Keywords: []types.KeywordIndex{ 285 | {"t1", 0, []int{6}}, 286 | {"t2", 0, []int{3, 19}}, 287 | {"t3", 0, []int{0, 22}}, 288 | }, 289 | }, 290 | make(chan<- bool), 291 | ) 292 | utils.Expect(t, "[0 10 [6 3 0]] ", 293 | indexedDocsToString(indexer.Lookup([]string{"t1", "t2", "t3"}, []string{}, nil, false))) 294 | } 295 | 296 | func TestLookupWithPartialLocations(t *testing.T) { 297 | var indexer Indexer 298 | indexer.Init(5, types.IndexerInitOptions{IndexType: types.LocationsIndex}) 299 | // doc0 = "token2 token4 token4 token2 token3 token4" + "label1"(不在文本中) 300 | indexer.AddDocument(&types.DocumentIndex{ 301 | DocId: 0, 302 | Keywords: []types.KeywordIndex{ 303 | {"token2", 0, []int{0, 21}}, 304 | {"token3", 0, []int{28}}, 305 | {"label1", 0, []int{}}, 306 | {"token4", 0, []int{7, 14, 35}}, 307 | }, 308 | }, 309 | make(chan<- bool), 310 | ) 311 | // doc1 = "token2 token4 token4 token2 token3 token4" 312 | indexer.AddDocument(&types.DocumentIndex{ 313 | DocId: 1, 314 | Keywords: []types.KeywordIndex{ 315 | {"token2", 0, []int{0, 21}}, 316 | {"token3", 0, []int{28}}, 317 | {"token4", 0, []int{7, 14, 35}}, 318 | }, 319 | }, 320 | make(chan<- bool), 321 | ) 322 | 323 | utils.Expect(t, "0 ", indicesToString(&indexer, "label1")) 324 | 325 | utils.Expect(t, "[0 1 [21 28]] ", 326 | indexedDocsToString(indexer.Lookup([]string{"token2", "token3"}, []string{"label1"}, nil, false))) 327 | } 328 | 329 | func TestLookupWithBM25(t *testing.T) { 330 | var indexer Indexer 331 | indexer.Init(31, types.IndexerInitOptions{ 332 | IndexType: types.FrequenciesIndex, 333 | BM25Parameters: &types.BM25Parameters{ 334 | K1: 1, 335 | B: 1, 336 | }, 337 | }) 338 | // doc0 = "token2 token4 token4 token2 token3 token4" 339 | indexer.AddDocument(&types.DocumentIndex{ 340 | DocId: 0, 341 | TokenLength: 6, 342 | Keywords: []types.KeywordIndex{ 343 | {"token2", 3, []int{0, 21}}, 344 | {"token3", 7, []int{28}}, 345 | {"token4", 15, []int{7, 14, 35}}, 346 | }, 347 | }, 348 | make(chan<- bool), 349 | ) 350 | // doc0 = "token6 token7" 351 | indexer.AddDocument(&types.DocumentIndex{ 352 | DocId: 1, 353 | TokenLength: 2, 354 | Keywords: []types.KeywordIndex{ 355 | {"token6", 3, []int{0}}, 356 | {"token7", 15, []int{7}}, 357 | }, 358 | }, 359 | make(chan<- bool), 360 | ) 361 | 362 | outputs, _ := indexer.Lookup([]string{"token2", "token3", "token4"}, []string{}, nil, false) 363 | 364 | // BM25 = log2(3) * (12/9 + 28/17 + 60/33) = 6.3433 365 | utils.Expect(t, "76055", int(outputs[0].BM25*10000)) 366 | } 367 | 368 | func TestLookupWithinDocIds(t *testing.T) { 369 | var indexer Indexer 370 | indexer.Init(2, types.IndexerInitOptions{IndexType: types.LocationsIndex}) 371 | // doc0 = "token2 token3" 372 | indexer.AddDocument(&types.DocumentIndex{ 373 | DocId: 0, 374 | Keywords: []types.KeywordIndex{ 375 | {"token2", 0, []int{0}}, 376 | {"token3", 0, []int{7}}, 377 | }, 378 | }, 379 | make(chan<- bool), 380 | ) 381 | // doc1 = "token1 token2 token3" 382 | indexer.AddDocument(&types.DocumentIndex{ 383 | DocId: 1, 384 | Keywords: []types.KeywordIndex{ 385 | {"token1", 0, []int{0}}, 386 | {"token2", 0, []int{7}}, 387 | {"token3", 0, []int{14}}, 388 | }, 389 | }, 390 | make(chan<- bool), 391 | ) 392 | // doc2 = "token1 token2" 393 | indexer.AddDocument(&types.DocumentIndex{ 394 | DocId: 2, 395 | Keywords: []types.KeywordIndex{ 396 | {"token1", 0, []int{0}}, 397 | {"token2", 0, []int{7}}, 398 | }, 399 | }, 400 | make(chan<- bool), 401 | ) 402 | // doc3 = "token2" 403 | indexer.AddDocument(&types.DocumentIndex{ 404 | DocId: 3, 405 | Keywords: []types.KeywordIndex{ 406 | {"token2", 0, []int{0}}, 407 | }, 408 | }, 409 | make(chan<- bool), 410 | ) 411 | 412 | docIds := make(map[uint64]bool) 413 | docIds[0] = true 414 | docIds[2] = true 415 | utils.Expect(t, "[2 0 [7]] [0 0 [0]] ", 416 | indexedDocsToString(indexer.Lookup([]string{"token2"}, []string{}, docIds, false))) 417 | } 418 | 419 | func TestLookupWithLocations(t *testing.T) { 420 | var indexer Indexer 421 | indexer.Init(0, types.IndexerInitOptions{IndexType: types.LocationsIndex}) 422 | // doc0 = "token2 token4 token4 token2 token3 token4" 423 | indexer.AddDocument(&types.DocumentIndex{ 424 | DocId: 0, 425 | Keywords: []types.KeywordIndex{ 426 | {"token2", 0, []int{0, 21}}, 427 | {"token3", 0, []int{28}}, 428 | {"token4", 0, []int{7, 14, 35}}, 429 | }, 430 | }, 431 | make(chan<- bool), 432 | ) 433 | 434 | docs, _ := indexer.Lookup([]string{"token2", "token3"}, []string{}, nil, false) 435 | utils.Expect(t, "[[0 21] [28]]", docs[0].TokenLocations) 436 | } 437 | -------------------------------------------------------------------------------- /core/ranker.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "github.com/huichen/wukong/types" 5 | "github.com/huichen/wukong/utils" 6 | "log" 7 | "sort" 8 | ) 9 | 10 | type Ranker struct { 11 | shard int 12 | // 文档信息 13 | *types.DocInfosShard 14 | 15 | initialized bool 16 | } 17 | 18 | func (ranker *Ranker) Init(shard int) { 19 | if ranker.initialized == true { 20 | log.Fatal("排序器不能初始化两次") 21 | } 22 | ranker.initialized = true 23 | 24 | ranker.shard = shard 25 | 26 | AddDocInfosShard(shard) 27 | ranker.DocInfosShard = DocInfoGroup[shard] 28 | } 29 | 30 | // 给某个文档添加评分字段 31 | func (ranker *Ranker) AddDoc(docId uint64, fields interface{}, dealDocInfoChan <-chan bool) *types.DocInfo { 32 | if ranker.initialized == false { 33 | log.Fatal("排序器尚未初始化") 34 | } 35 | 36 | <-dealDocInfoChan // 等待索引器处理完成 37 | 38 | ranker.DocInfosShard.Lock() 39 | defer ranker.DocInfosShard.Unlock() 40 | if _, found := ranker.DocInfosShard.DocInfos[docId]; !found { 41 | ranker.DocInfosShard.DocInfos[docId] = new(types.DocInfo) 42 | ranker.DocInfosShard.NumDocuments++ 43 | } 44 | ranker.DocInfosShard.DocInfos[docId].Fields = fields 45 | return ranker.DocInfosShard.DocInfos[docId] 46 | } 47 | 48 | // 删除某个文档的评分字段 49 | func (ranker *Ranker) RemoveDoc(docId uint64) { 50 | if ranker.initialized == false { 51 | log.Fatal("排序器尚未初始化") 52 | } 53 | 54 | ranker.DocInfosShard.Lock() 55 | delete(ranker.DocInfosShard.DocInfos, docId) 56 | ranker.DocInfosShard.NumDocuments-- 57 | ranker.DocInfosShard.Unlock() 58 | } 59 | 60 | // 给文档评分并排序 61 | func (ranker *Ranker) Rank( 62 | docs []types.IndexedDocument, options types.RankOptions, countDocsOnly bool) (types.ScoredDocuments, int) { 63 | if ranker.initialized == false { 64 | log.Fatal("排序器尚未初始化") 65 | } 66 | // 对每个文档评分 67 | var outputDocs types.ScoredDocuments 68 | numDocs := 0 69 | for _, d := range docs { 70 | ranker.DocInfosShard.RLock() 71 | // 判断doc是否存在 72 | if _, ok := ranker.DocInfosShard.DocInfos[d.DocId]; ok { 73 | fs := ranker.DocInfosShard.DocInfos[d.DocId].Fields 74 | ranker.DocInfosShard.RUnlock() 75 | // 计算评分并剔除没有分值的文档 76 | scores := options.ScoringCriteria.Score(d, fs) 77 | if len(scores) > 0 { 78 | if !countDocsOnly { 79 | outputDocs = append(outputDocs, types.ScoredDocument{ 80 | DocId: d.DocId, 81 | Scores: scores, 82 | TokenSnippetLocations: d.TokenSnippetLocations, 83 | TokenLocations: d.TokenLocations}) 84 | } 85 | numDocs++ 86 | } 87 | } else { 88 | ranker.DocInfosShard.RUnlock() 89 | } 90 | } 91 | 92 | // 排序 93 | if !countDocsOnly { 94 | if options.ReverseOrder { 95 | sort.Sort(sort.Reverse(outputDocs)) 96 | } else { 97 | sort.Sort(outputDocs) 98 | } 99 | // 当用户要求只返回部分结果时返回部分结果 100 | var start, end int 101 | if options.MaxOutputs != 0 { 102 | start = utils.MinInt(options.OutputOffset, len(outputDocs)) 103 | end = utils.MinInt(options.OutputOffset+options.MaxOutputs, len(outputDocs)) 104 | } else { 105 | start = utils.MinInt(options.OutputOffset, len(outputDocs)) 106 | end = len(outputDocs) 107 | } 108 | return outputDocs[start:end], numDocs 109 | } 110 | return outputDocs, numDocs 111 | } 112 | -------------------------------------------------------------------------------- /core/ranker_test.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "github.com/huichen/wukong/types" 5 | "github.com/huichen/wukong/utils" 6 | "reflect" 7 | "testing" 8 | ) 9 | 10 | type DummyScoringFields struct { 11 | label string 12 | counter int 13 | amount float32 14 | } 15 | 16 | type DummyScoringCriteria struct { 17 | Threshold float32 18 | } 19 | 20 | func (criteria DummyScoringCriteria) Score( 21 | doc types.IndexedDocument, fields interface{}) []float32 { 22 | if reflect.TypeOf(fields) == reflect.TypeOf(DummyScoringFields{}) { 23 | dsf := fields.(DummyScoringFields) 24 | value := float32(dsf.counter) + dsf.amount 25 | if value < criteria.Threshold { 26 | return []float32{} 27 | } 28 | return []float32{value} 29 | } 30 | return []float32{} 31 | } 32 | 33 | func TestRankDocument(t *testing.T) { 34 | var ranker Ranker 35 | ranker.Init(5) 36 | c1 := make(chan bool) 37 | c2 := make(chan bool) 38 | c3 := make(chan bool) 39 | close(c1) 40 | close(c2) 41 | close(c3) 42 | ranker.AddDoc(1, DummyScoringFields{}, 43 | c1, 44 | ) 45 | ranker.AddDoc(3, DummyScoringFields{}, 46 | c2, 47 | ) 48 | ranker.AddDoc(4, DummyScoringFields{}, 49 | c3, 50 | ) 51 | 52 | scoredDocs, _ := ranker.Rank([]types.IndexedDocument{ 53 | types.IndexedDocument{DocId: 1, BM25: 6}, 54 | types.IndexedDocument{DocId: 3, BM25: 24}, 55 | types.IndexedDocument{DocId: 4, BM25: 18}, 56 | }, types.RankOptions{ScoringCriteria: types.RankByBM25{}}, false) 57 | utils.Expect(t, "[3 [24000 ]] [4 [18000 ]] [1 [6000 ]] ", scoredDocsToString(scoredDocs)) 58 | 59 | scoredDocs, _ = ranker.Rank([]types.IndexedDocument{ 60 | types.IndexedDocument{DocId: 1, BM25: 6}, 61 | types.IndexedDocument{DocId: 3, BM25: 24}, 62 | types.IndexedDocument{DocId: 2, BM25: 0}, 63 | types.IndexedDocument{DocId: 4, BM25: 18}, 64 | }, types.RankOptions{ScoringCriteria: types.RankByBM25{}, ReverseOrder: true}, false) 65 | // doc0因为没有AddDoc所以没有添加进来 66 | utils.Expect(t, "[1 [6000 ]] [4 [18000 ]] [3 [24000 ]] ", scoredDocsToString(scoredDocs)) 67 | } 68 | 69 | func TestRankWithCriteria(t *testing.T) { 70 | var ranker Ranker 71 | ranker.Init(5) 72 | c1 := make(chan bool) 73 | c2 := make(chan bool) 74 | c3 := make(chan bool) 75 | close(c1) 76 | close(c2) 77 | close(c3) 78 | ranker.AddDoc(1, DummyScoringFields{ 79 | label: "label3", 80 | counter: 3, 81 | amount: 22.3, 82 | }, 83 | c1, 84 | ) 85 | ranker.AddDoc(2, DummyScoringFields{ 86 | label: "label4", 87 | counter: 1, 88 | amount: 2, 89 | }, 90 | c2, 91 | ) 92 | ranker.AddDoc(3, DummyScoringFields{ 93 | label: "label1", 94 | counter: 7, 95 | amount: 10.3, 96 | }, 97 | c3, 98 | ) 99 | ranker.AddDoc(4, DummyScoringFields{ 100 | label: "label1", 101 | counter: -1, 102 | amount: 2.3, 103 | }, 104 | c1, 105 | ) 106 | 107 | criteria := DummyScoringCriteria{} 108 | scoredDocs, _ := ranker.Rank([]types.IndexedDocument{ 109 | types.IndexedDocument{DocId: 1, TokenProximity: 6}, 110 | types.IndexedDocument{DocId: 2, TokenProximity: -1}, 111 | types.IndexedDocument{DocId: 3, TokenProximity: 24}, 112 | types.IndexedDocument{DocId: 4, TokenProximity: 18}, 113 | }, types.RankOptions{ScoringCriteria: criteria}, false) 114 | utils.Expect(t, "[1 [25300 ]] [3 [17300 ]] [2 [3000 ]] [4 [1300 ]] ", scoredDocsToString(scoredDocs)) 115 | 116 | criteria.Threshold = 4 117 | scoredDocs, _ = ranker.Rank([]types.IndexedDocument{ 118 | types.IndexedDocument{DocId: 1, TokenProximity: 6}, 119 | types.IndexedDocument{DocId: 2, TokenProximity: -1}, 120 | types.IndexedDocument{DocId: 3, TokenProximity: 24}, 121 | types.IndexedDocument{DocId: 4, TokenProximity: 18}, 122 | }, types.RankOptions{ScoringCriteria: criteria}, false) 123 | utils.Expect(t, "[1 [25300 ]] [3 [17300 ]] ", scoredDocsToString(scoredDocs)) 124 | } 125 | 126 | func TestRemoveDocument(t *testing.T) { 127 | var ranker Ranker 128 | ranker.Init(6) 129 | c := make(chan bool) 130 | close(c) 131 | ranker.AddDoc(1, DummyScoringFields{ 132 | label: "label3", 133 | counter: 3, 134 | amount: 22.3, 135 | }, 136 | c, 137 | ) 138 | ranker.AddDoc(2, DummyScoringFields{ 139 | label: "label4", 140 | counter: 1, 141 | amount: 2, 142 | }, 143 | c, 144 | ) 145 | ranker.AddDoc(3, DummyScoringFields{ 146 | label: "label1", 147 | counter: 7, 148 | amount: 10.3, 149 | }, 150 | c, 151 | ) 152 | ranker.RemoveDoc(3) 153 | 154 | criteria := DummyScoringCriteria{} 155 | scoredDocs, _ := ranker.Rank([]types.IndexedDocument{ 156 | types.IndexedDocument{DocId: 1, TokenProximity: 6}, 157 | types.IndexedDocument{DocId: 2, TokenProximity: -1}, 158 | types.IndexedDocument{DocId: 3, TokenProximity: 24}, 159 | types.IndexedDocument{DocId: 4, TokenProximity: 18}, 160 | }, types.RankOptions{ScoringCriteria: criteria}, false) 161 | utils.Expect(t, "[1 [25300 ]] [2 [3000 ]] ", scoredDocsToString(scoredDocs)) 162 | } 163 | -------------------------------------------------------------------------------- /core/test_utils.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "fmt" 5 | "github.com/huichen/wukong/types" 6 | ) 7 | 8 | func indicesToString(indexer *Indexer, token string) (output string) { 9 | indices := indexer.InvertedIndexShard.InvertedIndex[token] 10 | for i := 0; i < indexer.getIndexLength(indices); i++ { 11 | output += fmt.Sprintf("%d ", 12 | indexer.getDocId(indices, i)) 13 | } 14 | return 15 | } 16 | 17 | func indexedDocsToString(docs []types.IndexedDocument, numDocs int) (output string) { 18 | for _, doc := range docs { 19 | output += fmt.Sprintf("[%d %d %v] ", 20 | doc.DocId, doc.TokenProximity, doc.TokenSnippetLocations) 21 | } 22 | return 23 | } 24 | 25 | func scoredDocsToString(docs []types.ScoredDocument) (output string) { 26 | for _, doc := range docs { 27 | output += fmt.Sprintf("[%d [", doc.DocId) 28 | for _, score := range doc.Scores { 29 | output += fmt.Sprintf("%d ", int(score*1000)) 30 | } 31 | output += "]] " 32 | } 33 | return 34 | } 35 | -------------------------------------------------------------------------------- /data/README.md: -------------------------------------------------------------------------------- 1 | dictionary.txt 词典拷贝自 github.com/fxsjy/jieba 2 | 3 | stop_tokens.txt 停用词列表来自网络 4 | -------------------------------------------------------------------------------- /data/stop_tokens.txt: -------------------------------------------------------------------------------- 1 | , 2 | . 3 | ? 4 | ! 5 | " 6 | 7 | @ 8 | ， 9 | 。 10 | 、 11 | ？ 12 | ！ 13 | ： 14 | “ 15 | ” 16 | ； 17 | 　 18 | （ 19 | ） 20 | 《 21 | 》 22 | ~ 23 | * 24 | < 25 | > 26 | / 27 | \ 28 | | 29 | - 30 | _ 31 | + 32 | = 33 | & 34 | ^ 35 | % 36 | # 37 | ` 38 | ; 39 | $ 40 | ￥ 41 | ‘ 42 | ’ 43 | 〉 44 | 〈 45 | … 46 | ＞ 47 | ＜ 48 | ＠ 49 | ＃ 50 | ＄ 51 | ％ 52 | ︿ 53 | ＆ 54 | ＊ 55 | ＋ 56 | ～ 57 | ｜ 58 | ［ 59 | ］ 60 | ｛ 61 | ｝ 62 | 啊 63 | 阿 64 | 哎 65 | 哎呀 66 | 哎哟 67 | 唉 68 | 俺 69 | 俺们 70 | 按 71 | 按照 72 | 吧 73 | 吧哒 74 | 把 75 | 罢了 76 | 被 77 | 本 78 | 本着 79 | 比 80 | 比方 81 | 比如 82 | 鄙人 83 | 彼 84 | 彼此 85 | 边 86 | 别 87 | 别的 88 | 别说 89 | 并 90 | 并且 91 | 不比 92 | 不成 93 | 不单 94 | 不但 95 | 不独 96 | 不管 97 | 不光 98 | 不过 99 | 不仅 100 | 不拘 101 | 不论 102 | 不怕 103 | 不然 104 | 不如 105 | 不特 106 | 不惟 107 | 不问 108 | 不只 109 | 朝 110 | 朝着 111 | 趁 112 | 趁着 113 | 乘 114 | 冲 115 | 除 116 | 除此之外 117 | 除非 118 | 除了 119 | 此 120 | 此间 121 | 此外 122 | 从 123 | 从而 124 | 打 125 | 待 126 | 但 127 | 但是 128 | 当 129 | 当着 130 | 到 131 | 得 132 | 的 133 | 的话 134 | 等 135 | 等等 136 | 地 137 | 第 138 | 叮咚 139 | 对 140 | 对于 141 | 多 142 | 多少 143 | 而 144 | 而况 145 | 而且 146 | 而是 147 | 而外 148 | 而言 149 | 而已 150 | 尔后 151 | 反过来 152 | 反过来说 153 | 反之 154 | 非但 155 | 非徒 156 | 否则 157 | 嘎 158 | 嘎登 159 | 该 160 | 赶 161 | 个 162 | 各 163 | 各个 164 | 各位 165 | 各种 166 | 各自 167 | 给 168 | 根据 169 | 跟 170 | 故 171 | 故此 172 | 固然 173 | 关于 174 | 管 175 | 归 176 | 果然 177 | 果真 178 | 过 179 | 哈 180 | 哈哈 181 | 呵 182 | 和 183 | 何 184 | 何处 185 | 何况 186 | 何时 187 | 嘿 188 | 哼 189 | 哼唷 190 | 呼哧 191 | 乎 192 | 哗 193 | 还是 194 | 还有 195 | 换句话说 196 | 换言之 197 | 或 198 | 或是 199 | 或者 200 | 极了 201 | 及 202 | 及其 203 | 及至 204 | 即 205 | 即便 206 | 即或 207 | 即令 208 | 即若 209 | 即使 210 | 几 211 | 几时 212 | 己 213 | 既 214 | 既然 215 | 既是 216 | 继而 217 | 加之 218 | 假如 219 | 假若 220 | 假使 221 | 鉴于 222 | 将 223 | 较 224 | 较之 225 | 叫 226 | 接着 227 | 结果 228 | 借 229 | 紧接着 230 | 进而 231 | 尽 232 | 尽管 233 | 经 234 | 经过 235 | 就 236 | 就是 237 | 就是说 238 | 据 239 | 具体地说 240 | 具体说来 241 | 开始 242 | 开外 243 | 靠 244 | 咳 245 | 可 246 | 可见 247 | 可是 248 | 可以 249 | 况且 250 | 啦 251 | 来 252 | 来着 253 | 离 254 | 例如 255 | 哩 256 | 连 257 | 连同 258 | 两者 259 | 了 260 | 临 261 | 另 262 | 另外 263 | 另一方面 264 | 论 265 | 嘛 266 | 吗 267 | 慢说 268 | 漫说 269 | 冒 270 | 么 271 | 每 272 | 每当 273 | 们 274 | 莫若 275 | 某 276 | 某个 277 | 某些 278 | 拿 279 | 哪 280 | 哪边 281 | 哪儿 282 | 哪个 283 | 哪里 284 | 哪年 285 | 哪怕 286 | 哪天 287 | 哪些 288 | 哪样 289 | 那 290 | 那边 291 | 那儿 292 | 那个 293 | 那会儿 294 | 那里 295 | 那么 296 | 那么些 297 | 那么样 298 | 那时 299 | 那些 300 | 那样 301 | 乃 302 | 乃至 303 | 呢 304 | 能 305 | 你 306 | 你们 307 | 您 308 | 宁 309 | 宁可 310 | 宁肯 311 | 宁愿 312 | 哦 313 | 呕 314 | 啪达 315 | 旁人 316 | 呸 317 | 凭 318 | 凭借 319 | 其 320 | 其次 321 | 其二 322 | 其他 323 | 其它 324 | 其一 325 | 其余 326 | 其中 327 | 起 328 | 起见 329 | 岂但 330 | 恰恰相反 331 | 前后 332 | 前者 333 | 且 334 | 然而 335 | 然后 336 | 然则 337 | 让 338 | 人家 339 | 任 340 | 任何 341 | 任凭 342 | 如 343 | 如此 344 | 如果 345 | 如何 346 | 如其 347 | 如若 348 | 如上所述 349 | 若 350 | 若非 351 | 若是 352 | 啥 353 | 上下 354 | 尚且 355 | 设若 356 | 设使 357 | 甚而 358 | 甚么 359 | 甚至 360 | 省得 361 | 时候 362 | 什么 363 | 什么样 364 | 使得 365 | 是 366 | 是的 367 | 首先 368 | 谁 369 | 谁知 370 | 顺 371 | 顺着 372 | 似的 373 | 虽 374 | 虽然 375 | 虽说 376 | 虽则 377 | 随 378 | 随着 379 | 所 380 | 所以 381 | 他 382 | 他们 383 | 他人 384 | 它 385 | 它们 386 | 她 387 | 她们 388 | 倘 389 | 倘或 390 | 倘然 391 | 倘若 392 | 倘使 393 | 腾 394 | 替 395 | 通过 396 | 同 397 | 同时 398 | 哇 399 | 万一 400 | 往 401 | 望 402 | 为 403 | 为何 404 | 为了 405 | 为什么 406 | 为着 407 | 喂 408 | 嗡嗡 409 | 我 410 | 我们 411 | 呜 412 | 呜呼 413 | 乌乎 414 | 无论 415 | 无宁 416 | 毋宁 417 | 嘻 418 | 吓 419 | 相对而言 420 | 像 421 | 向 422 | 向着 423 | 嘘 424 | 呀 425 | 焉 426 | 沿 427 | 沿着 428 | 要 429 | 要不 430 | 要不然 431 | 要不是 432 | 要么 433 | 要是 434 | 也 435 | 也罢 436 | 也好 437 | 一 438 | 一般 439 | 一旦 440 | 一方面 441 | 一来 442 | 一切 443 | 一样 444 | 一则 445 | 依 446 | 依照 447 | 矣 448 | 以 449 | 以便 450 | 以及 451 | 以免 452 | 以至 453 | 以至于 454 | 以致 455 | 抑或 456 | 因 457 | 因此 458 | 因而 459 | 因为 460 | 哟 461 | 用 462 | 由 463 | 由此可见 464 | 由于 465 | 有 466 | 有的 467 | 有关 468 | 有些 469 | 又 470 | 于 471 | 于是 472 | 于是乎 473 | 与 474 | 与此同时 475 | 与否 476 | 与其 477 | 越是 478 | 云云 479 | 哉 480 | 再说 481 | 再者 482 | 在 483 | 在下 484 | 咱 485 | 咱们 486 | 则 487 | 怎 488 | 怎么 489 | 怎么办 490 | 怎么样 491 | 怎样 492 | 咋 493 | 照 494 | 照着 495 | 者 496 | 这 497 | 这边 498 | 这儿 499 | 这个 500 | 这会儿 501 | 这就是说 502 | 这里 503 | 这么 504 | 这么点儿 505 | 这么些 506 | 这么样 507 | 这时 508 | 这些 509 | 这样 510 | 正如 511 | 吱 512 | 之 513 | 之类 514 | 之所以 515 | 之一 516 | 只是 517 | 只限 518 | 只要 519 | 只有 520 | 至 521 | 至于 522 | 诸位 523 | 着 524 | 着呢 525 | 自 526 | 自从 527 | 自个儿 528 | 自各儿 529 | 自己 530 | 自家 531 | 自身 532 | 综上所述 533 | 总的来看 534 | 总的来说 535 | 总的说来 536 | 总而言之 537 | 总之 538 | 纵 539 | 纵令 540 | 纵然 541 | 纵使 542 | 遵照 543 | 作为 544 | 兮 545 | 呃 546 | 呗 547 | 咚 548 | 咦 549 | 喏 550 | 啐 551 | 喔唷 552 | 嗬 553 | 嗯 554 | 嗳 555 | 啊哈 556 | 啊呀 557 | 啊哟 558 | 挨次 559 | 挨个 560 | 挨家挨户 561 | 挨门挨户 562 | 挨门逐户 563 | 挨着 564 | 按理 565 | 按期 566 | 按时 567 | 按说 568 | 暗地里 569 | 暗中 570 | 暗自 571 | 昂然 572 | 八成 573 | 白白 574 | 半 575 | 梆 576 | 保管 577 | 保险 578 | 饱 579 | 背地里 580 | 背靠背 581 | 倍感 582 | 倍加 583 | 本人 584 | 本身 585 | 甭 586 | 比起 587 | 比如说 588 | 比照 589 | 毕竟 590 | 必 591 | 必定 592 | 必将 593 | 必须 594 | 便 595 | 别人 596 | 并非 597 | 并肩 598 | 并没 599 | 并没有 600 | 并排 601 | 并无 602 | 勃然 603 | 不 604 | 不必 605 | 不常 606 | 不大 607 | 不但...而且 608 | 不得 609 | 不得不 610 | 不得了 611 | 不得已 612 | 不迭 613 | 不定 614 | 不对 615 | 不妨 616 | 不管怎样 617 | 不会 618 | 不仅...而且 619 | 不仅仅 620 | 不仅仅是 621 | 不经意 622 | 不可开交 623 | 不可抗拒 624 | 不力 625 | 不了 626 | 不料 627 | 不满 628 | 不免 629 | 不能不 630 | 不起 631 | 不巧 632 | 不然的话 633 | 不日 634 | 不少 635 | 不胜 636 | 不时 637 | 不是 638 | 不同 639 | 不能 640 | 不要 641 | 不外 642 | 不外乎 643 | 不下 644 | 不限 645 | 不消 646 | 不已 647 | 不亦乐乎 648 | 不由得 649 | 不再 650 | 不择手段 651 | 不怎么 652 | 不曾 653 | 不知不觉 654 | 不止 655 | 不止一次 656 | 不至于 657 | 才 658 | 才能 659 | 策略地 660 | 差不多 661 | 差一点 662 | 常 663 | 常常 664 | 常言道 665 | 常言说 666 | 常言说得好 667 | 长此下去 668 | 长话短说 669 | 长期以来 670 | 长线 671 | 敞开儿 672 | 彻夜 673 | 陈年 674 | 趁便 675 | 趁机 676 | 趁热 677 | 趁势 678 | 趁早 679 | 成年 680 | 成年累月 681 | 成心 682 | 乘机 683 | 乘胜 684 | 乘势 685 | 乘隙 686 | 乘虚 687 | 诚然 688 | 迟早 689 | 充分 690 | 充其极 691 | 充其量 692 | 抽冷子 693 | 臭 694 | 初 695 | 出 696 | 出来 697 | 出去 698 | 除此 699 | 除此而外 700 | 除此以外 701 | 除开 702 | 除去 703 | 除却 704 | 除外 705 | 处处 706 | 川流不息 707 | 传 708 | 传说 709 | 传闻 710 | 串行 711 | 纯 712 | 纯粹 713 | 此后 714 | 此中 715 | 次第 716 | 匆匆 717 | 从不 718 | 从此 719 | 从此以后 720 | 从古到今 721 | 从古至今 722 | 从今以后 723 | 从宽 724 | 从来 725 | 从轻 726 | 从速 727 | 从头 728 | 从未 729 | 从无到有 730 | 从小 731 | 从新 732 | 从严 733 | 从优 734 | 从早到晚 735 | 从中 736 | 从重 737 | 凑巧 738 | 粗 739 | 存心 740 | 达旦 741 | 打从 742 | 打开天窗说亮话 743 | 大 744 | 大不了 745 | 大大 746 | 大抵 747 | 大都 748 | 大多 749 | 大凡 750 | 大概 751 | 大家 752 | 大举 753 | 大略 754 | 大面儿上 755 | 大事 756 | 大体 757 | 大体上 758 | 大约 759 | 大张旗鼓 760 | 大致 761 | 呆呆地 762 | 带 763 | 殆 764 | 待到 765 | 单 766 | 单纯 767 | 单单 768 | 但愿 769 | 弹指之间 770 | 当场 771 | 当儿 772 | 当即 773 | 当口儿 774 | 当然 775 | 当庭 776 | 当头 777 | 当下 778 | 当真 779 | 当中 780 | 倒不如 781 | 倒不如说 782 | 倒是 783 | 到处 784 | 到底 785 | 到了儿 786 | 到目前为止 787 | 到头 788 | 到头来 789 | 得起 790 | 得天独厚 791 | 的确 792 | 等到 793 | 叮当 794 | 顶多 795 | 定 796 | 动不动 797 | 动辄 798 | 陡然 799 | 都 800 | 独 801 | 独自 802 | 断然 803 | 顿时 804 | 多次 805 | 多多 806 | 多多少少 807 | 多多益善 808 | 多亏 809 | 多年来 810 | 多年前 811 | 而后 812 | 而论 813 | 而又 814 | 尔等 815 | 二话不说 816 | 二话没说 817 | 反倒 818 | 反倒是 819 | 反而 820 | 反手 821 | 反之亦然 822 | 反之则 823 | 方 824 | 方才 825 | 方能 826 | 放量 827 | 非常 828 | 非得 829 | 分期 830 | 分期分批 831 | 分头 832 | 奋勇 833 | 愤然 834 | 风雨无阻 835 | 逢 836 | 弗 837 | 甫 838 | 嘎嘎 839 | 该当 840 | 概 841 | 赶快 842 | 赶早不赶晚 843 | 敢 844 | 敢情 845 | 敢于 846 | 刚 847 | 刚才 848 | 刚好 849 | 刚巧 850 | 高低 851 | 格外 852 | 隔日 853 | 隔夜 854 | 个人 855 | 各式 856 | 更 857 | 更加 858 | 更进一步 859 | 更为 860 | 公然 861 | 共 862 | 共总 863 | 够瞧的 864 | 姑且 865 | 古来 866 | 故而 867 | 故意 868 | 固 869 | 怪 870 | 怪不得 871 | 惯常 872 | 光 873 | 光是 874 | 归根到底 875 | 归根结底 876 | 过于 877 | 毫不 878 | 毫无 879 | 毫无保留地 880 | 毫无例外 881 | 好在 882 | 何必 883 | 何尝 884 | 何妨 885 | 何苦 886 | 何乐而不为 887 | 何须 888 | 何止 889 | 很 890 | 很多 891 | 很少 892 | 轰然 893 | 后来 894 | 呼啦 895 | 忽地 896 | 忽然 897 | 互 898 | 互相 899 | 哗啦 900 | 话说 901 | 还 902 | 恍然 903 | 会 904 | 豁然 905 | 活 906 | 伙同 907 | 或多或少 908 | 或许 909 | 基本 910 | 基本上 911 | 基于 912 | 极 913 | 极大 914 | 极度 915 | 极端 916 | 极力 917 | 极其 918 | 极为 919 | 急匆匆 920 | 即将 921 | 即刻 922 | 即是说 923 | 几度 924 | 几番 925 | 几乎 926 | 几经 927 | 既...又 928 | 继之 929 | 加上 930 | 加以 931 | 间或 932 | 简而言之 933 | 简言之 934 | 简直 935 | 见 936 | 将才 937 | 将近 938 | 将要 939 | 交口 940 | 较比 941 | 较为 942 | 接连不断 943 | 接下来 944 | 皆可 945 | 截然 946 | 截至 947 | 藉以 948 | 借此 949 | 借以 950 | 届时 951 | 仅 952 | 仅仅 953 | 谨 954 | 进来 955 | 进去 956 | 近 957 | 近几年来 958 | 近来 959 | 近年来 960 | 尽管如此 961 | 尽可能 962 | 尽快 963 | 尽量 964 | 尽然 965 | 尽如人意 966 | 尽心竭力 967 | 尽心尽力 968 | 尽早 969 | 精光 970 | 经常 971 | 竟 972 | 竟然 973 | 究竟 974 | 就此 975 | 就地 976 | 就算 977 | 居然 978 | 局外 979 | 举凡 980 | 据称 981 | 据此 982 | 据实 983 | 据说 984 | 据我所知 985 | 据悉 986 | 具体来说 987 | 决不 988 | 决非 989 | 绝 990 | 绝不 991 | 绝顶 992 | 绝对 993 | 绝非 994 | 均 995 | 喀 996 | 看 997 | 看来 998 | 看起来 999 | 看上去 1000 | 看样子 1001 | 可好 1002 | 可能 1003 | 恐怕 1004 | 快 1005 | 快要 1006 | 来不及 1007 | 来得及 1008 | 来讲 1009 | 来看 1010 | 拦腰 1011 | 牢牢 1012 | 老 1013 | 老大 1014 | 老老实实 1015 | 老是 1016 | 累次 1017 | 累年 1018 | 理当 1019 | 理该 1020 | 理应 1021 | 历 1022 | 立 1023 | 立地 1024 | 立刻 1025 | 立马 1026 | 立时 1027 | 联袂 1028 | 连连 1029 | 连日 1030 | 连日来 1031 | 连声 1032 | 连袂 1033 | 临到 1034 | 另方面 1035 | 另行 1036 | 另一个 1037 | 路经 1038 | 屡 1039 | 屡次 1040 | 屡次三番 1041 | 屡屡 1042 | 缕缕 1043 | 率尔 1044 | 率然 1045 | 略 1046 | 略加 1047 | 略微 1048 | 略为 1049 | 论说 1050 | 马上 1051 | 蛮 1052 | 满 1053 | 没 1054 | 没有 1055 | 每逢 1056 | 每每 1057 | 每时每刻 1058 | 猛然 1059 | 猛然间 1060 | 莫 1061 | 莫不 1062 | 莫非 1063 | 莫如 1064 | 默默地 1065 | 默然 1066 | 呐 1067 | 那末 1068 | 奈 1069 | 难道 1070 | 难得 1071 | 难怪 1072 | 难说 1073 | 内 1074 | 年复一年 1075 | 凝神 1076 | 偶而 1077 | 偶尔 1078 | 怕 1079 | 砰 1080 | 碰巧 1081 | 譬如 1082 | 偏偏 1083 | 乒 1084 | 平素 1085 | 颇 1086 | 迫于 1087 | 扑通 1088 | 其后 1089 | 其实 1090 | 奇 1091 | 齐 1092 | 起初 1093 | 起来 1094 | 起首 1095 | 起头 1096 | 起先 1097 | 岂 1098 | 岂非 1099 | 岂止 1100 | 迄 1101 | 恰逢 1102 | 恰好 1103 | 恰恰 1104 | 恰巧 1105 | 恰如 1106 | 恰似 1107 | 千 1108 | 千万 1109 | 千万千万 1110 | 切 1111 | 切不可 1112 | 切莫 1113 | 切切 1114 | 切勿 1115 | 窃 1116 | 亲口 1117 | 亲身 1118 | 亲手 1119 | 亲眼 1120 | 亲自 1121 | 顷 1122 | 顷刻 1123 | 顷刻间 1124 | 顷刻之间 1125 | 请勿 1126 | 穷年累月 1127 | 取道 1128 | 去 1129 | 权时 1130 | 全都 1131 | 全力 1132 | 全年 1133 | 全然 1134 | 全身心 1135 | 然 1136 | 人人 1137 | 仍 1138 | 仍旧 1139 | 仍然 1140 | 日复一日 1141 | 日见 1142 | 日渐 1143 | 日益 1144 | 日臻 1145 | 如常 1146 | 如此等等 1147 | 如次 1148 | 如今 1149 | 如期 1150 | 如前所述 1151 | 如上 1152 | 如下 1153 | 汝 1154 | 三番两次 1155 | 三番五次 1156 | 三天两头 1157 | 瑟瑟 1158 | 沙沙 1159 | 上 1160 | 上来 1161 | 上去 1162 | -------------------------------------------------------------------------------- /docs/benchmarking.md: -------------------------------------------------------------------------------- 1 | 性能测试 2 | ==== 3 | 4 | 测试程序见 [examples/benchmark.go](/examples/benchmark.go) 5 | 6 | 测试数据为从52个微博账号里抓取的十万条微博（请从[这里](https://raw.githubusercontent.com/huichen/wukong/43f20b4c0921cc704cf41fe8653e66a3fcbb7e31/testdata/weibo_data.txt)下载，然后copy到testdata目录），通过benchmark.go中的-num_repeat_text参数（设为10）重复索引为一百万条，500M文本。测试环境Intel(R) Xeon(R) CPU E5-2650 v2 @ 2.60GHz 32核，128G内存。 7 | 8 | 改变测试程序中的NumShards变量可以改变数据单机裂分（sharding）的数目，裂分越多单请求的并发度越高延迟越小，但相应地每秒能处理的总请求数也会变少，比如： 9 | 10 | - 1个shard时：每秒索引1.3M个索引项，响应时间1.65毫秒，吞吐量每秒19.3K次搜索 11 | - 2个shard时：每秒索引1.8M个索引项，响应时间0.87毫秒，吞吐量每秒18.4K次搜索 12 | - 4个shard时：每秒索引1.9M个索引项，响应时间0.56毫秒，吞吐量每秒14.3K次搜索 13 | - 8个shard时：每秒索引2.0M个索引项，响应时间0.39毫秒，吞吐量每秒10.3K次搜索 14 | 15 | 这里的索引项是指一个不重复的“搜索键”-“文档”对，比如当一个文档中有N个不一样的搜索键时，该文档会产生N个索引项。 16 | 17 | 程序默认使用2个shard，你可以根据具体的需求在初始化引擎时改变这个值，见[types.EngineInitOptions.NumShards](/types/engine_init_options.go) 18 | 19 | # 性能分析 20 | 21 | benchmark.go也可以帮助你找到引擎的CPU和内存瓶颈在哪里。 22 | 23 | 分析性能瓶颈： 24 | ``` 25 | go build benchmark.go 26 | ./benchmark -cpuprofile=cpu.prof 27 | go tool pprof benchmark cpu.prof 28 | ``` 29 | 30 | 进入pprof终端后输入web命令可以生成类似下面的图，清晰地表示了每个组件消耗的CPU时间 31 | 32 | ![](https://raw.github.com/huichen/wukong/master/docs/cpu.png) 33 | 34 | 分析内存占用： 35 | ``` 36 | go build benchmark.go 37 | ./benchmark -memprofile=mem.prof 38 | go tool pprof benchmark mem.prof 39 | ``` 40 | 41 | pprof的使用见[这篇文章](http://blog.golang.org/profiling-go-programs)。 42 | -------------------------------------------------------------------------------- /docs/bm25.md: -------------------------------------------------------------------------------- 1 | # 定义 2 | 3 | BM25是搜索引擎的经典排序函数，用于衡量一组关键词和某文档的相关程度。BM25的定义为 4 | 5 | IDF * TF * (k1 + 1) 6 | BM25 = sum ---------------------------- 7 | TF + k1 * (1 - b + b * D / L) 8 | 9 | 其中sum对所有关键词求和，TF（term frequency）为某关键词在该文档中出现的词频，D为该文档的词数，L为所有文档的平均词数，k1和b为常数，在悟空里默认值为2.0和0.75，不过可以在引擎初始化的时候在[EngineInitOptions.IndexerInitOptions.BM25Parameters](/types/indexer_init_options.go)中修改。IDF（inverse document frequency）衡量关键词是否常见，悟空引擎使用带平滑的IDF公式 10 | 11 | 总文档数目 12 | IDF = log2( ------------------------ + 1 ) 13 | 出现该关键词的文档数目 14 | # 使用 15 | 16 | 索引器负责计算BM25，为了能计算文档的BM25值，必须保存文档中所有关键词的词频，这需要在引擎初始化时将[EngineInitOptions.IndexerInitOptions.IndexType](/types/indexer_init_options.go)至少设置为FrequenciesIndex（LocationsIndex也可计算BM25，但这种索引也保存词出现的位置，消耗更多内存）。 17 | 18 | 然后你可以在你[自定义的评分规则](/docs/custom_scoring_criteria.md)中调用IndexedDocument.BM25得到此值作为评分数据。如果你想完全依赖BM25评分，可以使用默认的评分规则，既RankByBM25。 19 | -------------------------------------------------------------------------------- /docs/codelab.md: -------------------------------------------------------------------------------- 1 | 悟空引擎入门 2 | ==== 3 | 4 | 在本篇的结束，你将学会用悟空引擎写一个简单的全文本微博搜索。 5 | 6 | 在阅读本篇之前，你需要对Go语言有基本了解，如果你还不会用Go，[这里](http://go-tour-zh.appspot.com/#1)有个教程。 7 | 8 | ## 引擎的原理 9 | 10 | 引擎中处理用户请求、分词、索引和排序分别由不同的协程（goroutines）完成。 11 | 12 | 1. 主协程，用于收发用户请求 13 | 2. 分词器（segmenter）协程，负责分词 14 | 3. 索引器（indexer）协程，负责建立和查找索引表 15 | 4. 排序器（ranker）协程，负责对文档评分排序 16 | 17 | ![](https://raw.github.com/huichen/wukong/master/docs/wukong.png) 18 | 19 | **索引流程** 20 | 21 | 当一个将文档（document）加入索引的请求进来以后，主协程会通过一个信道（channel）将要分词的文本发送给某个分词协程，该协程将文本分词后通过另一个信道发送给一个索引器协程。索引器协程建立从搜索键（search keyword）到文档的反向索引（inverted index），反向索引表保存在内存中方便快速调用。 22 | 23 | **搜索流程** 24 | 25 | 主协程接到用户请求，将请求短语在主协程内分词，然后通过信道发送给索引器，索引器查找每个搜索键对应的文档然后进行逻辑操作（归并求交集）得到一个精简的文档列表，此列表通过信道传递给排序器，排序器对文档进行评分（scoring）、筛选和排序，然后将排好序的文档通过指定的信道发送给主协程，主协程将结果返回给用户。 26 | 27 | 分词、索引和排序都有多个协程完成，中间结果保存在信道缓冲队列以避免阻塞。为了提高搜索的并发度降低延迟，悟空引擎将文档做了裂分（裂分数目可以由用户指定），索引和排序请求会发送到所有裂分上并行处理，结果在主协程进行二次归并排序。 28 | 29 | 上面就是悟空引擎的大致原理。任何完整的搜索系统都包括四个部分， **文档抓取** 、 **索引** 、 **搜索** 和 **显示** 。下面分别讲解这些部分是怎么实现的。 30 | 31 | ## 文档抓取 32 | 33 | 文档抓取的技术很多，多到可以单独拿出来写一篇文章。幸运的是微博抓取相对简单，可以通过新浪提供的API实现的，而且已经有[Go语言的SDK](http://github.com/huichen/gobo)可以并发抓取并且速度相当快。 34 | 35 | 我已经抓了大概十万篇微博放在了testdata/weibo_data.txt里(因为影响git clone的下载速度所以删除了，请从[这里](https://github.com/huichen/wukong/blob/43f20b4c0921cc704cf41fe8653e66a3fcbb7e31/testdata/weibo_data.txt?raw=true)下载)，所以你就不需要自己做了。文件中每行存储了一篇微博，格式如下 36 | 37 | <微博id>||||<时间戳>||||<用户id>||||<用户名>||||<转贴数>||||<评论数>||||<喜欢数>||||<小图片网址>||||<大图片网址>||||<正文> 38 | 39 | 微博保存在下面的结构体中方便查询，只载入了我们需要的数据： 40 | 41 | ```go 42 | type Weibo struct { 43 | Id uint64 44 | Timestamp uint64 45 | UserName string 46 | RepostsCount uint64 47 | Text string 48 | } 49 | ``` 50 | 51 | 如果你对抓取的细节感兴趣请见抓取程序[testdata/crawl_weibo_data.go](/testdata/crawl_weibo_data.go)。 52 | 53 | ## 索引 54 | 55 | 使用悟空引擎你需要import两个包 56 | 57 | ```go 58 | import ( 59 | "github.com/huichen/wukong/engine" 60 | "github.com/huichen/wukong/types" 61 | ) 62 | ``` 63 | 第一个包定义了引擎功能，第二个包定义了常用结构体。在使用引擎之前需要初始化，例如 64 | 65 | ```go 66 | var searcher engine.Engine 67 | searcher.Init(types.EngineInitOptions{ 68 | SegmenterDictionaries: "../../data/dictionary.txt", 69 | StopTokenFile: "../../data/stop_tokens.txt", 70 | IndexerInitOptions: &types.IndexerInitOptions{ 71 | IndexType: types.LocationsIndex, 72 | }, 73 | }) 74 | ``` 75 | [types.EngineInitOptions](/types/engine_init_options.go)定义了初始化引擎需要设定的参数，比如从何处载入分词字典文件，停用词列表，索引器类型，BM25参数等，以及默认的评分规则（见“搜索”一节）和输出分页选项。具体细节请阅读代码中结构体的注释。 76 | 77 | 特别需要强调的是请慎重选择IndexerInitOptions.IndexType的类型，共有三种不同类型的索引表： 78 | 79 | 1. DocIdsIndex，提供了最基本的索引，仅仅记录搜索键出现的文档docid。 80 | 2. FrequenciesIndex，除了记录docid外，还保存了搜索键在每个文档中出现的频率，如果你需要BM25那么FrequenciesIndex是你需要的。 81 | 3. LocationsIndex，这个不仅包括上两种索引的内容，还额外存储了关键词在文档中的具体位置，这用来[计算紧邻距离](/docs/token_proximity.md)。 82 | 83 | 这三种索引由上到下在提供更多计算能力的同时也消耗了更多的内存，特别是LocationsIndex，当文档很长时会占用大量内存。请根据需要平衡选择。 84 | 85 | 初始化好了以后就可以添加索引了，下面的例子将一条微博加入引擎 86 | 87 | ```go 88 | searcher.IndexDocument(docId, types.DocumentIndexData{ 89 | Content: weibo.Text, // Weibo结构体见上文的定义。必须是UTF-8格式。 90 | Fields: WeiboScoringFields{ 91 | Timestamp: weibo.Timestamp, 92 | RepostsCount: weibo.RepostsCount, 93 | }, 94 | }) 95 | ``` 96 | 97 | 文档的docId必须唯一，对微博来说可以直接用微博的ID。悟空引擎允许你加入三种索引数据： 98 | 99 | 1. 文档的正文（content），会被分词为关键词（tokens）加入索引。 100 | 2. 文档的关键词（tokens）。当正文为空的时候，允许用户绕过悟空内置的分词器直接输入文档关键词，这使得在引擎外部进行文档分词成为可能。 101 | 3. 文档的属性标签（labels），比如微博的作者，类别等。标签并不出现在正文中。 102 | 4. 自定义评分字段（scoring fields），这允许你给文档添加 **任意类型** 、 **任意结构** 的数据用于排序。“搜索”一节会进一步介绍自定义评分字段的用法。 103 | 104 | **特别注意的是** ，关键词（tokens）和标签（labels）组成了索引器中的搜索键（keywords），文档和代码中会反复出现这三个概念，请不要混淆。对正文的搜索就是在搜索键上的逻辑查询，比如一个文档正文中出现了“自行车”这个关键词，也有“健身”这样的分类标签，但“健身”这个词并不直接出现在正文中，当查询“自行车”+“健身”这样的搜索键组合时，这篇文章就会被查询到。设计标签的目的是为了方便从非字面意义的维度快速缩小查询范围。 105 | 106 | 引擎采用了非同步的索引方式，也就是说当IndexDocument返回时索引可能还没有加入索引表中，这方便你循环并发地加入索引。如果你需要等待索引添加完毕后再进行后续操作，请调用下面的函数 107 | 108 | ```go 109 | searcher.FlushIndex() 110 | ``` 111 | 112 | ## 搜索 113 | 114 | 搜索的过程分两步，第一步是在索引表中查找包含搜索键的文档，这在上一节已经介绍过。第二步是对所有索引到的文档进行排序。 115 | 116 | 排序的核心是对文档评分。悟空引擎允许你自定义任意的评分规则（scoring criteria）。在微博搜索例子中，我们定义的评分规则如下： 117 | 118 | 1. 首先按照关键词紧邻距离排序，比如搜索“自行车运动”，这个短语会被切分成两个关键词，“自行车”和“运动”，出现两个关键词紧邻的文章应该排在两个关键词分开的文章前面。 119 | 2. 然后按照微博的发布时间大致排序，每三天为一个梯队，较晚梯队的文章排在前面。 120 | 3. 最后给微博打分为 BM25*(1+转发数/10000) 121 | 122 | 这样的规则需要给每个文档保存一些评分数据，比如微博发布时间，微博的转发数等。这些数据保存在下面的结构体中 123 | ```go 124 | type WeiboScoringFields struct { 125 | Timestamp uint64 126 | RepostsCount uint64 127 | } 128 | ``` 129 | 你可能已经注意到了，这就是在上一节将文档加入索引时调用的IndexDocument函数传入的参数类型（实际上那个参数是interface{}类型的，因此可以传入任意类型的结构体）。这些数据保存在排序器的内存中等待调用。 130 | 131 | 有了这些数据，我们就可以评分了，代码如下： 132 | ```go 133 | type WeiboScoringCriteria struct { 134 | } 135 | 136 | func (criteria WeiboScoringCriteria) Score( 137 | doc types.IndexedDocument, fields interface{}) []float32 { 138 | if reflect.TypeOf(fields) != reflect.TypeOf(WeiboScoringFields{}) { 139 | return []float32{} 140 | } 141 | wsf := fields.(WeiboScoringFields) 142 | output := make([]float32, 3) 143 | if doc.TokenProximity > MaxTokenProximity { // 第一步 144 | output[0] = 1.0 / float32(doc.TokenProximity) 145 | } else { 146 | output[0] = 1.0 147 | } 148 | output[1] = float32(wsf.Timestamp / (SecondsInADay * 3)) // 第二步 149 | output[2] = float32(doc.BM25 * (1 + float32(wsf.RepostsCount)/10000)) // 第三步 150 | return output 151 | } 152 | ``` 153 | WeiboScoringCriteria实际上继承了types.ScoringCriteria接口，这个接口实现了Score函数。这个函数带有两个参数： 154 | 155 | 1. types.IndexedDocument参数传递了从索引器中得到的数据，比如词频，词的具体位置，BM25值，紧邻度等信息，具体见[types/index.go](/types/index.go)的注释。 156 | 2. 第二个参数是interface{}类型的，你可以把这个类型理解成C语言中的void指针，它可以指向任何数据类型。在我们的例子中指向的是WeiboScoringFields结构体，并通过反射机制检查是否是正确的类型。 157 | 158 | 有了自定义评分数据和自定义评分规则，我们就可以进行搜索了，见下面的代码 159 | 160 | ```go 161 | response := searcher.Search(types.SearchRequest{ 162 | Text: "自行车运动", 163 | RankOptions: &types.RankOptions{ 164 | ScoringCriteria: &WeiboScoringCriteria{}, 165 | OutputOffset: 0, 166 | MaxOutputs: 100, 167 | }, 168 | }) 169 | ``` 170 | 171 | 其中，Text是输入的搜索短语（必须是UTF-8格式），会被分词为关键词。和索引时相同，悟空引擎允许绕过内置的分词器直接输入关键词和文档标签，见types.SearchRequest结构体的注释。RankOptions定义了排序选项。WeiboScoringCriteria就是我们在上面定义的评分规则。另外你也可以通过OutputOffset和MaxOutputs参数控制分页输出。搜索结果保存在response变量中，具体内容见[types/search_response.go](/types/search_response.go)文件中定义的SearchResponse结构体，比如这个结构体返回了关键词出现在文档中的位置，可以用来生成文档的摘要。 172 | 173 | ## 显示 174 | 175 | 完成用户搜索的最后一步是将搜索结果呈现给用户。 通常的做法是将搜索引擎做成一个后台服务，然后让前端以JSON-RPC的方式调用它。前端并不属于悟空引擎本身因此就不多着墨了。 176 | 177 | ## 总结 178 | 179 | 读到这里，你应该对使用悟空引擎进行微博搜索有了基本了解，建议你自己动手将其完成。如果你没有耐心，可以看看已经完成的代码，见[examples/codelab/search_server.go](/examples/codelab/search_server.go)，总共不到200行。运行这个例子非常简单，进入examples/codelab目录后输入 180 | 181 | go run search_server.go 182 | 183 | 等待终端中出现“索引了xxx条微博”的输出后，在浏览器中打开[http://localhost:8080](http://localhost:8080) 即可进入搜索页面。 184 | 185 | 如果你想进一步了解悟空引擎，建议你直接阅读代码。代码的目录结构如下： 186 | 187 | /core 核心部件，包括索引器和排序器 188 | /data 字典文件和停用词文件 189 | /docs 文档 190 | /engine 引擎，包括主协程、分词协程、索引器协程和排序器协程的实现 191 | /examples 例子和性能测试程序 192 | /testdata 测试数据 193 | /types 常用结构体 194 | /utils 常用函数 195 | -------------------------------------------------------------------------------- /docs/cpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/andeya/wukong/d12b490f6e19b236b9f34e19259480e8fff0868b/docs/cpu.png -------------------------------------------------------------------------------- /docs/custom_scoring_criteria.md: -------------------------------------------------------------------------------- 1 | 自定义评分字段和评分规则 2 | === 3 | 4 | 悟空引擎支持在排序器内存中保存一些评分字段，并利用自定义评分规则给文档打分。例子: 5 | 6 | ```go 7 | // 自定义评分字段 8 | type MyScoringFields struct { 9 | // 加入一些文档排序用的数据，可以是任意类型，比如： 10 | label string 11 | counter int32 12 | someScore float32 13 | } 14 | 15 | // MyScoringCriteria实现了types.ScoringCriteria接口，也就是下面的Score函数 16 | type MyScoringCriteria struct { 17 | } 18 | func (criteria MyScoringCriteria) Score( 19 | doc types.IndexedDocument, fields interface{}) []float32 { 20 | // 首先检查评分字段是否为MyScoringFields类型的，如果不是则返回空切片，此文档将从结果中剔除 21 | if reflect.TypeOf(fields) != reflect.TypeOf(MySearchFields{}) { 22 | return []float32{} 23 | } 24 | 25 | // 匹配则进行类型转换 26 | myFields := fields.(MySearchFields) 27 | 28 | // 下面利用myFields中的数据给文档评分并返回分值 29 | } 30 | ``` 31 | 32 | 文档的MyScoringFields数据通过engine.Engine的IndexDocument函数传给排序器保存在内存中。然后通过Search函数的参数调用MyScoringCriteria进行查询。 33 | 34 | 当然，MyScoringCriteria的Score函数也可以通过docId从硬盘或数据库读取更多文档数据用于打分，但速度要比从内存中直接读慢许多，请在内存和速度之间合适取舍。 35 | 36 | [examples/custom_scoring_criteria.go](/examples/custom_scoring_criteria.go)中包含了一个利用自定义规则查询微博数据的例子。 37 | -------------------------------------------------------------------------------- /docs/distributed_indexing_and_search.md: -------------------------------------------------------------------------------- 1 | 分布式索引和搜索 2 | === 3 | 4 | 分布式搜索的原理如下： 5 | 6 | 当文档数量较多无法在一台机器内存中索引时，可以将文档按照文本内容的hash值裂分(sharding)，不同块交由不同服务器索引。在查找时同一请求分发到所有裂分服务器上，然后将所有服务器返回的结果归并重排序作为最终搜索结果输出。 7 | 8 | 为了保证裂分的均匀性，建议使用Go语言实现的Murmur3 hash函数: 9 | 10 | https://github.com/huichen/murmur 11 | 12 | 按照上面的原理很容易用悟空引擎实现分布式搜索（每个裂分服务器运行一个悟空引擎），但这样的分布式系统多数是高度定制的，比如任务的调度依赖于分布式环境，有时需要添加额外层的服务器以均衡负载，因此就不在这里实现了。 13 | -------------------------------------------------------------------------------- /docs/docker.md: -------------------------------------------------------------------------------- 1 | docker支持 2 | === 3 | 4 | ## 从源代码build docker镜像 5 | 6 | 1、请从[这里](https://github.com/huichen/wukong/blob/43f20b4c0921cc704cf41fe8653e66a3fcbb7e31/testdata/weibo_data.txt?raw=true)下载weibo_data.txt，放在testdata/目录下 7 | 8 | 2、进入examples/codelab目录 9 | 10 | 3、建立docker image 11 | 12 | ./build_docker_image.sh 13 | 14 | 4、运行docker container 15 | 16 | docker run -d -p 8080:8080 unmerged/wukong-codelab 17 | 18 | 在浏览器中打开 localhost:8080 即可打开搜索页面 19 | 20 | ## 直接从docker hub下载镜像 21 | 22 | 我已经建好了一个repo并上传到了docker hub，用下面的命令pull镜像 23 | 24 | docker pull unmerged/wukong-codelab 25 | 26 | 下载后运行镜像的方法和上面的相同。 27 | -------------------------------------------------------------------------------- /docs/feedback.md: -------------------------------------------------------------------------------- 1 | 如果你想使用这个引擎，除了需要遵守Apache License v2外没有任何限制，欢迎商业应用。 2 | 3 | 如果你想了解开发规划或者和我讨论这个引擎，请发信给我 unmerged [at] gmail [dot] com 4 | -------------------------------------------------------------------------------- /docs/persistent_storage.md: -------------------------------------------------------------------------------- 1 | 持久存储 2 | ==== 3 | 4 | 悟空引擎支持将搜索数据存入硬盘，并在当机重启动时从硬盘恢复数据。使用持久存储只需设置EngineInitOptions中的三个选项： 5 | 6 | ```go 7 | type EngineInitOptions struct { 8 | // 略过其他选项 9 | 10 | // 是否使用持久数据库，以及数据库文件保存的目录和裂分数目 11 | UsePersistentStorage bool 12 | PersistentStorageFolder string 13 | PersistentStorageShards int 14 | } 15 | ``` 16 | 17 | 当UsePersistentStorage为true时使用持久存储： 18 | 19 | 1. 在引擎启动时（engine.Init函数），引擎从PersistentStorageFolder指定的目录中读取 20 | 文档索引数据，重新计算索引表并给排序器注入排序数据。如果分词器的代码 21 | 或者词典有变化，这些变化会体现在启动后的引擎索引表中。 22 | 2. 在调用engine.IndexDocument时，引擎将索引数据写入到PersistentStorageFolder指定 23 | 的目录中。 24 | 3. PersistentStorageShards定义了数据库裂分数目，默认为8。为了得到最好的性能，请调整这个参数使得每个裂分文件小于100M。 25 | 4. 在调用engine.RemoveDocument删除一个文档后，该文档会从持久存储中剔除，下次启动 26 | 引擎时不会载入该文档。 27 | 28 | 29 | ### 必须注意事项 30 | 31 | 一、如果排序器使用[自定义评分字段](/docs/custom_scoring_criteria.md)，那么该类型必须在gob中注册，比如在左边的例子中需要在调用engine.Init前加入： 32 | ``` 33 | gob.Register(MyScoringFields{}) 34 | ``` 35 | 否则程序会崩溃。 36 | 37 | 二、在引擎退出时请使用engine.Close()来关闭数据库，如果数据库未关闭，数据库文件会被锁定， 38 | 这会导致引擎重启失败。解锁的方法是，进入PersistentStorageFolder指定的目录，删除所有以"."开头的文件即可。 39 | 40 | ### 性能测试 41 | 42 | [benchmark.go](/examples/benchmark.go)程序可用来测试持久存储的读写速度： 43 | 44 | ```go 45 | go run benchmark.go --num_repeat_text 1 --use_persistent 46 | ``` 47 | 48 | 不使用持久存储的结果： 49 | 50 | ``` 51 | 加入的索引总数3711375 52 | 建立索引花费时间 3.159781129s 53 | 建立索引速度每秒添加 1.174567 百万个索引 54 | 搜索平均响应时间 0.051146982400000006 毫秒 55 | 搜索吞吐量每秒 78205.98229466612 次查询 56 | ``` 57 | 58 | 使用持久存储： 59 | 60 | ``` 61 | 建立索引花费时间 13.068458511s 62 | 建立索引速度每秒添加 0.283995 百万个索引 63 | 搜索平均响应时间 0.05819595780000001 毫秒 64 | 搜索吞吐量每秒 68733.29611219149 次查询 65 | 从持久存储加入的索引总数3711375 66 | 从持久存储建立索引花费时间 6.406528378999999 67 | 从持久存储建立索引速度每秒添加 0.579311 百万个索引 68 | ``` 69 | 70 | 可以看出，和不使用持久存储相比： 71 | 72 | 1. 持久存储不影响搜索响应时间和吞吐量 73 | 2. 写入持久存储将索引时间延长了四倍 74 | 3. 从持久存储中导入数据将索引时间延长了一倍 75 | -------------------------------------------------------------------------------- /docs/realtime_indexing.md: -------------------------------------------------------------------------------- 1 | ## 动态修改索引表 2 | 3 | 悟空引擎支持搜索的同时添加索引（engine.IndexDocument函数），但由于添加索引时会对索引表进行写锁定，因此在添加索引的同时搜索性能会有所下降。请控制添加操作的频率，或者将大量添加操作转移到引擎比较空闲时进行。 4 | 5 | 删除一条文档（engine.RemoveDocument函数）也有同样的问题。但是删除操作不会对索引表进行修改，仅仅从排序器中删除该文档的自定义评分字段。因此，在悟空引擎上做大量的删除操作是内存低效的。当删除操作很频繁时，比如和添加操作的频率接近，建议周期性地重启引擎进行索引表重建。 6 | -------------------------------------------------------------------------------- /docs/token_proximity.md: -------------------------------------------------------------------------------- 1 | 关键词紧邻距离（Token Proximity） 2 | === 3 | 4 | 关键词紧邻距离用来衡量多个关键词在同一文档中是否相邻。比如用户搜索“中国足球”这一短语，包含“中国”和“足球”两个关键词，当这两个关键词按照同样顺序前后紧挨着出现在一个文档中时，紧邻距离为零，如果两词中间夹入很多词则紧邻距离较大。紧邻距离是一种衡量文档和多个关键词相关度的方法。紧邻距离虽然不应该作为给文档排序的唯一指标，但在一些情况下通过设定阈值可以过滤掉相当一部分无关的结果。 5 | 6 | N关键词的紧邻距离计算公式如下： 7 | 8 | 假定第i个关键词首字节出现在文本中的位置为P_i，长度L_i，紧邻距离为 9 | 10 | ArgMin(Sum(Abs(P_(i+1) - P_i - L_i))) 11 | 12 | 具体计算过程为先取定一个P_1，计算所有P_2的可能值中令Abs(P_2 - P_1 - L1)最小，然后固定P2后依照同样的方法选择P3，P4，等等。遍历所有可能的P_1得到最小值。 13 | 14 | 具体实现见[core/indexer.go](/core/indexer.go)文件中computeTokenProximity函数。 15 | 16 | 紧邻距离计算需要在索引器中保存每个分词的位置，这需要额外消耗内存，因此是默认关闭的，打开这一功能请在引擎初始化时设定 EngineInitOptions.IndexerInitOptions.IndexType为LocationsIndex。 17 | -------------------------------------------------------------------------------- /docs/why_wukong.md: -------------------------------------------------------------------------------- 1 | 创建悟空搜索引擎项目的初衷有两个。 2 | 3 | ### 知识应该容易获得 4 | 5 | 在网上浏览经常发现一些优秀的网站没有像样的内容搜索服务。网站的建设者投入巨大的时间和精力，为读者创造了丰富的阅读内容，却因为缺少好的搜索引擎极大降低了信息的流动性，新的内容源源不断被推送到首页，而旧的内容却被积压在N页之后变得默默无闻。这种基于时间轴的推荐，以及其它的种种推荐系统，其实都是在忽视和打压用户作为信息搜寻者的主观能动性。一个普遍现象是，很少有网站具备一个易用的、即时的、全文的和高质量的搜索引擎，久而久之，这让用户忘记了自己应该有这种主动搜寻信息的能力，于是用户不再在阅读时搜索，网站搜索缺乏搜索流量，这形成恶性循环。 6 | 7 | 通用的搜索引擎并不能解决这个问题：首先，不能有针对性的垂直搜索，对内容的抓取和索引简单粗暴，无法得到结构化数据，看不到隐藏的内容属性；其次，通用搜索引擎对内容的排序无法进行定制，实时性不够，无法成为网站社区的有机的一部分。 8 | 9 | 一个好的网络社区应该允许用户在阅读时搜索，搜索成为阅读的一部分；网站内部的流量应该是个性化的、有机的，而不应该由生硬的推荐系统主导用户阅读的方向。很可惜这样的网络社区很少，这种局面的造成一部分是因为缺乏一个开源的、容易定制的、高效的搜索引擎的存在。 10 | 11 | 这个项目就是为了实现这样理想的一个尝试。 12 | 13 | ### 搜索引擎技术应该普及 14 | 15 | 另一个目的就是工程上的好奇心。所有的软件工程师都应该对搜索引擎技术有好奇心，因为这是信息检索技术的核心，是将互联网的知识整合起来的关键，稍加变化可以用来解决较小规模的相关问题；所有致力于将信息提供给用户的工程师都应当使用搜索引擎技术，帮助普通人更容易地获得他们想要的知识；所有大数据领域的工程师都应该在职业生涯的某个时刻试着重写一个搜索引擎，了解一下其中的技术问题，并尝试去解决这些问题。 16 | 17 | 这个项目就是为了满足这样好奇心的一个实验。 18 | -------------------------------------------------------------------------------- /docs/wukong.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/andeya/wukong/d12b490f6e19b236b9f34e19259480e8fff0868b/docs/wukong.png -------------------------------------------------------------------------------- /engine/counters.go: -------------------------------------------------------------------------------- 1 | package engine 2 | 3 | func (engine *Engine) NumTokenIndexAdded() uint64 { 4 | return engine.numTokenIndexAdded 5 | } 6 | 7 | func (engine *Engine) NumDocumentsIndexed() uint64 { 8 | return engine.numDocumentsIndexed 9 | } 10 | -------------------------------------------------------------------------------- /engine/engine.go: -------------------------------------------------------------------------------- 1 | package engine 2 | 3 | import ( 4 | // "encoding/json" 5 | "fmt" 6 | "github.com/huichen/murmur" 7 | "github.com/huichen/sego" 8 | "github.com/huichen/wukong/core" 9 | "github.com/huichen/wukong/storage" 10 | "github.com/huichen/wukong/types" 11 | "github.com/huichen/wukong/utils" 12 | "log" 13 | "os" 14 | "runtime" 15 | "sort" 16 | "strconv" 17 | "sync/atomic" 18 | "time" 19 | ) 20 | 21 | const ( 22 | PersistentStorageFilePrefix = "wukong" 23 | ) 24 | 25 | type Engine struct { 26 | // 计数器，用来统计有多少文档被索引等信息 27 | numDocumentsIndexed uint64 28 | numIndexingRequests uint64 29 | numTokenIndexAdded uint64 30 | numDocumentsStored uint64 31 | 32 | // 记录初始化参数 33 | initOptions types.EngineInitOptions 34 | initialized bool 35 | 36 | indexers []core.Indexer 37 | rankers []core.Ranker 38 | segmenter sego.Segmenter 39 | stopTokens StopTokens 40 | // 数据库实例[shard][info/index]db 41 | dbs [][2]storage.Storage 42 | 43 | // 建立分词器使用的通信通道 44 | segmenterChannel chan segmenterRequest 45 | 46 | // 建立索引器使用的通信通道 47 | indexerAddDocumentChannels []chan indexerAddDocumentRequest 48 | indexerLookupChannels []chan indexerLookupRequest 49 | indexerRemoveDocChannels []chan indexerRemoveDocRequest 50 | 51 | // 建立排序器使用的通信通道 52 | rankerAddDocChannels []chan rankerAddDocRequest 53 | rankerRankChannels []chan rankerRankRequest 54 | rankerRemoveDocChannels []chan rankerRemoveDocRequest 55 | 56 | // 建立持久存储使用的通信通道 57 | persistentStorageIndexDocumentChannels []chan persistentStorageIndexDocumentRequest 58 | persistentStorageInitChannel chan bool 59 | } 60 | 61 | func (engine *Engine) Init(options types.EngineInitOptions) { 62 | // 初始化初始参数 63 | if engine.initialized { 64 | log.Fatal("请勿重复初始化引擎") 65 | } 66 | 67 | // 将线程数设置为CPU数 68 | runtime.GOMAXPROCS(runtime.NumCPU()) 69 | options.Init() 70 | engine.initOptions = options 71 | engine.initialized = true 72 | 73 | // 载入分词器词典 74 | engine.segmenter.LoadDictionary(options.SegmenterDictionaries) 75 | 76 | // 初始化停用词 77 | engine.stopTokens.Init(options.StopTokenFile) 78 | 79 | // 初始化索引器和排序器 80 | for shard := 0; shard < options.NumShards; shard++ { 81 | engine.indexers = append(engine.indexers, core.Indexer{}) 82 | engine.indexers[shard].Init(shard, *options.IndexerInitOptions) 83 | 84 | engine.rankers = append(engine.rankers, core.Ranker{}) 85 | engine.rankers[shard].Init(shard) 86 | } 87 | 88 | // 初始化分词器通道 89 | engine.segmenterChannel = make( 90 | chan segmenterRequest, options.NumSegmenterThreads) 91 | 92 | // 初始化索引器通道 93 | engine.indexerAddDocumentChannels = make( 94 | []chan indexerAddDocumentRequest, options.NumShards) 95 | engine.indexerRemoveDocChannels = make( 96 | []chan indexerRemoveDocRequest, options.NumShards) 97 | engine.indexerLookupChannels = make( 98 | []chan indexerLookupRequest, options.NumShards) 99 | for shard := 0; shard < options.NumShards; shard++ { 100 | engine.indexerAddDocumentChannels[shard] = make( 101 | chan indexerAddDocumentRequest, 102 | options.IndexerBufferLength) 103 | engine.indexerRemoveDocChannels[shard] = make( 104 | chan indexerRemoveDocRequest, 105 | options.IndexerBufferLength) 106 | engine.indexerLookupChannels[shard] = make( 107 | chan indexerLookupRequest, 108 | options.IndexerBufferLength) 109 | } 110 | 111 | // 初始化排序器通道 112 | engine.rankerAddDocChannels = make( 113 | []chan rankerAddDocRequest, options.NumShards) 114 | engine.rankerRankChannels = make( 115 | []chan rankerRankRequest, options.NumShards) 116 | engine.rankerRemoveDocChannels = make( 117 | []chan rankerRemoveDocRequest, options.NumShards) 118 | for shard := 0; shard < options.NumShards; shard++ { 119 | engine.rankerAddDocChannels[shard] = make( 120 | chan rankerAddDocRequest, 121 | options.RankerBufferLength) 122 | engine.rankerRankChannels[shard] = make( 123 | chan rankerRankRequest, 124 | options.RankerBufferLength) 125 | engine.rankerRemoveDocChannels[shard] = make( 126 | chan rankerRemoveDocRequest, 127 | options.RankerBufferLength) 128 | } 129 | 130 | // 初始化持久化存储通道 131 | if engine.initOptions.UsePersistentStorage { 132 | engine.persistentStorageIndexDocumentChannels = 133 | make([]chan persistentStorageIndexDocumentRequest, 134 | engine.initOptions.NumShards) 135 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 136 | engine.persistentStorageIndexDocumentChannels[shard] = make( 137 | chan persistentStorageIndexDocumentRequest) 138 | } 139 | engine.persistentStorageInitChannel = make( 140 | chan bool, engine.initOptions.NumShards) 141 | } 142 | 143 | // 启动分词器 144 | for iThread := 0; iThread < options.NumSegmenterThreads; iThread++ { 145 | go engine.segmenterWorker() 146 | } 147 | 148 | // 启动索引器和排序器 149 | for shard := 0; shard < options.NumShards; shard++ { 150 | go engine.indexerAddDocumentWorker(shard) 151 | go engine.indexerRemoveDocWorker(shard) 152 | go engine.rankerAddDocWorker(shard) 153 | go engine.rankerRemoveDocWorker(shard) 154 | 155 | for i := 0; i < options.NumIndexerThreadsPerShard; i++ { 156 | go engine.indexerLookupWorker(shard) 157 | } 158 | for i := 0; i < options.NumRankerThreadsPerShard; i++ { 159 | go engine.rankerRankWorker(shard) 160 | } 161 | } 162 | 163 | // 启动持久化存储工作协程 164 | if engine.initOptions.UsePersistentStorage { 165 | err := os.MkdirAll(engine.initOptions.PersistentStorageFolder, 0700) 166 | if err != nil { 167 | log.Fatal("无法创建目录", engine.initOptions.PersistentStorageFolder) 168 | } 169 | 170 | // 打开或者创建数据库 171 | engine.dbs = make([][2]storage.Storage, engine.initOptions.NumShards) 172 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 173 | dbPathInfo := engine.initOptions.PersistentStorageFolder + "/" + PersistentStorageFilePrefix + ".info." + strconv.Itoa(shard) 174 | dbInfo, err := storage.OpenStorage(dbPathInfo) 175 | if dbInfo == nil || err != nil { 176 | log.Fatal("无法打开数据库", dbPathInfo, ": ", err) 177 | } 178 | dbPathIndex := engine.initOptions.PersistentStorageFolder + "/" + PersistentStorageFilePrefix + ".index." + strconv.Itoa(shard) 179 | dbIndex, err := storage.OpenStorage(dbPathIndex) 180 | if dbIndex == nil || err != nil { 181 | log.Fatal("无法打开数据库", dbPathIndex, ": ", err) 182 | } 183 | engine.dbs[shard][getDB("info")] = dbInfo 184 | engine.dbs[shard][getDB("index")] = dbIndex 185 | 186 | } 187 | 188 | // 从数据库中恢复 189 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 190 | go engine.persistentStorageInitWorker(shard) 191 | } 192 | 193 | // 等待恢复完成 194 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 195 | <-engine.persistentStorageInitChannel 196 | } 197 | 198 | // 关闭并重新打开数据库 199 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 200 | engine.dbs[shard][0].Close() 201 | engine.dbs[shard][1].Close() 202 | dbPathInfo := engine.initOptions.PersistentStorageFolder + "/" + PersistentStorageFilePrefix + ".info." + strconv.Itoa(shard) 203 | dbInfo, err := storage.OpenStorage(dbPathInfo) 204 | if dbInfo == nil || err != nil { 205 | log.Fatal("无法打开数据库", dbPathInfo, ": ", err) 206 | } 207 | dbPathIndex := engine.initOptions.PersistentStorageFolder + "/" + PersistentStorageFilePrefix + ".index." + strconv.Itoa(shard) 208 | dbIndex, err := storage.OpenStorage(dbPathIndex) 209 | if dbIndex == nil || err != nil { 210 | log.Fatal("无法打开数据库", dbPathIndex, ": ", err) 211 | } 212 | engine.dbs[shard][getDB("info")] = dbInfo 213 | engine.dbs[shard][getDB("index")] = dbIndex 214 | } 215 | 216 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 217 | go engine.persistentStorageIndexDocumentWorker(shard) 218 | } 219 | } 220 | } 221 | 222 | // 将文档加入索引 223 | // 224 | // 输入参数： 225 | // docId 标识文档编号，必须唯一 226 | // data 见DocumentIndexData注释 227 | // 228 | // 注意： 229 | // 1. 这个函数是线程安全的，请尽可能并发调用以提高索引速度 230 | // 2. 这个函数调用是非同步的，也就是说在函数返回时有可能文档还没有加入索引中，因此 231 | // 如果立刻调用Search可能无法查询到这个文档。强制刷新索引请调用FlushIndex函数。 232 | func (engine *Engine) IndexDocument(docId uint64, data types.DocumentIndexData) { 233 | if !engine.initialized { 234 | log.Fatal("必须先初始化引擎") 235 | } 236 | atomic.AddUint64(&engine.numIndexingRequests, 1) 237 | shard := int(murmur.Murmur3([]byte(fmt.Sprint("%d", docId))) % uint32(engine.initOptions.NumShards)) 238 | engine.segmenterChannel <- segmenterRequest{ 239 | docId: docId, shard: shard, data: data} 240 | } 241 | 242 | // 只分词与过滤弃用词 243 | func (engine *Engine) Segment(content string) (keywords []string) { 244 | segments := engine.segmenter.Segment([]byte(content)) 245 | for _, segment := range segments { 246 | token := segment.Token().Text() 247 | if !engine.stopTokens.IsStopToken(token) { 248 | keywords = append(keywords, token) 249 | } 250 | } 251 | return 252 | } 253 | 254 | // 将文档从索引中删除 255 | // 256 | // 输入参数： 257 | // docId 标识文档编号，必须唯一 258 | // 259 | // 注意：这个函数仅从排序器中删除文档，索引器不会发生变化。 260 | func (engine *Engine) RemoveDocument(docId uint64) { 261 | if !engine.initialized { 262 | log.Fatal("必须先初始化引擎") 263 | } 264 | 265 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 266 | engine.indexerRemoveDocChannels[shard] <- indexerRemoveDocRequest{docId: docId} 267 | engine.rankerRemoveDocChannels[shard] <- rankerRemoveDocRequest{docId: docId} 268 | } 269 | 270 | if engine.initOptions.UsePersistentStorage { 271 | // 从数据库中删除 272 | shard := int(murmur.Murmur3([]byte(fmt.Sprint("%d", docId)))) % engine.initOptions.NumShards 273 | go engine.persistentStorageRemoveDocumentWorker(docId, shard) 274 | } 275 | } 276 | 277 | // 阻塞等待直到所有索引添加完毕 278 | func (engine *Engine) FlushIndex() { 279 | for { 280 | runtime.Gosched() 281 | if engine.numIndexingRequests == engine.numDocumentsIndexed && 282 | (!engine.initOptions.UsePersistentStorage || 283 | engine.numIndexingRequests == engine.numDocumentsStored) { 284 | return 285 | } 286 | } 287 | } 288 | 289 | // 查找满足搜索条件的文档，此函数线程安全 290 | func (engine *Engine) Search(request types.SearchRequest) (output types.SearchResponse) { 291 | if !engine.initialized { 292 | log.Fatal("必须先初始化引擎") 293 | } 294 | 295 | // for k, s := range core.DocInfoGroup { 296 | // log.Printf("DocInfo:%v,%v,%v\n", k, s.NumDocuments, s.DocInfos) 297 | // } 298 | // for k, s := range core.InvertedIndexGroup { 299 | // b, _ := json.Marshal(s.InvertedIndex) 300 | // log.Printf("InvertedIndex:%v,%v,%+v\n", k, s.TotalTokenLength, string(b)) 301 | // } 302 | 303 | // for k, s := range core.DocInfoGroup { 304 | // log.Printf("DocInfo:%v,%v\n", k, s.NumDocuments) 305 | // } 306 | // for k, s := range core.InvertedIndexGroup { 307 | // log.Printf("InvertedIndex:%#v,%v\n", k, s.TotalTokenLength) 308 | // } 309 | 310 | var rankOptions types.RankOptions 311 | if request.RankOptions == nil { 312 | rankOptions = *engine.initOptions.DefaultRankOptions 313 | } else { 314 | rankOptions = *request.RankOptions 315 | } 316 | if rankOptions.ScoringCriteria == nil { 317 | rankOptions.ScoringCriteria = engine.initOptions.DefaultRankOptions.ScoringCriteria 318 | } 319 | 320 | // 收集关键词 321 | tokens := []string{} 322 | if request.Text != "" { 323 | querySegments := engine.segmenter.Segment([]byte(request.Text)) 324 | for _, s := range querySegments { 325 | token := s.Token().Text() 326 | if !engine.stopTokens.IsStopToken(token) { 327 | tokens = append(tokens, s.Token().Text()) 328 | } 329 | } 330 | } else { 331 | for _, t := range request.Tokens { 332 | tokens = append(tokens, t) 333 | } 334 | } 335 | 336 | // 建立排序器返回的通信通道 337 | rankerReturnChannel := make( 338 | chan rankerReturnRequest, engine.initOptions.NumShards) 339 | 340 | // 生成查找请求 341 | lookupRequest := indexerLookupRequest{ 342 | countDocsOnly: request.CountDocsOnly, 343 | tokens: tokens, 344 | labels: request.Labels, 345 | docIds: request.DocIds, 346 | options: rankOptions, 347 | rankerReturnChannel: rankerReturnChannel, 348 | orderless: request.Orderless, 349 | } 350 | 351 | // 向索引器发送查找请求 352 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 353 | engine.indexerLookupChannels[shard] <- lookupRequest 354 | } 355 | 356 | // 从通信通道读取排序器的输出 357 | numDocs := 0 358 | rankOutput := types.ScoredDocuments{} 359 | timeout := request.Timeout 360 | isTimeout := false 361 | if timeout <= 0 { 362 | // 不设置超时 363 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 364 | rankerOutput := <-rankerReturnChannel 365 | if !request.CountDocsOnly { 366 | for _, doc := range rankerOutput.docs { 367 | rankOutput = append(rankOutput, doc) 368 | } 369 | } 370 | numDocs += rankerOutput.numDocs 371 | } 372 | } else { 373 | // 设置超时 374 | deadline := time.Now().Add(time.Millisecond * time.Duration(request.Timeout)) 375 | for shard := 0; shard < engine.initOptions.NumShards; shard++ { 376 | select { 377 | case rankerOutput := <-rankerReturnChannel: 378 | if !request.CountDocsOnly { 379 | for _, doc := range rankerOutput.docs { 380 | rankOutput = append(rankOutput, doc) 381 | } 382 | } 383 | numDocs += rankerOutput.numDocs 384 | case <-time.After(deadline.Sub(time.Now())): 385 | isTimeout = true 386 | break 387 | } 388 | } 389 | } 390 | 391 | // 再排序 392 | if !request.CountDocsOnly && !request.Orderless { 393 | if rankOptions.ReverseOrder { 394 | sort.Sort(sort.Reverse(rankOutput)) 395 | } else { 396 | sort.Sort(rankOutput) 397 | } 398 | } 399 | 400 | // 准备输出 401 | output.Tokens = tokens 402 | // 仅当CountDocsOnly为false时才充填output.Docs 403 | if !request.CountDocsOnly { 404 | if request.Orderless { 405 | // 无序状态无需对Offset截断 406 | output.Docs = rankOutput 407 | } else { 408 | var start, end int 409 | if rankOptions.MaxOutputs == 0 { 410 | start = utils.MinInt(rankOptions.OutputOffset, len(rankOutput)) 411 | end = len(rankOutput) 412 | } else { 413 | start = utils.MinInt(rankOptions.OutputOffset, len(rankOutput)) 414 | end = utils.MinInt(start+rankOptions.MaxOutputs, len(rankOutput)) 415 | } 416 | output.Docs = rankOutput[start:end] 417 | } 418 | } 419 | output.NumDocs = numDocs 420 | output.Timeout = isTimeout 421 | return 422 | } 423 | 424 | // 关闭引擎 425 | func (engine *Engine) Close() { 426 | engine.FlushIndex() 427 | core.DocInfoGroup = make(map[int]*types.DocInfosShard) 428 | core.InvertedIndexGroup = make(map[int]*types.InvertedIndexShard) 429 | if engine.initOptions.UsePersistentStorage { 430 | for _, db := range engine.dbs { 431 | db[0].Close() 432 | db[1].Close() 433 | } 434 | } 435 | } 436 | 437 | // 获取数据库类别索引 438 | func getDB(typ string) int { 439 | switch typ { 440 | case "info": 441 | return 0 442 | case "index": 443 | return 1 444 | } 445 | log.Fatal("数据库类别不正确") 446 | return 0 447 | } 448 | -------------------------------------------------------------------------------- /engine/engine_test.go: -------------------------------------------------------------------------------- 1 | package engine 2 | 3 | import ( 4 | "encoding/gob" 5 | "github.com/huichen/wukong/core" 6 | "github.com/huichen/wukong/types" 7 | "github.com/huichen/wukong/utils" 8 | "os" 9 | "reflect" 10 | "testing" 11 | ) 12 | 13 | type ScoringFields struct { 14 | A, B, C float32 15 | } 16 | 17 | func AddDocs(engine *Engine) { 18 | docId := uint64(0) 19 | engine.IndexDocument(docId, types.DocumentIndexData{ 20 | Content: "中国有十三亿人口人口", 21 | Fields: ScoringFields{1, 2, 3}, 22 | }) 23 | docId++ 24 | engine.IndexDocument(docId, types.DocumentIndexData{ 25 | Content: "中国人口", 26 | Fields: nil, 27 | }) 28 | docId++ 29 | engine.IndexDocument(docId, types.DocumentIndexData{ 30 | Content: "有人口", 31 | Fields: ScoringFields{2, 3, 1}, 32 | }) 33 | docId++ 34 | engine.IndexDocument(docId, types.DocumentIndexData{ 35 | Content: "有十三亿人口", 36 | Fields: ScoringFields{2, 3, 3}, 37 | }) 38 | docId++ 39 | engine.IndexDocument(docId, types.DocumentIndexData{ 40 | Content: "中国十三亿人口", 41 | Fields: ScoringFields{0, 9, 1}, 42 | }) 43 | 44 | engine.FlushIndex() 45 | } 46 | 47 | type RankByTokenProximity struct { 48 | } 49 | 50 | func (rule RankByTokenProximity) Score( 51 | doc types.IndexedDocument, fields interface{}) []float32 { 52 | if doc.TokenProximity < 0 { 53 | return []float32{} 54 | } 55 | return []float32{1.0 / (float32(doc.TokenProximity) + 1)} 56 | } 57 | 58 | func reset() { 59 | core.DocInfoGroup = make(map[int]*types.DocInfosShard) 60 | core.InvertedIndexGroup = make(map[int]*types.InvertedIndexShard) 61 | os.RemoveAll("wukong.persistent") 62 | } 63 | 64 | func TestEngineIndexDocument(t *testing.T) { 65 | reset() 66 | var engine Engine 67 | engine.Init(types.EngineInitOptions{ 68 | SegmenterDictionaries: "../testdata/test_dict.txt", 69 | DefaultRankOptions: &types.RankOptions{ 70 | OutputOffset: 0, 71 | MaxOutputs: 10, 72 | ScoringCriteria: &RankByTokenProximity{}, 73 | }, 74 | IndexerInitOptions: &types.IndexerInitOptions{ 75 | IndexType: types.LocationsIndex, 76 | }, 77 | }) 78 | 79 | AddDocs(&engine) 80 | 81 | outputs := engine.Search(types.SearchRequest{Text: "中国人口"}) 82 | utils.Expect(t, "2", len(outputs.Tokens)) 83 | utils.Expect(t, "中国", outputs.Tokens[0]) 84 | utils.Expect(t, "人口", outputs.Tokens[1]) 85 | utils.Expect(t, "3", len(outputs.Docs)) 86 | 87 | utils.Expect(t, "1", outputs.Docs[0].DocId) 88 | utils.Expect(t, "1000", int(outputs.Docs[0].Scores[0]*1000)) 89 | utils.Expect(t, "[0 6]", outputs.Docs[0].TokenSnippetLocations) 90 | 91 | utils.Expect(t, "4", outputs.Docs[1].DocId) 92 | utils.Expect(t, "100", int(outputs.Docs[1].Scores[0]*1000)) 93 | utils.Expect(t, "[0 15]", outputs.Docs[1].TokenSnippetLocations) 94 | 95 | utils.Expect(t, "0", outputs.Docs[2].DocId) 96 | utils.Expect(t, "76", int(outputs.Docs[2].Scores[0]*1000)) 97 | utils.Expect(t, "[0 18]", outputs.Docs[2].TokenSnippetLocations) 98 | } 99 | 100 | func TestReverseOrder(t *testing.T) { 101 | reset() 102 | var engine Engine 103 | engine.Init(types.EngineInitOptions{ 104 | SegmenterDictionaries: "../testdata/test_dict.txt", 105 | DefaultRankOptions: &types.RankOptions{ 106 | ReverseOrder: true, 107 | OutputOffset: 0, 108 | MaxOutputs: 10, 109 | ScoringCriteria: &RankByTokenProximity{}, 110 | }, 111 | IndexerInitOptions: &types.IndexerInitOptions{ 112 | IndexType: types.LocationsIndex, 113 | }, 114 | }) 115 | 116 | AddDocs(&engine) 117 | 118 | outputs := engine.Search(types.SearchRequest{Text: "中国人口"}) 119 | utils.Expect(t, "3", len(outputs.Docs)) 120 | 121 | utils.Expect(t, "0", outputs.Docs[0].DocId) 122 | utils.Expect(t, "4", outputs.Docs[1].DocId) 123 | utils.Expect(t, "1", outputs.Docs[2].DocId) 124 | } 125 | 126 | func TestOffsetAndMaxOutputs(t *testing.T) { 127 | reset() 128 | var engine Engine 129 | engine.Init(types.EngineInitOptions{ 130 | SegmenterDictionaries: "../testdata/test_dict.txt", 131 | DefaultRankOptions: &types.RankOptions{ 132 | ReverseOrder: true, 133 | OutputOffset: 1, 134 | MaxOutputs: 3, 135 | ScoringCriteria: &RankByTokenProximity{}, 136 | }, 137 | IndexerInitOptions: &types.IndexerInitOptions{ 138 | IndexType: types.LocationsIndex, 139 | }, 140 | }) 141 | 142 | AddDocs(&engine) 143 | 144 | outputs := engine.Search(types.SearchRequest{Text: "中国人口"}) 145 | utils.Expect(t, "2", len(outputs.Docs)) 146 | 147 | utils.Expect(t, "4", outputs.Docs[0].DocId) 148 | utils.Expect(t, "1", outputs.Docs[1].DocId) 149 | } 150 | 151 | type TestScoringCriteria struct { 152 | } 153 | 154 | func (criteria TestScoringCriteria) Score( 155 | doc types.IndexedDocument, fields interface{}) []float32 { 156 | if reflect.TypeOf(fields) != reflect.TypeOf(ScoringFields{}) { 157 | return []float32{} 158 | } 159 | fs := fields.(ScoringFields) 160 | return []float32{float32(doc.TokenProximity)*fs.A + fs.B*fs.C} 161 | } 162 | 163 | func TestSearchWithCriteria(t *testing.T) { 164 | reset() 165 | var engine Engine 166 | engine.Init(types.EngineInitOptions{ 167 | SegmenterDictionaries: "../testdata/test_dict.txt", 168 | DefaultRankOptions: &types.RankOptions{ 169 | ScoringCriteria: TestScoringCriteria{}, 170 | }, 171 | IndexerInitOptions: &types.IndexerInitOptions{ 172 | IndexType: types.LocationsIndex, 173 | }, 174 | }) 175 | 176 | AddDocs(&engine) 177 | 178 | outputs := engine.Search(types.SearchRequest{Text: "中国人口"}) 179 | utils.Expect(t, "2", len(outputs.Docs)) 180 | 181 | utils.Expect(t, "0", outputs.Docs[0].DocId) 182 | utils.Expect(t, "18000", int(outputs.Docs[0].Scores[0]*1000)) 183 | 184 | utils.Expect(t, "4", outputs.Docs[1].DocId) 185 | utils.Expect(t, "9000", int(outputs.Docs[1].Scores[0]*1000)) 186 | } 187 | 188 | func TestCompactIndex(t *testing.T) { 189 | reset() 190 | var engine Engine 191 | engine.Init(types.EngineInitOptions{ 192 | SegmenterDictionaries: "../testdata/test_dict.txt", 193 | DefaultRankOptions: &types.RankOptions{ 194 | ScoringCriteria: TestScoringCriteria{}, 195 | }, 196 | }) 197 | 198 | AddDocs(&engine) 199 | 200 | outputs := engine.Search(types.SearchRequest{Text: "中国人口"}) 201 | utils.Expect(t, "2", len(outputs.Docs)) 202 | 203 | utils.Expect(t, "4", outputs.Docs[0].DocId) 204 | utils.Expect(t, "9000", int(outputs.Docs[0].Scores[0]*1000)) 205 | 206 | utils.Expect(t, "0", outputs.Docs[1].DocId) 207 | utils.Expect(t, "6000", int(outputs.Docs[1].Scores[0]*1000)) 208 | } 209 | 210 | type BM25ScoringCriteria struct { 211 | } 212 | 213 | func (criteria BM25ScoringCriteria) Score( 214 | doc types.IndexedDocument, fields interface{}) []float32 { 215 | if reflect.TypeOf(fields) != reflect.TypeOf(ScoringFields{}) { 216 | return []float32{} 217 | } 218 | return []float32{doc.BM25} 219 | } 220 | 221 | func TestFrequenciesIndex(t *testing.T) { 222 | reset() 223 | var engine Engine 224 | engine.Init(types.EngineInitOptions{ 225 | SegmenterDictionaries: "../testdata/test_dict.txt", 226 | DefaultRankOptions: &types.RankOptions{ 227 | ScoringCriteria: BM25ScoringCriteria{}, 228 | }, 229 | IndexerInitOptions: &types.IndexerInitOptions{ 230 | IndexType: types.FrequenciesIndex, 231 | }, 232 | NumShards: 2, 233 | }) 234 | 235 | AddDocs(&engine) 236 | 237 | outputs := engine.Search(types.SearchRequest{Text: "中国人口"}) 238 | utils.Expect(t, "2", len(outputs.Docs)) 239 | t.Log(outputs.Docs) 240 | utils.Expect(t, "4", outputs.Docs[0].DocId) 241 | utils.Expect(t, "2285", int(outputs.Docs[0].Scores[0]*1000)) 242 | 243 | utils.Expect(t, "0", outputs.Docs[1].DocId) 244 | utils.Expect(t, "2260", int(outputs.Docs[1].Scores[0]*1000)) 245 | } 246 | 247 | func TestRemoveDocument(t *testing.T) { 248 | reset() 249 | var engine Engine 250 | engine.Init(types.EngineInitOptions{ 251 | SegmenterDictionaries: "../testdata/test_dict.txt", 252 | DefaultRankOptions: &types.RankOptions{ 253 | ScoringCriteria: TestScoringCriteria{}, 254 | }, 255 | NumShards: 2, 256 | }) 257 | 258 | AddDocs(&engine) 259 | engine.RemoveDocument(4) 260 | 261 | outputs := engine.Search(types.SearchRequest{Text: "中国人口"}) 262 | utils.Expect(t, "1", len(outputs.Docs)) 263 | 264 | utils.Expect(t, "0", outputs.Docs[0].DocId) 265 | utils.Expect(t, "6000", int(outputs.Docs[0].Scores[0]*1000)) 266 | } 267 | 268 | func TestEngineIndexDocumentWithTokens(t *testing.T) { 269 | reset() 270 | 271 | var engine Engine 272 | engine.Init(types.EngineInitOptions{ 273 | SegmenterDictionaries: "../testdata/test_dict.txt", 274 | DefaultRankOptions: &types.RankOptions{ 275 | OutputOffset: 0, 276 | MaxOutputs: 10, 277 | ScoringCriteria: &RankByTokenProximity{}, 278 | }, 279 | IndexerInitOptions: &types.IndexerInitOptions{ 280 | IndexType: types.LocationsIndex, 281 | }, 282 | }) 283 | 284 | docId := uint64(0) 285 | engine.IndexDocument(docId, types.DocumentIndexData{ 286 | Content: "", 287 | Tokens: []types.TokenData{ 288 | {"中国", []int{0}}, 289 | {"人口", []int{18, 24}}, 290 | }, 291 | Fields: ScoringFields{1, 2, 3}, 292 | }) 293 | docId++ 294 | engine.IndexDocument(docId, types.DocumentIndexData{ 295 | Content: "", 296 | Tokens: []types.TokenData{ 297 | {"中国", []int{0}}, 298 | {"人口", []int{6}}, 299 | }, 300 | Fields: ScoringFields{1, 2, 3}, 301 | }) 302 | docId++ 303 | engine.IndexDocument(docId, types.DocumentIndexData{ 304 | Content: "中国十三亿人口", 305 | Fields: ScoringFields{0, 9, 1}, 306 | }) 307 | 308 | engine.FlushIndex() 309 | 310 | outputs := engine.Search(types.SearchRequest{Text: "中国人口"}) 311 | utils.Expect(t, "2", len(outputs.Tokens)) 312 | utils.Expect(t, "中国", outputs.Tokens[0]) 313 | utils.Expect(t, "人口", outputs.Tokens[1]) 314 | utils.Expect(t, "3", len(outputs.Docs)) 315 | 316 | utils.Expect(t, "1", outputs.Docs[0].DocId) 317 | utils.Expect(t, "1000", int(outputs.Docs[0].Scores[0]*1000)) 318 | utils.Expect(t, "[0 6]", outputs.Docs[0].TokenSnippetLocations) 319 | 320 | utils.Expect(t, "2", outputs.Docs[1].DocId) 321 | utils.Expect(t, "100", int(outputs.Docs[1].Scores[0]*1000)) 322 | utils.Expect(t, "[0 15]", outputs.Docs[1].TokenSnippetLocations) 323 | 324 | utils.Expect(t, "0", outputs.Docs[2].DocId) 325 | utils.Expect(t, "76", int(outputs.Docs[2].Scores[0]*1000)) 326 | utils.Expect(t, "[0 18]", outputs.Docs[2].TokenSnippetLocations) 327 | } 328 | 329 | func TestEngineIndexDocumentWithPersistentStorage(t *testing.T) { 330 | reset() 331 | gob.Register(ScoringFields{}) 332 | var engine Engine 333 | engine.Init(types.EngineInitOptions{ 334 | SegmenterDictionaries: "../testdata/test_dict.txt", 335 | DefaultRankOptions: &types.RankOptions{ 336 | OutputOffset: 0, 337 | MaxOutputs: 10, 338 | ScoringCriteria: &RankByTokenProximity{}, 339 | }, 340 | IndexerInitOptions: &types.IndexerInitOptions{ 341 | IndexType: types.LocationsIndex, 342 | }, 343 | UsePersistentStorage: true, 344 | PersistentStorageFolder: "wukong.persistent", 345 | }) 346 | AddDocs(&engine) 347 | engine.RemoveDocument(4) 348 | engine.Close() 349 | 350 | var engine1 Engine 351 | engine1.Init(types.EngineInitOptions{ 352 | SegmenterDictionaries: "../testdata/test_dict.txt", 353 | DefaultRankOptions: &types.RankOptions{ 354 | OutputOffset: 0, 355 | MaxOutputs: 10, 356 | ScoringCriteria: &RankByTokenProximity{}, 357 | }, 358 | IndexerInitOptions: &types.IndexerInitOptions{ 359 | IndexType: types.LocationsIndex, 360 | }, 361 | UsePersistentStorage: true, 362 | PersistentStorageFolder: "wukong.persistent", 363 | }) 364 | 365 | outputs := engine1.Search(types.SearchRequest{Text: "中国人口"}) 366 | utils.Expect(t, "2", len(outputs.Tokens)) 367 | utils.Expect(t, "中国", outputs.Tokens[0]) 368 | utils.Expect(t, "人口", outputs.Tokens[1]) 369 | utils.Expect(t, "2", len(outputs.Docs)) 370 | 371 | utils.Expect(t, "1", outputs.Docs[0].DocId) 372 | utils.Expect(t, "1000", int(outputs.Docs[0].Scores[0]*1000)) 373 | utils.Expect(t, "[0 6]", outputs.Docs[0].TokenSnippetLocations) 374 | 375 | utils.Expect(t, "0", outputs.Docs[1].DocId) 376 | utils.Expect(t, "76", int(outputs.Docs[1].Scores[0]*1000)) 377 | utils.Expect(t, "[0 18]", outputs.Docs[1].TokenSnippetLocations) 378 | 379 | engine1.Close() 380 | os.RemoveAll("wukong.persistent") 381 | } 382 | 383 | func TestCountDocsOnly(t *testing.T) { 384 | reset() 385 | var engine Engine 386 | engine.Init(types.EngineInitOptions{ 387 | SegmenterDictionaries: "../testdata/test_dict.txt", 388 | DefaultRankOptions: &types.RankOptions{ 389 | ReverseOrder: true, 390 | OutputOffset: 0, 391 | MaxOutputs: 1, 392 | ScoringCriteria: &RankByTokenProximity{}, 393 | }, 394 | IndexerInitOptions: &types.IndexerInitOptions{ 395 | IndexType: types.LocationsIndex, 396 | }, 397 | NumShards: 2, 398 | }) 399 | 400 | AddDocs(&engine) 401 | engine.RemoveDocument(4) 402 | 403 | outputs := engine.Search(types.SearchRequest{Text: "中国人口", CountDocsOnly: true}) 404 | utils.Expect(t, "0", len(outputs.Docs)) 405 | utils.Expect(t, "2", len(outputs.Tokens)) 406 | utils.Expect(t, "2", outputs.NumDocs) 407 | } 408 | 409 | func TestSearchWithin(t *testing.T) { 410 | reset() 411 | var engine Engine 412 | engine.Init(types.EngineInitOptions{ 413 | SegmenterDictionaries: "../testdata/test_dict.txt", 414 | DefaultRankOptions: &types.RankOptions{ 415 | ReverseOrder: true, 416 | OutputOffset: 0, 417 | MaxOutputs: 10, 418 | ScoringCriteria: &RankByTokenProximity{}, 419 | }, 420 | IndexerInitOptions: &types.IndexerInitOptions{ 421 | IndexType: types.LocationsIndex, 422 | }, 423 | }) 424 | 425 | AddDocs(&engine) 426 | 427 | docIds := make(map[uint64]bool) 428 | docIds[4] = true 429 | docIds[0] = true 430 | outputs := engine.Search(types.SearchRequest{ 431 | Text: "中国人口", 432 | DocIds: docIds, 433 | }) 434 | utils.Expect(t, "2", len(outputs.Tokens)) 435 | utils.Expect(t, "中国", outputs.Tokens[0]) 436 | utils.Expect(t, "人口", outputs.Tokens[1]) 437 | utils.Expect(t, "2", len(outputs.Docs)) 438 | 439 | utils.Expect(t, "0", outputs.Docs[0].DocId) 440 | utils.Expect(t, "76", int(outputs.Docs[0].Scores[0]*1000)) 441 | utils.Expect(t, "[0 18]", outputs.Docs[0].TokenSnippetLocations) 442 | 443 | utils.Expect(t, "4", outputs.Docs[1].DocId) 444 | utils.Expect(t, "100", int(outputs.Docs[1].Scores[0]*1000)) 445 | utils.Expect(t, "[0 15]", outputs.Docs[1].TokenSnippetLocations) 446 | } 447 | -------------------------------------------------------------------------------- /engine/indexer_worker.go: -------------------------------------------------------------------------------- 1 | package engine 2 | 3 | import ( 4 | "github.com/huichen/wukong/types" 5 | "sync/atomic" 6 | ) 7 | 8 | type indexerAddDocumentRequest struct { 9 | document *types.DocumentIndex 10 | dealDocInfoChan chan<- bool 11 | } 12 | 13 | type indexerLookupRequest struct { 14 | countDocsOnly bool 15 | tokens []string 16 | labels []string 17 | docIds map[uint64]bool 18 | options types.RankOptions 19 | rankerReturnChannel chan rankerReturnRequest 20 | orderless bool 21 | } 22 | 23 | type indexerRemoveDocRequest struct { 24 | docId uint64 25 | } 26 | 27 | func (engine *Engine) indexerAddDocumentWorker(shard int) { 28 | for { 29 | request := <-engine.indexerAddDocumentChannels[shard] 30 | addInvertedIndex := engine.indexers[shard].AddDocument(request.document, request.dealDocInfoChan) 31 | // save 32 | if engine.initOptions.UsePersistentStorage { 33 | for k, v := range addInvertedIndex { 34 | engine.persistentStorageIndexDocumentChannels[shard] <- persistentStorageIndexDocumentRequest{ 35 | typ: "index", 36 | keyword: k, 37 | keywordIndices: v, 38 | } 39 | } 40 | } 41 | 42 | atomic.AddUint64(&engine.numTokenIndexAdded, 43 | uint64(len(request.document.Keywords))) 44 | atomic.AddUint64(&engine.numDocumentsIndexed, 1) 45 | } 46 | } 47 | 48 | func (engine *Engine) indexerLookupWorker(shard int) { 49 | for { 50 | request := <-engine.indexerLookupChannels[shard] 51 | 52 | var docs []types.IndexedDocument 53 | var numDocs int 54 | if request.docIds == nil { 55 | docs, numDocs = engine.indexers[shard].Lookup(request.tokens, request.labels, nil, request.countDocsOnly) 56 | } else { 57 | docs, numDocs = engine.indexers[shard].Lookup(request.tokens, request.labels, request.docIds, request.countDocsOnly) 58 | } 59 | 60 | if request.countDocsOnly { 61 | request.rankerReturnChannel <- rankerReturnRequest{numDocs: numDocs} 62 | continue 63 | } 64 | 65 | if len(docs) == 0 { 66 | request.rankerReturnChannel <- rankerReturnRequest{} 67 | continue 68 | } 69 | 70 | if request.orderless { 71 | var outputDocs []types.ScoredDocument 72 | for _, d := range docs { 73 | outputDocs = append(outputDocs, types.ScoredDocument{ 74 | DocId: d.DocId, 75 | TokenSnippetLocations: d.TokenSnippetLocations, 76 | TokenLocations: d.TokenLocations}) 77 | } 78 | request.rankerReturnChannel <- rankerReturnRequest{ 79 | docs: outputDocs, 80 | numDocs: len(outputDocs), 81 | } 82 | continue 83 | } 84 | 85 | rankerRequest := rankerRankRequest{ 86 | countDocsOnly: request.countDocsOnly, 87 | docs: docs, 88 | options: request.options, 89 | rankerReturnChannel: request.rankerReturnChannel, 90 | } 91 | engine.rankerRankChannels[shard] <- rankerRequest 92 | } 93 | } 94 | 95 | func (engine *Engine) indexerRemoveDocWorker(shard int) { 96 | for { 97 | request := <-engine.indexerRemoveDocChannels[shard] 98 | engine.indexers[shard].RemoveDoc(request.docId) 99 | } 100 | } 101 | -------------------------------------------------------------------------------- /engine/persistent_storage_worker.go: -------------------------------------------------------------------------------- 1 | package engine 2 | 3 | import ( 4 | "bytes" 5 | "encoding/binary" 6 | "encoding/gob" 7 | "github.com/huichen/wukong/core" 8 | "github.com/huichen/wukong/types" 9 | "sync" 10 | "sync/atomic" 11 | ) 12 | 13 | type persistentStorageIndexDocumentRequest struct { 14 | typ string //"info"or"index" 15 | 16 | // typ=="info"时，以下两个字段有效 17 | docId uint64 18 | docInfo *types.DocInfo 19 | 20 | // typ=="index"时，以下两个字段有效 21 | keyword string 22 | keywordIndices *types.KeywordIndices 23 | } 24 | 25 | func (engine *Engine) persistentStorageIndexDocumentWorker(shard int) { 26 | for { 27 | request := <-engine.persistentStorageIndexDocumentChannels[shard] 28 | switch request.typ { 29 | case "info": 30 | // 得到key 31 | b := make([]byte, 10) 32 | length := binary.PutUvarint(b, request.docId) 33 | 34 | // 得到value 35 | var buf bytes.Buffer 36 | enc := gob.NewEncoder(&buf) 37 | err := enc.Encode(request.docInfo) 38 | if err != nil { 39 | atomic.AddUint64(&engine.numDocumentsStored, 1) 40 | return 41 | } 42 | 43 | // 将key-value写入数据库 44 | engine.dbs[shard][getDB(request.typ)].Set(b[0:length], buf.Bytes()) 45 | atomic.AddUint64(&engine.numDocumentsStored, 1) 46 | 47 | case "index": 48 | // 得到key 49 | b := []byte(request.keyword) 50 | 51 | // 得到value 52 | var buf bytes.Buffer 53 | enc := gob.NewEncoder(&buf) 54 | err := enc.Encode(request.keywordIndices) 55 | if err != nil { 56 | return 57 | } 58 | 59 | // 将key-value写入数据库 60 | engine.dbs[shard][getDB(request.typ)].Set(b, buf.Bytes()) 61 | } 62 | } 63 | } 64 | 65 | func (engine *Engine) persistentStorageRemoveDocumentWorker(docId uint64, shard int) { 66 | // 得到key 67 | b := make([]byte, 10) 68 | length := binary.PutUvarint(b, docId) 69 | 70 | // 从数据库删除该key 71 | engine.dbs[shard][getDB("info")].Delete(b[0:length]) 72 | } 73 | 74 | func (engine *Engine) persistentStorageInitWorker(shard int) { 75 | var finish sync.WaitGroup 76 | finish.Add(2) 77 | // 恢复docInfo 78 | go func() { 79 | defer finish.Add(-1) 80 | engine.dbs[shard][getDB("info")].ForEach(func(k, v []byte) error { 81 | key, value := k, v 82 | // 得到docID 83 | docId, _ := binary.Uvarint(key) 84 | 85 | // 得到data 86 | buf := bytes.NewReader(value) 87 | dec := gob.NewDecoder(buf) 88 | var data types.DocInfo 89 | err := dec.Decode(&data) 90 | if err == nil { 91 | // 添加索引 92 | core.AddDocInfo(shard, docId, &data) 93 | } 94 | return nil 95 | }) 96 | }() 97 | 98 | // 恢复invertedIndex 99 | go func() { 100 | defer finish.Add(-1) 101 | engine.dbs[shard][getDB("index")].ForEach(func(k, v []byte) error { 102 | key, value := k, v 103 | // 得到keyword 104 | keyword := string(key) 105 | 106 | // 得到data 107 | buf := bytes.NewReader(value) 108 | dec := gob.NewDecoder(buf) 109 | var data types.KeywordIndices 110 | err := dec.Decode(&data) 111 | if err == nil { 112 | // 添加索引 113 | core.AddKeywordIndices(shard, keyword, &data) 114 | } 115 | return nil 116 | }) 117 | }() 118 | finish.Wait() 119 | engine.persistentStorageInitChannel <- true 120 | } 121 | -------------------------------------------------------------------------------- /engine/ranker_worker.go: -------------------------------------------------------------------------------- 1 | package engine 2 | 3 | import ( 4 | "github.com/huichen/wukong/types" 5 | ) 6 | 7 | type rankerAddDocRequest struct { 8 | docId uint64 9 | fields interface{} 10 | dealDocInfoChan <-chan bool 11 | } 12 | 13 | type rankerRankRequest struct { 14 | docs []types.IndexedDocument 15 | options types.RankOptions 16 | rankerReturnChannel chan rankerReturnRequest 17 | countDocsOnly bool 18 | } 19 | 20 | type rankerReturnRequest struct { 21 | docs types.ScoredDocuments 22 | numDocs int 23 | } 24 | 25 | type rankerRemoveDocRequest struct { 26 | docId uint64 27 | } 28 | 29 | func (engine *Engine) rankerAddDocWorker(shard int) { 30 | for { 31 | request := <-engine.rankerAddDocChannels[shard] 32 | docInfo := engine.rankers[shard].AddDoc(request.docId, request.fields, request.dealDocInfoChan) 33 | // save 34 | if engine.initOptions.UsePersistentStorage { 35 | engine.persistentStorageIndexDocumentChannels[shard] <- persistentStorageIndexDocumentRequest{ 36 | typ: "info", 37 | docId: request.docId, 38 | docInfo: docInfo, 39 | } 40 | } 41 | } 42 | } 43 | 44 | func (engine *Engine) rankerRankWorker(shard int) { 45 | for { 46 | request := <-engine.rankerRankChannels[shard] 47 | if request.options.MaxOutputs != 0 { 48 | request.options.MaxOutputs += request.options.OutputOffset 49 | } 50 | request.options.OutputOffset = 0 51 | outputDocs, numDocs := engine.rankers[shard].Rank(request.docs, request.options, request.countDocsOnly) 52 | request.rankerReturnChannel <- rankerReturnRequest{docs: outputDocs, numDocs: numDocs} 53 | } 54 | } 55 | 56 | func (engine *Engine) rankerRemoveDocWorker(shard int) { 57 | for { 58 | request := <-engine.rankerRemoveDocChannels[shard] 59 | engine.rankers[shard].RemoveDoc(request.docId) 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /engine/segmenter_worker.go: -------------------------------------------------------------------------------- 1 | package engine 2 | 3 | import ( 4 | "github.com/huichen/wukong/types" 5 | ) 6 | 7 | type segmenterRequest struct { 8 | docId uint64 9 | shard int 10 | data types.DocumentIndexData 11 | } 12 | 13 | func (engine *Engine) segmenterWorker() { 14 | for { 15 | request := <-engine.segmenterChannel 16 | 17 | tokensMap := make(map[string][]int) 18 | numTokens := 0 19 | if !engine.initOptions.NotUsingSegmenter && request.data.Content != "" { 20 | // 当文档正文不为空时，优先从内容分词中得到关键词 21 | segments := engine.segmenter.Segment([]byte(request.data.Content)) 22 | for _, segment := range segments { 23 | token := segment.Token().Text() 24 | if !engine.stopTokens.IsStopToken(token) { 25 | tokensMap[token] = append(tokensMap[token], segment.Start()) 26 | } 27 | } 28 | numTokens = len(segments) 29 | } else { 30 | // 否则载入用户输入的关键词 31 | for _, t := range request.data.Tokens { 32 | if !engine.stopTokens.IsStopToken(t.Text) { 33 | tokensMap[t.Text] = t.Locations 34 | } 35 | } 36 | numTokens = len(request.data.Tokens) 37 | } 38 | 39 | // 加入非分词的文档标签 40 | for _, label := range request.data.Labels { 41 | if !engine.initOptions.NotUsingSegmenter { 42 | if !engine.stopTokens.IsStopToken(label) { 43 | tokensMap[label] = []int{} 44 | } 45 | } else { 46 | tokensMap[label] = []int{} 47 | } 48 | } 49 | 50 | indexerRequest := indexerAddDocumentRequest{ 51 | document: &types.DocumentIndex{ 52 | DocId: request.docId, 53 | TokenLength: float32(numTokens), 54 | Keywords: make([]types.KeywordIndex, len(tokensMap)), 55 | }, 56 | } 57 | iTokens := 0 58 | for k, v := range tokensMap { 59 | indexerRequest.document.Keywords[iTokens] = types.KeywordIndex{ 60 | Text: k, 61 | // 非分词标注的词频设置为0，不参与tf-idf计算 62 | Frequency: float32(len(v)), 63 | Starts: v} 64 | iTokens++ 65 | } 66 | 67 | var dealDocInfoChan = make(chan bool, 1) 68 | 69 | indexerRequest.dealDocInfoChan = dealDocInfoChan 70 | engine.indexerAddDocumentChannels[request.shard] <- indexerRequest 71 | 72 | rankerRequest := rankerAddDocRequest{ 73 | docId: request.docId, 74 | fields: request.data.Fields, 75 | dealDocInfoChan: dealDocInfoChan, 76 | } 77 | engine.rankerAddDocChannels[request.shard] <- rankerRequest 78 | } 79 | } 80 | -------------------------------------------------------------------------------- /engine/stop_tokens.go: -------------------------------------------------------------------------------- 1 | package engine 2 | 3 | import ( 4 | "bufio" 5 | "log" 6 | "os" 7 | ) 8 | 9 | type StopTokens struct { 10 | stopTokens map[string]bool 11 | } 12 | 13 | // 从stopTokenFile中读入停用词，一个词一行 14 | // 文档索引建立时会跳过这些停用词 15 | func (st *StopTokens) Init(stopTokenFile string) { 16 | st.stopTokens = make(map[string]bool) 17 | if stopTokenFile == "" { 18 | return 19 | } 20 | 21 | file, err := os.Open(stopTokenFile) 22 | if err != nil { 23 | log.Fatal(err) 24 | } 25 | defer file.Close() 26 | 27 | scanner := bufio.NewScanner(file) 28 | for scanner.Scan() { 29 | text := scanner.Text() 30 | if text != "" { 31 | st.stopTokens[text] = true 32 | } 33 | } 34 | 35 | } 36 | 37 | func (st *StopTokens) IsStopToken(token string) bool { 38 | _, found := st.stopTokens[token] 39 | return found 40 | } 41 | -------------------------------------------------------------------------------- /examples/benchmark.go: -------------------------------------------------------------------------------- 1 | // 悟空性能测试 2 | package main 3 | 4 | import ( 5 | "bufio" 6 | "flag" 7 | "github.com/huichen/wukong/engine" 8 | "github.com/huichen/wukong/types" 9 | "log" 10 | "os" 11 | "runtime" 12 | "runtime/pprof" 13 | "strings" 14 | "time" 15 | ) 16 | 17 | const ( 18 | numRepeatQuery = 1000 19 | ) 20 | 21 | var ( 22 | weibo_data = flag.String( 23 | "weibo_data", 24 | "../testdata/weibo_data.txt", 25 | "微博数据") 26 | queries = flag.String( 27 | "queries", 28 | "女人母亲,你好中国,网络草根,热门微博,红十字会,"+ 29 | "鳄鱼表演,星座歧视,chinajoy,高帅富,假期计划", 30 | "待搜索的关键词") 31 | dictionaries = flag.String( 32 | "dictionaries", 33 | "../data/dictionary.txt", 34 | "分词字典文件") 35 | stop_token_file = flag.String( 36 | "stop_token_file", 37 | "../data/stop_tokens.txt", 38 | "停用词文件") 39 | cpuprofile = flag.String("cpuprofile", "", "处理器profile文件") 40 | memprofile = flag.String("memprofile", "", "内存profile文件") 41 | num_repeat_text = flag.Int("num_repeat_text", 10, "文本重复加入多少次") 42 | index_type = flag.Int("index_type", types.DocIdsIndex, "索引类型") 43 | use_persistent = flag.Bool("use_persistent", false, "是否使用持久存储") 44 | persistent_storage_folder = flag.String("persistent_storage_folder", "benchmark.persistent", "持久存储数据库保存的目录") 45 | persistent_storage_shards = flag.Int("persistent_storage_shards", 0, "持久数据库存储裂分数目") 46 | 47 | searcher = engine.Engine{} 48 | options = types.RankOptions{ 49 | OutputOffset: 0, 50 | MaxOutputs: 100, 51 | } 52 | searchQueries = []string{} 53 | 54 | NumShards = 2 55 | numQueryThreads = runtime.NumCPU() / NumShards 56 | ) 57 | 58 | func main() { 59 | // 解析命令行参数 60 | flag.Parse() 61 | searchQueries = strings.Split(*queries, ",") 62 | log.Printf("待搜索的关键词为\"%s\"", searchQueries) 63 | 64 | // 初始化 65 | tBeginInit := time.Now() 66 | searcher.Init(types.EngineInitOptions{ 67 | SegmenterDictionaries: *dictionaries, 68 | StopTokenFile: *stop_token_file, 69 | IndexerInitOptions: &types.IndexerInitOptions{ 70 | IndexType: *index_type, 71 | }, 72 | NumShards: NumShards, 73 | DefaultRankOptions: &options, 74 | UsePersistentStorage: *use_persistent, 75 | PersistentStorageFolder: *persistent_storage_folder, 76 | PersistentStorageShards: *persistent_storage_shards, 77 | }) 78 | tEndInit := time.Now() 79 | defer searcher.Close() 80 | 81 | // 打开将要搜索的文件 82 | file, err := os.Open(*weibo_data) 83 | if err != nil { 84 | log.Fatal(err) 85 | } 86 | defer file.Close() 87 | 88 | // 逐行读入 89 | log.Printf("读入文本 %s", *weibo_data) 90 | scanner := bufio.NewScanner(file) 91 | lines := []string{} 92 | size := 0 93 | for scanner.Scan() { 94 | var text string 95 | data := strings.Split(scanner.Text(), "||||") 96 | if len(data) != 10 { 97 | continue 98 | } 99 | text = data[9] 100 | if text != "" { 101 | size += len(text) * (*num_repeat_text) 102 | lines = append(lines, text) 103 | } 104 | } 105 | log.Print("文件行数", len(lines)) 106 | 107 | // 记录时间 108 | t0 := time.Now() 109 | 110 | // 打开处理器profile文件 111 | if *cpuprofile != "" { 112 | f, err := os.Create(*cpuprofile) 113 | if err != nil { 114 | log.Fatal(err) 115 | } 116 | pprof.StartCPUProfile(f) 117 | defer pprof.StopCPUProfile() 118 | } 119 | 120 | // 建索引 121 | log.Print("建索引 ... ") 122 | docId := uint64(1) 123 | for i := 0; i < *num_repeat_text; i++ { 124 | for _, line := range lines { 125 | searcher.IndexDocument(docId, types.DocumentIndexData{ 126 | Content: line}) 127 | docId++ 128 | if docId-docId/1000000*1000000 == 0 { 129 | log.Printf("已索引%d百万文档", docId/1000000) 130 | runtime.GC() 131 | } 132 | } 133 | } 134 | searcher.FlushIndex() 135 | log.Print("加入的索引总数", searcher.NumTokenIndexAdded()) 136 | 137 | // 记录时间 138 | t1 := time.Now() 139 | log.Printf("建立索引花费时间 %v", t1.Sub(t0)) 140 | log.Printf("建立索引速度每秒添加 %f 百万个索引", 141 | float64(searcher.NumTokenIndexAdded())/t1.Sub(t0).Seconds()/(1000000)) 142 | 143 | // 写入内存profile文件 144 | if *memprofile != "" { 145 | f, err := os.Create(*memprofile) 146 | if err != nil { 147 | log.Fatal(err) 148 | } 149 | pprof.WriteHeapProfile(f) 150 | defer f.Close() 151 | } 152 | 153 | // 记录时间 154 | t2 := time.Now() 155 | 156 | done := make(chan bool) 157 | for iThread := 0; iThread < numQueryThreads; iThread++ { 158 | go search(done) 159 | } 160 | for iThread := 0; iThread < numQueryThreads; iThread++ { 161 | <-done 162 | } 163 | 164 | // 记录时间并计算分词速度 165 | t3 := time.Now() 166 | log.Printf("搜索平均响应时间 %v 毫秒", 167 | t3.Sub(t2).Seconds()*1000/float64(numRepeatQuery*len(searchQueries))) 168 | log.Printf("搜索吞吐量每秒 %v 次查询", 169 | float64(numRepeatQuery*numQueryThreads*len(searchQueries))/ 170 | t3.Sub(t2).Seconds()) 171 | 172 | if *use_persistent { 173 | searcher.Close() 174 | t4 := time.Now() 175 | searcher1 := engine.Engine{} 176 | searcher1.Init(types.EngineInitOptions{ 177 | SegmenterDictionaries: *dictionaries, 178 | StopTokenFile: *stop_token_file, 179 | IndexerInitOptions: &types.IndexerInitOptions{ 180 | IndexType: *index_type, 181 | }, 182 | NumShards: NumShards, 183 | DefaultRankOptions: &options, 184 | UsePersistentStorage: *use_persistent, 185 | PersistentStorageFolder: *persistent_storage_folder, 186 | PersistentStorageShards: *persistent_storage_shards, 187 | }) 188 | defer searcher1.Close() 189 | t5 := time.Now() 190 | t := t5.Sub(t4).Seconds() - tEndInit.Sub(tBeginInit).Seconds() 191 | log.Print("从持久存储加入的索引总数", searcher1.NumTokenIndexAdded()) 192 | log.Printf("从持久存储建立索引花费时间 %v 秒", t) 193 | log.Printf("从持久存储建立索引速度每秒添加 %f 百万个索引", 194 | float64(searcher1.NumTokenIndexAdded())/t/(1000000)) 195 | 196 | } 197 | //os.RemoveAll(*persistent_storage_folder) 198 | } 199 | 200 | func search(ch chan bool) { 201 | for i := 0; i < numRepeatQuery; i++ { 202 | for _, query := range searchQueries { 203 | searcher.Search(types.SearchRequest{Text: query}) 204 | } 205 | } 206 | ch <- true 207 | } 208 | -------------------------------------------------------------------------------- /examples/codelab/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM busybox 2 | EXPOSE 8080 3 | ADD / / 4 | CMD ./search_server --weibo_data=weibo_data.txt --dict_file=dictionary.txt --stop_token_file=stop_tokens.txt --static_folder=static 5 | -------------------------------------------------------------------------------- /examples/codelab/build_docker_image.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | rm -rf docker 6 | mkdir docker 7 | CGO_ENABLED=0 GOOS=linux go build -a -installsuffix cgo -o docker/search_server 8 | cp Dockerfile docker/ 9 | cp static docker/ -r 10 | cp ../../data/dictionary.txt docker/ 11 | cp ../../data/stop_tokens.txt docker/ 12 | cp ../../testdata/weibo_data.txt docker/ 13 | 14 | docker build -t unmerged/wukong-codelab -f docker/Dockerfile docker/ 15 | -------------------------------------------------------------------------------- /examples/codelab/search_server.go: -------------------------------------------------------------------------------- 1 | // 一个微博搜索的例子。 2 | package main 3 | 4 | import ( 5 | "bufio" 6 | "encoding/gob" 7 | "encoding/json" 8 | "flag" 9 | "github.com/huichen/wukong/engine" 10 | "github.com/huichen/wukong/types" 11 | "io" 12 | "log" 13 | "net/http" 14 | "os" 15 | "os/signal" 16 | "reflect" 17 | "strconv" 18 | "strings" 19 | ) 20 | 21 | const ( 22 | SecondsInADay = 86400 23 | MaxTokenProximity = 2 24 | ) 25 | 26 | var ( 27 | searcher = engine.Engine{} 28 | wbs = map[uint64]Weibo{} 29 | weiboData = flag.String("weibo_data", "../../testdata/weibo_data.txt", "微博数据文件") 30 | dictFile = flag.String("dict_file", "../../data/dictionary.txt", "词典文件") 31 | stopTokenFile = flag.String("stop_token_file", "../../data/stop_tokens.txt", "停用词文件") 32 | staticFolder = flag.String("static_folder", "static", "静态文件目录") 33 | ) 34 | 35 | type Weibo struct { 36 | Id uint64 `json:"id"` 37 | Timestamp uint64 `json:"timestamp"` 38 | UserName string `json:"user_name"` 39 | RepostsCount uint64 `json:"reposts_count"` 40 | Text string `json:"text"` 41 | } 42 | 43 | /******************************************************************************* 44 | 索引 45 | *******************************************************************************/ 46 | func indexWeibo() { 47 | // 读入微博数据 48 | file, err := os.Open(*weiboData) 49 | if err != nil { 50 | log.Fatal(err) 51 | } 52 | defer file.Close() 53 | scanner := bufio.NewScanner(file) 54 | for scanner.Scan() { 55 | data := strings.Split(scanner.Text(), "||||") 56 | if len(data) != 10 { 57 | continue 58 | } 59 | wb := Weibo{} 60 | wb.Id, _ = strconv.ParseUint(data[0], 10, 64) 61 | wb.Timestamp, _ = strconv.ParseUint(data[1], 10, 64) 62 | wb.UserName = data[3] 63 | wb.RepostsCount, _ = strconv.ParseUint(data[4], 10, 64) 64 | wb.Text = data[9] 65 | wbs[wb.Id] = wb 66 | } 67 | 68 | log.Print("添加索引") 69 | for docId, weibo := range wbs { 70 | searcher.IndexDocument(docId, types.DocumentIndexData{ 71 | Content: weibo.Text, 72 | Fields: WeiboScoringFields{ 73 | Timestamp: weibo.Timestamp, 74 | RepostsCount: weibo.RepostsCount, 75 | }, 76 | }) 77 | } 78 | 79 | searcher.FlushIndex() 80 | log.Printf("索引了%d条微博\n", len(wbs)) 81 | } 82 | 83 | /******************************************************************************* 84 | 评分 85 | *******************************************************************************/ 86 | type WeiboScoringFields struct { 87 | Timestamp uint64 88 | RepostsCount uint64 89 | } 90 | 91 | type WeiboScoringCriteria struct { 92 | } 93 | 94 | func (criteria WeiboScoringCriteria) Score( 95 | doc types.IndexedDocument, fields interface{}) []float32 { 96 | if reflect.TypeOf(fields) != reflect.TypeOf(WeiboScoringFields{}) { 97 | return []float32{} 98 | } 99 | wsf := fields.(WeiboScoringFields) 100 | output := make([]float32, 3) 101 | if doc.TokenProximity > MaxTokenProximity { 102 | output[0] = 1.0 / float32(doc.TokenProximity) 103 | } else { 104 | output[0] = 1.0 105 | } 106 | output[1] = float32(wsf.Timestamp / (SecondsInADay * 3)) 107 | output[2] = float32(doc.BM25 * (1 + float32(wsf.RepostsCount)/10000)) 108 | return output 109 | } 110 | 111 | /******************************************************************************* 112 | JSON-RPC 113 | *******************************************************************************/ 114 | type JsonResponse struct { 115 | Docs []*Weibo `json:"docs"` 116 | } 117 | 118 | func JsonRpcServer(w http.ResponseWriter, req *http.Request) { 119 | query := req.URL.Query().Get("query") 120 | output := searcher.Search(types.SearchRequest{ 121 | Text: query, 122 | RankOptions: &types.RankOptions{ 123 | ScoringCriteria: &WeiboScoringCriteria{}, 124 | OutputOffset: 0, 125 | MaxOutputs: 100, 126 | }, 127 | }) 128 | 129 | // 整理为输出格式 130 | docs := []*Weibo{} 131 | for _, doc := range output.Docs { 132 | wb := wbs[doc.DocId] 133 | for _, t := range output.Tokens { 134 | wb.Text = strings.Replace(wb.Text, t, ""+t+"", -1) 135 | } 136 | docs = append(docs, &wb) 137 | } 138 | response, _ := json.Marshal(&JsonResponse{Docs: docs}) 139 | 140 | w.Header().Set("Content-Type", "application/json") 141 | io.WriteString(w, string(response)) 142 | } 143 | 144 | /******************************************************************************* 145 | 主函数 146 | *******************************************************************************/ 147 | func main() { 148 | // 解析命令行参数 149 | flag.Parse() 150 | 151 | // 初始化 152 | gob.Register(WeiboScoringFields{}) 153 | log.Print("引擎开始初始化") 154 | searcher.Init(types.EngineInitOptions{ 155 | NumShards: 10, 156 | SegmenterDictionaries: *dictFile, 157 | StopTokenFile: *stopTokenFile, 158 | IndexerInitOptions: &types.IndexerInitOptions{ 159 | IndexType: types.LocationsIndex, 160 | }, 161 | // 如果你希望使用持久存储，启用下面的选项 162 | // 默认使用boltdb持久化，如果你希望修改数据库类型 163 | // 请修改 WUKONG_STORAGE_ENGINE 环境变量 164 | UsePersistentStorage: true, 165 | PersistentStorageFolder: "weibo_search", 166 | }) 167 | log.Print("引擎初始化完毕") 168 | wbs = make(map[uint64]Weibo) 169 | 170 | // 索引 171 | log.Print("建索引开始") 172 | go indexWeibo() 173 | log.Print("建索引完毕") 174 | 175 | // 捕获ctrl-c 176 | c := make(chan os.Signal, 1) 177 | signal.Notify(c, os.Interrupt) 178 | go func() { 179 | for _ = range c { 180 | log.Print("捕获Ctrl-c，退出服务器") 181 | searcher.Close() 182 | os.Exit(0) 183 | } 184 | }() 185 | 186 | http.HandleFunc("/json", JsonRpcServer) 187 | http.Handle("/", http.FileServer(http.Dir(*staticFolder))) 188 | log.Print("服务器启动") 189 | log.Fatal(http.ListenAndServe(":8080", nil)) 190 | } 191 | -------------------------------------------------------------------------------- /examples/codelab/static/index.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 微博全文搜索 6 | 12 | 13 | 62 | 63 | 64 |

微博全文搜索

65 |

66 |
67 | 68 | 69 | -------------------------------------------------------------------------------- /examples/custom_scoring_criteria.go: -------------------------------------------------------------------------------- 1 | // 一个使用自定义评分规则搜索微博数据的例子 2 | // 3 | // 微博数据文件每行的格式是"||||||||||||||||" 4 | // , 和的文本长度做评分数据 5 | // 6 | // 自定义评分规则为： 7 | // 1. 首先排除关键词紧邻距离大于150个字节(五十个汉字)的微博 8 | // 2. 按照帖子距当前时间评分，精度为天，越晚的帖子评分越高 9 | // 3. 按照帖子BM25的整数部分排名 10 | // 4. 同一天的微博再按照转发数评分，转发越多的帖子评分越高 11 | // 5. 最后按照帖子长度评分，越长的帖子评分越高 12 | 13 | package main 14 | 15 | import ( 16 | "bufio" 17 | "encoding/gob" 18 | "flag" 19 | "fmt" 20 | "github.com/huichen/wukong/engine" 21 | "github.com/huichen/wukong/types" 22 | "log" 23 | "os" 24 | "reflect" 25 | "strconv" 26 | "strings" 27 | ) 28 | 29 | const ( 30 | SecondsInADay = 86400 31 | MaxTokenProximity = 150 32 | ) 33 | 34 | var ( 35 | weibo_data = flag.String( 36 | "weibo_data", 37 | "../testdata/weibo_data.txt", 38 | "索引的微博帖子，每行当作一个文档") 39 | query = flag.String( 40 | "query", 41 | "chinajoy游戏", 42 | "待搜索的短语") 43 | dictionaries = flag.String( 44 | "dictionaries", 45 | "../data/dictionary.txt", 46 | "分词字典文件") 47 | stop_token_file = flag.String( 48 | "stop_token_file", 49 | "../data/stop_tokens.txt", 50 | "停用词文件") 51 | 52 | searcher = engine.Engine{} 53 | options = types.RankOptions{ 54 | ScoringCriteria: WeiboScoringCriteria{}, 55 | OutputOffset: 0, 56 | MaxOutputs: 100, 57 | } 58 | searchQueries = []string{} 59 | ) 60 | 61 | // 微博评分字段 62 | type WeiboScoringFields struct { 63 | // 帖子的时间戳 64 | Timestamp uint32 65 | 66 | // 帖子的转发数 67 | RepostsCount uint32 68 | 69 | // 帖子的长度 70 | TextLength int 71 | } 72 | 73 | // 自定义的微博评分规则 74 | type WeiboScoringCriteria struct { 75 | } 76 | 77 | func (criteria WeiboScoringCriteria) Score( 78 | doc types.IndexedDocument, fields interface{}) []float32 { 79 | if doc.TokenProximity > MaxTokenProximity { // 评分第一步 80 | return []float32{} 81 | } 82 | if reflect.TypeOf(fields) != reflect.TypeOf(WeiboScoringFields{}) { 83 | return []float32{} 84 | } 85 | output := make([]float32, 4) 86 | wsf := fields.(WeiboScoringFields) 87 | output[0] = float32(wsf.Timestamp / SecondsInADay) // 评分第二步 88 | output[1] = float32(int(doc.BM25)) // 评分第三步 89 | output[2] = float32(wsf.RepostsCount) // 评分第四步 90 | output[3] = float32(wsf.TextLength) // 评分第五步 91 | return output 92 | } 93 | 94 | func main() { 95 | // 解析命令行参数 96 | flag.Parse() 97 | log.Printf("待搜索的短语为\"%s\"", *query) 98 | 99 | // 初始化 100 | gob.Register(WeiboScoringFields{}) 101 | searcher.Init(types.EngineInitOptions{ 102 | SegmenterDictionaries: *dictionaries, 103 | StopTokenFile: *stop_token_file, 104 | IndexerInitOptions: &types.IndexerInitOptions{ 105 | IndexType: types.LocationsIndex, 106 | }, 107 | DefaultRankOptions: &options, 108 | }) 109 | defer searcher.Close() 110 | 111 | // 读入微博数据 112 | file, err := os.Open(*weibo_data) 113 | if err != nil { 114 | log.Fatal(err) 115 | } 116 | defer file.Close() 117 | log.Printf("读入文本 %s", *weibo_data) 118 | scanner := bufio.NewScanner(file) 119 | lines := []string{} 120 | fieldsSlice := []WeiboScoringFields{} 121 | for scanner.Scan() { 122 | data := strings.Split(scanner.Text(), "||||") 123 | if len(data) != 10 { 124 | continue 125 | } 126 | timestamp, _ := strconv.ParseUint(data[1], 10, 32) 127 | repostsCount, _ := strconv.ParseUint(data[4], 10, 32) 128 | text := data[9] 129 | if text != "" { 130 | lines = append(lines, text) 131 | fields := WeiboScoringFields{ 132 | Timestamp: uint32(timestamp), 133 | RepostsCount: uint32(repostsCount), 134 | TextLength: len(text), 135 | } 136 | fieldsSlice = append(fieldsSlice, fields) 137 | } 138 | } 139 | log.Printf("读入%d条微博\n", len(lines)) 140 | 141 | // 建立索引 142 | log.Print("建立索引") 143 | for i, text := range lines { 144 | searcher.IndexDocument(uint64(i), 145 | types.DocumentIndexData{Content: text, Fields: fieldsSlice[i]}) 146 | } 147 | searcher.FlushIndex() 148 | log.Print("索引建立完毕") 149 | 150 | // 搜索 151 | log.Printf("开始查询") 152 | output := searcher.Search(types.SearchRequest{Text: *query}) 153 | 154 | // 显示 155 | fmt.Println() 156 | for _, doc := range output.Docs { 157 | fmt.Printf("%v %s\n\n", doc.Scores, lines[doc.DocId]) 158 | } 159 | log.Printf("查询完毕") 160 | } 161 | -------------------------------------------------------------------------------- /examples/simplest_example.go: -------------------------------------------------------------------------------- 1 | /* 2 | 3 | 没有比这个更简单的例子了。 4 | 5 | */ 6 | 7 | package main 8 | 9 | import ( 10 | "github.com/huichen/wukong/engine" 11 | "github.com/huichen/wukong/types" 12 | "log" 13 | ) 14 | 15 | var ( 16 | // searcher是线程安全的 17 | searcher = engine.Engine{} 18 | ) 19 | 20 | func main() { 21 | // 初始化 22 | searcher.Init(types.EngineInitOptions{ 23 | SegmenterDictionaries: "../data/dictionary.txt"}) 24 | defer searcher.Close() 25 | 26 | // 将文档加入索引 27 | searcher.IndexDocument(0, types.DocumentIndexData{Content: "此次百度收购将成中国互联网最大并购"}) 28 | searcher.IndexDocument(1, types.DocumentIndexData{Content: "百度宣布拟全资收购91无线业务"}) 29 | searcher.IndexDocument(2, types.DocumentIndexData{Content: "百度是中国最大的搜索引擎"}) 30 | 31 | // 强制索引刷新 32 | searcher.FlushIndex() 33 | 34 | // 搜索输出格式见types.SearchResponse结构体 35 | log.Print(searcher.Search(types.SearchRequest{Text: "百度中国"})) 36 | } 37 | -------------------------------------------------------------------------------- /license.txt: -------------------------------------------------------------------------------- 1 | Copyright 2013 Hui Chen 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | -------------------------------------------------------------------------------- /storage/bolt_storage.go: -------------------------------------------------------------------------------- 1 | package storage 2 | 3 | import ( 4 | "github.com/boltdb/bolt" 5 | "time" 6 | ) 7 | 8 | var wukong_documents = []byte("wukong_documents") 9 | 10 | type boltStorage struct { 11 | db *bolt.DB 12 | } 13 | 14 | func openBoltStorage(path string) (Storage, error) { 15 | db, err := bolt.Open(path, 0600, &bolt.Options{Timeout: 3600 * time.Second}) 16 | if err != nil { 17 | return nil, err 18 | } 19 | err = db.Update(func(tx *bolt.Tx) error { 20 | _, err := tx.CreateBucketIfNotExists(wukong_documents) 21 | return err 22 | }) 23 | if err != nil { 24 | db.Close() 25 | return nil, err 26 | } 27 | return &boltStorage{db}, nil 28 | } 29 | 30 | func (s *boltStorage) WALName() string { 31 | return s.db.Path() 32 | } 33 | 34 | func (s *boltStorage) Set(k []byte, v []byte) error { 35 | return s.db.Update(func(tx *bolt.Tx) error { 36 | return tx.Bucket(wukong_documents).Put(k, v) 37 | }) 38 | } 39 | 40 | func (s *boltStorage) Get(k []byte) (b []byte, err error) { 41 | err = s.db.View(func(tx *bolt.Tx) error { 42 | b = tx.Bucket(wukong_documents).Get(k) 43 | return nil 44 | }) 45 | return 46 | } 47 | 48 | func (s *boltStorage) Delete(k []byte) error { 49 | return s.db.Update(func(tx *bolt.Tx) error { 50 | return tx.Bucket(wukong_documents).Delete(k) 51 | }) 52 | } 53 | 54 | func (s *boltStorage) ForEach(fn func(k, v []byte) error) error { 55 | return s.db.View(func(tx *bolt.Tx) error { 56 | b := tx.Bucket(wukong_documents) 57 | c := b.Cursor() 58 | for k, v := c.First(); k != nil; k, v = c.Next() { 59 | if err := fn(k, v); err != nil { 60 | return err 61 | } 62 | } 63 | return nil 64 | }) 65 | } 66 | 67 | func (s *boltStorage) Close() error { 68 | return s.db.Close() 69 | } 70 | -------------------------------------------------------------------------------- /storage/bolt_storage_test.go: -------------------------------------------------------------------------------- 1 | package storage 2 | 3 | import ( 4 | "github.com/huichen/wukong/utils" 5 | "os" 6 | "testing" 7 | ) 8 | 9 | func TestOpenOrCreateBolt(t *testing.T) { 10 | db, err := openBoltStorage("bolt_test") 11 | utils.Expect(t, "", err) 12 | db.Close() 13 | 14 | db, err = openBoltStorage("bolt_test") 15 | utils.Expect(t, "", err) 16 | err = db.Set([]byte("key1"), []byte("value1")) 17 | utils.Expect(t, "", err) 18 | 19 | buffer := make([]byte, 100) 20 | buffer, err = db.Get([]byte("key1")) 21 | utils.Expect(t, "", err) 22 | utils.Expect(t, "value1", string(buffer)) 23 | 24 | walFile := db.WALName() 25 | db.Close() 26 | os.Remove(walFile) 27 | os.Remove("bolt_test") 28 | } 29 | -------------------------------------------------------------------------------- /storage/kv_storage.go: -------------------------------------------------------------------------------- 1 | package storage 2 | 3 | import ( 4 | "github.com/cznic/kv" 5 | "io" 6 | ) 7 | 8 | type kvStorage struct { 9 | db *kv.DB 10 | } 11 | 12 | func openKVStorage(path string) (Storage, error) { 13 | options := &kv.Options{} 14 | db, errOpen := kv.Open(path, options) 15 | if errOpen != nil { 16 | var errCreate error 17 | db, errCreate = kv.Create(path, options) 18 | if errCreate != nil { 19 | return &kvStorage{db}, errCreate 20 | } 21 | } 22 | return &kvStorage{db}, nil 23 | } 24 | 25 | func (s *kvStorage) WALName() string { 26 | return s.db.WALName() 27 | } 28 | 29 | func (s *kvStorage) Set(k []byte, v []byte) error { 30 | return s.db.Set(k, v) 31 | } 32 | 33 | func (s *kvStorage) Get(k []byte) ([]byte, error) { 34 | return s.db.Get(nil, k) 35 | } 36 | 37 | func (s *kvStorage) Delete(k []byte) error { 38 | return s.db.Delete(k) 39 | } 40 | 41 | func (s *kvStorage) ForEach(fn func(k, v []byte) error) error { 42 | iter, err := s.db.SeekFirst() 43 | if err == io.EOF { 44 | return nil 45 | } else if err != nil { 46 | return err 47 | } 48 | for { 49 | key, value, err := iter.Next() 50 | if err == io.EOF { 51 | break 52 | } else if err != nil { 53 | return err 54 | } 55 | if err := fn(key, value); err != nil { 56 | return err 57 | } 58 | } 59 | return nil 60 | } 61 | 62 | func (s *kvStorage) Close() error { 63 | return s.db.Close() 64 | } 65 | -------------------------------------------------------------------------------- /storage/kv_storage_test.go: -------------------------------------------------------------------------------- 1 | package storage 2 | 3 | import ( 4 | "github.com/huichen/wukong/utils" 5 | "os" 6 | "testing" 7 | ) 8 | 9 | func TestOpenOrCreateKv(t *testing.T) { 10 | db, err := openKVStorage("kv_test") 11 | utils.Expect(t, "", err) 12 | db.Close() 13 | 14 | db, err = openKVStorage("kv_test") 15 | utils.Expect(t, "", err) 16 | err = db.Set([]byte("key1"), []byte("value1")) 17 | utils.Expect(t, "", err) 18 | 19 | buffer := make([]byte, 100) 20 | buffer, err = db.Get([]byte("key1")) 21 | utils.Expect(t, "", err) 22 | utils.Expect(t, "value1", string(buffer)) 23 | 24 | walFile := db.WALName() 25 | db.Close() 26 | os.Remove(walFile) 27 | os.Remove("kv_test") 28 | } 29 | -------------------------------------------------------------------------------- /storage/ldb_storage.go: -------------------------------------------------------------------------------- 1 | package storage 2 | 3 | import ( 4 | "github.com/syndtr/goleveldb/leveldb" 5 | ) 6 | 7 | type leveldbStorage struct { 8 | db *leveldb.DB 9 | } 10 | 11 | func OpenLeveldbStorage(path string) (Storage, error) { 12 | 13 | if db, err := leveldb.OpenFile(path, nil); err != nil { 14 | return nil, err 15 | } else { 16 | return &leveldbStorage{db}, nil 17 | } 18 | 19 | } 20 | 21 | func (s *leveldbStorage) WALName() string { 22 | return "" //对于此数据库，本函数没什么卵用~ 23 | } 24 | 25 | func (s *leveldbStorage) Set(k, v []byte) error { 26 | return s.db.Put(k, v, nil) 27 | } 28 | 29 | func (s *leveldbStorage) Get(k []byte) ([]byte, error) { 30 | return s.db.Get(k, nil) 31 | } 32 | 33 | func (s *leveldbStorage) Delete(k []byte) error { 34 | return s.db.Delete(k, nil) 35 | } 36 | 37 | func (s *leveldbStorage) ForEach(fn func(k, v []byte) error) error { 38 | iter := s.db.NewIterator(nil, nil) 39 | for iter.Next() { 40 | // Remember that the contents of the returned slice should not be modified, and 41 | // only valid until the next call to Next. 42 | key := iter.Key() 43 | value := iter.Value() 44 | if err := fn(key, value); err != nil { 45 | return err 46 | } 47 | } 48 | iter.Release() 49 | return iter.Error() 50 | } 51 | 52 | func (s *leveldbStorage) Close() error { 53 | return s.db.Close() 54 | } 55 | -------------------------------------------------------------------------------- /storage/storage.go: -------------------------------------------------------------------------------- 1 | package storage 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | ) 7 | 8 | const DEFAULT_STORAGE_ENGINE = "ldb" 9 | 10 | var supportedStorage = map[string]func(path string) (Storage, error){ 11 | "kv": openKVStorage, 12 | "ldb": OpenLeveldbStorage, 13 | "bolt": openBoltStorage, 14 | } 15 | 16 | func RegisterStorageEngine(name string, fn func(path string) (Storage, error)) { 17 | supportedStorage[name] = fn 18 | } 19 | 20 | type Storage interface { 21 | Set(k, v []byte) error 22 | Get(k []byte) ([]byte, error) 23 | Delete(k []byte) error 24 | ForEach(fn func(k, v []byte) error) error 25 | Close() error 26 | WALName() string 27 | } 28 | 29 | func OpenStorage(path string) (Storage, error) { 30 | wse := os.Getenv("WUKONG_STORAGE_ENGINE") 31 | if wse == "" { 32 | wse = DEFAULT_STORAGE_ENGINE 33 | } 34 | if fn, has := supportedStorage[wse]; has { 35 | return fn(path) 36 | } 37 | return nil, fmt.Errorf("unsupported storage engine %v", wse) 38 | } 39 | -------------------------------------------------------------------------------- /testdata/crawl_weibo_data.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "flag" 5 | "fmt" 6 | "github.com/huichen/gobo" 7 | "github.com/huichen/gobo/contrib" 8 | "io/ioutil" 9 | "log" 10 | "os" 11 | "strings" 12 | "time" 13 | ) 14 | 15 | var ( 16 | access_token = flag.String("access_token", "", "用户的访问令牌") 17 | weibo = gobo.Weibo{} 18 | users_file = flag.String("users_file", "users.txt", "从该文件读入要下载的微博用户名，每个名字一行") 19 | output_file = flag.String("output_file", "weibo_data.txt", "将抓取的微博写入下面的文件") 20 | num_weibos = flag.Int("num_weibos", 2000, "从每个微博账号中抓取多少条微博") 21 | ) 22 | 23 | func main() { 24 | flag.Parse() 25 | 26 | // 读取用户名 27 | content, err := ioutil.ReadFile(*users_file) 28 | if err != nil { 29 | log.Fatal("无法读取-users_file") 30 | } 31 | users := strings.Split(string(content), "\n") 32 | 33 | outputFile, _ := os.Create(*output_file) 34 | defer outputFile.Close() 35 | 36 | // 抓微博 37 | for _, user := range users { 38 | if user == "" { 39 | continue 40 | } 41 | log.Printf("抓取 @%s 的微博", user) 42 | statuses, err := contrib.GetStatuses( 43 | &weibo, *access_token, user, 0, *num_weibos, 5000) // 超时5秒 44 | if err != nil { 45 | log.Print(err) 46 | continue 47 | } 48 | 49 | for _, status := range statuses { 50 | t, _ := time.Parse("Mon Jan 2 15:04:05 -0700 2006", status.Created_At) 51 | outputFile.WriteString(fmt.Sprintf( 52 | "%d||||%d||||%d||||%s||||%d||||%d||||%d||||%s||||%s||||%s\n", 53 | status.Id, uint32(t.Unix()), status.User.Id, status.User.Screen_Name, 54 | status.Reposts_Count, status.Comments_Count, status.Attitudes_Count, 55 | status.Thumbnail_Pic, status.Original_Pic, status.Text)) 56 | } 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /testdata/test_dict.txt: -------------------------------------------------------------------------------- 1 | 中 64 p1 2 | 国 64 p2 3 | 有 64 p3 4 | 三 64 p4 5 | 亿 64 p5 6 | 人 64 p6 7 | 口 64 p7 8 | 中国 32 p8 9 | 国有 8 p9 10 | 十三 16 p10 11 | 十三亿 4 p11 12 | 人口 16 p12 13 | -------------------------------------------------------------------------------- /testdata/users.txt: -------------------------------------------------------------------------------- 1 | 新浪房产 2 | 中国企业家杂志 3 | 头条博客 4 | 第一财经周刊 5 | 重庆晨报 6 | 华西都市报 7 | 羊城晚报 8 | 新浪军事 9 | 头条新闻 10 | 新浪体育 11 | 新浪财经 12 | 南方周末 13 | 南都周刊 14 | 财经网 15 | 中国新闻周刊 16 | 新浪证券 17 | 新浪娱乐 18 | 爱范儿 19 | 新京报 20 | 南方都市报 21 | 法制晚报 22 | 华尔街日报中文网 23 | 南方人物周刊 24 | 扬子晚报 25 | 新周刊 26 | Donews官方微博 27 | 现代快报 28 | 环球企业家 29 | 财新网 30 | 新发现杂志 31 | 环球网 32 | 新民晚报新民网 33 | 中国经济周刊 34 | 百度搜索风云榜 35 | 36氪 36 | 信息时报 37 | 新浪评论 38 | 新华社中国网事 39 | CSDN 40 | 谣言粉碎机 41 | 都市快报 42 | 果壳网 43 | 重庆商报 44 | 微天下 45 | 爱极客 46 | 外交小灵通 47 | 环球时报 48 | 人民网 49 | 新浪视野 50 | 人民日报 51 | 非诚勿扰 52 | 最体育 53 | -------------------------------------------------------------------------------- /types/doc_info.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | import ( 4 | "sync" 5 | ) 6 | 7 | // 文档信息[id]info 8 | type DocInfosShard struct { 9 | DocInfos map[uint64]*DocInfo 10 | NumDocuments uint64 // 这实际上是总文档数的一个近似 11 | sync.RWMutex 12 | } 13 | 14 | type DocInfo struct { 15 | Fields interface{} 16 | TokenLengths float32 17 | } 18 | -------------------------------------------------------------------------------- /types/document_index_data.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | type DocumentIndexData struct { 4 | // 文档全文（必须是UTF-8格式），用于生成待索引的关键词 5 | Content string 6 | 7 | // 文档的关键词 8 | // 当Content不为空的时候，优先从Content中分词得到关键词。 9 | // Tokens存在的意义在于绕过悟空内置的分词器，在引擎外部 10 | // 进行分词和预处理。 11 | Tokens []TokenData 12 | 13 | // 文档标签（必须是UTF-8格式），比如文档的类别属性等，这些标签并不出现在文档文本中 14 | Labels []string 15 | 16 | // 文档的评分字段，可以接纳任何类型的结构体 17 | Fields interface{} 18 | } 19 | 20 | // 文档的一个关键词 21 | type TokenData struct { 22 | // 关键词的字符串 23 | Text string 24 | 25 | // 关键词的首字节在文档中出现的位置 26 | Locations []int 27 | } 28 | -------------------------------------------------------------------------------- /types/engine_init_options.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | import ( 4 | "log" 5 | "runtime" 6 | ) 7 | 8 | var ( 9 | // EngineInitOptions的默认值 10 | defaultNumSegmenterThreads = runtime.NumCPU() 11 | defaultNumShards = 8 12 | defaultIndexerBufferLength = runtime.NumCPU() 13 | defaultNumIndexerThreadsPerShard = runtime.NumCPU() 14 | defaultRankerBufferLength = runtime.NumCPU() 15 | defaultNumRankerThreadsPerShard = runtime.NumCPU() 16 | defaultDefaultRankOptions = RankOptions{ 17 | ScoringCriteria: RankByBM25{}, 18 | } 19 | defaultIndexerInitOptions = IndexerInitOptions{ 20 | IndexType: FrequenciesIndex, 21 | BM25Parameters: &defaultBM25Parameters, 22 | } 23 | defaultBM25Parameters = BM25Parameters{ 24 | K1: 2.0, 25 | B: 0.75, 26 | } 27 | ) 28 | 29 | type EngineInitOptions struct { 30 | // 是否使用分词器 31 | // 默认使用，否则在启动阶段跳过SegmenterDictionaries和StopTokenFile设置 32 | // 如果你不需要在引擎内分词，可以将这个选项设为true 33 | // 注意，如果你不用分词器，那么在调用IndexDocument时DocumentIndexData中的Content会被忽略 34 | NotUsingSegmenter bool 35 | 36 | // 半角逗号分隔的字典文件，具体用法见 37 | // sego.Segmenter.LoadDictionary函数的注释 38 | SegmenterDictionaries string 39 | 40 | // 停用词文件 41 | StopTokenFile string 42 | 43 | // 分词器线程数 44 | NumSegmenterThreads int 45 | 46 | // 索引器/排序器/持久数据库的shard数目 47 | // 被检索/排序的文档会被均匀分配到各个shard中 48 | // 每个shard对应一对数据库文件（反向索引数据库和文档字段数据库） 49 | NumShards int 50 | 51 | // 索引器的信道缓冲长度 52 | IndexerBufferLength int 53 | 54 | // 索引器每个shard分配的线程数 55 | NumIndexerThreadsPerShard int 56 | 57 | // 排序器的信道缓冲长度 58 | RankerBufferLength int 59 | 60 | // 排序器每个shard分配的线程数 61 | NumRankerThreadsPerShard int 62 | 63 | // 索引器初始化选项 64 | IndexerInitOptions *IndexerInitOptions 65 | 66 | // 默认的搜索选项 67 | DefaultRankOptions *RankOptions 68 | 69 | // 是否使用持久数据库，以及数据库文件保存的目录 70 | UsePersistentStorage bool 71 | PersistentStorageFolder string 72 | } 73 | 74 | // 初始化EngineInitOptions，当用户未设定某个选项的值时用默认值取代 75 | func (options *EngineInitOptions) Init() { 76 | if !options.NotUsingSegmenter { 77 | if options.SegmenterDictionaries == "" { 78 | log.Fatal("字典文件不能为空") 79 | } 80 | } 81 | 82 | if options.NumSegmenterThreads == 0 { 83 | options.NumSegmenterThreads = defaultNumSegmenterThreads 84 | } 85 | 86 | if options.NumShards == 0 { 87 | options.NumShards = defaultNumShards 88 | } 89 | 90 | if options.IndexerBufferLength == 0 { 91 | options.IndexerBufferLength = defaultIndexerBufferLength 92 | } 93 | 94 | if options.NumIndexerThreadsPerShard == 0 { 95 | options.NumIndexerThreadsPerShard = defaultNumIndexerThreadsPerShard 96 | } 97 | 98 | if options.RankerBufferLength == 0 { 99 | options.RankerBufferLength = defaultRankerBufferLength 100 | } 101 | 102 | if options.NumRankerThreadsPerShard == 0 { 103 | options.NumRankerThreadsPerShard = defaultNumRankerThreadsPerShard 104 | } 105 | 106 | if options.IndexerInitOptions == nil { 107 | options.IndexerInitOptions = &defaultIndexerInitOptions 108 | } 109 | 110 | if options.IndexerInitOptions.BM25Parameters == nil { 111 | options.IndexerInitOptions.BM25Parameters = &defaultBM25Parameters 112 | } 113 | 114 | if options.DefaultRankOptions == nil { 115 | options.DefaultRankOptions = &defaultDefaultRankOptions 116 | } 117 | 118 | if options.DefaultRankOptions.ScoringCriteria == nil { 119 | options.DefaultRankOptions.ScoringCriteria = defaultDefaultRankOptions.ScoringCriteria 120 | } 121 | } 122 | -------------------------------------------------------------------------------- /types/index.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | type DocumentIndex struct { 4 | // 文本的DocId 5 | DocId uint64 6 | 7 | // 文本的关键词长 8 | TokenLength float32 9 | 10 | // 加入的索引键 11 | Keywords []KeywordIndex 12 | } 13 | 14 | // 反向索引项，这实际上标注了一个（搜索键，文档）对。 15 | type KeywordIndex struct { 16 | // 搜索键的UTF-8文本 17 | Text string 18 | 19 | // 搜索键词频 20 | Frequency float32 21 | 22 | // 搜索键在文档中的起始字节位置，按照升序排列 23 | Starts []int 24 | } 25 | 26 | // 索引器返回结果 27 | type IndexedDocument struct { 28 | DocId uint64 29 | 30 | // BM25，仅当索引类型为FrequenciesIndex或者LocationsIndex时返回有效值 31 | BM25 float32 32 | 33 | // 关键词在文档中的紧邻距离，紧邻距离的含义见computeTokenProximity的注释。 34 | // 仅当索引类型为LocationsIndex时返回有效值。 35 | TokenProximity int32 36 | 37 | // 紧邻距离计算得到的关键词位置，和Lookup函数输入tokens的长度一样且一一对应。 38 | // 仅当索引类型为LocationsIndex时返回有效值。 39 | TokenSnippetLocations []int 40 | 41 | // 关键词在文本中的具体位置。 42 | // 仅当索引类型为LocationsIndex时返回有效值。 43 | TokenLocations [][]int 44 | } 45 | -------------------------------------------------------------------------------- /types/indexer_init_options.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | // 这些常数定义了反向索引表存储的数据类型 4 | const ( 5 | // 仅存储文档的docId 6 | DocIdsIndex = 0 7 | 8 | // 存储关键词的词频，用于计算BM25 9 | FrequenciesIndex = 1 10 | 11 | // 存储关键词在文档中出现的具体字节位置（可能有多个） 12 | // 如果你希望得到关键词紧邻度数据，必须使用LocationsIndex类型的索引 13 | LocationsIndex = 2 14 | ) 15 | 16 | // 初始化索引器选项 17 | type IndexerInitOptions struct { 18 | // 索引表的类型，见上面的常数 19 | IndexType int 20 | 21 | // BM25参数 22 | BM25Parameters *BM25Parameters 23 | } 24 | 25 | // 见http://en.wikipedia.org/wiki/Okapi_BM25 26 | // 默认值见engine_init_options.go 27 | type BM25Parameters struct { 28 | K1 float32 29 | B float32 30 | } 31 | -------------------------------------------------------------------------------- /types/inverted_index.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | import ( 4 | "sync" 5 | ) 6 | 7 | // 反向索引表([关键词]反向索引表) 8 | type InvertedIndexShard struct { 9 | InvertedIndex map[string]*KeywordIndices 10 | TotalTokenLength float32 //总关键词数 11 | sync.RWMutex 12 | } 13 | 14 | // 反向索引表的一行，收集了一个搜索键出现的所有文档，按照DocId从小到大排序。 15 | type KeywordIndices struct { 16 | // 下面的切片是否为空，取决于初始化时IndexType的值 17 | DocIds []uint64 // 全部类型都有 18 | Frequencies []float32 // IndexType == FrequenciesIndex 19 | Locations [][]int // IndexType == LocationsIndex 20 | } 21 | -------------------------------------------------------------------------------- /types/scoring_criteria.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | // 评分规则通用接口 4 | type ScoringCriteria interface { 5 | // 给一个文档评分，文档排序时先用第一个分值比较，如果 6 | // 分值相同则转移到第二个分值，以此类推。 7 | // 返回空切片表明该文档应该从最终排序结果中剔除。 8 | Score(doc IndexedDocument, fields interface{}) []float32 9 | } 10 | 11 | // 一个简单的评分规则，文档分数为BM25 12 | type RankByBM25 struct { 13 | } 14 | 15 | func (rule RankByBM25) Score(doc IndexedDocument, fields interface{}) []float32 { 16 | return []float32{doc.BM25} 17 | } 18 | -------------------------------------------------------------------------------- /types/search_request.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | type SearchRequest struct { 4 | // 搜索的短语（必须是UTF-8格式），会被分词 5 | // 当值为空字符串时关键词会从下面的Tokens读入 6 | Text string 7 | 8 | // 关键词（必须是UTF-8格式），当Text不为空时优先使用Text 9 | // 通常你不需要自己指定关键词，除非你运行自己的分词程序 10 | Tokens []string 11 | 12 | // 文档标签（必须是UTF-8格式），标签不存在文档文本中，但也属于搜索键的一种 13 | Labels []string 14 | 15 | // 当不为nil时，仅从这些DocIds包含的键中搜索（忽略值） 16 | DocIds map[uint64]bool 17 | 18 | // 排序选项 19 | RankOptions *RankOptions 20 | 21 | // 超时，单位毫秒（千分之一秒）。此值小于等于零时不设超时。 22 | // 搜索超时的情况下仍有可能返回部分排序结果。 23 | Timeout int 24 | 25 | // 设为true时仅统计搜索到的文档个数，不返回具体的文档 26 | CountDocsOnly bool 27 | 28 | // 不排序，对于可在引擎外部（比如客户端）排序情况适用 29 | // 对返回文档很多的情况打开此选项可以有效节省时间 30 | Orderless bool 31 | } 32 | 33 | type RankOptions struct { 34 | // 文档的评分规则，值为nil时使用Engine初始化时设定的规则 35 | ScoringCriteria ScoringCriteria 36 | 37 | // 默认情况下（ReverseOrder=false）按照分数从大到小排序，否则从小到大排序 38 | ReverseOrder bool 39 | 40 | // 从第几条结果开始输出 41 | OutputOffset int 42 | 43 | // 最大输出的搜索结果数，为0时无限制 44 | MaxOutputs int 45 | } 46 | -------------------------------------------------------------------------------- /types/search_response.go: -------------------------------------------------------------------------------- 1 | package types 2 | 3 | import ( 4 | "github.com/huichen/wukong/utils" 5 | ) 6 | 7 | type SearchResponse struct { 8 | // 搜索用到的关键词 9 | Tokens []string 10 | 11 | // 搜索到的文档，已排序 12 | Docs []ScoredDocument 13 | 14 | // 搜索是否超时。超时的情况下也可能会返回部分结果 15 | Timeout bool 16 | 17 | // 搜索到的文档个数。注意这是全部文档中满足条件的个数，可能比返回的文档数要大 18 | NumDocs int 19 | } 20 | 21 | type ScoredDocument struct { 22 | DocId uint64 23 | 24 | // 文档的打分值 25 | // 搜索结果按照Scores的值排序，先按照第一个数排，如果相同则按照第二个数排序，依次类推。 26 | Scores []float32 27 | 28 | // 用于生成摘要的关键词在文本中的字节位置，该切片长度和SearchResponse.Tokens的长度一样 29 | // 只有当IndexType == LocationsIndex时不为空 30 | TokenSnippetLocations []int 31 | 32 | // 关键词出现的位置 33 | // 只有当IndexType == LocationsIndex时不为空 34 | TokenLocations [][]int 35 | } 36 | 37 | // 为了方便排序 38 | 39 | type ScoredDocuments []ScoredDocument 40 | 41 | func (docs ScoredDocuments) Len() int { 42 | return len(docs) 43 | } 44 | func (docs ScoredDocuments) Swap(i, j int) { 45 | docs[i], docs[j] = docs[j], docs[i] 46 | } 47 | func (docs ScoredDocuments) Less(i, j int) bool { 48 | // 为了从大到小排序，这实际上实现的是More的功能 49 | for iScore := 0; iScore < utils.MinInt(len(docs[i].Scores), len(docs[j].Scores)); iScore++ { 50 | if docs[i].Scores[iScore] > docs[j].Scores[iScore] { 51 | return true 52 | } else if docs[i].Scores[iScore] < docs[j].Scores[iScore] { 53 | return false 54 | } 55 | } 56 | return len(docs[i].Scores) > len(docs[j].Scores) 57 | } 58 | -------------------------------------------------------------------------------- /utils/test_utils.go: -------------------------------------------------------------------------------- 1 | package utils 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | ) 7 | 8 | func Expect(t *testing.T, expect string, actual interface{}) { 9 | actualString := fmt.Sprint(actual) 10 | if expect != actualString { 11 | t.Errorf("期待值=\"%s\", 实际=\"%s\"", expect, actualString) 12 | } 13 | } 14 | -------------------------------------------------------------------------------- /utils/utils.go: -------------------------------------------------------------------------------- 1 | package utils 2 | 3 | func AbsInt(a int) int { 4 | if a < 0 { 5 | return -a 6 | } 7 | return a 8 | } 9 | 10 | func MinInt(a, b int) int { 11 | if a < b { 12 | return a 13 | } 14 | return b 15 | } 16 | -------------------------------------------------------------------------------- /wukong.go: -------------------------------------------------------------------------------- 1 | package wukong 2 | 3 | import ( 4 | _ "github.com/boltdb/bolt" 5 | _ "github.com/cznic/kv" 6 | _ "github.com/huichen/murmur" 7 | _ "github.com/huichen/sego" 8 | ) 9 | --------------------------------------------------------------------------------