21107032990

Committed 18 Jan 2026 06:00AM UTC coverage: 88.74% (-0.6%) from 89.34%

Build # 21107032990

Build Type

push

github

Committed by

freeeve

Commit Message

Unify Batch API and clean up public interface

- Replace AddBatch/AddBatchN with consistent Batch() pattern across all types
- Index, BitmapFilter, and SortColumn now all use Batch()/BatchSize() -> Add() -> Flush()
- Make Document struct unexported (users use Batch().Add(id, text))
- Make BitmapFilter.Fields and SortColumn.Values unexported
- Add MemoryUsage() methods to BitmapFilter and SortColumn
- Update README with new API examples
- Apply gofmt -s

Run Details

157 of 182 new or added lines in 2 files covered. (86.26%)

61 existing lines in 2 files now uncovered.

1387 of 1563 relevant lines covered (88.74%)

240655.24 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

88.12

/fields.go

package roaringsearch

import (
        "cmp"
        "container/heap"
        "io"
        "os"
        "slices"
        "sync"
        "unsafe"

        "github.com/RoaringBitmap/roaring"
        "github.com/vmihailenco/msgpack/v5"
)

// BitmapFilter provides fast filtering by multiple category fields using bitmap indexes.
// Each field (e.g., "media_type", "language") can have multiple category values
// (e.g., "book", "movie"), and each category maps to a bitmap of document IDs.
//
// Example:
//
//        filter := NewBitmapFilter()
//        filter.Set(1, "media_type", "book")
//        filter.Set(1, "language", "english")
//        filter.Set(2, "media_type", "movie")
//
//        books := filter.Get("media_type", "book")           // bitmap of books
//        english := filter.Get("language", "english")        // bitmap of english
//        englishBooks := roaring.And(books, english)         // AND filter
type BitmapFilter struct {
        mu     sync.RWMutex
        fields map[string]map[string]*roaring.Bitmap
}

// NewBitmapFilter creates a new bitmap filter.
func NewBitmapFilter() *BitmapFilter {
        return &BitmapFilter{
                fields: make(map[string]map[string]*roaring.Bitmap),
        }
}

// Set assigns a document to a category within a field.
func (c *BitmapFilter) Set(docID uint32, field, category string) {
        c.mu.Lock()
        defer c.mu.Unlock()
        c.setLocked(docID, field, category)
}

func (c *BitmapFilter) setLocked(docID uint32, field, category string) {
        fieldMap, ok := c.fields[field]
        if !ok {
                fieldMap = make(map[string]*roaring.Bitmap)
                c.fields[field] = fieldMap
        }

        bm, ok := fieldMap[category]
        if !ok {
                bm = roaring.New()
                fieldMap[category] = bm
        }
        bm.Add(docID)
}

// FilterBatch accumulates entries for efficient batch insertion.
type FilterBatch struct {
        filter     *BitmapFilter
        field      string
        docIDs     []uint32
        categories []string
}

// Batch creates a new batch builder for the given field.
// Use BatchSize for better performance when you know the approximate count.
func (c *BitmapFilter) Batch(field string) *FilterBatch {
        return c.BatchSize(field, 1024)
}

// BatchSize creates a batch builder with pre-allocated capacity.
func (c *BitmapFilter) BatchSize(field string, size int) *FilterBatch {
        return &FilterBatch{
                filter:     c,
                field:      field,
                docIDs:     make([]uint32, 0, size),
                categories: make([]string, 0, size),
        }
}

// Add adds a document with a category to the batch.
func (b *FilterBatch) Add(docID uint32, category string) {
        b.docIDs = append(b.docIDs, docID)
        b.categories = append(b.categories, category)
}

// Flush commits all accumulated entries to the filter.
func (b *FilterBatch) Flush() {
        if len(b.docIDs) == 0 {
                return
        }

        n := len(b.docIDs)

        // Pass 1: build integer index array using linear search (fast for few categories)
        categoryList := make([]string, 0, 16)
        indices := make([]int, n)

        for i, cat := range b.categories {
                // Linear search - faster than map for small category counts
                idx := -1
                for j, existing := range categoryList {
                        if existing == cat {
                                idx = j
                                break
                        }
                }
                if idx == -1 {
                        idx = len(categoryList)
                        categoryList = append(categoryList, cat)
                }
                indices[i] = idx
        }

        numCats := len(categoryList)

        // Pass 2: count per category (fast integer array access)
        counts := make([]int, numCats)
        for _, idx := range indices {
                counts[idx]++
        }

        // Pre-allocate exact-sized groups
        groups := make([][]uint32, numCats)
        for i, count := range counts {
                groups[i] = make([]uint32, 0, count)
        }

        // Pass 3: fill groups (fast integer array access, no reallocation)
        for i, idx := range indices {
                groups[idx] = append(groups[idx], b.docIDs[i])
        }

        b.filter.mu.Lock()
        defer b.filter.mu.Unlock()

        fieldMap, ok := b.filter.fields[b.field]
        if !ok {
                fieldMap = make(map[string]*roaring.Bitmap)
                b.filter.fields[b.field] = fieldMap
        }

        // Create bitmaps
        bitmaps := make([]*roaring.Bitmap, numCats)
        for idx := range groups {
                cat := categoryList[idx]
                bm, ok := fieldMap[cat]
                if !ok {
                        bm = roaring.New()
                        fieldMap[cat] = bm
                }
                bitmaps[idx] = bm
        }

        // Parallel AddMany if enough categories
        if numCats >= 4 {
                var wg sync.WaitGroup
                for idx, ids := range groups {
                        wg.Add(1)
                        go func(bm *roaring.Bitmap, docIDs []uint32) {
                                defer wg.Done()
                                bm.AddMany(docIDs)
                        }(bitmaps[idx], ids)
                }
                wg.Wait()
        } else {
                for idx, ids := range groups {
                        bitmaps[idx].AddMany(ids)
                }
        }

        // Clear for reuse
        b.docIDs = b.docIDs[:0]
        b.categories = b.categories[:0]
}

