github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/lsmkv/bucket_roaring_set.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 lsmkv
    13  
    14  import (
    15  	"errors"
    16  	"fmt"
    17  
    18  	"github.com/weaviate/sroar"
    19  	"github.com/weaviate/weaviate/entities/lsmkv"
    20  )
    21  
    22  func (b *Bucket) RoaringSetAddOne(key []byte, value uint64) error {
    23  	if err := checkStrategyRoaringSet(b.strategy); err != nil {
    24  		return err
    25  	}
    26  
    27  	b.flushLock.RLock()
    28  	defer b.flushLock.RUnlock()
    29  
    30  	return b.active.roaringSetAddOne(key, value)
    31  }
    32  
    33  func (b *Bucket) RoaringSetRemoveOne(key []byte, value uint64) error {
    34  	if err := checkStrategyRoaringSet(b.strategy); err != nil {
    35  		return err
    36  	}
    37  
    38  	b.flushLock.RLock()
    39  	defer b.flushLock.RUnlock()
    40  
    41  	return b.active.roaringSetRemoveOne(key, value)
    42  }
    43  
    44  func (b *Bucket) RoaringSetAddList(key []byte, values []uint64) error {
    45  	if err := checkStrategyRoaringSet(b.strategy); err != nil {
    46  		return err
    47  	}
    48  
    49  	b.flushLock.RLock()
    50  	defer b.flushLock.RUnlock()
    51  
    52  	return b.active.roaringSetAddList(key, values)
    53  }
    54  
    55  func (b *Bucket) RoaringSetAddBitmap(key []byte, bm *sroar.Bitmap) error {
    56  	if err := checkStrategyRoaringSet(b.strategy); err != nil {
    57  		return err
    58  	}
    59  
    60  	b.flushLock.RLock()
    61  	defer b.flushLock.RUnlock()
    62  
    63  	return b.active.roaringSetAddBitmap(key, bm)
    64  }
    65  
    66  func (b *Bucket) RoaringSetGet(key []byte) (*sroar.Bitmap, error) {
    67  	if err := checkStrategyRoaringSet(b.strategy); err != nil {
    68  		return nil, err
    69  	}
    70  
    71  	b.flushLock.RLock()
    72  	defer b.flushLock.RUnlock()
    73  
    74  	segments, err := b.disk.roaringSetGet(key)
    75  	if err != nil {
    76  		return nil, err
    77  	}
    78  
    79  	if b.flushing != nil {
    80  		flushing, err := b.flushing.roaringSetGet(key)
    81  		if err != nil {
    82  			if !errors.Is(err, lsmkv.NotFound) {
    83  				return nil, err
    84  			}
    85  		} else {
    86  			segments = append(segments, flushing)
    87  		}
    88  	}
    89  
    90  	memtable, err := b.active.roaringSetGet(key)
    91  	if err != nil {
    92  		if !errors.Is(err, lsmkv.NotFound) {
    93  			return nil, err
    94  		}
    95  	} else {
    96  		segments = append(segments, memtable)
    97  	}
    98  
    99  	return segments.Flatten(), nil
   100  }
   101  
   102  func checkStrategyRoaringSet(bucketStrat string) error {
   103  	if bucketStrat == StrategyRoaringSet {
   104  		return nil
   105  	}
   106  
   107  	return fmt.Errorf("this method requires a roaring set strategy, got: %s",
   108  		bucketStrat)
   109  }