Merge #154074

154074: vecbench: add flag to control index creation timing r=mw5h a=mw5h

Previously, the vecbench SQL provider would create its test table with a
vector index. This patch gives us the option to retain that behavior
(the default) or to import into an empty table without a vector index
and then use CREATE INDEX to build the index after data is inserted.

Add the --index-after flag to provide choice between two index creation
workflows:
- Default: Create table with index, then import data
- Flag set: Create table, import data, then create index

Flag validation prevents use with --memstore since this only applies
to the SQL provider.

Also adds progress monitoring interface with CheckIndexCreationStatus()
method for tracking async index creation progress.

Informs: #146691
Release note: none

Co-authored-by: Matt White <[email protected]>

Author: craig[bot]

pkg/cmd/vecbench/main.go

@@ -48,6 +48,7 @@ const (
 )
 
 const (
+	ClearLine  = "\033[2K"
 	HideCursor = "\033[?25l"
 	ShowCursor = "\033[?25h"
 )
@@ -78,6 +79,8 @@ var flagDisableAdaptiveSearch = flag.Bool(
 var flagMemStore = flag.Bool("memstore", false, "Use in-memory store instead of CockroachDB")
 var flagDBConnStr = flag.String("db", "postgresql://root@localhost:26257",
 	"Database connection string (when not using --memstore)")
+var flagCreateIndexAfterImport = flag.Bool("index-after", false,
+	"Create vector index after data import instead of during table creation (SQL provider only)")
 
 // vecbench benchmarks vector index in-memory build and search performance on a
 // variety of datasets. Datasets are downloaded from the
@@ -125,7 +128,7 @@ func main() {
 	// Hide the cursor, but ensure it's restored on exit, including Ctrl+C.
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
-	stopper.RunAsyncTask(ctx, "Ctrl+C", func(context.Context) {
+	err := stopper.RunAsyncTask(ctx, "Ctrl+C", func(context.Context) {
 		select {
 		case <-c:
 			fmt.Print(ShowCursor)
@@ -136,12 +139,17 @@ func main() {
 			break
 		}
 	})
+	if err != nil {
+		fmt.Printf("Failed to install Ctrl-C handler: %v\n", err)
+	}
 	fmt.Print(HideCursor)
 	defer fmt.Print(ShowCursor)
 
-	// Start pprof server at http://localhost:8080/debug/pprof/
+	// Start pprof server at http://localhost:8080/debug/pprof/.
 	go func() {
-		http.ListenAndServe("localhost:8080", nil)
+		if err := http.ListenAndServe("localhost:8080", nil); err != nil {
+			fmt.Printf("Failed to start pprof server: %v\n", err)
+		}
 	}()
 
 	switch flag.Arg(0) {
@@ -340,6 +348,14 @@ func (vb *vectorBench) BuildIndex() {
 		panic(err)
 	}
 
+	// Create index immediately if flag is not set (default behavior).
+	if !*flagCreateIndexAfterImport {
+		err = vb.provider.CreateIndex(vb.ctx)
+		if err != nil {
+			panic(err)
+		}
+	}
+
 	// Compute percentile latencies.
 	estimator := NewPercentileEstimator(1000)
 
@@ -368,7 +384,7 @@ func (vb *vectorBench) BuildIndex() {
 	var lastInserted int
 	batchSize := *flagBatchSize
 
-	// Reset the dataset to start from the beginning
+	// Reset the dataset to start from the beginning.
 	vb.data.Reset()
 
 	for {
@@ -384,7 +400,7 @@ func (vb *vectorBench) BuildIndex() {
 		trainBatch := vb.data.Train
 		insertedBefore := int(insertCount.Load())
 
-		// Create primary keys for this batch
+		// Create primary keys for this batch.
 		primaryKeys := make([]cspann.KeyBytes, trainBatch.Count)
 		keyBuf := make(cspann.KeyBytes, trainBatch.Count*4)
 		for i := range trainBatch.Count {
@@ -456,6 +472,50 @@ func (vb *vectorBench) BuildIndex() {
 		}
 	}
 
+	// Create index after import if flag is set.
+	if *flagCreateIndexAfterImport {
+		if !*flagHideProgress {
+			fmt.Println()
+		}
+		startIndexTime := timeutil.Now()
+
+		// Start index creation in a goroutine to track progress.
+		done := make(chan error, 1)
+		go func() {
+			done <- vb.provider.CreateIndex(vb.ctx)
+		}()
+
+		// Track progress until CreateIndex returns.
+		ticker := time.NewTicker(time.Second)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case err := <-done:
+				// CreateIndex returned, stop tracking.
+				if err != nil {
+					panic(err)
+				}
+				fmt.Printf(ClearLine+White+"\rVector index creation completed in %v\n"+Reset,
+					roundDuration(timeutil.Since(startIndexTime)))
+				goto indexCreated
+			case <-ticker.C:
+				if *flagHideProgress {
+					break
+				}
+				// Check progress.
+				progress, err := vb.provider.CheckIndexCreationStatus(vb.ctx)
+				if err != nil {
+					fmt.Printf(Red+"Error checking index status: %v\n"+Reset, err)
+				} else {
+					fmt.Printf(ClearLine+White+"\rIndex creation progress: %.1f%% in %v"+Reset,
+						progress*100, timeutil.Since(startIndexTime).Truncate(time.Second))
+				}
+			}
+		}
+	indexCreated:
+	}
+
 	fmt.Printf(White+"\nBuilt index in %v\n"+Reset, roundDuration(startAt.Elapsed()))
 
 	// Ensure that index is persisted so it can be reused.
@@ -494,6 +554,11 @@ func (vb *vectorBench) ensureDataset(ctx context.Context) {
 func newVectorProvider(
 	stopper *stop.Stopper, datasetName string, dims int, distanceMetric vecpb.DistanceMetric,
 ) (VectorProvider, error) {
+	// Validate flag compatibility.
+	if *flagCreateIndexAfterImport && *flagMemStore {
+		return nil, errors.New("--create-index-after-import flag cannot be used with --memstore")
+	}
+
 	options := cspann.IndexOptions{
 		MinPartitionSize:      minPartitionSize,
 		MaxPartitionSize:      maxPartitionSize,

pkg/cmd/vecbench/mem_provider.go

@@ -110,7 +110,7 @@ func (m *MemProvider) Load(ctx context.Context) (bool, error) {
 	return true, nil
 }
 
-// New implements the VectorProvider interface
+// New implements the VectorProvider interface.
 func (m *MemProvider) New(ctx context.Context) error {
 	// Clear any existing state.
 	m.Close()
@@ -149,6 +149,18 @@ func (m *MemProvider) InsertVectors(
 	})
 }
 
+// CreateIndex implements the VectorProvider interface.
+func (m *MemProvider) CreateIndex(ctx context.Context) error {
+	// No-op for in-memory provider as index is built incrementally during insertion.
+	return nil
+}
+
+// CheckIndexCreationStatus implements the VectorProvider interface.
+func (m *MemProvider) CheckIndexCreationStatus(ctx context.Context) (float64, error) {
+	// Always return 100% complete for in-memory provider since index is built incrementally.
+	return 1.0, nil
+}
+
 // SetupSearch implements the VectorProvider interface.
 func (m *MemProvider) SetupSearch(
 	ctx context.Context, maxResults int, beamSize int,

pkg/cmd/vecbench/sql_provider.go

@@ -48,6 +48,7 @@ type SQLProvider struct {
 	options     cspann.IndexOptions
 	pool        *pgxpool.Pool
 	tableName   string
+	indexName   string
 	retryCount  atomic.Uint64
 }
 
@@ -85,6 +86,7 @@ func NewSQLProvider(
 
 	// Create sanitized table and index names.
 	tableName := fmt.Sprintf("vecbench_%s", sanitizeIdentifier(datasetName))
+	indexName := fmt.Sprintf("vecbench_%s_embedding_idx", sanitizeIdentifier(datasetName))
 
 	return &SQLProvider{
 		datasetName: datasetName,
@@ -93,6 +95,7 @@ func NewSQLProvider(
 		options:     options,
 		pool:        pool,
 		tableName:   tableName,
+		indexName:   indexName,
 	}, nil
 }
 
@@ -136,6 +139,20 @@ func (s *SQLProvider) New(ctx context.Context) error {
 		return errors.Wrap(err, "dropping table")
 	}
 
+	// Enable vector indexes if not already enabled.
+	_, err = s.pool.Exec(ctx, "SET CLUSTER SETTING feature.vector_index.enabled = true")
+	if err != nil {
+		return errors.Wrap(err, "enabling vector indexes")
+	}
+
+	// Create table without vector index. Index creation is handled separately.
+	_, err = s.pool.Exec(ctx, fmt.Sprintf(`CREATE TABLE %s (
+		id BYTES PRIMARY KEY,
+		embedding VECTOR(%d))`, s.tableName, s.dims))
+	return errors.Wrap(err, "creating table")
+}
+
+func (s *SQLProvider) CreateIndex(ctx context.Context) error {
 	var opClass string
 	switch s.distMetric {
 	case vecpb.CosineDistance:
@@ -144,13 +161,52 @@ func (s *SQLProvider) New(ctx context.Context) error {
 		opClass = " vector_ip_ops"
 	}
 
-	_, err = s.pool.Exec(ctx, fmt.Sprintf(`
-		CREATE TABLE %s (
-			id BYTES PRIMARY KEY,
-			embedding VECTOR(%d),
-			VECTOR INDEX (embedding%s)
-		)`, s.tableName, s.dims, opClass))
-	return errors.Wrap(err, "creating table")
+	_, err := s.pool.Exec(ctx, fmt.Sprintf(
+		"CREATE VECTOR INDEX %s ON %s (embedding%s)",
+		s.indexName,
+		s.tableName,
+		opClass,
+	))
+	return errors.Wrap(err, "creating index")
+}
+
+// CheckIndexCreationStatus implements the VectorProvider interface.
+func (s *SQLProvider) CheckIndexCreationStatus(ctx context.Context) (float64, error) {
+	var status string
+	var fractionCompleted float64
+
+	// Query for jobs related to our index creation.
+	rows, err := s.pool.Query(ctx, `
+		SELECT status, fraction_completed
+		FROM [SHOW JOBS]
+		WHERE description LIKE '%%vecbench%%'
+		AND job_type = 'NEW SCHEMA CHANGE'
+		ORDER BY created DESC
+		LIMIT 1`)
+	if err != nil {
+		return 0, errors.Wrap(err, "querying job progress")
+	}
+	defer rows.Close()
+
+	if !rows.Next() {
+		if err := rows.Err(); err != nil {
+			return 0, errors.Wrap(err, "querying job progress")
+		}
+		return 0.0, nil
+	}
+
+	if err := rows.Scan(&status, &fractionCompleted); err != nil {
+		return 0, errors.Wrap(err, "scanning job progress")
+	}
+
+	if status == "succeeded" {
+		return 1.0, nil
+	} else if status == "failed" || status == "canceled" {
+		return fractionCompleted, errors.Newf("index creation job failed with status: %s", status)
+	}
+
+	// Job is still running.
+	return fractionCompleted, nil
 }
 
 // InsertVectors implements the VectorProvider interface.
@@ -296,7 +352,7 @@ func (s *SQLProvider) FormatStats() string {
 	return ""
 }
 
-// sanitizeIdentifier makes a string safe to use as a SQL identifier
+// sanitizeIdentifier makes a string safe to use as a SQL identifier.
 func sanitizeIdentifier(s string) string {
 	// Replace non-alphanumeric characters with underscores.
 	return strings.Map(func(r rune) rune {

pkg/cmd/vecbench/vector_provider.go

@@ -54,6 +54,15 @@ type VectorProvider interface {
 	// identified by a key.
 	InsertVectors(ctx context.Context, keys []cspann.KeyBytes, vectors vector.Set) error
 
+	// CreateIndex creates a vector index on the data. This is called after
+	// table/store creation when the index should be created before data import.
+	CreateIndex(ctx context.Context) error
+
+	// CheckIndexCreationStatus returns the percentage complete (0.0-1.0) of index
+	// creation and any error. For providers that don't support async index creation,
+	// this should return 1.0, nil.
+	CheckIndexCreationStatus(ctx context.Context) (float64, error)
+
 	// SetupSearch allows the provider to perform expensive up-front steps in
 	// preparation for many calls to Search. It returns provider-specific state
 	// that will be passed to Search.