vecstore: search all PKs for get full vector index

When we're changing primary keys, it's possible that the PrimaryKey()
isn't actually keyed by the keys embedded as suffixes in the vector
index. Since we need these indexes to be able to execute
vecstore.GetFullVectors(), build the GetFullVectorsFetchSpec using the
first primary key encoded index that has the keys that we have embedded
in our key.

This change addresses a bug discovered during construction of the schema
changer DML injection tests.

Release note: None

Author: mw5h

pkg/sql/vecindex/vecstore/store_txn.go

@@ -353,20 +353,38 @@ func InitGetFullVectorsFetchSpec(
 	spec.FamilyIDs = splitter.FamilyIDs()
 
 	// We need to decode the columns for the PK from the vector index. If the
-	// vector index is being mutated, there's a chance the primary key is changing.
-	// The vector index needs to point into the NEW primary key, so find that here.
+	// vector index is being mutated or the primary key is being changed, this
+	// can get a bit complicated. We need to find the PK for which the vector
+	// index was encoded, so that the vector index has the PK's key columns In most
+	// cases, this will just be the primary key, which is the first index returned
+	// by AllIndexes().
 	var primaryKey catalog.Index
-	primaryKey = tableDesc.GetPrimaryIndex()
-	if indexDesc.IsMutation() {
-		for _, idx := range tableDesc.NonPrimaryIndexes() {
-			if idx.GetEncodingType() == catenumpb.PrimaryIndexEncoding && !idx.IsTemporaryIndexForBackfill() {
-				primaryKey = idx
-				break
-			}
+	haveCols := indexDesc.CollectKeyColumnIDs()
+	haveCols.UnionWith(indexDesc.CollectKeySuffixColumnIDs())
+	for _, idx := range tableDesc.AllIndexes() {
+		if idx.GetEncodingType() != catenumpb.PrimaryIndexEncoding {
+			continue
+		}
+
+		if idx.IsTemporaryIndexForBackfill() {
+			continue
+		}
+
+		if !idx.CollectKeyColumnIDs().SubsetOf(haveCols) {
+			continue
 		}
+
+		primaryKey = idx
+		break
 	}
+	if primaryKey == nil {
+		return errors.AssertionFailedf("No primary key of '%s' has column '%s' and is keyed by %v",
+			tableDesc.GetName(), vectorCol.GetName(), haveCols)
+	}
+
 	keyCols := primaryKey.CollectKeyColumnIDs().Ordered()
-	return rowenc.InitIndexFetchSpec(&spec.ExtractPKFetchSpec, evalCtx.Codec, tableDesc, indexDesc, keyCols)
+	return rowenc.InitIndexFetchSpec(&spec.ExtractPKFetchSpec, evalCtx.Codec, tableDesc,
+		indexDesc, keyCols)
 }
 
 // PKDecoder is used to extract the primary key from a vector index.