Merge #141399

141399: tabledesc: implement vector index retrieval in the catalog r=mw5h a=mw5h

Add a VectorIndexes() method to the catalog's index cache. This patch also addes ForEachVectorIndex() and FindVectorIndex(), primarily so that we can hook up the index_test to ensure that VectorIndexes() is working properly.

Epic: CRDB-42943

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

Author: craig[bot]

pkg/sql/catalog/descriptor.go

@@ -461,6 +461,11 @@ type TableDescriptor interface {
 	// produced by AllIndexes and removing indexes with empty expressions.
 	PartialIndexes() []Index
 
+	// VectorIndexes returns a slice of all vector indexes in the underlying
+	// proto, in their canonical order. This is equivalent to taking the slice
+	// produced by AllIndexes and removing indexes which are not vector indexes.
+	VectorIndexes() []Index
+
 	// PublicNonPrimaryIndexes returns a slice of all active secondary indexes,
 	// in their canonical order. This is equivalent to the Indexes array in the
 	// proto.

pkg/sql/catalog/table_elements.go

@@ -780,6 +780,11 @@ func ForEachPartialIndex(desc TableDescriptor, f func(idx Index) error) error {
 	return forEachIndex(desc.PartialIndexes(), f)
 }
 
+// ForEachVectorIndex is like ForEachIndex over VectorIndexes().
+func ForEachVectorIndex(desc TableDescriptor, f func(idx Index) error) error {
+	return forEachIndex(desc.VectorIndexes(), f)
+}
+
 // ForEachNonPrimaryIndex is like ForEachIndex over
 // NonPrimaryIndexes().
 func ForEachNonPrimaryIndex(desc TableDescriptor, f func(idx Index) error) error {
@@ -846,12 +851,18 @@ func FindNonDropIndex(desc TableDescriptor, test func(idx Index) bool) Index {
 	return findIndex(desc.NonDropIndexes(), test)
 }
 
-// FindPartialIndex returns the first index in PartialIndex() for which test
+// FindPartialIndex returns the first index in PartialIndexes() for which test
 // returns true.
 func FindPartialIndex(desc TableDescriptor, test func(idx Index) bool) Index {
 	return findIndex(desc.PartialIndexes(), test)
 }
 
+// FindVectorIndex returns the first index in VectorIndexes() for which test
+// returns true.
+func FindVectorIndex(desc TableDescriptor, test func(idx Index) bool) Index {
+	return findIndex(desc.VectorIndexes(), test)
+}
+
 // FindPublicNonPrimaryIndex returns the first index in PublicNonPrimaryIndex()
 // for which test returns true.
 func FindPublicNonPrimaryIndex(desc TableDescriptor, test func(idx Index) bool) Index {

pkg/sql/catalog/tabledesc/index.go

@@ -664,6 +664,7 @@ type indexCache struct {
 	deletableNonPrimary  []catalog.Index
 	deleteOnlyNonPrimary []catalog.Index
 	partial              []catalog.Index
+	vector               []catalog.Index
 }
 
 // newIndexCache returns a fresh fully-populated indexCache struct for the
@@ -713,13 +714,13 @@ func newIndexCache(desc *descpb.TableDescriptor, mutations *mutationCache) *inde
 		if !idx.Dropped() && (!idx.Primary() || desc.IsPhysicalTable()) {
 			lazyAllocAppendIndex(&c.nonDrop, idx, len(c.all))
 		}
-		// TODO(ssd): We exclude backfilling indexes from
-		// IsPartial() for the unprincipled reason of not
-		// wanting to modify all of the code that assumes
-		// these are always at least delete-only.
+		// Include only deletable indexes.
 		if idx.IsPartial() && !idx.Backfilling() {
 			lazyAllocAppendIndex(&c.partial, idx, len(c.all))
 		}
+		if idx.GetType() == idxtype.VECTOR && !idx.Backfilling() {
+			lazyAllocAppendIndex(&c.vector, idx, len(c.all))
+		}
 	}
 	return &c
 }

pkg/sql/catalog/tabledesc/index_test.go

@@ -180,6 +180,12 @@ func TestIndexInterface(t *testing.T) {
 			catalog.ForEachPartialIndex,
 			catalog.FindPartialIndex,
 		},
+		{"VectorIndex",
+			[]string{"s7"},
+			catalog.TableDescriptor.VectorIndexes,
+			catalog.ForEachVectorIndex,
+			catalog.FindVectorIndex,
+		},
 		{
 			"PublicNonPrimaryIndex",
 			indexNames[1:],

pkg/sql/catalog/tabledesc/table_desc.go

@@ -352,11 +352,24 @@ func (desc *wrapper) NonDropIndexes() []catalog.Index {
 
 // PartialIndexes returns a slice of all partial indexes in the underlying
 // proto, in their canonical order. This is equivalent to taking the slice
-// produced by AllIndexes and filtering indexes with non-empty expressions.
+// produced by DeletableNonPrimaryIndexes and filtering indexes with non-empty
+// expressions.
+//
+// Backfilling indexes are excluded.
 func (desc *wrapper) PartialIndexes() []catalog.Index {
 	return desc.getExistingOrNewIndexCache().partial
 }
 
+// VectorIndexes returns a slice of all vector indexes in the underlying
+// proto, in their canonical order. This is equivalent to taking the slice
+// produced by DeletableNonPrimaryIndexes and filtering indexes that are not
+// vector indexes.
+//
+// Backfilling indexes are excluded.
+func (desc *wrapper) VectorIndexes() []catalog.Index {
+	return desc.getExistingOrNewIndexCache().vector
+}
+
 // NonPrimaryIndexes returns a slice of all non-primary indexes, in
 // their canonical order. This is equivalent to taking the slice
 // produced by AllIndexes and removing the primary index.