// Remove removes a document from all categories across all fields.
func (c *BitmapFilter) Remove(docID uint32) {
        c.mu.Lock()
        defer c.mu.Unlock()

        for _, fieldMap := range c.fields {
                for _, bm := range fieldMap {
                        bm.Remove(docID)
                }
        }
}

// Get returns a bitmap of documents in the given category for a field.
// Returns nil if field or category doesn't exist.
func (c *BitmapFilter) Get(field, category string) *roaring.Bitmap {
        c.mu.RLock()
        defer c.mu.RUnlock()

        fieldMap, ok := c.fields[field]
        if !ok {
                return nil
        }
        return fieldMap[category]
}

// GetAny returns a bitmap of documents in ANY of the given categories (OR).
func (c *BitmapFilter) GetAny(field string, categories []string) *roaring.Bitmap {
        c.mu.RLock()
        defer c.mu.RUnlock()

        fieldMap, ok := c.fields[field]
        if !ok {
                return roaring.New()
        }

        result := roaring.New()
        for _, cat := range categories {
                if bm, ok := fieldMap[cat]; ok {
                        result.Or(bm)
                }
        }
        return result
}

// Categories returns all category values for a given field.
func (c *BitmapFilter) Categories(field string) []string {
        c.mu.RLock()
        defer c.mu.RUnlock()

        fieldMap, ok := c.fields[field]
        if !ok {
                return nil
        }

        cats := make([]string, 0, len(fieldMap))
        for cat := range fieldMap {
                cats = append(cats, cat)
        }
        return cats
}

// Counts returns the number of documents in each category for a field.
func (c *BitmapFilter) Counts(field string) map[string]uint64 {
        c.mu.RLock()
        defer c.mu.RUnlock()

        fieldMap, ok := c.fields[field]
        if !ok {
                return nil
        }

        counts := make(map[string]uint64, len(fieldMap))
        for cat, bm := range fieldMap {
                counts[cat] = bm.GetCardinality()
        }
        return counts
}

// AllCounts returns counts for all fields and categories.
func (c *BitmapFilter) AllCounts() map[string]map[string]uint64 {
        c.mu.RLock()
        defer c.mu.RUnlock()

        result := make(map[string]map[string]uint64, len(c.fields))
        for field, fieldMap := range c.fields {
                counts := make(map[string]uint64, len(fieldMap))
                for cat, bm := range fieldMap {
                        counts[cat] = bm.GetCardinality()
                }
                result[field] = counts
        }
        return result
}

// MemoryUsage returns the total memory used by all bitmaps in bytes.
func (c *BitmapFilter) MemoryUsage() uint64 {
        c.mu.RLock()
        defer c.mu.RUnlock()

        var total uint64
        for _, fieldMap := range c.fields {
                for _, bm := range fieldMap {
                        total += bm.GetSizeInBytes()
                }
        }
        return total
}

// bitmapFilterData is the serializable representation.
type bitmapFilterData struct {
        Fields map[string]map[string][]byte `msgpack:"fields"`
}

// SaveToFile saves the bitmap filter to a file.
func (c *BitmapFilter) SaveToFile(path string) error {
        file, err := os.Create(path)
        if err != nil {
                return err
        }
        defer file.Close()
        return c.Encode(file)
}

// Encode writes the bitmap filter to a writer.
func (c *BitmapFilter) Encode(w io.Writer) error {
        c.mu.RLock()
        defer c.mu.RUnlock()

        data := bitmapFilterData{
                Fields: make(map[string]map[string][]byte, len(c.fields)),
        }

        for field, fieldMap := range c.fields {
                data.Fields[field] = make(map[string][]byte, len(fieldMap))
                for cat, bm := range fieldMap {
                        bytes, err := bm.ToBytes()
                        if err != nil {
                                return err
                        }
                        data.Fields[field][cat] = bytes
                }
        }

        return msgpack.NewEncoder(w).Encode(data)
}

// LoadBitmapFilter loads a bitmap filter from a file.
func LoadBitmapFilter(path string) (*BitmapFilter, error) {
        file, err := os.Open(path)
        if err != nil {
                return nil, err
        }
        defer file.Close()
        return ReadBitmapFilter(file)
}

// ReadBitmapFilter reads a bitmap filter from a reader.
func ReadBitmapFilter(r io.Reader) (*BitmapFilter, error) {
        var data bitmapFilterData
        if err := msgpack.NewDecoder(r).Decode(&data); err != nil {
                return nil, err
        }

        c := &BitmapFilter{
                fields: make(map[string]map[string]*roaring.Bitmap, len(data.Fields)),
        }

        for field, fieldMap := range data.Fields {
                c.fields[field] = make(map[string]*roaring.Bitmap, len(fieldMap))
                for cat, bytes := range fieldMap {
                        bm := roaring.New()
                        if err := bm.UnmarshalBinary(bytes); err != nil {
                                return nil, err
                        }
                        c.fields[field][cat] = bm
                }
        }

        return c, nil
}

// SortColumn provides a typed columnar array for sorting documents by a value.
// Uses heap-based partial sort for efficient top-K queries.
//
// Example:
//
//        ratings := NewSortColumn[uint16]()
//        ratings.Set(1, 85)
//        ratings.Set(2, 92)
//
//        // Sort all docs
//        results := ratings.Sort([]uint32{1, 2}, false, 10)
//
//        // Sort filtered docs from a bitmap
//        results := ratings.SortBitmapDesc(filteredBitmap, 100)
type SortColumn[T cmp.Ordered] struct {
        mu       sync.RWMutex
        values   []T
        maxDocID uint32
}

// SortedResult holds a document ID and its sort value.
type SortedResult[T cmp.Ordered] struct {
        DocID uint32
        Value T
}

// NewSortColumn creates a new typed sort column.
func NewSortColumn[T cmp.Ordered]() *SortColumn[T] {
        return &SortColumn[T]{
                values: make([]T, 0),
        }
}

