Repository: blevesearch/bleve
Pull request #2002 (2024-03-25 08:09:17)

MB-61134: Consider segment fileSize (when with vectors) during merge plan computation
Existing merge policy use a segment's total and live docs count as the measure for determining the merge tasks. 
A doc may contain multiple vector fields ( even chunked vectors ) , of varying vector dimensions, this usually means that index size will be greater than docs size by orders of magnitude.

Thus, docs count as merge policy measure, can easily lead us to formation of huge segments. That's a problem during merge and search time. 

Thus, this PR intends to utilize segment fileSize as an additional limiting check, to ensure that we don't end up creating huge segments.


example:
2M docs, each with a vector field of 2048 dimensions
size := 2000000 * 2048 * 4 bytes = 16GB



abhinavdangeti: @moshaad7 would you update the PR description above to explain everything the code here intends for.
The code in itself is looking mostly good - I'd love to hear if you've gotten this perf tested to see if we get the same gains as that updated merge policy got us a couple weeks ago.
abhinavdangeti: @moshaad7 I think you talked about this but for the record - you'll be getting us details on how much did we see file size grow to say with 25-50million documents (with vector content) - yes?

Patch:
diff --git a/index/scorch/mergeplan/merge_plan.go b/index/scorch/mergeplan/merge_plan.go
index 75235066..b4e7a372 100644
--- a/index/scorch/mergeplan/merge_plan.go
+++ b/index/scorch/mergeplan/merge_plan.go
@@ -37,6 +37,11 @@ type Segment interface {
 	// Size of the live data of the segment; i.e., FullSize() minus
 	// any logical deletions.
 	LiveSize() int64
+
+	HasVector() bool
+
+	// Size of the persisted segment file.
+	FileSize() int64
 }
 
 // Plan() will functionally compute a merge plan.  A segment will be
