cspann: change TryGetPartitionMetadata to handle batches

Previously, TryGetPartitionMetadata fetched metadata for a single partition.
This commit changes that to instead fetch metadata for a batch of multiple
partitions in one call. This will be used by a future commit.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

pkg/sql/vecindex/cspann/commontest/storetests.go

@@ -76,10 +76,11 @@ type MakeStoreFunc func(quantizer quantize.Quantizer) TestStore
 type StoreTestSuite struct {
 	suite.Suite
 
-	ctx           context.Context
-	makeStore     MakeStoreFunc
-	rootQuantizer quantize.Quantizer
-	quantizer     quantize.Quantizer
+	ctx              context.Context
+	makeStore        MakeStoreFunc
+	rootQuantizer    quantize.Quantizer
+	quantizer        quantize.Quantizer
+	nextPartitionKey cspann.PartitionKey
 }
 
 // NewStoreTestSuite constructs a new suite of tests that run against
@@ -88,10 +89,12 @@ type StoreTestSuite struct {
 // tests.
 func NewStoreTestSuite(ctx context.Context, makeStore MakeStoreFunc) *StoreTestSuite {
 	return &StoreTestSuite{
-		ctx:           ctx,
-		makeStore:     makeStore,
-		rootQuantizer: quantize.NewUnQuantizer(2, vecpb.L2SquaredDistance),
-		quantizer:     quantize.NewRaBitQuantizer(2, 42, vecpb.L2SquaredDistance)}
+		ctx:              ctx,
+		makeStore:        makeStore,
+		rootQuantizer:    quantize.NewUnQuantizer(2, vecpb.L2SquaredDistance),
+		quantizer:        quantize.NewRaBitQuantizer(2, 42, vecpb.L2SquaredDistance),
+		nextPartitionKey: cspann.RootKey + 1,
+	}
 }
 
 func (suite *StoreTestSuite) TestRunTransaction() {
@@ -700,31 +703,38 @@ func (suite *StoreTestSuite) TestTryGetPartitionMetadata() {
 
 	doTest := func(treeID int) {
 		treeKey := store.MakeTreeKey(suite.T(), treeID)
-		partitionKey := cspann.PartitionKey(10)
-
-		// Partition does not yet exist.
-		_, err := store.TryGetPartitionMetadata(suite.ctx, treeKey, partitionKey)
-		suite.ErrorIs(err, cspann.ErrPartitionNotFound)
 
-		// Create partition with some vectors in it.
-		partitionKey, partition := suite.createTestPartition(store, treeKey)
+		// Create two partition with vectors in them.
+		partitionKey1, partition1 := suite.createTestPartition(store, treeKey)
+		partitionKey2, partition2 := suite.createTestPartition(store, treeKey)
 
-		// Fetch back only the metadata and validate it.
-		partitionMetadata, err := store.TryGetPartitionMetadata(suite.ctx, treeKey, partitionKey)
+		// Fetch metadata for the partitions, along with one that doesn't exist.
+		toGet := []cspann.PartitionMetadataToGet{
+			{Key: partitionKey1},
+			{Key: cspann.PartitionKey(9999)},
+			{Key: partitionKey2},
+		}
+		err := store.TryGetPartitionMetadata(suite.ctx, treeKey, toGet)
 		suite.NoError(err)
-		suite.True(partitionMetadata.Equal(partition.Metadata()))
+
+		// Validate that partition 9999 does not exist.
+		suite.Equal(cspann.PartitionMetadata{}, toGet[1].Metadata)
+
+		// Validate metadata for other partitions.
+		suite.True(partition1.Metadata().Equal(&toGet[0].Metadata))
+		suite.True(partition2.Metadata().Equal(&toGet[2].Metadata))
 
 		// Update the metadata and verify we get the updated values.
-		expected := *partition.Metadata()
+		expected := toGet[0].Metadata
 		metadata := expected
 		metadata.StateDetails.MakeUpdating(30)
 		suite.NoError(store.TryUpdatePartitionMetadata(
-			suite.ctx, treeKey, partitionKey, metadata, expected))
+			suite.ctx, treeKey, partitionKey1, metadata, expected))
 
 		// Fetch updated metadata and validate.
-		partitionMetadata, err = store.TryGetPartitionMetadata(suite.ctx, treeKey, partitionKey)
+		err = store.TryGetPartitionMetadata(suite.ctx, treeKey, toGet[:1])
 		suite.NoError(err)
-		suite.True(partitionMetadata.Equal(&metadata))
+		suite.True(toGet[0].Metadata.Equal(&metadata))
 	}
 
 	suite.Run("default tree", func() {
@@ -1043,7 +1053,8 @@ func (suite *StoreTestSuite) TestTryClearPartition() {
 func (suite *StoreTestSuite) createTestPartition(
 	store TestStore, treeKey cspann.TreeKey,
 ) (cspann.PartitionKey, *cspann.Partition) {
-	partitionKey := cspann.PartitionKey(10)
+	partitionKey := suite.nextPartitionKey
+	suite.nextPartitionKey++
 	metadata := cspann.MakeReadyPartitionMetadata(cspann.SecondLevel, vector.T{4, 3})
 	suite.NoError(store.TryCreateEmptyPartition(suite.ctx, treeKey, partitionKey, metadata))
 	vectors := vector.MakeSet(2)

pkg/sql/vecindex/cspann/fixup_split.go

@@ -209,6 +209,7 @@ func (fw *fixupWorker) splitPartition(
 			return errors.Wrapf(errFixupAborted, "reloading partition, metadata timestamp changed")
 		}
 	} else if metadata.StateDetails.State != MissingState {
+		// Fetch metadata for already created left and right partitions.
 		leftMetadata, err = fw.getPartitionMetadata(ctx, leftPartitionKey)
 		if err != nil {
 			return err
@@ -363,16 +364,15 @@ func (fw *fixupWorker) getPartition(
 // getPartitionMetadata returns the up-to-date metadata of the partition with
 // the given key.
 func (fw *fixupWorker) getPartitionMetadata(
-	ctx context.Context, partitionKey PartitionKey,
+       ctx context.Context, partitionKey PartitionKey,
 ) (PartitionMetadata, error) {
-	metadata, err := fw.index.store.TryGetPartitionMetadata(ctx, fw.treeKey, partitionKey)
+	fw.tempMetadataToGet = ensureSliceLen(fw.tempMetadataToGet, 1)
+	fw.tempMetadataToGet[0].Key = partitionKey
+	err := fw.index.store.TryGetPartitionMetadata(ctx, fw.treeKey, fw.tempMetadataToGet)
 	if err != nil {
-		metadata, err = suppressRaceErrors(err)
-		if err != nil {
-			return PartitionMetadata{}, errors.Wrapf(err, "getting metadata for partition %d", partitionKey)
-		}
+		return PartitionMetadata{}, errors.Wrapf(err, "getting metadata for partition %d", partitionKey)
 	}
-	return metadata, nil
+	return fw.tempMetadataToGet[0].Metadata, nil
 }
 
 // updateMetadata updates the given partition's metadata record, on the
@@ -586,8 +586,7 @@ func (fw *fixupWorker) createSplitSubPartition(
 		// Load parent metadata to verify that it's in a state that allows inserts.
 		parentMetadata, err := fw.getPartitionMetadata(ctx, parentPartitionKey)
 		if err != nil {
-			return PartitionMetadata{}, errors.Wrapf(err,
-				"getting parent partition %d metadata", parentPartitionKey)
+			return PartitionMetadata{}, err
 		}
 
 		parentLevel := sourceMetadata.Level + 1

pkg/sql/vecindex/cspann/fixup_worker.go

@@ -95,6 +95,7 @@ type fixupWorker struct {
 	tempVectorsWithKeys []VectorWithKey
 	tempChildKey        [1]ChildKey
 	tempValueBytes      [1]ValueBytes
+	tempMetadataToGet   []PartitionMetadataToGet
 }
 
 // ewFixupWorker returns a new worker for the given processor.

pkg/sql/vecindex/cspann/memstore/memstore.go

@@ -377,17 +377,26 @@ func (s *Store) TryGetPartition(
 
 // TryGetPartitionMetadata implements the Store interface.
 func (s *Store) TryGetPartitionMetadata(
-	ctx context.Context, treeKey cspann.TreeKey, partitionKey cspann.PartitionKey,
-) (cspann.PartitionMetadata, error) {
-	memPart := s.lockPartition(treeKey, partitionKey, uniqueOwner, false /* isExclusive */)
-	if memPart == nil {
-		// Partition does not exist.
-		return cspann.PartitionMetadata{}, cspann.ErrPartitionNotFound
+	ctx context.Context, treeKey cspann.TreeKey, toGet []cspann.PartitionMetadataToGet,
+) error {
+	for i := range toGet {
+		item := &toGet[i]
+
+		func() {
+			memPart := s.lockPartition(treeKey, item.Key, uniqueOwner, false /* isExclusive */)
+			if memPart == nil {
+				// Partition does not exist, so map it to Missing.
+				item.Metadata = cspann.PartitionMetadata{}
+				return
+			}
+			defer memPart.lock.ReleaseShared()
+
+			// Return a copy of the metadata.
+			item.Metadata = *memPart.lock.partition.Metadata()
+		}()
 	}
-	defer memPart.lock.ReleaseShared()
 
-	// Return a copy of the metadata.
-	return *memPart.lock.partition.Metadata(), nil
+	return nil
 }
 
 // TryUpdatePartitionMetadata implements the Store interface.

pkg/sql/vecindex/cspann/memstore/memstore_test.go

@@ -197,10 +197,11 @@ func TestInMemoryStoreUpdateStats(t *testing.T) {
 		require.Equal(t, []cspann.CVStats{}, stats.CVStats)
 
 		// Increase root partition level and check stats.
-		metadata, err := store.TryGetPartitionMetadata(ctx, treeKey, cspann.RootKey)
+		toGet := []cspann.PartitionMetadataToGet{{Key: cspann.RootKey}}
+		err = store.TryGetPartitionMetadata(ctx, treeKey, toGet)
 		require.NoError(t, err)
 
-		expected := metadata
+		expected := toGet[0].Metadata
 		metadata.Level = 3
 		err = store.TryUpdatePartitionMetadata(ctx, treeKey, cspann.RootKey, metadata, expected)
 		require.NoError(t, err)

pkg/sql/vecindex/cspann/store.go

@@ -44,6 +44,15 @@ type PartitionToSearch struct {
 	Count int
 }
 
+// PartitionMetadataToGet contains information about partition metadata to be
+// fetched by the TryGetPartitionMetadata method.
+type PartitionMetadataToGet struct {
+	// Key specifies which partition's metadata to fetch and is set by the caller.
+	Key PartitionKey
+	// Metadata is the metadata for the partition and is set by the Store.
+	Metadata PartitionMetadata
+}
+
 // Store encapsulates the component that's actually storing the vectors, whether
 // that's in a CRDB cluster for production or in memory for testing and
 // benchmarking. Callers can use Store to start and commit transactions against
@@ -92,13 +101,14 @@ type Store interface {
 		ctx context.Context, treeKey TreeKey, partitionKey PartitionKey,
 	) (*Partition, error)
 
-	// TryGetPartitionMetadata returns just the metadata of the requested
-	// partition, if it exists, else it returns ErrPartitionNotFound. This is
-	// more efficient than loading the entire partition when only metadata is
-	// needed.
+	// TryGetPartitionMetadata returns the metadata of the requested partitions.
+	// If a partition does not exist, its state is set to Missing.
+	//
+	// NOTE: The caller owns the "toGet" memory. The Store implementation should
+	// not try to use it after returning.
 	TryGetPartitionMetadata(
-		ctx context.Context, treeKey TreeKey, partitionKey PartitionKey,
-	) (PartitionMetadata, error)
+		ctx context.Context, treeKey TreeKey, toGet []PartitionMetadataToGet,
+	) error
 
 	// TryUpdatePartitionMetadata updates the partition's metadata only if it's
 	// equal to the expected value, else it returns a ConditionFailedError. If

pkg/sql/vecindex/vecstore/store.go

@@ -287,18 +287,36 @@ func (s *Store) TryGetPartition(
 }
 
 // TryGetPartitionMetadata is part of the cspann.Store interface. It returns the
-// metadata of an existing partition.
+// metadata for a batch of partitions.
 func (s *Store) TryGetPartitionMetadata(
-	ctx context.Context, treeKey cspann.TreeKey, partitionKey cspann.PartitionKey,
-) (cspann.PartitionMetadata, error) {
+	ctx context.Context, treeKey cspann.TreeKey, toGet []cspann.PartitionMetadataToGet,
+) error {
+	// Construct a batch with one Get request per partition.
 	b := s.kv.NewBatch()
-	metadataKey := vecencoding.EncodeMetadataKey(s.prefix, treeKey, partitionKey)
-	b.Get(metadataKey)
-	if err := s.kv.Run(ctx, b); err != nil {
-		return cspann.PartitionMetadata{},
-			errors.Wrapf(err, "getting partition metadata for %d", partitionKey)
+	for i := range toGet {
+		metadataKey := vecencoding.EncodeMetadataKey(s.prefix, treeKey, toGet[i].Key)
+		b.Get(metadataKey)
+	}
+
+	// Run the batch and return results.
+	var err error
+	if err = s.kv.Run(ctx, b); err != nil {
+		return errors.Wrapf(err, "getting partition metadata for %d partitions", len(toGet))
 	}
-	return s.getMetadataFromKVResult(partitionKey, &b.Results[0])
+
+	for i := range toGet {
+		item := &toGet[i]
+		item.Metadata, err = s.getMetadataFromKVResult(item.Key, &b.Results[i])
+
+		// If partition is missing, just return Missing metadata.
+		if err != nil {
+			if !errors.Is(err, cspann.ErrPartitionNotFound) {
+				return err
+			}
+		}
+	}
+
+	return nil
 }
 
 // TryUpdatePartitionMetadata is part of the cspann.Store interface. It updates