Merge #143517 #143965

143517: split: add SplitStatistics function to Decider r=angles-n-daemons a=angles-n-daemons

split: add SplitStatistics function to Decider

So that external callers can be aware of the internal statistics on the load based splitters, this PR adds a fairly straightforward "SplitStatistics" function to the replica split deciders.

Fixes: #138760
Epic: CRDB-43150

Release note: None

143965: cspann: fix partition data sorting bug r=drewkimball a=andy-kimball

When sorting partition data into left and right groupings, vectors can be sorted in a different order than associated child keys and value bytes. This commit updates the logic to operate on all the partition data at once, not just the vectors.

Epic: CRDB-42943

Release note: None

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

Author: 3 people

pkg/kv/kvserver/split/decider.go

@@ -28,6 +28,11 @@ type PopularKey struct {
 	Frequency float64
 }
 
+type SplitStatistics struct {
+	AccessDirection float64
+	PopularKey      PopularKey
+}
+
 type LoadBasedSplitter interface {
 	redact.SafeFormatter
 	// Record informs the LoadBasedSplitter about where the span lies with regard
@@ -359,6 +364,21 @@ func (d *Decider) MaybeSplitKey(ctx context.Context, now time.Time) roachpb.Key
 	return key
 }
 
+// SplitStatistics gets the split stats of the current replica if load-based
+// splitting has been engaged.
+func (d *Decider) SplitStatistics() *SplitStatistics {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	if d.mu.splitFinder != nil {
+		return &SplitStatistics{
+			AccessDirection: d.mu.splitFinder.AccessDirection(),
+			PopularKey:      d.mu.splitFinder.PopularKey(),
+		}
+	}
+	return nil
+}
+
 // Reset deactivates any current attempt at determining a split key. The method
 // also discards any historical stat tracking information.
 func (d *Decider) Reset(now time.Time) {

pkg/kv/kvserver/split/decider_test.go

@@ -374,6 +374,110 @@ func TestMaxStatTracker(t *testing.T) {
 	require.Equal(t, 1, mt.curIdx)
 }
 
+func TestSplitStatisticsGeneral(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	for _, test := range []struct {
+		name        string
+		useWeighted bool
+		expected    *SplitStatistics
+	}{
+		{"unweighted", false, &SplitStatistics{
+			AccessDirection: 0.4945791444904396,
+			PopularKey: PopularKey{
+				Key:       keys.SystemSQLCodec.TablePrefix(uint32(52)),
+				Frequency: 0.05,
+			},
+		}},
+		{"weighted", true, &SplitStatistics{
+			AccessDirection: 0.3885786802030457,
+			PopularKey: PopularKey{
+				Key:       keys.SystemSQLCodec.TablePrefix(uint32(111)),
+				Frequency: 0.05,
+			},
+		}},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			rand := rand.New(rand.NewPCG(11, 11))
+			timeStart := 1000
+
+			var decider Decider
+			loadSplitConfig := testLoadSplitConfig{
+				randSource:    rand,
+				useWeighted:   test.useWeighted,
+				statRetention: time.Second,
+				statThreshold: 1,
+			}
+
+			Init(&decider, &loadSplitConfig, &LoadSplitterMetrics{
+				PopularKeyCount: metric.NewCounter(metric.Metadata{}),
+				NoSplitKeyCount: metric.NewCounter(metric.Metadata{}),
+			}, SplitCPU)
+
+			for i := 1; i <= 1000; i++ {
+				k := i
+				if i > 500 {
+					k = 500 - i
+				}
+				decider.Record(context.Background(), ms(timeStart+i*50), ld(1), func() roachpb.Span {
+					return roachpb.Span{Key: keys.SystemSQLCodec.TablePrefix(uint32(k))}
+				})
+			}
+
+			assert.Equal(t, decider.SplitStatistics(), test.expected)
+		})
+	}
+}
+
+func TestSplitStatisticsPopularKey(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	for _, test := range []struct {
+		name        string
+		useWeighted bool
+		expected    *SplitStatistics
+	}{
+		{"unweighted", false, &SplitStatistics{
+			AccessDirection: 1,
+			PopularKey: PopularKey{
+				Key:       keys.SystemSQLCodec.TablePrefix(uint32(100)),
+				Frequency: 1,
+			},
+		}},
+		{"weighted", true, &SplitStatistics{
+			AccessDirection: 1,
+			PopularKey: PopularKey{
+				Key:       keys.SystemSQLCodec.TablePrefix(uint32(100)),
+				Frequency: 1,
+			},
+		}},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			rand := rand.New(rand.NewPCG(11, 11))
+			timeStart := 1000
+
+			var decider Decider
+			loadSplitConfig := testLoadSplitConfig{
+				randSource:    rand,
+				useWeighted:   test.useWeighted,
+				statRetention: time.Second,
+				statThreshold: 1,
+			}
+
+			Init(&decider, &loadSplitConfig, &LoadSplitterMetrics{
+				PopularKeyCount: metric.NewCounter(metric.Metadata{}),
+				NoSplitKeyCount: metric.NewCounter(metric.Metadata{}),
+			}, SplitCPU)
+
+			for i := 1; i <= 1000; i++ {
+				decider.Record(context.Background(), ms(timeStart+i*50), ld(1), func() roachpb.Span {
+					return roachpb.Span{Key: keys.SystemSQLCodec.TablePrefix(uint32(100))}
+				})
+			}
+
+			assert.Equal(t, decider.SplitStatistics(), test.expected)
+		})
+	}
+}
+
 func TestDeciderMetrics(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	rng := rand.New(rand.NewPCG(11, 11))

pkg/sql/vecindex/cspann/BUILD.bazel

@@ -53,6 +53,7 @@ go_test(
         "childkey_dedup_test.go",
         "cspannpb_test.go",
         "fixup_processor_test.go",
+        "fixup_split_test.go",
         "fixup_worker_test.go",
         "index_stats_test.go",
         "index_test.go",

pkg/sql/vecindex/cspann/fixup_split.go

@@ -139,7 +139,7 @@ func (fw *fixupWorker) splitPartition(
 	}
 
 	log.VEventf(ctx, 2, "splitting partition %d with %d vectors (parent=%d, state=%s)",
-		partitionKey, parentPartitionKey, partition.Count(), metadata.StateDetails.String())
+		partitionKey, partition.Count(), parentPartitionKey, metadata.StateDetails.String())
 
 	// Update partition's state to Splitting.
 	if metadata.StateDetails.State == ReadyState {
@@ -454,13 +454,13 @@ func (fw *fixupWorker) createSplitSubPartition(
 	partitionKey PartitionKey,
 	centroid vector.T,
 ) (targetMetadata PartitionMetadata, err error) {
-	const format = "creating split sub-partition %d, with parent %d"
+	const format = "creating split sub-partition %d (source=%d, parent=%d)"
 
 	defer func() {
-		err = errors.Wrapf(err, format, partitionKey, parentPartitionKey)
+		err = errors.Wrapf(err, format, partitionKey, sourcePartitionKey, parentPartitionKey)
 	}()
 
-	log.VEventf(ctx, 2, format, partitionKey, parentPartitionKey)
+	log.VEventf(ctx, 2, format, partitionKey, sourcePartitionKey, parentPartitionKey)
 
 	// Create an empty partition in the Updating state.
 	targetMetadata = PartitionMetadata{
@@ -548,15 +548,15 @@ func (fw *fixupWorker) addToParentPartition(
 func (fw *fixupWorker) deletePartition(
 	ctx context.Context, parentPartitionKey, partitionKey PartitionKey,
 ) (err error) {
-	const format = "deleting partition %d (parent=%d, state=%d)"
+	const format = "deleting partition %d (parent=%d, state=%s)"
 	var parentMetadata PartitionMetadata
 
 	defer func() {
 		err = errors.Wrapf(err, format,
-			partitionKey, parentPartitionKey, parentMetadata.StateDetails.State)
+			partitionKey, parentPartitionKey, parentMetadata.StateDetails.String())
 	}()
 
-	// Load parent partition to verify that it's in a state that allows inserts.
+	// Load parent partition to verify that it's in a state that allows removes.
 	var parentPartition *Partition
 	parentPartition, err = fw.getPartition(ctx, parentPartitionKey)
 	if err != nil {
@@ -566,7 +566,8 @@ func (fw *fixupWorker) deletePartition(
 		parentMetadata = *parentPartition.Metadata()
 	}
 
-	log.VEventf(ctx, 2, format, partitionKey, parentPartitionKey, parentMetadata.StateDetails.State)
+	log.VEventf(ctx, 2, format,
+		partitionKey, parentPartitionKey, parentMetadata.StateDetails.String())
 
 	if !parentMetadata.StateDetails.State.AllowAddOrRemove() {
 		// Child could not be removed from the parent because it doesn't exist or
@@ -640,21 +641,16 @@ func (fw *fixupWorker) copyToSplitSubPartitions(
 	leftOffsets, rightOffsets = kmeans.AssignPartitions(
 		vectors, leftMetadata.Centroid, rightMetadata.Centroid, tempOffsets)
 
-	// Sort vectors into contiguous left and right groupings.
-	sortVectors(&fw.workspace, vectors, leftOffsets, rightOffsets)
+	// 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)
 	leftVectors := vectors
 	rightVectors := leftVectors.SplitAt(len(leftOffsets))
-
-	childKeys := make([]ChildKey, vectors.Count)
-	valueBytes := make([]ValueBytes, vectors.Count)
-	leftChildKeys := copyByOffsets(
-		sourcePartition.ChildKeys(), childKeys[:len(leftOffsets)], leftOffsets)
-	rightChildKeys := copyByOffsets(
-		sourcePartition.ChildKeys(), childKeys[len(leftOffsets):], rightOffsets)
-	leftValueBytes := copyByOffsets(
-		sourcePartition.ValueBytes(), valueBytes[:len(leftOffsets)], leftOffsets)
-	rightValueBytes := copyByOffsets(
-		sourcePartition.ValueBytes(), valueBytes[len(leftOffsets):], rightOffsets)
+	leftChildKeys := childKeys[:len(leftOffsets)]
+	rightChildKeys := childKeys[len(leftOffsets):]
+	leftValueBytes := valueBytes[:len(leftOffsets)]
+	rightValueBytes := valueBytes[len(leftOffsets):]
 
 	log.VEventf(ctx, 2, format,
 		len(leftOffsets), sourceState.Target1, len(rightOffsets), sourceState.Target2)
@@ -740,63 +736,63 @@ func suppressRaceErrors(err error) (PartitionMetadata, error) {
 	return PartitionMetadata{}, err
 }
 
-// sortVectors sorts the input vectors in-place, according to the provided left
-// and right offsets, which reference vectors by position. Vectors at left
-// offsets are sorted at the beginning of the slice, followed by vectors at
-// right offsets. The internal ordering among left and right vectors is not
-// defined.
+// 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.
 //
-// NOTE: The left and right offsets are modified in-place with the updated
-// positions of the vectors.
-func sortVectors(w *workspace.T, vectors vector.Set, leftOffsets, rightOffsets []uint64) {
+// 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.
+func splitPartitionData(
+	w *workspace.T,
+	vectors vector.Set,
+	childKeys []ChildKey,
+	valueBytes []ValueBytes,
+	leftOffsets, rightOffsets []uint64,
+) {
 	tempVector := w.AllocFloats(vectors.Dims)
 	defer w.FreeFloats(tempVector)
 
-	// Sort left and right offsets.
-	slices.Sort(leftOffsets)
-	slices.Sort(rightOffsets)
-
-	// Any left offsets that point beyond the end of the left list indicate that
-	// a vector needs to be moved from the right half of vectors to the left half.
-	// The reverse is true for right offsets. Because the left and right offsets
-	// are in sorted order, out-of-bounds offsets must be at the end of the left
-	// list and the beginning of the right list. Therefore, the algorithm just
-	// needs to iterate over those out-of-bounds offsets and swap the positions
-	// of the referenced vectors.
-	li := len(leftOffsets) - 1
-	ri := 0
-
-	var rightToLeft, leftToRight vector.T
-	for li >= 0 {
-		left := int(leftOffsets[li])
-		if left < len(leftOffsets) {
-			break
+	left := 0
+	right := 0
+	for {
+		// Find a misplaced "right" element from the left side.
+		var leftOffset int
+		for {
+			if left >= len(leftOffsets) {
+				return
+			}
+			leftOffset = int(leftOffsets[left])
+			left++
+			if leftOffset >= len(leftOffsets) {
+				break
+			}
 		}
 
-		right := int(rightOffsets[ri])
-		if right >= len(leftOffsets) {
-			panic(errors.AssertionFailedf(
-				"expected equal number of left and right offsets that need to be swapped"))
+		// 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
+			}
 		}
 
-		// Swap vectors.
-		rightToLeft = vectors.At(left)
-		leftToRight = vectors.At(right)
+		// Swap the two elements.
+		rightToLeft := vectors.At(leftOffset)
+		leftToRight := vectors.At(rightOffset)
 		copy(tempVector, rightToLeft)
 		copy(rightToLeft, leftToRight)
 		copy(leftToRight, tempVector)
 
-		leftOffsets[li] = uint64(left)
-		rightOffsets[ri] = uint64(right)
-
-		li--
-		ri++
-	}
-}
+		tempChildKey := childKeys[leftOffset]
+		childKeys[leftOffset] = childKeys[rightOffset]
+		childKeys[rightOffset] = tempChildKey
 
-func copyByOffsets[T any](source, target []T, offsets []uint64) []T {
-	for i := range offsets {
-		target[i] = source[offsets[i]]
+		tempValueBytes := valueBytes[leftOffset]
+		valueBytes[leftOffset] = valueBytes[rightOffset]
+		valueBytes[rightOffset] = tempValueBytes
 	}
-	return target
 }