@@ -76,6 +81,11 @@ type MergePlanOptions struct {
 	// planner’s predicted sizes.
 	MaxSegmentSize int64
 
+	// Max size (in bytes) of the persisted segment file that contains the
+	// vectors.  This is used to prevent merging of segments that
+	// contain vectors that are too large.
+	MaxSegmentFileSize int64
+
 	// The growth factor for each tier in a staircase of idealized
 	// segments computed by CalcBudget().
 	TierGrowth float64
@@ -128,6 +138,7 @@ var ErrMaxSegmentSizeTooLarge = errors.New("MaxSegmentSize exceeds the size limi
 var DefaultMergePlanOptions = MergePlanOptions{
 	MaxSegmentsPerTier:   10,
 	MaxSegmentSize:       5000000,
+	MaxSegmentFileSize:   4000000000, // 4GB
 	TierGrowth:           10.0,
 	SegmentsPerMergeTask: 10,
 	FloorSegmentSize:     2000,
@@ -170,8 +181,17 @@ func plan(segmentsIn []Segment, o *MergePlanOptions) (*MergePlan, error) {
 			minLiveSize = segment.LiveSize()
 		}
 
+		isEligible := segment.LiveSize() < o.MaxSegmentSize/2
+		// An eligible segment (based on #documents) may be too large
+		// and thus need a stricter check based on the file size.
+		// This is particularly important for segments that contain
+		// vectors.
+		if isEligible && segment.HasVector() && o.MaxSegmentFileSize > 0 {
+			isEligible = segment.FileSize() < o.MaxSegmentFileSize/2
+		}
+
 		// Only small-enough segments are eligible.
-		if segment.LiveSize() < o.MaxSegmentSize/2 {
+		if isEligible {
 			eligibles = append(eligibles, segment)
 			eligiblesLiveSize += segment.LiveSize()
 		}
@@ -215,14 +235,25 @@ func plan(segmentsIn []Segment, o *MergePlanOptions) (*MergePlan, error) {
 		for startIdx := 0; startIdx < len(eligibles); startIdx++ {
 			var roster []Segment
 			var rosterLiveSize int64
+			var rosterFileSize int64 // useful for segments with vectors
 
 			for idx := startIdx; idx < len(eligibles) && len(roster) < o.SegmentsPerMergeTask; idx++ {
 				eligible := eligibles[idx]
 
-				if rosterLiveSize+eligible.LiveSize() < o.MaxSegmentSize {
-					roster = append(roster, eligible)
-					rosterLiveSize += eligible.LiveSize()
+				if rosterLiveSize+eligible.LiveSize() >= o.MaxSegmentSize {
+					continue
 				}
+
+				if eligible.HasVector() {
+					efs := eligible.FileSize()
+					if rosterFileSize+efs >= o.MaxSegmentFileSize {
+						continue
+					}
+					rosterFileSize += efs
+				}
+
+				roster = append(roster, eligible)
+				rosterLiveSize += eligible.LiveSize()
 			}
 
 			if len(roster) > 0 {
diff --git a/index/scorch/scorch.go b/index/scorch/scorch.go
index 2e6435ee..60256520 100644
--- a/index/scorch/scorch.go
+++ b/index/scorch/scorch.go
@@ -822,6 +822,10 @@ func (fs *fieldStats) Aggregate(stats segment.FieldStats) {
 
 // Returns the stats map
 func (fs *fieldStats) Fetch() map[string]map[string]uint64 {
+	if fs == nil {
+		return nil
+	}
+
 	return fs.statMap
 }
 
diff --git a/index/scorch/snapshot_segment.go b/index/scorch/snapshot_segment.go
index 1c14af72..96e59a31 100644
--- a/index/scorch/snapshot_segment.go
+++ b/index/scorch/snapshot_segment.go
@@ -16,6 +16,7 @@ package scorch
 
 import (
 	"bytes"
+	"os"
 	"sync"
 	"sync/atomic"
 
@@ -66,6 +67,31 @@ func (s *SegmentSnapshot) LiveSize() int64 {
 	return int64(s.Count())
 }
 
+func (s *SegmentSnapshot) HasVector() bool {
+	// number of vectors, for each vector field in the segment
+	numVecs := s.stats.Fetch()["num_vectors"]
+	return len(numVecs) > 0
+}
+
+func (s *SegmentSnapshot) FileSize() int64 {
+	ps, ok := s.segment.(segment.PersistedSegment)
+	if !ok {
+		return 0
+	}
+
+	path := ps.Path()
+	if path == "" {
+		return 0
+	}
+
+	fi, err := os.Stat(path)
+	if err != nil {
+		return 0
+	}
+
+	return fi.Size()
+}
+
 func (s *SegmentSnapshot) Close() error {
 	return s.segment.Close()
 }


Test patch:
diff --git a/index/scorch/mergeplan/merge_plan_test.go b/index/scorch/mergeplan/merge_plan_test.go
index 2791d09c..4a577109 100644
--- a/index/scorch/mergeplan/merge_plan_test.go
+++ b/index/scorch/mergeplan/merge_plan_test.go
@@ -30,11 +30,16 @@ type segment struct {
 	MyId       uint64
 	MyFullSize int64
 	MyLiveSize int64
+
+	MyHasVector bool
+	MyFileSize  int64
 }
 
 func (s *segment) Id() uint64      { return s.MyId }
 func (s *segment) FullSize() int64 { return s.MyFullSize }
 func (s *segment) LiveSize() int64 { return s.MyLiveSize }
+func (s *segment) HasVector() bool { return s.MyHasVector }
+func (s *segment) FileSize() int64 { return s.MyFileSize }
 
 func makeLinearSegments(n int) (rv []Segment) {
 	for i := 0; i < n; i++ {
@@ -589,3 +594,102 @@ func emit(descrip string, cycle int, step int, segments []Segment, plan *MergePl
 	fmt.Printf("%s %d.%d ---------- %s\n", descrip, cycle, step, suffix)
 	fmt.Printf("%s\n", ToBarChart(descrip, 100, segments, plan))
 }
+
+// -----------------------------------------------------------------------------
+// Test Vector Segment Merging
+
+func TestPlanMaxSegmentFileSize(t *testing.T) {
+	tests := []struct {
+		segments []Segment
+		o        *MergePlanOptions
+
+		expectedTasks [][]uint64
+	}{
+		{
+			[]Segment{
+				&segment{ // ineligible
+					MyId:       1,
+					MyFullSize: 4000,
+					MyLiveSize: 3900,
+
+					MyHasVector: true,
+					MyFileSize:  3900 * 1000 * 4, // > 2MB
+				},
+				&segment{ // ineligible
+					MyId:       2,
+					MyFullSize: 6000,
+					MyLiveSize: 5500, // > 5000
+
+					MyHasVector: true,
+					MyFileSize:  5500 * 1000 * 4, // > 2MB
+				},
+				&segment{ // eligible
+					MyId:       3,
+					MyFullSize: 500,
+					MyLiveSize: 490,
+
+					MyHasVector: true,
+					MyFileSize:  490 * 1000 * 4,
+				},
+				&segment{ // eligible
+					MyId:       4,
+					MyFullSize: 500,
+					MyLiveSize: 480,
+
+					MyHasVector: true,
+					MyFileSize:  480 * 1000 * 4,
+				},
+				&segment{ // eligible
+					MyId:       5,
+					MyFullSize: 500,
+					MyLiveSize: 300,
+
+					MyHasVector: true,
+					MyFileSize:  300 * 1000 * 4,
+				},
+				&segment{ // eligible
+					MyId:       6,
+					MyFullSize: 500,
+					MyLiveSize: 400,
+
+					MyHasVector: true,
+					MyFileSize:  400 * 1000 * 4,
+				},
+			},
+			&MergePlanOptions{
+				MaxSegmentSize: 5000, // number of documents
+				// considering vector dimension as 1000
+				// vectorBytes = 5000 * 1000 * 4 = 20MB, which is too large
+				// So, let's set the fileSize limit to 4MB
+				MaxSegmentFileSize:   4000000, // 4MB
+				MaxSegmentsPerTier:   1,
+				SegmentsPerMergeTask: 2,
+				TierGrowth:           2.0,
+				FloorSegmentSize:     1,
+			},
+			[][]uint64{
+				{3, 4},
+			},
+		},
+	}
+
+	for testi, test := range tests {
+		t.Run(fmt.Sprintf("Test-%d", testi), func(t *testing.T) {
+			plans, err := Plan(test.segments, test.o)
+			if err != nil {
+				t.Fatalf("Plan failed, err: %v", err)
+			}
+
+			for i, task := range plans.Tasks {
+				var segIDs []uint64
+				for _, seg := range task.Segments {
+					segIDs = append(segIDs, seg.Id())
+				}
+
+				if !reflect.DeepEqual(segIDs, test.expectedTasks[0]) {
+					t.Errorf("expected task segments: %v, got: %v", test.expectedTasks[i], segIDs)
+				}
+			}
+		})
+	}
+}