sql: add execution support for vector search with multiple prefix keys

This commit changes the vector-search operator to propagate a constraint
set rather than a single list of values to be used in constraining index
prefix columns. Execution has been modified as well to handle multiple
prefixes. This change will allow the vector-search optimizer rule to
handle cases where prefix columns can take on multiple possible values.
Note that modifying the vector search rule to take advantage of this
change is left for a following commit.

Epic: CRDB-42943

Release note: None

Author: DrewKimball

pkg/sql/distsql_physical_planner.go

@@ -4569,7 +4569,7 @@ func (dsp *DistSQLPlanner) planVectorSearch(
 
 	colTypes := getTypesFromResultColumns(planInfo.columns)
 	spec := &execinfrapb.VectorSearchSpec{
-		PrefixKey:           planInfo.prefixKey,
+		PrefixKeys:          planInfo.prefixKeys,
 		QueryVector:         queryVector,
 		TargetNeighborCount: planInfo.targetNeighborCount,
 	}

pkg/sql/distsql_spec_exec_factory.go

@@ -1452,7 +1452,7 @@ func (e *distSQLSpecExecFactory) ConstructVectorSearch(
 	table cat.Table,
 	index cat.Index,
 	outCols exec.TableColumnOrdinalSet,
-	prefixKey constraint.Key,
+	prefixConstraint *constraint.Constraint,
 	queryVector tree.TypedExpr,
 	targetNeighborCount uint64,
 ) (exec.Node, error) {
@@ -1461,17 +1461,17 @@ func (e *distSQLSpecExecFactory) ConstructVectorSearch(
 	cols := makeColList(table, outCols)
 	resultCols := colinfo.ResultColumnsFromColumns(tabDesc.GetID(), cols)
 
-	// Encode the prefix values as a roachpb.Key.
+	// Encode the prefix constraint as a list of roachpb.Keys.
 	var sb span.Builder
 	sb.Init(e.planner.EvalContext(), e.planner.ExecCfg().Codec, tabDesc, indexDesc)
-	encPrefixKey, _, err := sb.EncodeConstraintKey(prefixKey)
+	prefixKeys, err := sb.KeysFromVectorPrefixConstraint(e.ctx, prefixConstraint)
 	if err != nil {
 		return nil, err
 	}
 	planInfo := &vectorSearchPlanningInfo{
 		table:               tabDesc,
 		index:               indexDesc,
-		prefixKey:           encPrefixKey,
+		prefixKeys:          prefixKeys,
 		queryVector:         queryVector,
 		targetNeighborCount: targetNeighborCount,
 		cols:                cols,

pkg/sql/execinfrapb/flow_diagram.go

@@ -103,9 +103,15 @@ func (v *VectorSearchSpec) summary() (string, []string) {
 		fmt.Sprintf("Nearest Neighbor Target Count: %d", v.TargetNeighborCount),
 		fmt.Sprintf("Query Vector: %s", vector.T(v.QueryVector).String()),
 	}
-	if len(v.PrefixKey) > 0 {
-		vals, _ := encoding.PrettyPrintValuesWithTypes(nil /* valDirs */, v.PrefixKey)
-		details = append(details, fmt.Sprintf("Prefix Vals: %s", strings.Join(vals, "/")))
+	if len(v.PrefixKeys) > 0 {
+		// Only show the first prefix key.
+		var spanStr strings.Builder
+		vals, _ := encoding.PrettyPrintValuesWithTypes(nil /* valDirs */, v.PrefixKeys[0])
+		spanStr.WriteString(fmt.Sprintf("Prefix Vals: %s", strings.Join(vals, "/")))
+		if len(v.PrefixKeys) > 1 {
+			spanStr.WriteString(fmt.Sprintf(" and %d more", len(v.PrefixKeys)-1))
+		}
+		details = append(details, spanStr.String())
 	}
 	return "VectorSearch", details
 }

pkg/sql/execinfrapb/processors_sql.proto

@@ -1102,9 +1102,9 @@ message InsertSpec {
 message VectorSearchSpec {
   optional sqlbase.IndexFetchSpec fetch_spec = 1 [(gogoproto.nullable) = false];
 
-  // PrefixKey constrains the prefix index columns to a single value. It is
-  // empty for an index without prefix columns.
-  optional bytes prefix_key = 2 [(gogoproto.casttype) = "github.com/cockroachdb/cockroach/pkg/roachpb.Key"];
+  // PrefixKeys, if set, contains keys that each constrain every index prefix
+  // column to a single value. It is set IFF the index has prefix columns.
+  repeated bytes prefix_keys = 2 [(gogoproto.casttype) = "github.com/cockroachdb/cockroach/pkg/roachpb.Key"];
 
   // QueryVector is the vector to search for.
   repeated float query_vector = 3;

pkg/sql/opt/exec/execbuilder/relational.go

@@ -3949,25 +3949,6 @@ func (b *Builder) buildVectorSearch(
 				"vector search output column %d is not a primary key column", col)
 		}
 	}
-	// Evaluate the prefix expressions.
-	var prefixKey constraint.Key
-	if len(search.PrefixVals) > 0 {
-		values := make([]tree.Datum, len(search.PrefixVals))
-		for i, expr := range search.PrefixVals {
-			// The expression is either a placeholder or a constant.
-			if p, ok := expr.(*memo.PlaceholderExpr); ok {
-				val, err := eval.Expr(b.ctx, b.evalCtx, p.Value)
-				if err != nil {
-					return execPlan{}, colOrdMap{}, err
-				}
-				values[i] = val
-			} else {
-				values[i] = memo.ExtractConstDatum(expr)
-			}
-		}
-		prefixKey = constraint.MakeCompositeKey(values...)
-	}
-
 	outColOrds, outColMap := b.getColumns(search.Cols, search.Table)
 	ctx := buildScalarCtx{}
 	queryVector, err := b.buildScalar(&ctx, search.QueryVector)
@@ -3978,7 +3959,7 @@ func (b *Builder) buildVectorSearch(
 
 	var res execPlan
 	res.root, err = b.factory.ConstructVectorSearch(
-		table, index, outColOrds, prefixKey, queryVector, targetNeighborCount,
+		table, index, outColOrds, search.PrefixConstraint, queryVector, targetNeighborCount,
 	)
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err

pkg/sql/opt/exec/explain/emit.go

@@ -942,10 +942,15 @@ func (e *emitter) emitNodeAttributes(ctx context.Context, evalCtx *eval.Context,
 		a := n.args.(*vectorSearchArgs)
 		e.emitTableAndIndex("table", a.Table, a.Index, "" /* suffix */)
 		ob.Attr("target count", a.TargetNeighborCount)
-		if ob.flags.Verbose {
-			if !a.PrefixKey.IsEmpty() {
-				ob.Attr("prefix key", a.PrefixKey)
+		if a.PrefixConstraint != nil {
+			params := exec.ScanParams{
+				NeededCols:      a.OutCols,
+				IndexConstraint: a.PrefixConstraint,
 			}
+			e.emitSpans("prefix spans", a.Table, a.Index, params)
+		}
+		if ob.flags.Verbose {
+			// Vectors can have many dimensions, so don't print them unless verbose.
 			ob.Expr("query vector", a.QueryVector, nil /* varColumns */)
 		}
 

pkg/sql/opt/exec/factory.opt

@@ -860,7 +860,7 @@ define VectorSearch {
     Table cat.Table
     Index cat.Index
     OutCols exec.TableColumnOrdinalSet
-    PrefixKey constraint.Key
+    PrefixConstraint *constraint.Constraint
     QueryVector tree.TypedExpr
     TargetNeighborCount uint64
 }

pkg/sql/opt/memo/expr_format.go

@@ -747,6 +747,13 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 
 	case *VectorSearchExpr:
 		tp.Childf("target nearest neighbors: %d", t.TargetNeighborCount)
+		if t.PrefixConstraint != nil {
+			n := tp.Childf("prefix constraint: %s", t.PrefixConstraint.Columns.String())
+			for i := 0; i < t.PrefixConstraint.Spans.Count(); i++ {
+				spanString := t.PrefixConstraint.Spans.Get(i).String()
+				n.Child(cat.MaybeMarkRedactable(spanString, f.RedactableValues))
+			}
+		}
 
 	case *VectorMutationSearchExpr:
 		if t.IsIndexPut {

pkg/sql/opt/ops/relational.opt

@@ -1439,11 +1439,6 @@ define RecursiveCTEPrivate {
 # distances to the query vector, and applies the NN limit.
 [Relational, Telemetry]
 define VectorSearch {
-    # PrefixVals is set iff the vector index has a prefix of non-vector columns.
-    # It contains a list of constant values or placeholders that map 1-1 to the
-    # prefix columns.
-    PrefixVals ScalarListExpr
-
     # QueryVector is the scalar query vector. It is either a constant or a
     # placeholder.
     QueryVector ScalarExpr
@@ -1460,6 +1455,10 @@ define VectorSearchPrivate {
     # cat.Table.Index() method in order to fetch the cat.Index metadata.
     Index IndexOrdinal
 
+    # PrefixConstraint, if set, contains a set of spans that each constrain
+    # every index prefix column to a single value.
+    PrefixConstraint Constraint
+
     # Cols is the set of columns produced by the vector search operator. This
     # is currently always the set of primary key columns.
     Cols ColSet

pkg/sql/opt/xform/limit_funcs.go

@@ -469,7 +469,6 @@ func (c *CustomFuncs) TryGenerateVectorSearch(
 			// This index is for a different vector column.
 			return
 		}
-		var prefixVals memo.ScalarListExpr
 		if index.PrefixColumnCount() > 0 {
 			// TODO(drewk, mw5h): support multi-column vector indexes.
 			return
@@ -479,7 +478,7 @@ func (c *CustomFuncs) TryGenerateVectorSearch(
 		limitInt := int64(*limit.(*tree.DInt))
 		indexCols := c.PrimaryKeyCols(sp.Table)
 		vectorSearch := c.e.f.ConstructVectorSearch(
-			prefixVals, queryVector,
+			queryVector,
 			&memo.VectorSearchPrivate{
 				Table:               sp.Table,
 				Index:               index.Ordinal(),