// Set sets the value for a document.
func (col *SortColumn[T]) Set(docID uint32, value T) {
        col.mu.Lock()
        defer col.mu.Unlock()
        col.setLocked(docID, value)
}

func (col *SortColumn[T]) setLocked(docID uint32, value T) {
        // Grow array if needed
        if docID >= uint32(len(col.values)) {
                newSize := docID + 1
                if newSize < uint32(len(col.values)*5/4) {
                        newSize = uint32(len(col.values) * 5 / 4)
                }
                if newSize < 1024 {
                        newSize = 1024
                }
                newValues := make([]T, newSize)
                copy(newValues, col.values)
                col.values = newValues
        }

        col.values[docID] = value

        if docID > col.maxDocID {
                col.maxDocID = docID
        }
}

// SortColumnBatch accumulates entries for efficient batch insertion.
type SortColumnBatch[T cmp.Ordered] struct {
        col    *SortColumn[T]
        docIDs []uint32
        values []T
}

// Batch creates a new batch builder for this column.
// Use BatchSize for better performance when you know the approximate count.
func (col *SortColumn[T]) Batch() *SortColumnBatch[T] {
        return col.BatchSize(1024)
}

// BatchSize creates a batch builder with pre-allocated capacity.
func (col *SortColumn[T]) BatchSize(size int) *SortColumnBatch[T] {
        return &SortColumnBatch[T]{
                col:    col,
                docIDs: make([]uint32, 0, size),
                values: make([]T, 0, size),
        }
}

// Add adds a document with a value to the batch.
func (b *SortColumnBatch[T]) Add(docID uint32, value T) {
        b.docIDs = append(b.docIDs, docID)
        b.values = append(b.values, value)
}

// Flush commits all accumulated entries to the column.
func (b *SortColumnBatch[T]) Flush() {
        if len(b.docIDs) == 0 {
                return
        }

        // Find max docID to pre-allocate
        var maxID uint32
        for _, id := range b.docIDs {
                if id > maxID {
                        maxID = id
                }
        }

        b.col.mu.Lock()
        defer b.col.mu.Unlock()

        // Pre-allocate if needed
        if maxID >= uint32(len(b.col.values)) {
                newValues := make([]T, maxID+1)
                copy(newValues, b.col.values)
                b.col.values = newValues
        }

        // Set all values
        for i, id := range b.docIDs {
                b.col.values[id] = b.values[i]
                if id > b.col.maxDocID {
                        b.col.maxDocID = id
                }
        }

        // Clear for reuse
        b.docIDs = b.docIDs[:0]
        b.values = b.values[:0]
}

// Get returns the value for a document.
func (col *SortColumn[T]) Get(docID uint32) T {
        col.mu.RLock()
        defer col.mu.RUnlock()

        var zero T
        if docID >= uint32(len(col.values)) {
                return zero
        }
        return col.values[docID]
}

// MemoryUsage returns the memory used by the values array in bytes.
func (col *SortColumn[T]) MemoryUsage() uint64 {
        col.mu.RLock()
        defer col.mu.RUnlock()

        var zero T
        return uint64(len(col.values)) * uint64(unsafe.Sizeof(zero))
}

// Sort sorts document IDs by their value.
// Uses heap-based partial sort when limit is small relative to input.
func (col *SortColumn[T]) Sort(docIDs []uint32, asc bool, limit int) []SortedResult[T] {
        col.mu.RLock()
        defer col.mu.RUnlock()

        return col.sortLocked(docIDs, asc, limit)
}

// SortDesc is a convenience method for descending sort.
func (col *SortColumn[T]) SortDesc(docIDs []uint32, limit int) []SortedResult[T] {
        return col.Sort(docIDs, false, limit)
}

// SortBitmap sorts documents from a bitmap by their value.
func (col *SortColumn[T]) SortBitmap(bm *roaring.Bitmap, asc bool, limit int) []SortedResult[T] {
        if bm == nil || bm.IsEmpty() {
                return nil
        }

        col.mu.RLock()
        defer col.mu.RUnlock()

        return col.sortLocked(bm.ToArray(), asc, limit)
}

// SortBitmapDesc is a convenience method for descending bitmap sort.
func (col *SortColumn[T]) SortBitmapDesc(bm *roaring.Bitmap, limit int) []SortedResult[T] {
        return col.SortBitmap(bm, false, limit)
}

func (col *SortColumn[T]) sortLocked(docIDs []uint32, asc bool, limit int) []SortedResult[T] {
        if len(docIDs) == 0 {
                return nil
        }

        values := col.values

        // Use heap for partial sort when limit is small relative to input
        if limit > 0 && limit < len(docIDs)/4 {
                return col.heapSort(docIDs, values, asc, limit)
        }

        // Full sort
        results := make([]SortedResult[T], len(docIDs))
        for i, docID := range docIDs {
                var value T
                if docID < uint32(len(values)) {
                        value = values[docID]
                }
                results[i] = SortedResult[T]{DocID: docID, Value: value}
        }

        if asc {
                slices.SortFunc(results, func(a, b SortedResult[T]) int {
                        return cmp.Compare(a.Value, b.Value)
                })
        } else {
                slices.SortFunc(results, func(a, b SortedResult[T]) int {
                        return cmp.Compare(b.Value, a.Value)
                })
        }

        if limit > 0 && limit < len(results) {
                results = results[:limit]
        }

        return results
}

func (col *SortColumn[T]) heapSort(docIDs []uint32, values []T, asc bool, limit int) []SortedResult[T] {
        h := &resultHeap[T]{
                items: make([]SortedResult[T], 0, limit),
                asc:   asc,
        }

        for _, docID := range docIDs {
                var value T
                if docID < uint32(len(values)) {
                        value = values[docID]
                }

                if h.Len() < limit {
                        h.items = append(h.items, SortedResult[T]{DocID: docID, Value: value})
                        if h.Len() == limit {
                                heap.Init(h)
                        }
                } else {
                        top := h.items[0]
                        better := (asc && value < top.Value) || (!asc && value > top.Value)
                        if better {
                                h.items[0] = SortedResult[T]{DocID: docID, Value: value}
                                heap.Fix(h, 0)
                        }
                }
        }

        if h.Len() < limit && h.Len() > 0 {
                heap.Init(h)
        }

        results := make([]SortedResult[T], h.Len())
        for i := len(results) - 1; i >= 0; i-- {
                results[i] = heap.Pop(h).(SortedResult[T])
        }

        return results
}

