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.

Author: mw5h

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/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/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:

pkg/sql/catalog/tabledesc/structured.go

@@ -2797,3 +2797,17 @@ func (desc *Mutable) BumpExternalAsOf(timestamp hlc.Timestamp) error {
 func (desc *Mutable) ForceModificationTime(modificationTime hlc.Timestamp) {
 	desc.ModificationTime = modificationTime
 }
+
+func UpdateVectorIndexPrefixColDirections(
+	vecIndexDesc *descpb.IndexDescriptor, primaryIndexDesc *descpb.IndexDescriptor,
+) {
+	numPrefixCols := len(vecIndexDesc.KeyColumnIDs) - 1
+	for i := 0; i < numPrefixCols; i++ {
+		for j, pkColID := range primaryIndexDesc.KeyColumnIDs {
+			if vecIndexDesc.KeyColumnIDs[i] == pkColID {
+				vecIndexDesc.KeyColumnDirections[i] = primaryIndexDesc.KeyColumnDirections[j]
+				break
+			}
+		}
+	}
+}

pkg/sql/create_index.go

@@ -267,6 +267,8 @@ func makeIndexDescriptor(
 		if err != nil {
 			return nil, err
 		}
+
+		tabledesc.UpdateVectorIndexPrefixColDirections(&indexDesc, tableDesc.GetPrimaryIndex().IndexDesc())
 	}
 
 	if n.Sharded != nil {

pkg/sql/create_table.go

@@ -2849,6 +2849,9 @@ func replaceLikeTableOpts(n *tree.CreateTable, params runParams) (tree.TableDefs
 						Column:    tree.Name(name),
 						Direction: tree.Ascending,
 					}
+					if !indexDef.Type.HasScannablePrefix() {
+						elem.Direction = tree.DefaultDirection
+					}
 					col, err := catalog.MustFindColumnByID(td, idx.GetKeyColumnID(j))
 					if err != nil {
 						return nil, err

pkg/sql/logictest/testdata/logic_test/expression_index

@@ -10,7 +10,7 @@ CREATE TABLE t (
   INDEX t_a_plus_b_idx ((a + b)),
   INDEX t_lower_c_a_plus_b_idx (lower(c), (a + b)),
   INVERTED INDEX t_b_j_a_asc (b DESC, (j->'a') ASC),
-  VECTOR INDEX t_vec (a DESC, (v::VECTOR(3))),
+  VECTOR INDEX t_vec (a, (v::VECTOR(3))),
   FAMILY (k, a, b, c, j, v)
 )
 
@@ -30,7 +30,7 @@ CREATE TABLE public.t (
   INDEX t_a_plus_b_idx ((a + b) ASC),
   INDEX t_lower_c_a_plus_b_idx (lower(c) ASC, (a + b) ASC),
   INVERTED INDEX t_b_j_a_asc (b DESC, (j->'a':::STRING)),
-  VECTOR INDEX t_vec (a DESC, (v::VECTOR(3)) vector_l2_ops),
+  VECTOR INDEX t_vec (a, (v::VECTOR(3)) vector_l2_ops),
   FAMILY fam_0_k_a_b_c_j_v (k, a, b, c, j, v)
 );
 
@@ -50,7 +50,7 @@ CREATE TABLE public.t (
   INDEX t_a_plus_b_idx ((a + b) ASC),
   INDEX t_lower_c_a_plus_b_idx (lower(c) ASC, (a + b) ASC),
   INVERTED INDEX t_b_j_a_asc (b DESC, (j->'a':::STRING)),
-  VECTOR INDEX t_vec (a DESC, (v::VECTOR(3)) vector_l2_ops),
+  VECTOR INDEX t_vec (a, (v::VECTOR(3)) vector_l2_ops),
   FAMILY fam_0_k_a_b_c_j_v (k, a, b, c, j, v)
 ) WITH (schema_locked = true);
 
@@ -70,7 +70,7 @@ CREATE TABLE public.t (
   INDEX t_a_plus_b_idx ((a + b) ASC),
   INDEX t_lower_c_a_plus_b_idx (lower(c) ASC, (a + b) ASC),
   INVERTED INDEX t_b_j_a_asc (b DESC, (j->‹×›:::STRING)),
-  VECTOR INDEX t_vec (a DESC, (v::VECTOR(3)) vector_l2_ops),
+  VECTOR INDEX t_vec (a, (v::VECTOR(3)) vector_l2_ops),
   FAMILY fam_0_k_a_b_c_j_v (k, a, b, c, j, v)
 );
 
@@ -90,7 +90,7 @@ CREATE TABLE public.t (
   INDEX t_a_plus_b_idx ((a + b) ASC),
   INDEX t_lower_c_a_plus_b_idx (lower(c) ASC, (a + b) ASC),
   INVERTED INDEX t_b_j_a_asc (b DESC, (j->‹×›:::STRING)),
-  VECTOR INDEX t_vec (a DESC, (v::VECTOR(3)) vector_l2_ops),
+  VECTOR INDEX t_vec (a, (v::VECTOR(3)) vector_l2_ops),
   FAMILY fam_0_k_a_b_c_j_v (k, a, b, c, j, v)
 ) WITH (schema_locked = true);
 
@@ -178,7 +178,7 @@ CREATE TABLE err (a INT, b INT, VECTOR INDEX (a, (a + b)))
 statement error column \(vec1::VECTOR\(3\)\) has type vector, which is only allowed as the last column in a vector index
 CREATE TABLE err (a INT, vec1 VECTOR, VECTOR INDEX ((vec1::VECTOR(3)), a))
 
-statement error the last column in a vector index cannot have the DESC option
+statement error pq: a vector index does not support the DESC option
 CREATE TABLE err (a INT, vec1 VECTOR, VECTOR INDEX (a, (vec1::VECTOR(3)) DESC))
 
 statement error vector column \(vec1::VECTOR\) does not have a fixed number of dimensions, so it cannot be indexed\nDETAIL: specify the number of dimensions in the type, like VECTOR\(128\) for 128 dimensions
@@ -476,7 +476,7 @@ CREATE VECTOR INDEX err ON t (a, (a + b))
 statement error column \(v::VECTOR\(3\)\) has type vector, which is only allowed as the last column in a vector index
 CREATE VECTOR INDEX err ON t ((v::VECTOR(3)), a)
 
-statement error the last column in a vector index cannot have the DESC option
+statement error pq: a vector index does not support the DESC option
 CREATE VECTOR INDEX err ON t (a, (v::VECTOR(3)) DESC)
 
 statement error vector column \(v::VECTOR\) does not have a fixed number of dimensions, so it cannot be indexed\nDETAIL: specify the number of dimensions in the type, like VECTOR\(128\) for 128 dimensions
@@ -640,7 +640,7 @@ CREATE TABLE public.copy_indexes (
   INDEX named_idx ((a + 1:::INT8) ASC),
   UNIQUE INDEX src_expr_key ((a + 10:::INT8) ASC),
   INVERTED INDEX src_expr_expr1_idx ((a + b) ASC, (j->'a':::STRING)),
-  VECTOR INDEX src_b_expr_idx (b ASC, (v::VECTOR(3)) vector_l2_ops)
+  VECTOR INDEX src_b_expr_idx (b, (v::VECTOR(3)) vector_l2_ops)
 );
 
 skipif config schema-locked-disabled
@@ -659,7 +659,7 @@ CREATE TABLE public.copy_indexes (
   INDEX named_idx ((a + 1:::INT8) ASC),
   UNIQUE INDEX src_expr_key ((a + 10:::INT8) ASC),
   INVERTED INDEX src_expr_expr1_idx ((a + b) ASC, (j->'a':::STRING)),
-  VECTOR INDEX src_b_expr_idx (b ASC, (v::VECTOR(3)) vector_l2_ops)
+  VECTOR INDEX src_b_expr_idx (b, (v::VECTOR(3)) vector_l2_ops)
 ) WITH (schema_locked = true);
 
 # Inaccessible expression index columns should be copied if the indexes are
@@ -689,7 +689,7 @@ CREATE TABLE public.copy_all (
   INDEX named_idx ((a + 1:::INT8) ASC),
   UNIQUE INDEX src_expr_key ((a + 10:::INT8) ASC),
   INVERTED INDEX src_expr_expr1_idx ((a + b) ASC, (j->'a':::STRING)),
-  VECTOR INDEX src_b_expr_idx (b ASC, (v::VECTOR(3)) vector_l2_ops)
+  VECTOR INDEX src_b_expr_idx (b, (v::VECTOR(3)) vector_l2_ops)
 );
 
 skipif config schema-locked-disabled
@@ -708,7 +708,7 @@ CREATE TABLE public.copy_all (
   INDEX named_idx ((a + 1:::INT8) ASC),
   UNIQUE INDEX src_expr_key ((a + 10:::INT8) ASC),
   INVERTED INDEX src_expr_expr1_idx ((a + b) ASC, (j->'a':::STRING)),
-  VECTOR INDEX src_b_expr_idx (b ASC, (v::VECTOR(3)) vector_l2_ops)
+  VECTOR INDEX src_b_expr_idx (b, (v::VECTOR(3)) vector_l2_ops)
 ) WITH (schema_locked = true);
 
 # Inaccessible expression index columns should be copied if the indexes are