Merge #143403 #143446

143403: vecindex: improve Index.Format method r=drewkimball a=andy-kimball

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

143446: sql: optimize query to populate `crdb_internal.table_row_statistics` r=mgartner a=mgartner

The query to populate the `crdb_internal.table_row_statistics` virtual
table has been rewritten to avoid scanning the `system.table_statistics`
table twice.

Before this commit the query plan was:

    • group (hash)
    │ group by: tableID
    │
    └── • hash join
        │ actual row count: 515,822
        │ equality: (tableID, createdAt) = (tableID, max)
        │
        ├── • scan
        │     table: table_statistics@primary
        │     spans: FULL SCAN
        │
        └── • group (streaming)
            │ group by: tableID
            │ ordered: +"tableID"
            │
            └── • scan
                  table: table_statistics@primary
                  spans: FULL SCAN

Now it is:

    • distinct
    │ distinct on: tableID
    │ order key: tableID
    │
    └── • sort
        │ order: +"tableID",-"createdAt",-"rowCount"
        │ already ordered: +"tableID"
        │
        └── • scan
              table: table_statistics@primary
              spans: FULL SCAN

The `crdb_internal.table_row_statistics` table is used to populate the
`estimated_row_count` column in the output of `SHOW TABLES`. In
pathological cases where there a many rows in `system.table_statistics`,
this new query makes `SHOW TABLES` significantly faster. In a test of
mine with ~1.1 million rows in `system.table_statistics` setup with
[this](https://gist.github.com/mgartner/b72b39901be0d942d5a026054e688a8c),
I observed the latency of `SHOW TABLES` drop by ~50% with the new query,
from ~720ms to ~350ms.

Informs #143438

Release note (performance improvement): `SHOW TABLES` is now faster,
especially in cases where there are many tables, both live and
previously dropped.


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

Author: 3 people

pkg/sql/crdb_internal.go

@@ -768,15 +768,11 @@ CREATE TABLE crdb_internal.table_row_statistics (
 		// contention on the stats table. We pass a nil transaction so that the AS
 		// OF clause can be independent of any outer query.
 		query := fmt.Sprintf(`
-           SELECT s."tableID", max(s."rowCount")
-             FROM system.table_statistics AS s
-             JOIN (
-                    SELECT "tableID", max("createdAt") AS last_dt
-                      FROM system.table_statistics
-                     GROUP BY "tableID"
-                  ) AS l ON l."tableID" = s."tableID" AND l.last_dt = s."createdAt"
-            AS OF SYSTEM TIME '%s'
-            GROUP BY s."tableID"`, statsAsOfTimeClusterMode.String(&p.ExecCfg().Settings.SV))
+			SELECT DISTINCT ON ("tableID") "tableID", "rowCount"
+			FROM system.table_statistics
+			AS OF SYSTEM TIME '%s'
+			ORDER BY "tableID", "createdAt" DESC, "rowCount" DESC`,
+			statsAsOfTimeClusterMode.String(&p.ExecCfg().Settings.SV))
 		statRows, err := p.ExtendedEvalContext().ExecCfg.InternalDB.Executor().QueryBufferedEx(
 			ctx, "crdb-internal-statistics-table", nil,
 			sessiondata.NodeUserSessionDataOverride,

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