cspann: do not duplicate vectors to avoid an empty partition

Previously, the split fixup would sometimes add an interior centroid
vector to two different parent partitions. In the past, this was
necessary in order to avoid creating an empty partition that could
unbalance the K-means tree. But this is no longer necessary; empty
partitions are now allowed. So this commit removes that code.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

pkg/sql/vecindex/cspann/fixup_split.go

@@ -1053,15 +1053,6 @@ func (fw *fixupWorker) copyToSplitSubPartitions(
 		return err
 	}
 
-	if sourcePartition.Level() != LeafLevel && vectors.Count == 1 {
-		// This should have been a merge, not a split, but we're too far into the
-		// split operation to back out now, so avoid an empty non-root partition by
-		// duplicating the last remaining vector in both partitions.
-		rightVectors = leftVectors
-		rightChildKeys = leftChildKeys
-		rightValueBytes = leftValueBytes
-	}
-
 	rightPartitionKey := sourceState.Target2
 	added, err = fw.addToPartition(ctx,
 		rightPartitionKey, rightVectors, rightChildKeys, rightValueBytes, rightMetadata)

pkg/sql/vecindex/cspann/index_test.go

@@ -17,6 +17,7 @@ import (
 	"sync"
 	"testing"
 	"time"
+	"unicode/utf8"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann/commontest"
@@ -162,7 +163,30 @@ func (s *testState) LoadIndex(d *datadriven.TestData) string {
 	s.makeNewIndex(d)
 
 	lines := strings.Split(d.Input, "\n")
-	s.loadIndexFromFormat(s.TreeKey, lines, 0)
+
+	// Determine the root level by determining max indent.
+	maxIndent := 0
+	for _, line := range lines {
+		// Lines ending with the "│" rune don't increase the level
+		if strings.HasSuffix(line, "│") {
+			continue
+		}
+
+		// Leaf vector lines don't increase the level.
+		idx := strings.Index(line, "• ")
+		if idx != -1 {
+			// Note that the '•' rune is 3 UTF-8 bytes.
+			firstChar := line[idx+4]
+			if firstChar < '0' || firstChar > '9' {
+				// This is a leaf vector.
+				continue
+			}
+		}
+
+		maxIndent = max(maxIndent, getIndexFormatIndent(line))
+	}
+
+	s.loadIndexFromFormat(s.TreeKey, lines, cspann.Level(maxIndent+1))
 
 	return fmt.Sprintf("Loaded %d vectors.\n", len(s.MemStore.GetAllVectors()))
 }
@@ -848,8 +872,8 @@ func (s *testState) parseKeyAndVector(line string) (cspann.KeyBytes, vector.T) {
 }
 
 func (s *testState) loadIndexFromFormat(
-	treeKey cspann.TreeKey, lines []string, indent int,
-) (remaining []string, level cspann.Level, centroid vector.T, childKey cspann.ChildKey) {
+	treeKey cspann.TreeKey, lines []string, level cspann.Level,
+) (remaining []string, centroid vector.T, childKey cspann.ChildKey) {
 	// Ensure line contains "• ", note that the '•' rune is 3 UTF-8 bytes.
 	idx := strings.Index(lines[0], "• ")
 	require.NotEqual(s.T, -1, idx)
@@ -863,7 +887,7 @@ func (s *testState) loadIndexFromFormat(
 		vec := s.parseVector(line)
 		s.MemStore.InsertVector(keyBytes, vec)
 		randomized := s.Index.TransformVector(vec, make(vector.T, len(vec)))
-		return lines[1:], cspann.LeafLevel, randomized, cspann.ChildKey{KeyBytes: keyBytes}
+		return lines[1:], randomized, cspann.ChildKey{KeyBytes: keyBytes}
 	}
 
 	// This is an interior partition.
@@ -885,33 +909,33 @@ func (s *testState) loadIndexFromFormat(
 	centroid = s.parseVector(line)
 	centroid = s.Index.TransformVector(centroid, make(vector.T, len(centroid)))
 
+	// Parse any children.
+	parentIndent := getIndexFormatIndent(lines[0])
 	lines = lines[1:]
 
-	childLevel := cspann.LeafLevel
 	childVectors := vector.MakeSet(len(centroid))
 	childKeys := []cspann.ChildKey(nil)
 
-	if len(lines) > 0 && strings.HasSuffix(lines[0], "│") {
-		// There are children, so loop over them.
-		childIndent := len(lines[0]) - len("│")
-		for len(lines) > 0 && len(lines[0]) > childIndent {
-			remainder := lines[0][childIndent:]
-			if remainder == "│" {
-				// Skip line.
-				lines = lines[1:]
-				continue
-			} else if strings.HasPrefix(remainder, "├") || strings.HasPrefix(remainder, "└") {
-				var childVector vector.T
-				lines, childLevel, childVector, childKey = s.loadIndexFromFormat(treeKey, lines, childIndent)
-				childVectors.Add(childVector)
-				childKeys = append(childKeys, childKey)
-			}
+	for len(lines) > 0 {
+		childIndent := getIndexFormatIndent(lines[0])
+		if childIndent <= parentIndent {
+			break
+		}
+		if strings.HasSuffix(lines[0], "│") {
+			// Skip line.
+			lines = lines[1:]
+			continue
+		} else {
+			var childVector vector.T
+			lines, childVector, childKey = s.loadIndexFromFormat(treeKey, lines, level-1)
+			childVectors.Add(childVector)
+			childKeys = append(childKeys, childKey)
 		}
 	}
 
 	// Always create partition in Ready state so that adds are allowed.
 	metadata := cspann.PartitionMetadata{
-		Level:    childLevel,
+		Level:    level,
 		Centroid: centroid,
 	}
 	metadata.StateDetails.MakeReady()
@@ -934,7 +958,7 @@ func (s *testState) loadIndexFromFormat(
 		require.NoError(s.T, err)
 	}
 
-	return lines, childLevel + 1, centroid, cspann.ChildKey{PartitionKey: partitionKey}
+	return lines, centroid, cspann.ChildKey{PartitionKey: partitionKey}
 }
 
 // parsePartitionStateDetails parses a partition state details string in this
@@ -1267,3 +1291,22 @@ func validateIndex(
 	})
 	require.Empty(t, missingKeys)
 }
+
+// getIndexFormatIndent returns the number of indentation levels in the given
+// formatted line. For example, for this snippet:
+//
+// • 1 (0, 0)
+// │
+// ├───• 2 (-9, 18)
+//
+// Line 1 would have indent of 0 and lines 2 and 3 would have indent of 1.
+func getIndexFormatIndent(line string) int {
+	// Compute level for line containing the "•" rune.
+	idx := strings.Index(line, "• ")
+	if idx != -1 {
+		return utf8.RuneCountInString(line[:idx]) / 4
+	}
+
+	// Otherwise, this must be a line ending with the "│" rune.
+	return utf8.RuneCountInString(line)/4 + 1
+}

pkg/sql/vecindex/cspann/testdata/split-non-root-step.ddt

@@ -558,43 +558,12 @@ force-split partition-key=2 parent-partition-key=1 steps=7
 │
 └───• 5 (5, 6) [Updating:2]
 
-# Next steps should duplicate the last remaining vector in sub-partition #5
-# rather than leave it empty.
-force-split partition-key=2 parent-partition-key=1 steps=2
-----
-• 1 (3, 6)
-│
-├───• 2 (3, 6) [DrainingForSplit:4,5]
-├───• 4 (5, 6) [Updating:2]
-│   │
-│   └───• 3 (5, 6)
-│       │
-│       ├───• vec2 (3, 5)
-│       ├───• vec3 (6, 8)
-│       └───• vec4 (6, 5)
-│
-└───• 5 (5, 6) [Updating:2]
-    │
-    └───• 3 (5, 6)
-        │
-        ├───• vec2 (3, 5)
-        ├───• vec3 (6, 8)
-        └───• vec4 (6, 5)
-
-# Finish split. Partition #3 should be referenced by both partition #4 and
-# partition #5.
+# Target partition #5 should be left empty.
 force-split partition-key=2 parent-partition-key=1
 ----
 • 1 (3, 6)
 │
 ├───• 5 (5, 6)
-│   │
-│   └───• 3 (5, 6)
-│       │
-│       ├───• vec2 (3, 5)
-│       ├───• vec3 (6, 8)
-│       └───• vec4 (6, 5)
-│
 └───• 4 (5, 6)
     │
     └───• 3 (5, 6)

pkg/sql/vecindex/cspann/testdata/split-root-step.ddt

@@ -158,19 +158,12 @@ force-split partition-key=1 root=3 steps=5
     ├───• vec1 (1, 2)
     └───• vec2 (4, 3)
 
-# Next steps should duplicate the last remaining vector in partition #1 and also
-# add it to partition #4. This will prevent partition #4 from being empty, which
-# would be a violation of the constraint that the K-means tree must be balanced.
+# Partition #4 should be left empty.
 force-split partition-key=1 root=4 steps=1
 ----
 • 4 (2.5, 2.5) [Updating:1]
-│
-└───• 2 (2.5, 2.5)
-    │
-    ├───• vec1 (1, 2)
-    └───• vec2 (4, 3)
 
-# Finish the split, with partition #2 a child of both sub-partitions.
+# Finish the split.
 force-split partition-key=1
 ----
 • 1 (6.8, 4.2)
@@ -183,11 +176,6 @@ force-split partition-key=1
 │       └───• vec2 (4, 3)
 │
 └───• 4 (2.5, 2.5)
-    │
-    └───• 2 (2.5, 2.5)
-        │
-        ├───• vec1 (1, 2)
-        └───• vec2 (4, 3)
 
 # ----------------------------------------------------------------------
 # Search the tree when the root is in splitting states.

pkg/sql/vecindex/cspann/testdata/split.ddt

@@ -821,3 +821,47 @@ force-split partition-key=3 parent-partition-key=1
     │
     ├───• vec5 (5, 3)
     └───• vec4 (6, 2)
+
+# ----------------------------------------------------------------------
+# Empty split target partition. This regresses a bug caused by code that
+# was trying to insert partition 4 into both partition 5 and 6, in order
+# to avoid creating an empty partition. Instead, partition 4 wasn't
+# inserted into either partition, becoming a "lost update". The fix was
+# to remove the code entirely, as it's no longer a problem for empty
+# partitions to exist.
+# ----------------------------------------------------------------------
+
+load-index min-partition-size=1 max-partition-size=4 beam-size=2
+• 1 (0, 0)
+│
+├───• 2 (-9, 18)
+│   │
+│   ├───• 3 (-9, 18) [Splitting:5,6]
+│   │   │
+│   │   └───• 4 (-9, 18)
+│   │       │
+│   │       ├───• vec39 (-4, 13)
+│   │       └───• vec60 (-14, 23)
+│   │
+│   ├───• 5 (-9, 18) [Updating:3]
+│   └───• 6 (4, 3) [Updating:3]
+│
+└───• 7 (5, 8)
+----
+Loaded 2 vectors.
+
+force-split partition-key=3 parent-partition-key=2
+----
+• 1 (0, 0)
+│
+├───• 2 (-9, 18)
+│   │
+│   ├───• 6 (4, 3)
+│   │   │
+│   │   └───• 4 (-9, 18)
+│   │       │
+│   │       ├───• vec39 (-4, 13)
+│   │       └───• vec60 (-14, 23)
+│   │
+│   └───• 5 (-9, 18)
+└───• 7 (5, 8)