Merge #147307

147307: fix pulling full vectors from the primary key when column families are in use r=mw5h a=mw5h

#### vecindex: delete Manager.GetWithDesc() method

This method was added because backfilling indexes are not considered
deletable. In retrospect, there's no reason why we need to restrict an
internal interface this way and deleting it removes a fair bit of
complexity from upcoming changes to the vector store.

Informs: #146046
Release note: None

#### fetcher: return spanID from NextRowDecoded()

Previously, Fetcher.NextRowDecoded() discarded the spanID it got from
Fetcher.NextRow(). This information is needed for a future use of
NextRowDecoded(), so this patch passes it through. Existing users ignore
this return value.

Informs: #146046
Release note: None

#### 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.

#### vecstore: correctly handle composite pk columns

Previously, the vector store would blindly use primary key column bytes
encoded in a vector index key to rebuild the primary key when retrieving
full vector values. SpanBuilder will attempt to reused already encoded
datums when building a span, but it can't do so if the encoded datum was
encoded in a different direction than required for the span being built.
For non-composite columns this was not an issue because they have all
the data needed to re-encode in the key value itself. For composite
column, the data encoded in the key is a key analogue and can't be used
to reconstruct the original column value, which is stored as a suffix
column.

For vector search, the search processor has access to the ValueBytes
array as well as the key, so it's able to rebuild column values as
needed. However, the fixup processor only has access to the key bytes,
which requires a bit of trickery.

In this patch, we disallow assigning directionality to vector index key
columns. The mechanics of similarity search preclude scannign prefixes,
so this was always a nonsensical thing to do. We now disallow it
entirely and take the further step to mechanically switch the
directionality of vector index prefix columns to match the primary key,
ensuring that the encoded value of these shared columns is always
directly usable.

We also move the logic for decoding and rebuilding the primary key into
the vector store so that it can be shared between fixups and the vector
search processor. Note that there is an additional composite key decoder
available for the vector search processor because it does indeed need to
decode the primary key bytes back to column values.

Along the way, we:
*  Remove ASC/DESC from vector index columns.
*  Introduce a new index type method to describe indexes with
   non-directional prefixes.
*  Alter SHOW commands to no longer show vector index column
   directionality.
*  Introduce a bunch of test-gated code to ensure everything I said
   above is in fact true and remains so in the future.

Informs: #146046
Release note (sql change): It is no longer permitted to assign
directionality to any vector index column. Prefix columns (those columns
before the vector column in a compound vector index key) are not
scannable in a vector index, so directionality isn't relevant to them.

Co-authored-by: Matt White <[email protected]>

Author: craig[bot]

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/ccl/logictestccl/testdata/logic_test/partitioning_index

@@ -211,11 +211,11 @@ SHOW CREATE TABLE vector
             b INT8 NULL,
             v VECTOR(3) NULL,
             CONSTRAINT vector_pkey PRIMARY KEY (a ASC),
-            VECTOR INDEX vector_b_v_idx (b ASC, v vector_l2_ops) PARTITION BY LIST (b) (
+            VECTOR INDEX vector_b_v_idx (b, v vector_l2_ops) PARTITION BY LIST (b) (
               PARTITION p1 VALUES IN ((1)),
               PARTITION pu VALUES IN ((NULL))
             ),
-            VECTOR INDEX vector_idx (a ASC, b ASC, v vector_l2_ops) PARTITION BY LIST (a, b) (
+            VECTOR INDEX vector_idx (a, b, v vector_l2_ops) PARTITION BY LIST (a, b) (
               PARTITION p1 VALUES IN ((1, 1), (2, 2))
             )
           ) WITH (schema_locked = true)

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/alter_primary_key.go

@@ -851,6 +851,11 @@ func setKeySuffixAndStoredColumnIDsFromPrimary(
 				"indexed vector column cannot be part of the primary key")
 		}
 	}
+
+	if vecIdx {
+		tabledesc.UpdateVectorIndexPrefixColDirections(toAdd, primary)
+	}
+
 	// Finally, add all the stored columns if it is not already a key or key suffix column.
 	toAddOldStoredColumnIDs := toAdd.StoreColumnIDs
 	toAddOldStoredColumnNames := toAdd.StoreColumnNames

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)
 		}
 	}
@@ -891,18 +904,29 @@ func (ib *IndexBackfiller) initIndexes(
 			return err
 		}
 
-		vecIndex, err := vecIndexManager.GetWithDesc(ctx, desc, idx)
+		vecIndex, err := vecIndexManager.Get(ctx, desc.GetID(), idx.GetID())
 		if err != nil {
 			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/catalog/catformat/index.go

@@ -274,10 +274,14 @@ func FormatIndexElements(
 				}
 			}
 		}
-		// The last column of an inverted or vector index cannot have a DESC
-		// direction because it does not have a linear ordering. Since the default
-		// direction is ASC, we omit the direction entirely for inverted/vector
-		// index columns.
+		// Vector indexes do not support ASC/DESC modifiers.
+		if !index.Type.HasScannablePrefix() {
+			continue
+		}
+		// The last column of an inverted index cannot have a DESC direction
+		// because it does not have a linear ordering. Since the default
+		// direction is ASC, we omit the direction entirely for inverted index
+		// columns.
 		if i < n-1 || index.Type.HasLinearOrdering() {
 			f.WriteByte(' ')
 			f.WriteString(index.KeyColumnDirections[i].String())

pkg/sql/catalog/descpb/index.go

@@ -9,6 +9,8 @@ import (
 	"fmt"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/idxtype"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
@@ -47,6 +49,11 @@ func (desc *IndexDescriptor) FillColumns(elems tree.IndexElemList) error {
 		if c.Expr != nil {
 			return errors.AssertionFailedf("index elem expression should have been replaced with a column")
 		}
+		// Vector index prefix columns don't have a direction, nor do the vector columns themselves.
+		if !desc.Type.HasScannablePrefix() && c.Direction != tree.DefaultDirection {
+			return pgerror.Newf(pgcode.FeatureNotSupported,
+				"%s does not support the %s option", idxtype.ErrorText(desc.Type), c.Direction)
+		}
 		desc.KeyColumnNames = append(desc.KeyColumnNames, string(c.Column))
 		switch c.Direction {
 		case tree.Ascending, tree.DefaultDirection: