Merge #149652 #150437

149652: pkg/util/log: reducing allocs in otlp sink r=TheComputerM a=TheComputerM

this aims to increase the performance and decrease the number of
allocations made in the otlp sink.

Benchmark results:
```
name         old time/op    new time/op    delta
OTLPSink-12     536µs ± 1%     535µs ± 0%     ~     (p=0.156 n=10+9)

name         old alloc/op   new alloc/op   delta
OTLPSink-12    18.5kB ±25%    15.7kB ±11%  -14.97%  (p=0.004 n=9+9)

name         old allocs/op  new allocs/op  delta
OTLPSink-12       132 ± 0%       123 ± 0%   -6.82%  (p=0.000 n=10+10)
```

benchmark was done in a e2e environment (crdb->otel-collector->datadog)
and with gzip compression turned on.

part of: CRDB-51667

Release note: none

150437: vecbench: update vecbench and vecann to handle segmented datasets r=drewkimball a=andy-kimball

A recent change updated vector datasets to segment huge sets of training vectors across multiple files. This commit updates vecbench and vecann to consume these multi-file datasets. Vecbench will separately read each file and insert its vectors into the store, so that memory usage stays low. Vecann will load all files into memory so that it can randomly choose vectors to insert and search.

Epic: CRDB-42943
Release note: None

Release note (cli change):

Co-authored-by: Mudit Somani <[email protected]>
Co-authored-by: Andrew Kimball <[email protected]>

Author: 3 people

pkg/cmd/vecbench/main.go

@@ -188,8 +188,7 @@ type vectorBench struct {
 	datasetName    string
 	distanceMetric vecpb.DistanceMetric
 	provider       VectorProvider
-	buildData      vecann.Dataset
-	searchData     vecann.SearchDataset
+	data           vecann.Dataset
 }
 
 // newVectorBench creates a new VecBench instance for the given dataset.
@@ -218,11 +217,11 @@ func newVectorBench(ctx context.Context, stopper *stop.Stopper, datasetName stri
 // SearchIndex downloads, builds, and searches an index for the given dataset.
 func (vb *vectorBench) SearchIndex() {
 	// Ensure the data needed for the search is available.
-	vb.ensureDataset(vb.ctx, true /* forSearch */)
+	vb.ensureDataset(vb.ctx)
 
 	var err error
 	vb.provider, err = newVectorProvider(
-		vb.stopper, vb.datasetName, vb.searchData.Test.Dims, vb.distanceMetric)
+		vb.stopper, vb.datasetName, vb.data.Dims, vb.distanceMetric)
 	if err != nil {
 		panic(err)
 	}
@@ -249,10 +248,10 @@ func (vb *vectorBench) SearchIndex() {
 
 		// Search for test vectors.
 		var sumRecall, sumVectors, sumLeafVectors, sumFullVectors, sumPartitions float64
-		count := vb.searchData.Test.Count
+		count := vb.data.Test.Count
 		for i := range count {
 			// Calculate truth set for the vector.
-			queryVector := vb.searchData.Test.At(i)
+			queryVector := vb.data.Test.At(i)
 
 			var stats cspann.SearchStats
 			prediction, err := vb.provider.Search(vb.ctx, state, queryVector, &stats)
@@ -264,7 +263,7 @@ func (vb *vectorBench) SearchIndex() {
 			truth := make([]cspann.KeyBytes, maxResults)
 			for neighbor := range maxResults {
 				primaryKey := primaryKeys[neighbor*4 : neighbor*4+4]
-				binary.BigEndian.PutUint32(primaryKey, uint32(vb.searchData.Neighbors[i][neighbor]))
+				binary.BigEndian.PutUint32(primaryKey, uint32(vb.data.Neighbors[i][neighbor]))
 				truth[neighbor] = primaryKey
 			}
 
@@ -287,7 +286,7 @@ func (vb *vectorBench) SearchIndex() {
 	fmt.Printf(White+"%s\n"+Reset, vb.datasetName)
 	fmt.Printf(
 		White+"%d train vectors, %d test vectors, %d dimensions, %d/%d min/max partitions, build beam size %d\n"+Reset,
-		vb.searchData.Count, vb.searchData.Test.Count, vb.searchData.Test.Dims,
+		vb.data.TrainCount, vb.data.Test.Count, vb.data.Dims,
 		minPartitionSize, maxPartitionSize, *flagBeamSize)
 	fmt.Println(vb.provider.FormatStats())
 
@@ -308,12 +307,12 @@ func (vb *vectorBench) SearchIndex() {
 // it is rebuilt from scratch.
 func (vb *vectorBench) BuildIndex() {
 	// Ensure dataset is downloaded and cached.
-	vb.ensureDataset(vb.ctx, false /* forSearch */)
+	vb.ensureDataset(vb.ctx)
 
 	// Construct the vector provider.
 	var err error
 	vb.provider, err = newVectorProvider(
-		vb.stopper, vb.datasetName, vb.buildData.Train.Dims, vb.distanceMetric)
+		vb.stopper, vb.datasetName, vb.data.Dims, vb.distanceMetric)
 	if err != nil {
 		panic(err)
 	}
@@ -324,14 +323,6 @@ func (vb *vectorBench) BuildIndex() {
 		panic(err)
 	}
 
-	// Create unique primary key for each vector.
-	primaryKeys := make([]cspann.KeyBytes, vb.buildData.Train.Count)
-	keyBuf := make(cspann.KeyBytes, vb.buildData.Train.Count*4)
-	for i := range vb.buildData.Train.Count {
-		primaryKeys[i] = keyBuf[i*4 : i*4+4]
-		binary.BigEndian.PutUint32(primaryKeys[i], uint32(i))
-	}
-
 	// Compute percentile latencies.
 	estimator := NewPercentileEstimator(1000)
 
@@ -355,71 +346,96 @@ func (vb *vectorBench) BuildIndex() {
 	fmt.Printf(White+"Building index for dataset: %s\n"+Reset, vb.datasetName)
 	startAt := crtime.NowMono()
 
-	// Insert vectors into the provider on multiple goroutines.
+	// Insert vectors into the provider using batches of training vectors.
 	var insertCount atomic.Uint64
-	procs := runtime.GOMAXPROCS(-1)
-	countPerProc := (vb.buildData.Train.Count + procs) / procs
+	var lastInserted int
 	batchSize := *flagBatchSize
-	for i := 0; i < vb.buildData.Train.Count; i += countPerProc {
-		end := min(i+countPerProc, vb.buildData.Train.Count)
-		go func(start, end int) {
-			// Break vector group into batches that each insert a batch of vectors.
-			for j := start; j < end; j += batchSize {
-				startMono := crtime.NowMono()
-				vectors := vb.buildData.Train.Slice(j, min(j+batchSize, end)-j)
-				err := vb.provider.InsertVectors(vb.ctx, primaryKeys[j:j+vectors.Count], vectors)
-				if err != nil {
-					panic(err)
-				}
-				estimator.Add(startMono.Elapsed().Seconds() / float64(vectors.Count))
-				insertCount.Add(uint64(vectors.Count))
-			}
-		}(i, end)
-	}
 
-	// Compute ops per second.
-	var lastInserted int
+	// Reset the dataset to start from the beginning
+	vb.data.Reset()
 
-	// Update progress every second.
-	lastProgressAt := startAt
 	for {
-		time.Sleep(time.Second)
-
-		// Calculate exactly how long it's been since last progress update.
-		now := crtime.NowMono()
-		sinceProgress := now.Sub(lastProgressAt)
-		lastProgressAt = now
-
-		// Calculate ops per second over the last second.
-		totalInserted := int(insertCount.Load())
-		opsPerSec := float64(totalInserted-lastInserted) / sinceProgress.Seconds()
-		lastInserted = totalInserted
-
-		cp.AddSample(throughput, opsPerSec)
-		cp.AddSample(p50, estimator.Estimate(0.50)*1000)
-		cp.AddSample(p90, estimator.Estimate(0.90)*1000)
-		cp.AddSample(p99, estimator.Estimate(0.99)*1000)
-
-		// Add provider-specific metric samples.
-		metrics, err := vb.provider.GetMetrics()
+		// Get next batch of train vectors.
+		hasMore, err := vb.data.Next()
 		if err != nil {
 			panic(err)
 		}
-		for i, metric := range metrics {
-			cp.AddSample(metricIds[i], metric.Value)
+		if !hasMore {
+			// No more batches.
+			break
 		}
-		cp.Plot()
-
-		if !*flagHideProgress {
-			sinceStart := now.Sub(startAt)
-			fmt.Printf(White+"\rInserted %d / %d vectors (%.2f%%) in %v"+Reset,
-				totalInserted, vb.buildData.Train.Count,
-				(float64(totalInserted)/float64(vb.buildData.Train.Count))*100,
-				sinceStart.Truncate(time.Second))
+		trainBatch := vb.data.Train
+		insertedBefore := int(insertCount.Load())
+
+		// Create primary keys for this batch
+		primaryKeys := make([]cspann.KeyBytes, trainBatch.Count)
+		keyBuf := make(cspann.KeyBytes, trainBatch.Count*4)
+		for i := range trainBatch.Count {
+			primaryKeys[i] = keyBuf[i*4 : i*4+4]
+			binary.BigEndian.PutUint32(primaryKeys[i], uint32(insertedBefore+i))
 		}
 
-		if totalInserted >= vb.buildData.Train.Count {
-			break
+		procs := runtime.GOMAXPROCS(-1)
+		countPerProc := (vb.data.Train.Count + procs) / procs
+		for i := 0; i < vb.data.Train.Count; i += countPerProc {
+			end := min(i+countPerProc, vb.data.Train.Count)
+			go func(start, end int) {
+				// Break vector group into batches that each insert a batch of vectors.
+				for j := start; j < end; j += batchSize {
+					startMono := crtime.NowMono()
+					vectors := vb.data.Train.Slice(j, min(j+batchSize, end)-j)
+					err := vb.provider.InsertVectors(vb.ctx, primaryKeys[j:j+vectors.Count], vectors)
+					if err != nil {
+						panic(err)
+					}
+					estimator.Add(startMono.Elapsed().Seconds() / float64(vectors.Count))
+					insertCount.Add(uint64(vectors.Count))
+				}
+			}(i, end)
+		}
+
+		// Update progress every second.
+		lastProgressAt := startAt
+		for {
+			time.Sleep(time.Second)
+
+			// Calculate exactly how long it's been since last progress update.
+			now := crtime.NowMono()
+			sinceProgress := now.Sub(lastProgressAt)
+			lastProgressAt = now
+
+			// Calculate ops per second over the last second.
+			totalInserted := int(insertCount.Load())
+			opsPerSec := float64(totalInserted-lastInserted) / sinceProgress.Seconds()
+			lastInserted = totalInserted
+
+			cp.AddSample(throughput, opsPerSec)
+			cp.AddSample(p50, estimator.Estimate(0.50)*1000)
+			cp.AddSample(p90, estimator.Estimate(0.90)*1000)
+			cp.AddSample(p99, estimator.Estimate(0.99)*1000)
+
+			// Add provider-specific metric samples.
+			metrics, err := vb.provider.GetMetrics()
+			if err != nil {
+				panic(err)
+			}
+			for i, metric := range metrics {
+				cp.AddSample(metricIds[i], metric.Value)
+			}
+			cp.Plot()
+
+			if !*flagHideProgress {
+				sinceStart := now.Sub(startAt)
+				fmt.Printf(White+"\rInserted %d / %d vectors (%.2f%%) in %v"+Reset,
+					totalInserted, vb.data.TrainCount,
+					(float64(totalInserted)/float64(vb.data.TrainCount))*100,
+					sinceStart.Truncate(time.Second))
+			}
+
+			// Check if we've inserted all vectors in the batch.
+			if (totalInserted - insertedBefore) >= vb.data.Train.Count {
+				break
+			}
 		}
 	}
 
@@ -432,9 +448,8 @@ func (vb *vectorBench) BuildIndex() {
 }
 
 // ensureDataset ensures that the dataset has been downloaded and cached to
-// disk. It also loads the data into memory. If "forSearch" is true, then only
-// test vectors are loaded, not train vectors.
-func (vb *vectorBench) ensureDataset(ctx context.Context, forSearch bool) {
+// disk. It also loads the data into memory.
+func (vb *vectorBench) ensureDataset(ctx context.Context) {
 	loader := vecann.DatasetLoader{
 		DatasetName: vb.datasetName,
 		OnProgress: func(ctx context.Context, format string, args ...any) {
@@ -450,20 +465,11 @@ func (vb *vectorBench) ensureDataset(ctx context.Context, forSearch bool) {
 		},
 	}
 
-	if forSearch {
-		// Only need to load the search data, not the insert data. Loading
-		// the insert data can take 10+ seconds for large datasets.
-		if err := loader.LoadForSearch(ctx); err != nil {
-			panic(err)
-		}
-	} else {
-		if err := loader.Load(ctx); err != nil {
-			panic(err)
-		}
+	if err := loader.Load(ctx); err != nil {
+		panic(err)
 	}
 
-	vb.buildData = loader.Data
-	vb.searchData = loader.SearchData
+	vb.data = loader.Data
 }
 
 // newVectorProvider creates a new in-memory or SQL based vector provider that