vecindex: improve Index.Format method

Now that we have a non-transactional Store.TryGetPartition method, use
it in Index.Format. The Format method is used in testing/debugging, and
doesn't need to be transactional. Removing its dependence on GetPartition
will eventually allow us to remove that method entirely.

Also begin showing partition state if it is not Ready. This will be
needed as we begin using other states. Update all the places where we
create partition metadata to set the state so that Format shows expected
results.

Epic: CRDB-42943

Release note: None

Author: andy-kimball

pkg/sql/vecindex/cspann/BUILD.bazel

@@ -37,6 +37,7 @@ go_library(
         "//pkg/util/num32",
         "//pkg/util/stop",
         "//pkg/util/syncutil",
+        "//pkg/util/timeutil",
         "//pkg/util/vector",
         "@com_github_cockroachdb_crlib//crtime",
         "@com_github_cockroachdb_errors//:errors",

pkg/sql/vecindex/cspann/commontest/BUILD.bazel

@@ -14,7 +14,6 @@ go_library(
         "//pkg/sql/vecindex/cspann/testutils",
         "//pkg/sql/vecindex/cspann/workspace",
         "//pkg/util/num32",
-        "//pkg/util/timeutil",
         "//pkg/util/vector",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",

pkg/sql/vecindex/cspann/commontest/storetests.go

@@ -14,7 +14,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann/testutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann/workspace"
 	"github.com/cockroachdb/cockroach/pkg/util/num32"
-	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/vector"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/suite"
@@ -229,7 +228,10 @@ func (suite *StoreTestSuite) TestNonRootPartition() {
 		childKeys := []cspann.ChildKey{partitionKey1, partitionKey2}
 		valueBytes := []cspann.ValueBytes{valueBytes1, valueBytes2}
 		metadata := cspann.PartitionMetadata{
-			Level: cspann.SecondLevel, Centroid: quantizedSet.GetCentroid()}
+			Level:        cspann.SecondLevel,
+			Centroid:     quantizedSet.GetCentroid(),
+			StateDetails: cspann.MakeReadyDetails(),
+		}
 		partition := cspann.NewPartition(
 			metadata, suite.quantizer, quantizedSet, childKeys, valueBytes)
 
@@ -400,18 +402,12 @@ func (suite *StoreTestSuite) TestTryCreateEmptyPartition() {
 		treeKey := store.MakeTreeKey(suite.T(), treeID)
 		partitionKey := cspann.PartitionKey(10)
 		centroid := vector.T{4, 3}
-		timestamp := timeutil.Now()
 
 		// Create empty partition.
 		metadata := cspann.PartitionMetadata{
-			Level:    cspann.SecondLevel,
-			Centroid: centroid,
-			StateDetails: cspann.PartitionStateDetails{
-				State:     cspann.SplittingState,
-				Target1:   20,
-				Target2:   30,
-				Timestamp: timestamp,
-			},
+			Level:        cspann.SecondLevel,
+			Centroid:     centroid,
+			StateDetails: cspann.MakeSplittingDetails(20, 30),
 		}
 		suite.NoError(store.TryCreateEmptyPartition(suite.ctx, treeKey, partitionKey, metadata))
 
@@ -500,21 +496,16 @@ func (suite *StoreTestSuite) TestTryGetPartition() {
 		treeKey := store.MakeTreeKey(suite.T(), treeID)
 		partitionKey := cspann.PartitionKey(10)
 		centroid := vector.T{4, 3}
-		timestamp := timeutil.Now()
 
 		// Partition does not yet exist.
 		_, err := store.TryGetPartition(suite.ctx, treeKey, partitionKey)
 		suite.ErrorIs(err, cspann.ErrPartitionNotFound)
 
 		// Create partition with some vectors in it.
 		metadata := cspann.PartitionMetadata{
-			Level:    cspann.LeafLevel,
-			Centroid: centroid,
-			StateDetails: cspann.PartitionStateDetails{
-				State:     cspann.UpdatingState,
-				Source:    20,
-				Timestamp: timestamp,
-			},
+			Level:        cspann.LeafLevel,
+			Centroid:     centroid,
+			StateDetails: cspann.MakeUpdatingDetails(20),
 		}
 		suite.NoError(store.TryCreateEmptyPartition(suite.ctx, treeKey, partitionKey, metadata))
 		vectors := vector.MakeSet(2)
@@ -576,11 +567,7 @@ func (suite *StoreTestSuite) TestTryGetPartitionMetadata() {
 		// Update the metadata and verify we get the updated values.
 		expected := *partition.Metadata()
 		metadata := expected
-		metadata.StateDetails = cspann.PartitionStateDetails{
-			State:     cspann.UpdatingState,
-			Source:    30,
-			Timestamp: timeutil.Now(),
-		}
+		metadata.StateDetails = cspann.MakeUpdatingDetails(30)
 		suite.NoError(store.TryUpdatePartitionMetadata(
 			suite.ctx, treeKey, partitionKey, metadata, expected))
 
@@ -661,17 +648,12 @@ func (suite *StoreTestSuite) TestTryAddToPartition() {
 		treeKey := store.MakeTreeKey(suite.T(), treeID)
 		partitionKey := cspann.PartitionKey(10)
 		centroid := vector.T{4, 3}
-		timestamp := timeutil.Now()
 
 		// Partition does not yet exist.
 		metadata := cspann.PartitionMetadata{
-			Level:    cspann.LeafLevel,
-			Centroid: centroid,
-			StateDetails: cspann.PartitionStateDetails{
-				State:     cspann.UpdatingState,
-				Source:    20,
-				Timestamp: timestamp,
-			},
+			Level:        cspann.LeafLevel,
+			Centroid:     centroid,
+			StateDetails: cspann.MakeUpdatingDetails(20),
 		}
 		addVectors := vector.MakeSet(2)
 		addVectors.Add(vec1)
@@ -854,12 +836,9 @@ func (suite *StoreTestSuite) createTestPartition(
 ) (cspann.PartitionKey, *cspann.Partition) {
 	partitionKey := cspann.PartitionKey(10)
 	metadata := cspann.PartitionMetadata{
-		Level:    cspann.SecondLevel,
-		Centroid: vector.T{4, 3},
-		StateDetails: cspann.PartitionStateDetails{
-			State:     cspann.ReadyState,
-			Timestamp: timeutil.Now(),
-		},
+		Level:        cspann.SecondLevel,
+		Centroid:     vector.T{4, 3},
+		StateDetails: cspann.MakeReadyDetails(),
 	}
 	suite.NoError(store.TryCreateEmptyPartition(suite.ctx, treeKey, partitionKey, metadata))
 	vectors := vector.MakeSet(2)
@@ -973,7 +952,10 @@ func (suite *StoreTestSuite) insertLeafPartition(store TestStore, treeID int) cs
 	childKeys := []cspann.ChildKey{primaryKey1, primaryKey2, primaryKey3}
 	valueBytes := []cspann.ValueBytes{valueBytes1, valueBytes2, valueBytes3}
 	metadata := cspann.PartitionMetadata{
-		Level: cspann.LeafLevel, Centroid: quantizedSet.GetCentroid()}
+		Level:        cspann.LeafLevel,
+		Centroid:     quantizedSet.GetCentroid(),
+		StateDetails: cspann.MakeReadyDetails(),
+	}
 	partition := cspann.NewPartition(
 		metadata, suite.quantizer, quantizedSet, childKeys, valueBytes)
 
@@ -1117,7 +1099,10 @@ func (suite *StoreTestSuite) setRootPartition(store TestStore, treeID int) {
 	childKeys := []cspann.ChildKey{partitionKey1, partitionKey2}
 	valueBytes := []cspann.ValueBytes{valueBytes1, valueBytes2}
 	metadata := cspann.PartitionMetadata{
-		Level: cspann.SecondLevel, Centroid: quantizedSet.GetCentroid()}
+		Level:        cspann.SecondLevel,
+		Centroid:     quantizedSet.GetCentroid(),
+		StateDetails: cspann.MakeReadyDetails(),
+	}
 	newRoot := cspann.NewPartition(
 		metadata, suite.rootQuantizer, quantizedSet, childKeys, valueBytes)
 

pkg/sql/vecindex/cspann/fixup_worker.go

@@ -309,7 +309,10 @@ func (fw *fixupWorker) splitPartition(
 		}
 		valueBytes := make([]ValueBytes, 2)
 		rootMetadata := PartitionMetadata{
-			Level: partition.Level() + 1, Centroid: quantizedSet.GetCentroid()}
+			Level:        partition.Level() + 1,
+			Centroid:     quantizedSet.GetCentroid(),
+			StateDetails: MakeReadyDetails(),
+		}
 		rootPartition := NewPartition(
 			rootMetadata, fw.index.rootQuantizer, quantizedSet, childKeys, valueBytes)
 		if err = fw.txn.SetRootPartition(ctx, fw.treeKey, rootPartition); err != nil {
@@ -623,7 +626,11 @@ func (fw *fixupWorker) mergePartition(
 			return errors.AssertionFailedf("only root partition can have zero vectors")
 		}
 		quantizedSet := fw.index.rootQuantizer.Quantize(&fw.workspace, vectors)
-		metadata := PartitionMetadata{Level: partition.Level(), Centroid: quantizedSet.GetCentroid()}
+		metadata := PartitionMetadata{
+			Level:        partition.Level(),
+			Centroid:     quantizedSet.GetCentroid(),
+			StateDetails: MakeReadyDetails(),
+		}
 		rootPartition := NewPartition(
 			metadata, fw.index.rootQuantizer, quantizedSet, partition.ChildKeys(), partition.ValueBytes())
 		if err = fw.txn.SetRootPartition(ctx, fw.treeKey, rootPartition); err != nil {

pkg/sql/vecindex/cspann/index.go

@@ -918,6 +918,10 @@ type FormatOptions struct {
 	// PrimaryKeyStrings, if true, indicates that primary key bytes should be
 	// interpreted as strings. This is used for testing scenarios.
 	PrimaryKeyStrings bool
+	// RootPartitionKey is the key of the partition that is used as the root of
+	// the formatted output. If this is InvalidKey, then RootKey is used by
+	// default.
+	RootPartitionKey PartitionKey
 }
 
 // Format formats the vector index as a tree-formatted string similar to this,
@@ -933,7 +937,7 @@ type FormatOptions struct {
 // values are rounded to 4 decimal places. Centroids are printed next to
 // partition keys.
 func (vi *Index) Format(
-	ctx context.Context, idxCtx *Context, treeKey TreeKey, options FormatOptions,
+	ctx context.Context, treeKey TreeKey, options FormatOptions,
 ) (str string, err error) {
 	// Write formatted bytes to this buffer.
 	var buf bytes.Buffer
@@ -953,15 +957,24 @@ func (vi *Index) Format(
 
 	var helper func(partitionKey PartitionKey, parentPrefix string, childPrefix string) error
 	helper = func(partitionKey PartitionKey, parentPrefix string, childPrefix string) error {
-		partition, err := idxCtx.txn.GetPartition(ctx, treeKey, partitionKey)
+		partition, err := vi.store.TryGetPartition(ctx, treeKey, partitionKey)
 		if err != nil {
 			if errors.Is(err, ErrPartitionNotFound) {
-				// This should never happen, and indicates a bug in the vector index
-				// implementation. Fallback to showing MISSING.
+				// Something else might be modifying the tree as we're trying to
+				// print it. Fallback to showing MISSING.
 				buf.WriteString(parentPrefix)
 				buf.WriteString("• ")
 				buf.WriteString(strconv.FormatInt(int64(partitionKey), 10))
-				buf.WriteString(" (MISSING)\n")
+
+				// If this is the root partition, then show synthesized empty
+				// partition.
+				if partitionKey == RootKey {
+					centroid := make(vector.T, vi.quantizer.GetDims())
+					buf.WriteByte(' ')
+					utils.WriteVector(&buf, centroid, 4)
+				} else {
+					buf.WriteString(" (MISSING)\n")
+				}
 				return nil
 			}
 			return err
@@ -976,6 +989,12 @@ func (vi *Index) Format(
 		buf.WriteString(strconv.FormatInt(int64(partitionKey), 10))
 		buf.WriteByte(' ')
 		utils.WriteVector(&buf, original, 4)
+		details := partition.Metadata().StateDetails
+		if details.State != ReadyState {
+			buf.WriteString(" [")
+			buf.WriteString(details.String())
+			buf.WriteByte(']')
+		}
 		buf.WriteByte('\n')
 
 		if partition.Count() == 0 {
@@ -995,7 +1014,10 @@ func (vi *Index) Format(
 
 			if partition.Level() == LeafLevel {
 				refs := []VectorWithKey{{Key: childKey}}
-				if err = idxCtx.txn.GetFullVectors(ctx, treeKey, refs); err != nil {
+				err = vi.store.RunTransaction(ctx, func(tx Txn) error {
+					return tx.GetFullVectors(ctx, treeKey, refs)
+				})
+				if err != nil {
 					return err
 				}
 				buf.WriteString(parentPrefix)
@@ -1029,7 +1051,11 @@ func (vi *Index) Format(
 		return nil
 	}
 
-	if err = helper(RootKey, "", ""); err != nil {
+	rootKey := options.RootPartitionKey
+	if rootKey == InvalidKey {
+		rootKey = RootKey
+	}
+	if err = helper(rootKey, "", ""); err != nil {
 		return "", err
 	}
 	return buf.String(), nil