github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/helpers/allow_list.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package helpers
    13  
    14  import (
    15  	"github.com/weaviate/sroar"
    16  	"github.com/weaviate/weaviate/adapters/repos/db/roaringset"
    17  )
    18  
    19  type AllowList interface {
    20  	Insert(ids ...uint64)
    21  	Contains(id uint64) bool
    22  	DeepCopy() AllowList
    23  	Slice() []uint64
    24  
    25  	IsEmpty() bool
    26  	Len() int
    27  	Min() uint64
    28  	Max() uint64
    29  	Size() uint64
    30  	Truncate(uint64) AllowList
    31  
    32  	Iterator() AllowListIterator
    33  	LimitedIterator(limit int) AllowListIterator
    34  }
    35  
    36  type AllowListIterator interface {
    37  	Next() (uint64, bool)
    38  	Len() int
    39  }
    40  
    41  func NewAllowList(ids ...uint64) AllowList {
    42  	return NewAllowListFromBitmap(roaringset.NewBitmap(ids...))
    43  }
    44  
    45  func NewAllowListFromBitmap(bm *sroar.Bitmap) AllowList {
    46  	return &bitmapAllowList{bm: bm}
    47  }
    48  
    49  func NewAllowListFromBitmapDeepCopy(bm *sroar.Bitmap) AllowList {
    50  	return NewAllowListFromBitmap(bm.Clone())
    51  }
    52  
    53  type bitmapAllowList struct {
    54  	bm *sroar.Bitmap
    55  }
    56  
    57  func (al *bitmapAllowList) Insert(ids ...uint64) {
    58  	al.bm.SetMany(ids)
    59  }
    60  
    61  func (al *bitmapAllowList) Contains(id uint64) bool {
    62  	return al.bm.Contains(id)
    63  }
    64  
    65  func (al *bitmapAllowList) DeepCopy() AllowList {
    66  	return NewAllowListFromBitmapDeepCopy(al.bm)
    67  }
    68  
    69  func (al *bitmapAllowList) Slice() []uint64 {
    70  	return al.bm.ToArray()
    71  }
    72  
    73  func (al *bitmapAllowList) IsEmpty() bool {
    74  	return al.bm.IsEmpty()
    75  }
    76  
    77  func (al *bitmapAllowList) Len() int {
    78  	return al.bm.GetCardinality()
    79  }
    80  
    81  func (al *bitmapAllowList) Min() uint64 {
    82  	return al.bm.Minimum()
    83  }
    84  
    85  func (al *bitmapAllowList) Max() uint64 {
    86  	return al.bm.Maximum()
    87  }
    88  
    89  func (al *bitmapAllowList) Size() uint64 {
    90  	// TODO provide better size estimation
    91  	return uint64(1.5 * float64(len(al.bm.ToBuffer())))
    92  }
    93  
    94  func (al *bitmapAllowList) Truncate(upTo uint64) AllowList {
    95  	card := al.bm.GetCardinality()
    96  	if upTo < uint64(card) {
    97  		al.bm.RemoveRange(upTo, uint64(al.bm.GetCardinality()+1))
    98  	}
    99  	return al
   100  }
   101  
   102  func (al *bitmapAllowList) Iterator() AllowListIterator {
   103  	return al.LimitedIterator(0)
   104  }
   105  
   106  func (al *bitmapAllowList) LimitedIterator(limit int) AllowListIterator {
   107  	return newBitmapAllowListIterator(al.bm, limit)
   108  }
   109  
   110  type bitmapAllowListIterator struct {
   111  	len     int
   112  	counter int
   113  	it      *sroar.Iterator
   114  }
   115  
   116  func newBitmapAllowListIterator(bm *sroar.Bitmap, limit int) AllowListIterator {
   117  	len := bm.GetCardinality()
   118  	if limit > 0 && limit < len {
   119  		len = limit
   120  	}
   121  
   122  	return &bitmapAllowListIterator{
   123  		len:     len,
   124  		counter: 0,
   125  		it:      bm.NewIterator(),
   126  	}
   127  }
   128  
   129  func (i *bitmapAllowListIterator) Next() (uint64, bool) {
   130  	if i.counter >= i.len {
   131  		return 0, false
   132  	}
   133  	i.counter++
   134  	return i.it.Next(), true
   135  }
   136  
   137  func (i *bitmapAllowListIterator) Len() int {
   138  	return i.len
   139  }