// resultHeap implements heap.Interface for SortedResult.
type resultHeap[T cmp.Ordered] struct {
        items []SortedResult[T]
        asc   bool
}

func (h *resultHeap[T]) Len() int { return len(h.items) }

func (h *resultHeap[T]) Less(i, j int) bool {
        if h.asc {
                return h.items[i].Value > h.items[j].Value // max-heap for ascending
        }
        return h.items[i].Value < h.items[j].Value // min-heap for descending
}

func (h *resultHeap[T]) Swap(i, j int) { h.items[i], h.items[j] = h.items[j], h.items[i] }

func (h *resultHeap[T]) Push(x any) {
        h.items = append(h.items, x.(SortedResult[T]))
}

func (h *resultHeap[T]) Pop() any {
        n := len(h.items)
        x := h.items[n-1]
        h.items = h.items[:n-1]
        return x
}

// sortColumnData is the serializable representation.
type sortColumnData[T cmp.Ordered] struct {
        Values   []T    `msgpack:"values"`
        MaxDocID uint32 `msgpack:"max_doc_id"`
}

// SaveToFile saves the sort column to a file.
func (col *SortColumn[T]) SaveToFile(path string) error {
        file, err := os.Create(path)
        if err != nil {
                return err
        }
        defer file.Close()
        return col.Encode(file)
}

// Encode writes the sort column to a writer.
func (col *SortColumn[T]) Encode(w io.Writer) error {
        col.mu.RLock()
        defer col.mu.RUnlock()

        data := sortColumnData[T]{
                Values:   col.values[:col.maxDocID+1],
                MaxDocID: col.maxDocID,
        }

        return msgpack.NewEncoder(w).Encode(data)
}

// LoadSortColumn loads a sort column from a file.
func LoadSortColumn[T cmp.Ordered](path string) (*SortColumn[T], error) {
        file, err := os.Open(path)
        if err != nil {
                return nil, err
        }
        defer file.Close()
        return ReadSortColumn[T](file)
}

// ReadSortColumn reads a sort column from a reader.
func ReadSortColumn[T cmp.Ordered](r io.Reader) (*SortColumn[T], error) {
        var data sortColumnData[T]
        if err := msgpack.NewDecoder(r).Decode(&data); err != nil {
                return nil, err
        }

        return &SortColumn[T]{
                values:   data.Values,
                maxDocID: data.MaxDocID,
        }, nil
}

