Merge #147881

147881: cspann: simplify partition assignment code r=drewkimball a=andy-kimball

Previously, partition assignment worked by assigning vectors by their offsets in the vector set. AssignPartitions returned leftOffsets and rightOffsets sets containing the offsets for each partition. The offsets are awkward to handle and there's a TODO to switch this to use an assignments slice instead.

This commit switchs AssignPartitions to fill out an assignments slice, where 0 indicates a vector has been assigned to the left partition, or 1 to the right. It also updates all callers to use assignments rather than offsets.

Epic: CRDB-42943
Release note: None

Co-authored-by: Andrew Kimball <[email protected]>

Author: craig[bot]

pkg/sql/vecindex/cspann/fixup_split.go

@@ -711,33 +711,33 @@ func (fw *fixupWorker) copyToSplitSubPartitions(
 	vectors vector.Set,
 	leftMetadata, rightMetadata PartitionMetadata,
 ) (err error) {
-	var leftOffsets, rightOffsets []uint64
+	var leftCount int
 	sourceState := sourcePartition.Metadata().StateDetails
 
 	defer func() {
 		err = errors.Wrapf(err,
 			"assigning %d vectors to left partition %d and %d vectors to right partition %d",
-			len(leftOffsets), sourceState.Target1, len(rightOffsets), sourceState.Target2)
+			leftCount, sourceState.Target1, vectors.Count-leftCount, sourceState.Target2)
 	}()
 
-	tempOffsets := fw.workspace.AllocUint64s(vectors.Count)
-	defer fw.workspace.FreeUint64s(tempOffsets)
+	tempAssignments := fw.workspace.AllocUint64s(vectors.Count)
+	defer fw.workspace.FreeUint64s(tempAssignments)
 
 	// Assign vectors to the partition with the nearest centroid.
 	kmeans := BalancedKmeans{Workspace: &fw.workspace, Rand: fw.rng}
-	leftOffsets, rightOffsets = kmeans.AssignPartitions(
-		vectors, leftMetadata.Centroid, rightMetadata.Centroid, tempOffsets)
+	leftCount = kmeans.AssignPartitions(
+		vectors, leftMetadata.Centroid, rightMetadata.Centroid, tempAssignments)
 
 	// Assign vectors and associated keys and values into contiguous left and right groupings.
 	childKeys := slices.Clone(sourcePartition.ChildKeys())
 	valueBytes := slices.Clone(sourcePartition.ValueBytes())
-	splitPartitionData(&fw.workspace, vectors, childKeys, valueBytes, leftOffsets, rightOffsets)
+	splitPartitionData(&fw.workspace, vectors, childKeys, valueBytes, tempAssignments)
 	leftVectors := vectors
-	rightVectors := leftVectors.SplitAt(len(leftOffsets))
-	leftChildKeys := childKeys[:len(leftOffsets)]
-	rightChildKeys := childKeys[len(leftOffsets):]
-	leftValueBytes := valueBytes[:len(leftOffsets)]
-	rightValueBytes := valueBytes[len(leftOffsets):]
+	rightVectors := leftVectors.SplitAt(leftCount)
+	leftChildKeys := childKeys[:leftCount]
+	rightChildKeys := childKeys[leftCount:]
+	leftValueBytes := valueBytes[:leftCount]
+	rightValueBytes := valueBytes[leftCount:]
 
 	// Add vectors to left and right sub-partitions. Note that this may not be
 	// transactional; if an error occurs, any vectors already added may not be
