使用roaring bitmap创建倒排索引

package main

import (
    "fmt"
    "github.com/RoaringBitmap/roaring"
    "sync"
    "testing"
)

// InvertedIndex 倒排索引（线程安全）
type InvertedIndex struct {
    index sync.Map // key: term string, value: *roaring.Bitmap
}

func NewInvertedIndex() *InvertedIndex {
    return &InvertedIndex{}
}

// AddDocument 将文档内容加入倒排索引
// docID: 文档唯一标识（需保证唯一性）
// terms: 文档分词后的词项列表
func (idx *InvertedIndex) AddDocument(docID uint32, terms []string) {
    for _, term := range terms {
        // 获取或创建该词项对应的位图
        bm, _ := idx.index.LoadOrStore(term, roaring.New())
        bitmap := bm.(*roaring.Bitmap)
        bitmap.Add(docID)
    }
}

// Search 查找包含指定词项的文档ID集合
func (idx *InvertedIndex) Search(term string) *roaring.Bitmap {
    if bm, ok := idx.index.Load(term); ok {
        // 返回克隆，避免外部修改影响内部数据
        return bm.(*roaring.Bitmap).Clone()
    }
    return roaring.New() // 返回空集合
}

// SearchAND 查找同时包含多个词项的文档（交集）
func (idx *InvertedIndex) SearchAND(terms ...string) *roaring.Bitmap {
    if len(terms) == 0 {
        return roaring.New()
    }

    result := idx.Search(terms[0])
    for _, term := range terms[1:] {
        bm := idx.Search(term)
        result.And(bm) // 位图交集操作
    }
    return result
}

// SearchOR 查找包含任意词项的文档（并集）
func (idx *InvertedIndex) SearchOR(terms ...string) *roaring.Bitmap {
    result := roaring.New()
    for _, term := range terms {
        bm := idx.Search(term)
        result.Or(bm) // 位图并集操作
    }
    return result
}

func TestBitmap(t *testing.T) {
    idx := NewInvertedIndex()

    // 添加文档（假设文档ID为 1, 2, 3）
    idx.AddDocument(1, []string{"apple", "fruit", "red"})
    idx.AddDocument(2, []string{"banana", "fruit", "yellow"})
    idx.AddDocument(3, []string{"apple", "pie", "sweet"})

    // 查询单个词项
    fmt.Println("包含 'apple' 的文档:", idx.Search("apple").ToArray())
    // 输出: [1 3]

    // 多词交集查询（AND）
    fmt.Println("同时包含 'apple' 和 'fruit' 的文档:", idx.SearchAND("apple", "fruit").ToArray())
    // 输出: [1]

    // 多词并集查询（OR）
    fmt.Println("包含 'banana' 或 'pie' 的文档:", idx.SearchOR("banana", "pie").ToArray())
    // 输出: [2 3]
}