vecindex: suppress partition not found error in SearchPartitions

The SearchPartitions method is spec'd to return InvalidLevel for any partition
that is not found in the tree. However, the vecstore was instead returning
ErrPartitionNotFound, which was causing occasional insertion failures under
stress. This commit updates the vecstore to suppress ErrPartitionNotFound,
mapping it to a PartitionToSearch value of InvalidLevel, MissingState and
Count=0.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

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

@@ -346,14 +346,14 @@ func (suite *StoreTestSuite) TestSearchPartitions() {
 		treeKey := store.MakeTreeKey(suite.T(), treeID)
 
 		// Create non-root partition with some vectors in it.
-		testPartitionKey, _ := suite.createTestPartition(store, treeKey)
+		testPartitionKey, testPartition := suite.createTestPartition(store, treeKey)
 
 		// Create another partition.
 		testPartitionKey2 := cspann.PartitionKey(20)
 		metadata := cspann.PartitionMetadata{
 			Level:        cspann.SecondLevel,
 			Centroid:     vector.T{2, 4},
-			StateDetails: cspann.MakeReadyDetails(),
+			StateDetails: cspann.MakeSplittingDetails(10, 20),
 		}
 		suite.NoError(store.TryCreateEmptyPartition(suite.ctx, treeKey, testPartitionKey2, metadata))
 		vectors := vector.MakeSet(2)
@@ -368,20 +368,36 @@ func (suite *StoreTestSuite) TestSearchPartitions() {
 
 		suite.NoError(store.RunTransaction(suite.ctx, func(txn cspann.Txn) error {
 			searchSet := cspann.SearchSet{MaxResults: 2}
-			toSearch := []cspann.PartitionToSearch{{Key: testPartitionKey}, {Key: testPartitionKey2}}
+			toSearch := []cspann.PartitionToSearch{
+				{Key: testPartitionKey},
+				{Key: cspann.PartitionKey(99)}, // Partition does not exist.
+				{Key: testPartitionKey2},
+			}
 			err := txn.SearchPartitions(suite.ctx, treeKey, toSearch, vector.T{6, 1}, &searchSet)
 			suite.NoError(err)
+
+			// Validate partition info.
+			suite.Equal(cspann.SecondLevel, toSearch[0].Level)
+			suite.Equal(testPartition.Metadata().StateDetails, toSearch[0].StateDetails)
+			suite.Equal(3, toSearch[0].Count)
+
+			suite.Equal(cspann.InvalidLevel, toSearch[1].Level)
+			suite.Equal(cspann.PartitionStateDetails{}, toSearch[1].StateDetails)
+			suite.Equal(0, toSearch[1].Count)
+
+			suite.Equal(cspann.SecondLevel, toSearch[2].Level)
+			suite.Equal(metadata.StateDetails, toSearch[2].StateDetails)
+			suite.Equal(2, toSearch[2].Count)
+
+			// Validate search results.
 			result1 := cspann.SearchResult{
 				QuerySquaredDistance: 4.2, ErrorBound: 50.99, CentroidDistance: 7.21,
 				ParentPartitionKey: testPartitionKey2, ChildKey: partitionKey4, ValueBytes: valueBytes4}
 			result2 := cspann.SearchResult{
 				QuerySquaredDistance: 8, ErrorBound: 0, CentroidDistance: 0,
 				ParentPartitionKey: testPartitionKey, ChildKey: partitionKey3, ValueBytes: valueBytes3}
 			suite.Equal(cspann.SearchResults{result1, result2}, RoundResults(searchSet.PopResults(), 2))
-			suite.Equal(cspann.SecondLevel, toSearch[0].Level)
-			suite.Equal(3, toSearch[0].Count)
-			suite.Equal(cspann.SecondLevel, toSearch[1].Level)
-			suite.Equal(2, toSearch[1].Count)
+
 			return nil
 		}))
 	}

pkg/sql/vecindex/vecstore/store_txn.go

@@ -270,13 +270,21 @@ func (tx *Txn) SearchPartitions(
 		partition, err := tx.store.decodePartition(
 			treeKey, toSearch[i].Key, &tx.codec, &b.Results[i*2], &b.Results[i*2+1])
 		if err != nil {
-			return err
+			if errors.Is(err, cspann.ErrPartitionNotFound) {
+				// Partition not found, so return InvalidLevel, MissingState, and
+				// Count=0.
+				toSearch[i].Level = cspann.InvalidLevel
+				toSearch[i].StateDetails = cspann.PartitionStateDetails{}
+				toSearch[i].Count = 0
+			} else {
+				return err
+			}
+		} else {
+			toSearch[i].Level = partition.Level()
+			toSearch[i].StateDetails = partition.Metadata().StateDetails
+			toSearch[i].Count = partition.Search(
+				&tx.codec.workspace, toSearch[i].Key, queryVector, searchSet)
 		}
-
-		toSearch[i].Level = partition.Level()
-		toSearch[i].StateDetails = partition.Metadata().StateDetails
-		toSearch[i].Count = partition.Search(
-			&tx.codec.workspace, toSearch[i].Key, queryVector, searchSet)
 	}
 
 	return nil