vecstore: use span.Builder in implementation of GetFullVectors

Previously, the vector store stored the table and index descriptors for
the vector index it was associated with. Since the store itself did not
have a lease on this descriptor and would outlive the transaction that
caused the store to be created, we could end up using a stale table
descriptor for vector operations.

This patch moves table information into the vector transaction in the
form of an index fetch spec created by a leaseholder on the table. The
vecindex transaction is expected to have a lifetime shorter than the
lease held by the caller, so the descriptors that generate this fetch
spec will not go stale.

Addtionally, we generate the needed family IDs for a span.Splitter(), to
avoid reading families that we don't require. This is also included with
the fetch spec.

This patch also:
* Introduces a proto buffer to hold the index fetch spec and family IDs.
* Introduces a function to initialize this proto buffer.
* Introduces a testing API into the vecstore so that tests can provide
  fake family descriptors to vecstore.

Fixes: #146046
Release note (bug fix): A bug where using column families on tables
with vector indexes would cause the index to fail to return results has
been fixed.

Author: mw5h

pkg/BUILD.bazel

@@ -2379,6 +2379,7 @@ GO_TARGETS = [
     "//pkg/sql/vecindex/vecencoding:vecencoding_test",
     "//pkg/sql/vecindex/vecpb:vecpb",
     "//pkg/sql/vecindex/vecpb:vecpb_test",
+    "//pkg/sql/vecindex/vecstore/vecstorepb:vecstorepb",
     "//pkg/sql/vecindex/vecstore:vecstore",
     "//pkg/sql/vecindex/vecstore:vecstore_test",
     "//pkg/sql/vecindex:vecindex",

pkg/gen/protobuf.bzl

@@ -79,6 +79,7 @@ PROTOBUF_SRCS = [
     "//pkg/sql/vecindex/cspann/quantize:quantize_go_proto",
     "//pkg/sql/vecindex/cspann:cspann_go_proto",
     "//pkg/sql/vecindex/vecpb:vecpb_go_proto",
+    "//pkg/sql/vecindex/vecstore/vecstorepb:vecstorepb_go_proto",
     "//pkg/storage/enginepb:enginepb_go_proto",
     "//pkg/testutils/grpcutils:grpcutils_go_proto",
     "//pkg/ts/catalog:catalog_go_proto",

pkg/sql/BUILD.bazel

@@ -544,6 +544,7 @@ go_library(
         "//pkg/sql/types",
         "//pkg/sql/vecindex",
         "//pkg/sql/vecindex/vecpb",
+        "//pkg/sql/vecindex/vecstore",
         "//pkg/sql/vtable",
         "//pkg/storage",
         "//pkg/storage/enginepb",

pkg/sql/backfill/BUILD.bazel

@@ -45,6 +45,8 @@ go_library(
         "//pkg/sql/vecindex",
         "//pkg/sql/vecindex/cspann",
         "//pkg/sql/vecindex/vecencoding",
+        "//pkg/sql/vecindex/vecstore",
+        "//pkg/sql/vecindex/vecstore/vecstorepb",
         "//pkg/util/admission/admissionpb",
         "//pkg/util/ctxgroup",
         "//pkg/util/hlc",
@@ -81,6 +83,7 @@ go_test(
         "//pkg/sql/catalog/tabledesc",
         "//pkg/sql/execinfra",
         "//pkg/sql/sem/catid",
+        "//pkg/sql/sem/eval",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",

pkg/sql/backfill/backfill.go

@@ -12,7 +12,6 @@ import (
 	"time"
 	"unsafe"
 
-	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
@@ -39,6 +38,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecencoding"
+	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecstore"
+	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecstore/vecstorepb"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -477,8 +478,12 @@ type VectorIndexHelper struct {
 	numPrefixCols int
 	// vecIndex is the vector index retrieved from the vector index manager.
 	vecIndex *cspann.Index
+	// fullVecFetchSpec is the fetch spec for the full vector.
+	fullVecFetchSpec vecstorepb.GetFullVectorsFetchSpec
 	// indexPrefix are the prefix bytes for this index (/Tenant/Table/Index).
 	indexPrefix []byte
+	// evalCtx is the evaluation context used for vector operations.
+	evalCtx *eval.Context
 }
 
 // ReEncodeVector takes a rowenc.indexEntry, extracts the key values, unquantized
@@ -509,7 +514,8 @@ func (vih *VectorIndexHelper) ReEncodeVector(
 
 	// Locate a new partition for the key and re-encode the vector.
 	var searcher vecindex.MutationSearcher
-	searcher.Init(vih.vecIndex, txn)
+	searcher.Init(vih.evalCtx, vih.vecIndex, txn, &vih.fullVecFetchSpec)
+
 	if err := searcher.SearchForInsert(ctx, roachpb.Key(key.Prefix), vec); err != nil {
 		return &rowenc.IndexEntry{}, err
 	}
@@ -598,7 +604,7 @@ func (ib *IndexBackfiller) InitForLocalUse(
 ) error {
 
 	// Initialize ib.added.
-	if err := ib.initIndexes(ctx, evalCtx.Codec, desc, nil /* allowList */, 0 /*sourceIndex*/, nil); err != nil {
+	if err := ib.initIndexes(ctx, evalCtx, desc, nil /* allowList */, 0 /*sourceIndex*/, nil); err != nil {
 		return err
 	}
 
@@ -744,7 +750,14 @@ func (ib *IndexBackfiller) InitForDistributedUse(
 	evalCtx := flowCtx.NewEvalCtx()
 
 	// Initialize ib.added.
-	if err := ib.initIndexes(ctx, evalCtx.Codec, desc, allowList, sourceIndexID, flowCtx.Cfg.VecIndexManager.(*vecindex.Manager)); err != nil {
+	if err := ib.initIndexes(
+		ctx,
+		evalCtx,
+		desc,
+		allowList,
+		sourceIndexID,
+		flowCtx.Cfg.VecIndexManager.(*vecindex.Manager),
+	); err != nil {
 		return err
 	}
 
@@ -838,7 +851,7 @@ func (ib *IndexBackfiller) initCols(desc catalog.TableDescriptor) (err error) {
 // If `allowList` is nil, we add all adding index mutations.
 func (ib *IndexBackfiller) initIndexes(
 	ctx context.Context,
-	codec keys.SQLCodec,
+	evalCtx *eval.Context,
 	desc catalog.TableDescriptor,
 	allowList []catid.IndexID,
 	sourceIndexID catid.IndexID,
@@ -867,7 +880,7 @@ func (ib *IndexBackfiller) initIndexes(
 			(allowListAsSet.Empty() || allowListAsSet.Contains(m.AsIndex().GetID())) {
 			idx := m.AsIndex()
 			ib.added = append(ib.added, idx)
-			keyPrefix := rowenc.MakeIndexKeyPrefix(codec, desc.GetID(), idx.GetID())
+			keyPrefix := rowenc.MakeIndexKeyPrefix(evalCtx.Codec, desc.GetID(), idx.GetID())
 			ib.keyPrefixes = append(ib.keyPrefixes, keyPrefix)
 		}
 	}
@@ -896,13 +909,24 @@ func (ib *IndexBackfiller) initIndexes(
 			return err
 		}
 
-		ib.VectorIndexes[idx.GetID()] = VectorIndexHelper{
+		helper := VectorIndexHelper{
 			vectorOrd:     vectorCol.Ordinal(),
 			centroid:      make(vector.T, idx.GetVecConfig().Dims),
 			numPrefixCols: idx.NumKeyColumns() - 1,
 			vecIndex:      vecIndex,
-			indexPrefix:   rowenc.MakeIndexKeyPrefix(codec, desc.GetID(), idx.GetID()),
+			indexPrefix:   rowenc.MakeIndexKeyPrefix(evalCtx.Codec, desc.GetID(), idx.GetID()),
+			evalCtx:       ib.evalCtx,
+		}
+		if err := vecstore.InitGetFullVectorsFetchSpec(
+			&helper.fullVecFetchSpec,
+			evalCtx,
+			desc,
+			idx,
+			ib.sourceIndex,
+		); err != nil {
+			return err
 		}
+		ib.VectorIndexes[idx.GetID()] = helper
 	}
 
 	return nil

pkg/sql/backfill/index_backfiller_cols_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/catid"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/stretchr/testify/require"
 )
@@ -418,7 +419,7 @@ func TestInitIndexesAllowList(t *testing.T) {
 		ib := &IndexBackfiller{}
 		err := ib.initIndexes(
 			context.Background(),
-			keys.SystemSQLCodec,
+			&eval.Context{Codec: keys.SystemSQLCodec},
 			desc,
 			nil, /* allowList */
 			0,   /* sourceIndexID */
@@ -434,7 +435,7 @@ func TestInitIndexesAllowList(t *testing.T) {
 		ib := &IndexBackfiller{}
 		err := ib.initIndexes(
 			context.Background(),
-			keys.SystemSQLCodec,
+			&eval.Context{Codec: keys.SystemSQLCodec},
 			desc,
 			[]catid.IndexID{3}, /* allowList */
 			0,                  /* sourceIndexID */

pkg/sql/distsql_physical_planner.go

@@ -52,6 +52,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlinstance"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecstore"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -4602,6 +4603,16 @@ func (dsp *DistSQLPlanner) planVectorSearch(
 		return err
 	}
 
+	if err := vecstore.InitGetFullVectorsFetchSpec(
+		&spec.GetFullVectorsFetchSpec,
+		planCtx.EvalContext(),
+		planInfo.table,
+		planInfo.index,
+		planInfo.table.GetPrimaryIndex(),
+	); err != nil {
+		return err
+	}
+
 	// Execute the vector search on the gateway node.
 	corePlacement := []physicalplan.ProcessorCorePlacement{{
 		SQLInstanceID: dsp.gatewaySQLInstanceID,
@@ -4675,6 +4686,16 @@ func (dsp *DistSQLPlanner) planVectorMutationSearch(
 		return err
 	}
 
+	if err := vecstore.InitGetFullVectorsFetchSpec(
+		&spec.GetFullVectorsFetchSpec,
+		planCtx.EvalContext(),
+		planInfo.table,
+		planInfo.index,
+		planInfo.table.GetPrimaryIndex(),
+	); err != nil {
+		return err
+	}
+
 	// The vector mutation search can be conducted for each row independently, so
 	// it's fine to instantiate one instance for each stream.
 	pSpec := execinfrapb.ProcessorCoreUnion{VectorMutationSearch: spec}

pkg/sql/execinfrapb/BUILD.bazel

@@ -100,6 +100,7 @@ proto_library(
         "//pkg/sql/inverted:inverted_proto",
         "//pkg/sql/sessiondatapb:sessiondatapb_proto",
         "//pkg/sql/types:types_proto",
+        "//pkg/sql/vecindex/vecstore/vecstorepb:vecstorepb_proto",
         "//pkg/util/hlc:hlc_proto",
         "//pkg/util/optional:optional_proto",
         "//pkg/util/tracing/tracingpb:tracingpb_proto",
@@ -133,6 +134,7 @@ go_proto_library(
         "//pkg/sql/inverted",
         "//pkg/sql/sessiondatapb",
         "//pkg/sql/types",
+        "//pkg/sql/vecindex/vecstore/vecstorepb",
         "//pkg/util/hlc",
         "//pkg/util/optional",
         "//pkg/util/tracing/tracingpb",

pkg/sql/execinfrapb/processors_sql.proto

@@ -24,6 +24,7 @@ import "sql/types/types.proto";
 import "sql/execinfrapb/data.proto";
 import "sql/execinfrapb/processors_base.proto";
 import "sql/inverted/span_expression.proto";
+import "sql/vecindex/vecstore/vecstorepb/fullvecfetchspec.proto";
 import "util/hlc/timestamp.proto";
 
 // ValuesCoreSpec is the core of a processor that has no inputs and generates
@@ -1099,6 +1100,8 @@ message VectorSearchSpec {
 
   // TargetNeighborCount is the number of nearest neighbors to search for.
   optional uint64 target_neighbor_count = 4 [(gogoproto.nullable) = false];
+
+  optional vecindex.vecstore.vecstorepb.GetFullVectorsFetchSpec get_full_vectors_fetch_spec = 5 [(gogoproto.nullable) = false];
 }
 
 // VectorMutationSearchSpec is the specification for a vector-mutation-search
@@ -1128,4 +1131,6 @@ message VectorMutationSearchSpec {
   // IsIndexPut is true if the search is being conducted for an index put,
   // instead of an index del.
   optional bool is_index_put = 7 [(gogoproto.nullable) = false];
+
+  optional vecindex.vecstore.vecstorepb.GetFullVectorsFetchSpec get_full_vectors_fetch_spec = 8 [(gogoproto.nullable) = false];
 }