@@ -748,7 +748,7 @@ func (fw *fixupWorker) copyToSplitSubPartitions(
 		leftPartitionKey, leftVectors, leftChildKeys, leftValueBytes, leftMetadata)
 	if added {
 		log.VEventf(ctx, 2, "assigned %d vectors to left partition %d (level=%d, state=%s)",
-			len(leftOffsets), leftPartitionKey, leftMetadata.Level, leftMetadata.StateDetails.String())
+			leftCount, leftPartitionKey, leftMetadata.Level, leftMetadata.StateDetails.String())
 	}
 	if err != nil {
 		return err
@@ -768,7 +768,7 @@ func (fw *fixupWorker) copyToSplitSubPartitions(
 		rightPartitionKey, rightVectors, rightChildKeys, rightValueBytes, rightMetadata)
 	if added {
 		log.VEventf(ctx, 2, "assigned %d vectors to right partition %d (level=%d, state=%s)",
-			len(rightOffsets), rightPartitionKey,
+			vectors.Count-leftCount, rightPartitionKey,
 			rightMetadata.Level, rightMetadata.StateDetails.String())
 	}
 	if err != nil {
@@ -798,62 +798,55 @@ func suppressRaceErrors(err error) (PartitionMetadata, error) {
 }
 
 // splitPartitionData groups the provided partition data according to the left
-// and right offsets. All data referenced by left offsets will be moved to the
-// left of each set or slice. All data referenced by right offsets will be moved
-// to the right. The internal ordering of elements on each side is not defined.
-//
-// TODO(andyk): Passing in left and right offsets makes this overly complex. It
-// would be better to pass an assignments slice of the same length as the
-// partition data, where 0=left and 1=right.
+// and right offsets. The assignments slice specifies which partition the data
+// will be moved into: 0 for left and 1 for right. The internal ordering of
+// elements on each side is not defined.
 func splitPartitionData(
 	w *workspace.T,
 	vectors vector.Set,
 	childKeys []ChildKey,
 	valueBytes []ValueBytes,
-	leftOffsets, rightOffsets []uint64,
+	assignments []uint64,
 ) {
 	tempVector := w.AllocFloats(vectors.Dims)
 	defer w.FreeFloats(tempVector)
 
+	// Use a two-pointer approach to partition the data. left points to the next
+	// position where a left element should go. right points to the next position
+	// where a right element should go (from the end).
 	left := 0
-	right := 0
+	right := len(assignments) - 1
+
 	for {
-		// Find a misplaced "right" element from the left side.
-		var leftOffset int
-		for {
-			if left >= len(leftOffsets) {
-				return
-			}
-			leftOffset = int(leftOffsets[left])
+		// Find a misplaced element on the left side (should be 0 but is 1).
+		for left < right && assignments[left] == 0 {
 			left++
-			if leftOffset >= len(leftOffsets) {
-				break
-			}
 		}
 
-		// There must be a misplaced "left" element from the right side.
-		var rightOffset int
-		for {
-			rightOffset = int(rightOffsets[right])
-			right++
-			if rightOffset < len(leftOffsets) {
-				break
-			}
+		// Find a misplaced element on the right side (should be 1 but is 0).
+		for left < right && assignments[right] == 1 {
+			right--
+		}
+
+		if left >= right {
+			// No more misplaced elements, so break.
+			break
 		}
 
-		// Swap the two elements.
-		rightToLeft := vectors.At(leftOffset)
-		leftToRight := vectors.At(rightOffset)
-		copy(tempVector, rightToLeft)
-		copy(rightToLeft, leftToRight)
-		copy(leftToRight, tempVector)
+		// Swap vectors.
+		leftVector := vectors.At(left)
+		rightVector := vectors.At(right)
+		copy(tempVector, leftVector)
+		copy(leftVector, rightVector)
+		copy(rightVector, tempVector)
+
+		// Swap child keys.
+		childKeys[left], childKeys[right] = childKeys[right], childKeys[left]
 
-		tempChildKey := childKeys[leftOffset]
-		childKeys[leftOffset] = childKeys[rightOffset]
-		childKeys[rightOffset] = tempChildKey
+		// Swap value bytes.
+		valueBytes[left], valueBytes[right] = valueBytes[right], valueBytes[left]
 
-		tempValueBytes := valueBytes[leftOffset]
-		valueBytes[leftOffset] = valueBytes[rightOffset]
-		valueBytes[rightOffset] = tempValueBytes
+		left++
+		right--
 	}
 }

pkg/sql/vecindex/cspann/fixup_split_test.go

@@ -45,78 +45,54 @@ func TestSplitPartitionData(t *testing.T) {
 	}
 
 	testCases := []struct {
-		desc          string
-		leftOffsets   []uint64
-		rightOffsets  []uint64
-		expectedLeft  []uint64
-		expectedRight []uint64
+		desc        string
+		assignments []uint64
+		expected    []int
 	}{
 		{
-			desc:         "no reordering",
-			leftOffsets:  []uint64{0, 1, 2, 3},
-			rightOffsets: []uint64{4, 5, 6},
+			desc:        "no reordering",
+			assignments: []uint64{0, 0, 0, 0, 1, 1, 1},
+			expected:    []int{0, 1, 2, 3, 4, 5, 6},
 		},
 		{
-			desc:         "only one on left",
-			leftOffsets:  []uint64{1},
-			rightOffsets: []uint64{0, 2, 3, 4, 5, 6},
+			desc:        "only one on left",
+			assignments: []uint64{1, 0, 1, 1, 1, 1, 1},
+			expected:    []int{1, 0, 2, 3, 4, 5, 6},
 		},
 		{
-			desc:         "only one on right",
-			leftOffsets:  []uint64{0, 1, 2, 4, 5, 6},
-			rightOffsets: []uint64{3},
+			desc:        "only one on right",
+			assignments: []uint64{0, 0, 0, 1, 0, 0, 0},
+			expected:    []int{0, 1, 2, 6, 4, 5, 3},
 		},
 		{
-			desc:         "interleaved",
-			leftOffsets:  []uint64{0, 2, 4, 6},
-			rightOffsets: []uint64{1, 3, 5},
+			desc:        "interleaved",
+			assignments: []uint64{0, 1, 0, 1, 0, 1, 0},
+			expected:    []int{0, 6, 2, 4, 3, 5, 1},
 		},
 		{
-			desc:         "another interleaved",
-			leftOffsets:  []uint64{1, 4, 5},
-			rightOffsets: []uint64{0, 2, 3, 6},
+			desc:        "another interleaved",
+			assignments: []uint64{1, 0, 1, 1, 0, 0, 1},
+			expected:    []int{5, 1, 4, 3, 2, 0, 6},
 		},
 		{
-			desc:         "reversed",
-			leftOffsets:  []uint64{4, 5, 6},
-			rightOffsets: []uint64{0, 1, 2, 3},
-		},
-		{
-			desc:         "out of order",
-			leftOffsets:  []uint64{5, 4, 6},
-			rightOffsets: []uint64{3, 0, 1, 2},
+			desc:        "reversed",
+			assignments: []uint64{1, 1, 1, 1, 0, 0, 0},
+			expected:    []int{6, 5, 4, 3, 2, 1, 0},
 		},
 	}
 
-	findKey := func(allKeys []ChildKey, toFind ChildKey) int {
-		for i, key := range allKeys {
-			if key.Equal(toFind) {
-				return i
-			}
-		}
-		return -1
-	}
-
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			tempVectors := vectors.Clone()
 			tempChildKeys := slices.Clone(childKeys)
 			tempValueBytes := slices.Clone(valueBytes)
-			splitPartitionData(&workspace,
-				tempVectors, tempChildKeys, tempValueBytes, tc.leftOffsets, tc.rightOffsets)
+			splitPartitionData(&workspace, tempVectors, tempChildKeys, tempValueBytes, tc.assignments)
 
 			// Ensure that partition data is on the correct side.
-			for originalOffset := range childKeys {
-				newOffset := findKey(tempChildKeys, childKeys[originalOffset])
-
-				if newOffset < len(tc.leftOffsets) {
-					require.Contains(t, tc.leftOffsets, uint64(originalOffset))
-				} else {
-					require.Contains(t, tc.rightOffsets, uint64(originalOffset))
-				}
-
-				require.Equal(t, tempVectors.At(newOffset), vectors.At(originalOffset))
-				require.Equal(t, tempValueBytes[newOffset], valueBytes[originalOffset])
+			for i := range tc.expected {
+				require.Equal(t, tempVectors.At(tc.expected[i]), vectors.At(i))
+				require.Equal(t, tempChildKeys[tc.expected[i]], childKeys[i])
+				require.Equal(t, tempValueBytes[tc.expected[i]], valueBytes[i])
 			}
 		})
 	}