cspann: do not overwrite root level when searching targets

In the case where the root partition is in a non-Ready state, we need
to search its target partitions. Previously, we were overwriting the
root partition's level, which should be the same as the root partition's
level. However, if there's a bug, it can be different, and overriding
the root level causes cascading problems downstream, including a panic
that can take down the server. This commit adds code to check the level
and assert that it's equal to the root partition level rather than
silently overwriting it.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

pkg/sql/vecindex/cspann/index.go

@@ -1107,17 +1107,6 @@ func (vi *Index) Format(
 	return buf.String(), nil
 }
 
-// ensureSliceCap returns a slice of length = 0, of the given capacity and
-// generic type. If the existing slice has enough capacity, that slice is
-// returned after setting its length to zero. Otherwise, a new, larger slice is
-// allocated.
-func ensureSliceCap[T any](s []T, c int) []T {
-	if cap(s) < c {
-		return make([]T, 0, c)
-	}
-	return s[:0]
-}
-
 // ensureSliceLen returns a slice of the given length and generic type. If the
 // existing slice has enough capacity, that slice is returned after adjusting
 // its length. Otherwise, a new, larger slice is allocated.

pkg/sql/vecindex/cspann/searcher.go

@@ -93,7 +93,7 @@ func (s *searcher) Next(ctx context.Context) (ok bool, err error) {
 			// Next must have been called by an insert operation. This code path
 			// should only be hit when the insert needs to go into the root partition.
 			if root.Level() != s.idxCtx.level-1 {
-				panic(errors.AssertionFailedf("caller passed invalid level %d", s.idxCtx.level))
+				return false, errors.AssertionFailedf("caller passed invalid level %d", s.idxCtx.level)
 			}
 			s.searchSet.Add(&SearchResult{
 				ChildKey: ChildKey{PartitionKey: RootKey},
@@ -479,11 +479,15 @@ func (s *levelSearcher) searchChildPartitions(
 		// In the DrainingForSplit state, the target partitions are still at
 		// the same level as the root partition, so merge their contents into
 		// the search set (which will remove any duplicates).
-		s.idxCtx.tempToSearch = ensureSliceCap(s.idxCtx.tempToSearch, 2)
-		level, err = s.searchChildPartitions(ctx, s.tempResults[:2])
+		targetLevel, err := s.searchChildPartitions(ctx, s.tempResults[:2])
 		if err != nil {
 			return InvalidLevel, err
 		}
+		if targetLevel != level {
+			return InvalidLevel, errors.AssertionFailedf(
+				"target partitions cannot have level %d when DrainingForSplit root has level %d",
+				targetLevel, level)
+		}
 	} else if rootState.State == AddingLevelState {
 		// In the AddingLevel state, the target partitions should be treated as
 		// children of the root partition. Add their partition keys to the search

pkg/sql/vecindex/cspann/testdata/search-for-insert.ddt

@@ -108,3 +108,74 @@ format-tree
     ├───• vec2 (7, 4)
     ├───• vec4 (5, 5)
     └───• vec5 (8, 11)
+
+# ----------------------------------------------------------------------
+# Search for insert when tree is empty in various states.
+# ----------------------------------------------------------------------
+
+# Splitting state.
+load-index min-partition-size=1 max-partition-size=3 beam-size=2
+• 1 (0, 0) [Splitting:2,3]
+----
+Loaded 0 vectors.
+
+search-for-insert
+(1, 1)
+----
+partition 1, centroid=(0, 0), sqdist=0
+
+# DrainingForSplit state.
+load-index min-partition-size=1 max-partition-size=3 beam-size=2
+• 1 (0, 0)
+----
+Loaded 0 vectors.
+
+force-split partition-key=1 steps=4
+----
+• 1 (0, 0) [DrainingForSplit:2,3]
+
+search-for-insert
+(1, 1)
+----
+partition 2, centroid=(0, 0), sqdist=2
+
+# AddingLevel state before sub-partitions have been added.
+load-index min-partition-size=1 max-partition-size=3 beam-size=2
+• 1 (0, 0) [AddingLevel:2,3]
+----
+Loaded 0 vectors.
+
+search-for-insert
+(1, 1)
+----
+partition 1, centroid=(0, 0), sqdist=0
+
+# AddingLevel state after sub-partitions have been added.
+load-index min-partition-size=1 max-partition-size=3 beam-size=2
+• 1 (0, 0) [AddingLevel:2,3]
+│
+├───• 2 (1, 1)
+│
+└───• 3 (2, 2)
+----
+Loaded 0 vectors.
+
+search-for-insert
+(1, 1)
+----
+partition 2, centroid=(1, 1), sqdist=0
+
+# DeletingForSplit state.
+load-index min-partition-size=1 max-partition-size=3 beam-size=2
+• 1 (0, 0)
+│
+├───• 2 (1, 1) [DeletingForSplit:2,3]
+│
+└───• 3 (2, 2)
+----
+Loaded 0 vectors.
+
+search-for-insert
+(1, 1)
+----
+partition 3, centroid=(2, 2), sqdist=2