github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/lsmkv/memtable_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  	"github.com/pkg/errors"
    16  	"github.com/weaviate/sroar"
    17  	"github.com/weaviate/weaviate/adapters/repos/db/roaringset"
    18  )
    19  
    20  func (m *Memtable) roaringSetAddOne(key []byte, value uint64) error {
    21  	return m.roaringSetAddList(key, []uint64{value})
    22  }
    23  
    24  func (m *Memtable) roaringSetAddList(key []byte, values []uint64) error {
    25  	if err := checkStrategyRoaringSet(m.strategy); err != nil {
    26  		return err
    27  	}
    28  
    29  	m.Lock()
    30  	defer m.Unlock()
    31  
    32  	if err := m.roaringSetAddCommitLog(key, roaringset.NewBitmap(values...), roaringset.NewBitmap()); err != nil {
    33  		return err
    34  	}
    35  
    36  	m.roaringSet.Insert(key, roaringset.Insert{Additions: values})
    37  
    38  	m.roaringSetAdjustMeta(len(values))
    39  	return nil
    40  }
    41  
    42  func (m *Memtable) roaringSetAddBitmap(key []byte, bm *sroar.Bitmap) error {
    43  	if err := checkStrategyRoaringSet(m.strategy); err != nil {
    44  		return err
    45  	}
    46  
    47  	m.Lock()
    48  	defer m.Unlock()
    49  
    50  	if err := m.roaringSetAddCommitLog(key, bm, roaringset.NewBitmap()); err != nil {
    51  		return err
    52  	}
    53  
    54  	m.roaringSet.Insert(key, roaringset.Insert{Additions: bm.ToArray()})
    55  
    56  	m.roaringSetAdjustMeta(bm.GetCardinality())
    57  	return nil
    58  }
    59  
    60  func (m *Memtable) roaringSetRemoveOne(key []byte, value uint64) error {
    61  	return m.roaringSetRemoveList(key, []uint64{value})
    62  }
    63  
    64  func (m *Memtable) roaringSetRemoveList(key []byte, values []uint64) error {
    65  	if err := checkStrategyRoaringSet(m.strategy); err != nil {
    66  		return err
    67  	}
    68  
    69  	m.Lock()
    70  	defer m.Unlock()
    71  
    72  	if err := m.roaringSetAddCommitLog(key, roaringset.NewBitmap(), roaringset.NewBitmap(values...)); err != nil {
    73  		return err
    74  	}
    75  
    76  	m.roaringSet.Insert(key, roaringset.Insert{Deletions: values})
    77  
    78  	m.roaringSetAdjustMeta(len(values))
    79  	return nil
    80  }
    81  
    82  func (m *Memtable) roaringSetRemoveBitmap(key []byte, bm *sroar.Bitmap) error {
    83  	if err := checkStrategyRoaringSet(m.strategy); err != nil {
    84  		return err
    85  	}
    86  
    87  	m.Lock()
    88  	defer m.Unlock()
    89  
    90  	if err := m.roaringSetAddCommitLog(key, roaringset.NewBitmap(), bm); err != nil {
    91  		return err
    92  	}
    93  
    94  	m.roaringSet.Insert(key, roaringset.Insert{Deletions: bm.ToArray()})
    95  
    96  	m.roaringSetAdjustMeta(bm.GetCardinality())
    97  	return nil
    98  }
    99  
   100  func (m *Memtable) roaringSetAddRemoveBitmaps(key []byte, additions *sroar.Bitmap, deletions *sroar.Bitmap) error {
   101  	if err := checkStrategyRoaringSet(m.strategy); err != nil {
   102  		return err
   103  	}
   104  
   105  	m.Lock()
   106  	defer m.Unlock()
   107  
   108  	if err := m.roaringSetAddCommitLog(key, additions, deletions); err != nil {
   109  		return err
   110  	}
   111  
   112  	m.roaringSet.Insert(key, roaringset.Insert{
   113  		Additions: additions.ToArray(),
   114  		Deletions: deletions.ToArray(),
   115  	})
   116  
   117  	m.roaringSetAdjustMeta(additions.GetCardinality() + deletions.GetCardinality())
   118  	return nil
   119  }
   120  
   121  func (m *Memtable) roaringSetGet(key []byte) (roaringset.BitmapLayer, error) {
   122  	if err := checkStrategyRoaringSet(m.strategy); err != nil {
   123  		return roaringset.BitmapLayer{}, err
   124  	}
   125  
   126  	m.RLock()
   127  	defer m.RUnlock()
   128  
   129  	return m.roaringSet.Get(key)
   130  }
   131  
   132  func (m *Memtable) roaringSetAdjustMeta(entriesChanged int) {
   133  	// in the worst case roaring bitmaps take 2 bytes per entry. A reasonable
   134  	// estimation is therefore to take the changed entries and multiply them by
   135  	// 2.
   136  	m.size += uint64(entriesChanged * 2)
   137  	m.metrics.size(m.size)
   138  	m.updateDirtyAt()
   139  }
   140  
   141  func (m *Memtable) roaringSetAddCommitLog(key []byte, additions *sroar.Bitmap, deletions *sroar.Bitmap) error {
   142  	if node, err := roaringset.NewSegmentNode(key, additions, deletions); err != nil {
   143  		return errors.Wrap(err, "create node for commit log")
   144  	} else if err := m.commitlog.add(node); err != nil {
   145  		return errors.Wrap(err, "add node to commit log")
   146  	}
   147  	return nil
   148  }