Merge pull request weaviate#7801 from weaviate/rangeable_index_in_memory

Feature: rangeable index in memory

Author: antas-marcin

adapters/handlers/rest/configure_api.go

@@ -373,6 +373,7 @@ func MakeAppState(ctx context.Context, options *swag.CommandLineOptionsGroup) *s
 		SeparateObjectsCompactions:          appState.ServerConfig.Config.Persistence.LSMSeparateObjectsCompactions,
 		MaxSegmentSize:                      appState.ServerConfig.Config.Persistence.LSMMaxSegmentSize,
 		CycleManagerRoutinesFactor:          appState.ServerConfig.Config.Persistence.LSMCycleManagerRoutinesFactor,
+		IndexRangeableInMemory:              appState.ServerConfig.Config.Persistence.IndexRangeableInMemory,
 		HNSWMaxLogSize:                      appState.ServerConfig.Config.Persistence.HNSWMaxLogSize,
 		HNSWWaitForCachePrefill:             appState.ServerConfig.Config.HNSWStartupWaitForVectorCache,
 		HNSWFlatSearchConcurrency:           appState.ServerConfig.Config.HNSWFlatSearchConcurrency,

adapters/repos/db/index.go

@@ -637,6 +637,7 @@ type IndexConfig struct {
 	SegmentsCleanupIntervalSeconds      int
 	SeparateObjectsCompactions          bool
 	CycleManagerRoutinesFactor          int
+	IndexRangeableInMemory              bool
 	MaxSegmentSize                      int64
 	HNSWMaxLogSize                      int64
 	HNSWWaitForCachePrefill             bool

adapters/repos/db/init.go

@@ -93,6 +93,7 @@ func (db *DB) init(ctx context.Context) error {
 				SegmentsCleanupIntervalSeconds:      db.config.SegmentsCleanupIntervalSeconds,
 				SeparateObjectsCompactions:          db.config.SeparateObjectsCompactions,
 				CycleManagerRoutinesFactor:          db.config.CycleManagerRoutinesFactor,
+				IndexRangeableInMemory:              db.config.IndexRangeableInMemory,
 				MaxSegmentSize:                      db.config.MaxSegmentSize,
 				HNSWMaxLogSize:                      db.config.HNSWMaxLogSize,
 				HNSWWaitForCachePrefill:             db.config.HNSWWaitForCachePrefill,

adapters/repos/db/inverted/searcher_doc_bitmap.go

@@ -30,6 +30,7 @@ func (s *Searcher) docBitmap(ctx context.Context, b *lsmkv.Bucket, limit int,
 	pv *propValuePair,
 ) (bm docBitmap, err error) {
 	before := time.Now()
+	strategy := "geo"
 	defer func() {
 		took := time.Since(before)
 		vals := map[string]any{
@@ -39,6 +40,7 @@ func (s *Searcher) docBitmap(ctx context.Context, b *lsmkv.Bucket, limit int,
 			"took_string": took.String(),
 			"value":       pv.value,
 			"count":       bm.count(),
+			"strategy":    strategy,
 		}
 
 		helpers.AnnotateSlowQueryLogAppend(ctx, "build_allow_list_doc_bitmap", vals)
@@ -51,6 +53,7 @@ func (s *Searcher) docBitmap(ctx context.Context, b *lsmkv.Bucket, limit int,
 		bm, err = s.docBitmapGeo(ctx, pv)
 		return
 	}
+	strategy = b.Strategy()
 
 	// all other operators perform operations on the inverted index which we
 	// can serve directly

adapters/repos/db/lsmkv/bucket.go

@@ -142,6 +142,10 @@ type Bucket struct {
 	// introduces latency of segment availability, for the tradeoff of
 	// ensuring segment files have integrity before reading them.
 	enableChecksumValidation bool
+
+	// keep segments in memory for more performant search
+	// (currently used by roaringsetrange inverted indexes)
+	keepSegmentsInMemory bool
 }
 
 func NewBucketCreator() *Bucket { return &Bucket{} }
@@ -205,6 +209,7 @@ func (*Bucket) NewBucket(ctx context.Context, dir, rootDir string, logger logrus
 			maxSegmentSize:           b.maxSegmentSize,
 			cleanupInterval:          b.segmentsCleanupInterval,
 			enableChecksumValidation: b.enableChecksumValidation,
+			keepSegmentsInMemory:     b.keepSegmentsInMemory,
 		}, b.allocChecker)
 	if err != nil {
 		return nil, fmt.Errorf("init disk segments: %w", err)
@@ -1343,23 +1348,34 @@ func (b *Bucket) FlushAndSwitch() error {
 		return fmt.Errorf("precompute metadata: %w", err)
 	}
 
+	flushing := b.flushing
 	if err := b.atomicallyAddDiskSegmentAndRemoveFlushing(segment); err != nil {
 		return fmt.Errorf("add segment and remove flushing: %w", err)
 	}
 
-	if b.strategy == StrategyInverted && !tombstones.IsEmpty() {
-		if err = func() error {
-			b.disk.maintenanceLock.RLock()
-			defer b.disk.maintenanceLock.RUnlock()
-			// add flushing memtable tombstones to all segments
-			for _, seg := range b.disk.segments {
-				if _, err := seg.MergeTombstones(tombstones); err != nil {
-					return fmt.Errorf("merge tombstones: %w", err)
+	switch b.strategy {
+	case StrategyInverted:
+		if !tombstones.IsEmpty() {
+			if err = func() error {
+				b.disk.maintenanceLock.RLock()
+				defer b.disk.maintenanceLock.RUnlock()
+				// add flushing memtable tombstones to all segments
+				for _, seg := range b.disk.segments {
+					if _, err := seg.MergeTombstones(tombstones); err != nil {
+						return fmt.Errorf("merge tombstones: %w", err)
+					}
 				}
+				return nil
+			}(); err != nil {
+				return fmt.Errorf("add tombstones: %w", err)
+			}
+		}
+
+	case StrategyRoaringSetRange:
+		if b.keepSegmentsInMemory {
+			if err := b.disk.roaringSetRangeSegmentInMemory.MergeMemtable(flushing.roaringSetRange); err != nil {
+				return fmt.Errorf("merge roaringsetrange memtable to segment-in-memory: %w", err)
 			}
-			return nil
-		}(); err != nil {
-			return fmt.Errorf("add tombstones: %w", err)
 		}
 	}
 

adapters/repos/db/lsmkv/bucket_options.go

@@ -185,3 +185,10 @@ func WithForceCompaction(opt bool) BucketOption {
 		return nil
 	}
 }
+
+func WithKeepSegmentsInMemory(keep bool) BucketOption {
+	return func(b *Bucket) error {
+		b.keepSegmentsInMemory = keep
+		return nil
+	}
+}

adapters/repos/db/lsmkv/bucket_roaring_set_range.go

@@ -52,7 +52,14 @@ func (b *Bucket) ReaderRoaringSetRange() ReaderRoaringSetRange {
 
 	b.flushLock.RLock()
 
-	readers, releaseSegmentGroup := b.disk.newRoaringSetRangeReaders()
+	var release func()
+	var readers []roaringsetrange.InnerReader
+	if b.keepSegmentsInMemory {
+		reader, releaseInt := b.disk.roaringSetRangeSegmentInMemory.Reader()
+		readers, release = []roaringsetrange.InnerReader{reader}, releaseInt
+	} else {
+		readers, release = b.disk.newRoaringSetRangeReaders()
+	}
 
 	// we have a flush-RLock, so we have the guarantee that the flushing state
 	// will not change for the lifetime of the cursor, thus there can only be two
@@ -63,7 +70,7 @@ func (b *Bucket) ReaderRoaringSetRange() ReaderRoaringSetRange {
 	readers = append(readers, b.active.newRoaringSetRangeReader())
 
 	return roaringsetrange.NewCombinedReader(readers, func() {
-		releaseSegmentGroup()
+		release()
 		b.flushLock.RUnlock()
 	}, concurrency.SROAR_MERGE, b.logger)
 }

adapters/repos/db/lsmkv/segment_group.go

@@ -25,6 +25,7 @@ import (
 	"github.com/sirupsen/logrus"
 	"github.com/weaviate/weaviate/adapters/repos/db/lsmkv/segmentindex"
 	"github.com/weaviate/weaviate/adapters/repos/db/roaringset"
+	"github.com/weaviate/weaviate/adapters/repos/db/roaringsetrange"
 	"github.com/weaviate/weaviate/entities/cyclemanager"
 	"github.com/weaviate/weaviate/entities/diskio"
 	"github.com/weaviate/weaviate/entities/lsmkv"
@@ -83,6 +84,8 @@ type SegmentGroup struct {
 	cleanupInterval    time.Duration
 	lastCleanupCall    time.Time
 	lastCompactionCall time.Time
+
+	roaringSetRangeSegmentInMemory *roaringsetrange.SegmentInMemory
 }
 
 type sgConfig struct {
@@ -98,6 +101,7 @@ type sgConfig struct {
 	maxSegmentSize           int64
 	cleanupInterval          time.Duration
 	enableChecksumValidation bool
+	keepSegmentsInMemory     bool
 }
 
 func newSegmentGroup(logger logrus.FieldLogger, metrics *Metrics,
@@ -359,7 +363,8 @@ func newSegmentGroup(logger logrus.FieldLogger, metrics *Metrics,
 		sg.strategy = StrategyMapCollection
 	}
 
-	if sg.strategy == StrategyInverted {
+	switch sg.strategy {
+	case StrategyInverted:
 		// start with last but one segment, as the last one doesn't need tombstones for now
 		for i := len(sg.segments) - 2; i >= 0; i-- {
 			// avoid crashing if segment has no tombstones
@@ -371,6 +376,17 @@ func newSegmentGroup(logger logrus.FieldLogger, metrics *Metrics,
 				return nil, fmt.Errorf("init segment %s: merge tombstones %w", sg.segments[i].path, err)
 			}
 		}
+
+	case StrategyRoaringSetRange:
+		if cfg.keepSegmentsInMemory {
+			sg.roaringSetRangeSegmentInMemory = roaringsetrange.NewSegmentInMemory()
+			for _, seg := range sg.segments {
+				cursor := seg.newRoaringSetRangeCursor()
+				if err := sg.roaringSetRangeSegmentInMemory.MergeSegmentByCursor(cursor); err != nil {
+					return nil, fmt.Errorf("build segment-in-memory of strategy '%s': %w", sg.strategy, err)
+				}
+			}
+		}
 	}
 
 	return sg, nil

adapters/repos/db/migrator.go

@@ -119,6 +119,7 @@ func (m *Migrator) AddClass(ctx context.Context, class *models.Class,
 			SegmentsCleanupIntervalSeconds:      m.db.config.SegmentsCleanupIntervalSeconds,
 			SeparateObjectsCompactions:          m.db.config.SeparateObjectsCompactions,
 			CycleManagerRoutinesFactor:          m.db.config.CycleManagerRoutinesFactor,
+			IndexRangeableInMemory:              m.db.config.IndexRangeableInMemory,
 			MaxSegmentSize:                      m.db.config.MaxSegmentSize,
 			HNSWMaxLogSize:                      m.db.config.HNSWMaxLogSize,
 			HNSWWaitForCachePrefill:             m.db.config.HNSWWaitForCachePrefill,

adapters/repos/db/repo.go

@@ -219,6 +219,7 @@ type Config struct {
 	Replication                         replication.GlobalConfig
 	MaximumConcurrentShardLoads         int
 	CycleManagerRoutinesFactor          int
+	IndexRangeableInMemory              bool
 }
 
 // GetIndex returns the index if it exists or nil if it doesn't