1	package roaringsearch
2
3	import (
4	"cmp"
5	"container/heap"
6	"io"
7	"os"
8	"slices"
9	"sync"
10	"unsafe"
11
12	"github.com/RoaringBitmap/roaring"
13	"github.com/vmihailenco/msgpack/v5"
14	)
15
16	// BitmapFilter provides fast filtering by multiple category fields using bitmap indexes.
17	// Each field (e.g., "media_type", "language") can have multiple category values
18	// (e.g., "book", "movie"), and each category maps to a bitmap of document IDs.
19	//
20	// Example:
21	//
22	// filter := NewBitmapFilter()
23	// filter.Set(1, "media_type", "book")
24	// filter.Set(1, "language", "english")
25	// filter.Set(2, "media_type", "movie")
26	//
27	// books := filter.Get("media_type", "book") // bitmap of books
28	// english := filter.Get("language", "english") // bitmap of english
29	// englishBooks := roaring.And(books, english) // AND filter
30	type BitmapFilter struct {
31	mu sync.RWMutex
32	fields map[string]map[string]*roaring.Bitmap
33	}
34
35	// NewBitmapFilter creates a new bitmap filter.
36	func NewBitmapFilter() *BitmapFilter {	11✔
37	return &BitmapFilter{	11✔
38	fields: make(map[string]map[string]*roaring.Bitmap),	11✔
39	}	11✔
40	}	11✔
41
42	// Set assigns a document to a category within a field.
43	func (c *BitmapFilter) Set(docID uint32, field, category string) {	47✔
44	c.mu.Lock()	47✔
45	defer c.mu.Unlock()	47✔
46	c.setLocked(docID, field, category)	47✔
47	}	47✔
48
49	func (c *BitmapFilter) setLocked(docID uint32, field, category string) {	47✔
50	fieldMap, ok := c.fields[field]	47✔
51	if !ok {	61✔
52	fieldMap = make(map[string]*roaring.Bitmap)	14✔
53	c.fields[field] = fieldMap	14✔
54	}	14✔
55
56	bm, ok := fieldMap[category]	47✔
57	if !ok {	78✔
58	bm = roaring.New()	31✔
59	fieldMap[category] = bm	31✔
60	}	31✔
61	bm.Add(docID)	47✔
62	}
63
64	// FilterBatch accumulates entries for efficient batch insertion.
65	type FilterBatch struct {
66	filter *BitmapFilter
67	field string
68	docIDs []uint32
69	categories []string
70	}
71
72	// Batch creates a new batch builder for the given field.
73	// Use BatchSize for better performance when you know the approximate count.
74	func (c BitmapFilter) Batch(field string) FilterBatch {	1✔
75	return c.BatchSize(field, 1024)	1✔
76	}	1✔
77
78	// BatchSize creates a batch builder with pre-allocated capacity.
79	func (c BitmapFilter) BatchSize(field string, size int) FilterBatch {	1✔
80	return &FilterBatch{	1✔
81	filter: c,	1✔
82	field: field,	1✔
83	docIDs: make([]uint32, 0, size),	1✔
84	categories: make([]string, 0, size),	1✔
85	}	1✔
86	}	1✔
87
88	// Add adds a document with a category to the batch.
89	func (b *FilterBatch) Add(docID uint32, category string) {	6✔
90	b.docIDs = append(b.docIDs, docID)	6✔
91	b.categories = append(b.categories, category)	6✔
92	}	6✔
93
94	// Flush commits all accumulated entries to the filter.
95	func (b *FilterBatch) Flush() {	3✔
96	if len(b.docIDs) == 0 {	4✔
97	return	1✔
98	}	1✔
99
100	n := len(b.docIDs)	2✔
101		2✔
102	// Pass 1: build integer index array using linear search (fast for few categories)	2✔
103	categoryList := make([]string, 0, 16)	2✔
104	indices := make([]int, n)	2✔
105		2✔
106	for i, cat := range b.categories {	8✔
107	// Linear search - faster than map for small category counts	6✔
108	idx := -1	6✔
109	for j, existing := range categoryList {	11✔
110	if existing == cat {	8✔
111	idx = j	3✔
112	break	3✔
113	}
114	}
115	if idx == -1 {	9✔
116	idx = len(categoryList)	3✔
117	categoryList = append(categoryList, cat)	3✔
118	}	3✔
119	indices[i] = idx	6✔
120	}
121
122	numCats := len(categoryList)	2✔
123		2✔
124	// Pass 2: count per category (fast integer array access)	2✔
125	counts := make([]int, numCats)	2✔
126	for _, idx := range indices {	8✔
127	counts[idx]++	6✔
128	}	6✔
129
130	// Pre-allocate exact-sized groups
131	groups := make([][]uint32, numCats)	2✔
132	for i, count := range counts {	5✔
133	groups[i] = make([]uint32, 0, count)	3✔
134	}	3✔
135
136	// Pass 3: fill groups (fast integer array access, no reallocation)
137	for i, idx := range indices {	8✔
138	groups[idx] = append(groups[idx], b.docIDs[i])	6✔
139	}	6✔
140
141	b.filter.mu.Lock()	2✔
142	defer b.filter.mu.Unlock()	2✔
143		2✔
144	fieldMap, ok := b.filter.fields[b.field]	2✔
145	if !ok {	3✔
146	fieldMap = make(map[string]*roaring.Bitmap)	1✔
147	b.filter.fields[b.field] = fieldMap	1✔
148	}	1✔
149
150	// Create bitmaps
151	bitmaps := make([]*roaring.Bitmap, numCats)	2✔
152	for idx := range groups {	5✔
153	cat := categoryList[idx]	3✔
154	bm, ok := fieldMap[cat]	3✔
155	if !ok {	6✔
156	bm = roaring.New()	3✔
157	fieldMap[cat] = bm	3✔
158	}	3✔
159	bitmaps[idx] = bm	3✔
160	}
161
162	// Parallel AddMany if enough categories
163	if numCats >= 4 {	2✔
NEW 164	var wg sync.WaitGroup	×
NEW 165	for idx, ids := range groups {	×
NEW 166	wg.Add(1)	×
NEW 167	go func(bm *roaring.Bitmap, docIDs []uint32) {	×
NEW 168	defer wg.Done()	×
NEW 169	bm.AddMany(docIDs)	×
NEW 170	}(bitmaps[idx], ids)	×
171	}
NEW 172	wg.Wait()	×
173	} else {	2✔
174	for idx, ids := range groups {	5✔
175	bitmaps[idx].AddMany(ids)	3✔
176	}	3✔
177	}
178
179	// Clear for reuse
180	b.docIDs = b.docIDs[:0]	2✔
181	b.categories = b.categories[:0]	2✔
182	}
183
184	// Remove removes a document from all categories across all fields.
185	func (c *BitmapFilter) Remove(docID uint32) {	1✔
186	c.mu.Lock()	1✔
187	defer c.mu.Unlock()	1✔
188		1✔
189	for _, fieldMap := range c.fields {	2✔
190	for _, bm := range fieldMap {	2✔
191	bm.Remove(docID)	1✔
192	}	1✔
193	}
194	}
195
196	// Get returns a bitmap of documents in the given category for a field.
197	// Returns nil if field or category doesn't exist.
198	func (c BitmapFilter) Get(field, category string) roaring.Bitmap {	19✔
199	c.mu.RLock()	19✔
200	defer c.mu.RUnlock()	19✔
201		19✔
202	fieldMap, ok := c.fields[field]	19✔
203	if !ok {	20✔
204	return nil	1✔
205	}	1✔
206	return fieldMap[category]	18✔
207	}
208
209	// GetAny returns a bitmap of documents in ANY of the given categories (OR).
210	func (c BitmapFilter) GetAny(field string, categories []string) roaring.Bitmap {	1✔
211	c.mu.RLock()	1✔
212	defer c.mu.RUnlock()	1✔
213		1✔
214	fieldMap, ok := c.fields[field]	1✔
215	if !ok {	1✔
UNCOV 216	return roaring.New()	×
UNCOV 217	}	×
218
219	result := roaring.New()	1✔
220	for _, cat := range categories {	3✔
221	if bm, ok := fieldMap[cat]; ok {	4✔
222	result.Or(bm)	2✔
223	}	2✔
224	}
225	return result	1✔
226	}
227
228	// Categories returns all category values for a given field.
229	func (c *BitmapFilter) Categories(field string) []string {	2✔
230	c.mu.RLock()	2✔
231	defer c.mu.RUnlock()	2✔
232		2✔
233	fieldMap, ok := c.fields[field]	2✔
234	if !ok {	3✔
235	return nil	1✔
236	}	1✔
237
238	cats := make([]string, 0, len(fieldMap))	1✔
239	for cat := range fieldMap {	4✔
240	cats = append(cats, cat)	3✔
241	}	3✔
242	return cats	1✔
243	}
244
245	// Counts returns the number of documents in each category for a field.
246	func (c *BitmapFilter) Counts(field string) map[string]uint64 {	1✔
247	c.mu.RLock()	1✔
248	defer c.mu.RUnlock()	1✔
249		1✔
250	fieldMap, ok := c.fields[field]	1✔
251	if !ok {	1✔
UNCOV 252	return nil	×
UNCOV 253	}	×
254
255	counts := make(map[string]uint64, len(fieldMap))	1✔
256	for cat, bm := range fieldMap {	3✔
257	counts[cat] = bm.GetCardinality()	2✔
258	}	2✔
259	return counts	1✔
260	}
261
262	// AllCounts returns counts for all fields and categories.
263	func (c *BitmapFilter) AllCounts() map[string]map[string]uint64 {	1✔
264	c.mu.RLock()	1✔
265	defer c.mu.RUnlock()	1✔
266		1✔
267	result := make(map[string]map[string]uint64, len(c.fields))	1✔
268	for field, fieldMap := range c.fields {	3✔
269	counts := make(map[string]uint64, len(fieldMap))	2✔
270	for cat, bm := range fieldMap {	6✔
271	counts[cat] = bm.GetCardinality()	4✔
272	}	4✔
273	result[field] = counts	2✔
274	}
275	return result	1✔
276	}
277
278	// MemoryUsage returns the total memory used by all bitmaps in bytes.
NEW 279	func (c *BitmapFilter) MemoryUsage() uint64 {	×
NEW 280	c.mu.RLock()	×
NEW 281	defer c.mu.RUnlock()	×
NEW 282		×
NEW 283	var total uint64	×
NEW 284	for _, fieldMap := range c.fields {	×
NEW 285	for _, bm := range fieldMap {	×
NEW 286	total += bm.GetSizeInBytes()	×
NEW 287	}	×
288	}
NEW 289	return total	×
290	}
291
292	// bitmapFilterData is the serializable representation.
293	type bitmapFilterData struct {
294	Fields map[string]map[string][]byte `msgpack:"fields"`
295	}
296
297	// SaveToFile saves the bitmap filter to a file.
298	func (c *BitmapFilter) SaveToFile(path string) error {	1✔
299	file, err := os.Create(path)	1✔
300	if err != nil {	1✔
UNCOV 301	return err	×
UNCOV 302	}	×
303	defer file.Close()	1✔
304	return c.Encode(file)	1✔
305	}
306
307	// Encode writes the bitmap filter to a writer.
308	func (c *BitmapFilter) Encode(w io.Writer) error {	1✔
309	c.mu.RLock()	1✔
310	defer c.mu.RUnlock()	1✔
311		1✔
312	data := bitmapFilterData{	1✔
313	Fields: make(map[string]map[string][]byte, len(c.fields)),	1✔
314	}	1✔
315		1✔
316	for field, fieldMap := range c.fields {	3✔
317	data.Fields[field] = make(map[string][]byte, len(fieldMap))	2✔
318	for cat, bm := range fieldMap {	7✔
319	bytes, err := bm.ToBytes()	5✔
320	if err != nil {	5✔
321	return err	×
UNCOV 322	}	×
323	data.Fields[field][cat] = bytes	5✔
324	}
325	}
326
327	return msgpack.NewEncoder(w).Encode(data)	1✔
328	}
329
330	// LoadBitmapFilter loads a bitmap filter from a file.
331	func LoadBitmapFilter(path string) (*BitmapFilter, error) {	1✔
332	file, err := os.Open(path)	1✔
333	if err != nil {	1✔
UNCOV 334	return nil, err	×
UNCOV 335	}	×
336	defer file.Close()	1✔
337	return ReadBitmapFilter(file)	1✔
338	}
339
340	// ReadBitmapFilter reads a bitmap filter from a reader.
341	func ReadBitmapFilter(r io.Reader) (*BitmapFilter, error) {	1✔
342	var data bitmapFilterData	1✔
343	if err := msgpack.NewDecoder(r).Decode(&data); err != nil {	1✔
UNCOV 344	return nil, err	×
UNCOV 345	}	×
346
347	c := &BitmapFilter{	1✔
348	fields: make(map[string]map[string]*roaring.Bitmap, len(data.Fields)),	1✔
349	}	1✔
350		1✔
351	for field, fieldMap := range data.Fields {	3✔
352	c.fields[field] = make(map[string]*roaring.Bitmap, len(fieldMap))	2✔
353	for cat, bytes := range fieldMap {	7✔
354	bm := roaring.New()	5✔
355	if err := bm.UnmarshalBinary(bytes); err != nil {	5✔
UNCOV 356	return nil, err	×
UNCOV 357	}	×
358	c.fields[field][cat] = bm	5✔
359	}
360	}
361
362	return c, nil	1✔
363	}
364
365	// SortColumn provides a typed columnar array for sorting documents by a value.
366	// Uses heap-based partial sort for efficient top-K queries.
367	//
368	// Example:
369	//
370	// ratings := NewSortColumn[uint16]()
371	// ratings.Set(1, 85)
372	// ratings.Set(2, 92)
373	//
374	// // Sort all docs
375	// results := ratings.Sort([]uint32{1, 2}, false, 10)
376	//
377	// // Sort filtered docs from a bitmap
378	// results := ratings.SortBitmapDesc(filteredBitmap, 100)
379	type SortColumn[T cmp.Ordered] struct {
380	mu sync.RWMutex
381	values []T
382	maxDocID uint32
383	}
384
385	// SortedResult holds a document ID and its sort value.
386	type SortedResult[T cmp.Ordered] struct {
387	DocID uint32
388	Value T
389	}
390
391	// NewSortColumn creates a new typed sort column.
392	func NewSortColumn[T cmp.Ordered]() *SortColumn[T] {	13✔
393	return &SortColumn[T]{	13✔
394	values: make([]T, 0),	13✔
395	}	13✔
396	}	13✔
397
398	// Set sets the value for a document.
399	func (col *SortColumn[T]) Set(docID uint32, value T) {	153✔
400	col.mu.Lock()	153✔
401	defer col.mu.Unlock()	153✔
402	col.setLocked(docID, value)	153✔
403	}	153✔
404
405	func (col *SortColumn[T]) setLocked(docID uint32, value T) {	153✔
406	// Grow array if needed	153✔
407	if docID >= uint32(len(col.values)) {	165✔
408	newSize := docID + 1	12✔
409	if newSize < uint32(len(col.values)*5/4) {	12✔
NEW 410	newSize = uint32(len(col.values) * 5 / 4)	×
UNCOV 411	}	×
412	if newSize < 1024 {	24✔
413	newSize = 1024	12✔
414	}	12✔
415	newValues := make([]T, newSize)	12✔
416	copy(newValues, col.values)	12✔
417	col.values = newValues	12✔
418	}
419
420	col.values[docID] = value	153✔
421		153✔
422	if docID > col.maxDocID {	306✔
423	col.maxDocID = docID	153✔
424	}	153✔
425	}
426
427	// SortColumnBatch accumulates entries for efficient batch insertion.
428	type SortColumnBatch[T cmp.Ordered] struct {
429	col *SortColumn[T]
430	docIDs []uint32
431	values []T
432	}
433
434	// Batch creates a new batch builder for this column.
435	// Use BatchSize for better performance when you know the approximate count.
436	func (col SortColumn[T]) Batch() SortColumnBatch[T] {	1✔
437	return col.BatchSize(1024)	1✔
438	}	1✔
439
440	// BatchSize creates a batch builder with pre-allocated capacity.
441	func (col SortColumn[T]) BatchSize(size int) SortColumnBatch[T] {	1✔
442	return &SortColumnBatch[T]{	1✔
443	col: col,	1✔
444	docIDs: make([]uint32, 0, size),	1✔
445	values: make([]T, 0, size),	1✔
446	}	1✔
447	}	1✔
448
449	// Add adds a document with a value to the batch.
450	func (b *SortColumnBatch[T]) Add(docID uint32, value T) {	5✔
451	b.docIDs = append(b.docIDs, docID)	5✔
452	b.values = append(b.values, value)	5✔
453	}	5✔
454
455	// Flush commits all accumulated entries to the column.
456	func (b *SortColumnBatch[T]) Flush() {	3✔
457	if len(b.docIDs) == 0 {	4✔
458	return	1✔
459	}	1✔
460
461	// Find max docID to pre-allocate
462	var maxID uint32	2✔
463	for _, id := range b.docIDs {	7✔
464	if id > maxID {	10✔
465	maxID = id	5✔
466	}	5✔
467	}
468
469	b.col.mu.Lock()	2✔
470	defer b.col.mu.Unlock()	2✔
471		2✔
472	// Pre-allocate if needed	2✔
473	if maxID >= uint32(len(b.col.values)) {	4✔
474	newValues := make([]T, maxID+1)	2✔
475	copy(newValues, b.col.values)	2✔
476	b.col.values = newValues	2✔
477	}	2✔
478
479	// Set all values
480	for i, id := range b.docIDs {	7✔
481	b.col.values[id] = b.values[i]	5✔
482	if id > b.col.maxDocID {	10✔
483	b.col.maxDocID = id	5✔
484	}	5✔
485	}
486
487	// Clear for reuse
488	b.docIDs = b.docIDs[:0]	2✔
489	b.values = b.values[:0]	2✔
490	}
491
492	// Get returns the value for a document.
493	func (col *SortColumn[T]) Get(docID uint32) T {	9✔
494	col.mu.RLock()	9✔
495	defer col.mu.RUnlock()	9✔
496		9✔
497	var zero T	9✔
498	if docID >= uint32(len(col.values)) {	9✔
UNCOV 499	return zero	×
UNCOV 500	}	×
501	return col.values[docID]	9✔
502	}
503
504	// MemoryUsage returns the memory used by the values array in bytes.
NEW 505	func (col *SortColumn[T]) MemoryUsage() uint64 {	×
NEW 506	col.mu.RLock()	×
NEW 507	defer col.mu.RUnlock()	×
NEW 508		×
NEW 509	var zero T	×
NEW 510	return uint64(len(col.values)) * uint64(unsafe.Sizeof(zero))	×
UNCOV 511	}	×
512
513	// Sort sorts document IDs by their value.
514	// Uses heap-based partial sort when limit is small relative to input.
515	func (col *SortColumn[T]) Sort(docIDs []uint32, asc bool, limit int) []SortedResult[T] {	12✔
516	col.mu.RLock()	12✔
517	defer col.mu.RUnlock()	12✔
518		12✔
519	return col.sortLocked(docIDs, asc, limit)	12✔
520	}	12✔
521
522	// SortDesc is a convenience method for descending sort.
523	func (col *SortColumn[T]) SortDesc(docIDs []uint32, limit int) []SortedResult[T] {	4✔
524	return col.Sort(docIDs, false, limit)	4✔
525	}	4✔
526
527	// SortBitmap sorts documents from a bitmap by their value.
528	func (col SortColumn[T]) SortBitmap(bm roaring.Bitmap, asc bool, limit int) []SortedResult[T] {	4✔
529	if bm == nil \|\| bm.IsEmpty() {	6✔
530	return nil	2✔
531	}	2✔
532
533	col.mu.RLock()	2✔
534	defer col.mu.RUnlock()	2✔
535		2✔
536	return col.sortLocked(bm.ToArray(), asc, limit)	2✔
537	}
538
539	// SortBitmapDesc is a convenience method for descending bitmap sort.
540	func (col SortColumn[T]) SortBitmapDesc(bm roaring.Bitmap, limit int) []SortedResult[T] {	2✔
541	return col.SortBitmap(bm, false, limit)	2✔
542	}	2✔
543
544	func (col *SortColumn[T]) sortLocked(docIDs []uint32, asc bool, limit int) []SortedResult[T] {	14✔
545	if len(docIDs) == 0 {	14✔
UNCOV 546	return nil	×
UNCOV 547	}	×
548
549	values := col.values	14✔
550		14✔
551	// Use heap for partial sort when limit is small relative to input	14✔
552	if limit > 0 && limit < len(docIDs)/4 {	17✔
553	return col.heapSort(docIDs, values, asc, limit)	3✔
554	}	3✔
555
556	// Full sort
557	results := make([]SortedResult[T], len(docIDs))	11✔
558	for i, docID := range docIDs {	47✔
559	var value T	36✔
560	if docID < uint32(len(values)) {	72✔
561	value = values[docID]	36✔
562	}	36✔
563	results[i] = SortedResult[T]{DocID: docID, Value: value}	36✔
564	}
565
566	if asc {	14✔
567	slices.SortFunc(results, func(a, b SortedResult[T]) int {	12✔
568	return cmp.Compare(a.Value, b.Value)	9✔
569	})	9✔
570	} else {	8✔
571	slices.SortFunc(results, func(a, b SortedResult[T]) int {	40✔
572	return cmp.Compare(b.Value, a.Value)	32✔
573	})	32✔
574	}
575
576	if limit > 0 && limit < len(results) {	13✔
577	results = results[:limit]	2✔
578	}	2✔
579
580	return results	11✔
581	}
582
583	func (col *SortColumn[T]) heapSort(docIDs []uint32, values []T, asc bool, limit int) []SortedResult[T] {	3✔
584	h := &resultHeap[T]{	3✔
585	items: make([]SortedResult[T], 0, limit),	3✔
586	asc: asc,	3✔
587	}	3✔
588		3✔
589	for _, docID := range docIDs {	223✔
590	var value T	220✔
591	if docID < uint32(len(values)) {	440✔
592	value = values[docID]	220✔
593	}	220✔
594
595	if h.Len() < limit {	244✔
596	h.items = append(h.items, SortedResult[T]{DocID: docID, Value: value})	24✔
597	if h.Len() == limit {	27✔
598	heap.Init(h)	3✔
599	}	3✔
600	} else {	196✔
601	top := h.items[0]	196✔
602	better := (asc && value < top.Value) \|\| (!asc && value > top.Value)	196✔
603	if better {	302✔
604	h.items[0] = SortedResult[T]{DocID: docID, Value: value}	106✔
605	heap.Fix(h, 0)	106✔
606	}	106✔
607	}
608	}
609
610	if h.Len() < limit && h.Len() > 0 {	3✔
UNCOV 611	heap.Init(h)	×
UNCOV 612	}	×
613
614	results := make([]SortedResult[T], h.Len())	3✔
615	for i := len(results) - 1; i >= 0; i-- {	27✔
616	results[i] = heap.Pop(h).(SortedResult[T])	24✔
617	}	24✔
618
619	return results	3✔
620	}
621
622	// resultHeap implements heap.Interface for SortedResult.
623	type resultHeap[T cmp.Ordered] struct {
624	items []SortedResult[T]
625	asc bool
626	}
627
628	func (h *resultHeap[T]) Len() int { return len(h.items) }	388✔
629
630	func (h *resultHeap[T]) Less(i, j int) bool {	584✔
631	if h.asc {	625✔
632	return h.items[i].Value > h.items[j].Value // max-heap for ascending	41✔
633	}	41✔
634	return h.items[i].Value < h.items[j].Value // min-heap for descending	543✔
635	}
636
637	func (h *resultHeap[T]) Swap(i, j int) { h.items[i], h.items[j] = h.items[j], h.items[i] }	328✔
638
639	func (h *resultHeap[T]) Push(x any) {	3✔
640	h.items = append(h.items, x.(SortedResult[T]))	3✔
641	}	3✔
642
643	func (h *resultHeap[T]) Pop() any {	25✔
644	n := len(h.items)	25✔
645	x := h.items[n-1]	25✔
646	h.items = h.items[:n-1]	25✔
647	return x	25✔
648	}	25✔
649
650	// sortColumnData is the serializable representation.
651	type sortColumnData[T cmp.Ordered] struct {
652	Values []T `msgpack:"values"`
653	MaxDocID uint32 `msgpack:"max_doc_id"`
654	}
655
656	// SaveToFile saves the sort column to a file.
657	func (col *SortColumn[T]) SaveToFile(path string) error {	1✔
658	file, err := os.Create(path)	1✔
659	if err != nil {	1✔
UNCOV 660	return err	×
UNCOV 661	}	×
662	defer file.Close()	1✔
663	return col.Encode(file)	1✔
664	}
665
666	// Encode writes the sort column to a writer.
667	func (col *SortColumn[T]) Encode(w io.Writer) error {	1✔
668	col.mu.RLock()	1✔
669	defer col.mu.RUnlock()	1✔
670		1✔
671	data := sortColumnData[T]{	1✔
672	Values: col.values[:col.maxDocID+1],	1✔
673	MaxDocID: col.maxDocID,	1✔
674	}	1✔
675		1✔
676	return msgpack.NewEncoder(w).Encode(data)	1✔
677	}	1✔
678
679	// LoadSortColumn loads a sort column from a file.
680	func LoadSortColumn[T cmp.Ordered](path string) (*SortColumn[T], error) {	1✔
681	file, err := os.Open(path)	1✔
682	if err != nil {	1✔
UNCOV 683	return nil, err	×
UNCOV 684	}	×
685	defer file.Close()	1✔
686	return ReadSortColumn[T](file)	1✔
687	}
688
689	// ReadSortColumn reads a sort column from a reader.
690	func ReadSortColumn[T cmp.Ordered](r io.Reader) (*SortColumn[T], error) {	1✔
691	var data sortColumnData[T]	1✔
692	if err := msgpack.NewDecoder(r).Decode(&data); err != nil {	1✔
UNCOV 693	return nil, err	×
UNCOV 694	}	×
695
696	return &SortColumn[T]{	1✔
697	values: data.Values,	1✔
698	maxDocID: data.MaxDocID,	1✔
699	}, nil	1✔
700	}

freeeve / roaringsearch / 21107032990

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous