Merge #143628

143628: sql: add support for prefix columns to vector search r=michae2,andy-kimball a=DrewKimball

#### opt: use lookup-join instead of index-join for vector search

Index joins expect all input rows to find a match in the primary index,
excluding locked rows when in skip locked mode. Vector indexes can have
"dangling" vectors after a delete operation fails to locate a vector in
the index. Accordingly, this commit changes the vector-search optimizer
rule to use a lookup-join instead of an index-join to avoid violating
assumptions made elsewhere in the optimizer.

Epic: CRDB-42943

Release note: None

#### opt: do not plan vector search for a locking scan

This commit disallows vector-search when the original scan is performing
locking. Even if we did support this, it is unclear which rows should be
locked, so it seems best to disallow it entirely.

Epic: CRDB-42943

Release note: None

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

#### opt: add support for prefix columns to vector search rule

This commit adds support for prefix columns to the `GenerateVectorSearch`
optimizer rule. Select filters over the matched scan as well as check
constraint filters can be used if they constrain every prefix column to
single-key spans, like the following examples:
```
prefix1 = 1 AND prefix2 = 2

(prefix1, prefix2) IN ((1, 2), (3, 4))
```
Currently, additional filters beyond those used to constrain the
prefix are not yet supported. This will have to wait until the vector
index is able to stream results in approximate order of distance.

Fixes #143206

Release note: None

Co-authored-by: Drew Kimball <[email protected]>

Author: craig[bot]

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/execbuilder/testdata/not_visible_index

@@ -576,8 +576,10 @@ vectorized: false
 │
 └── • render
     │
-    └── • index join
+    └── • lookup join
         │ table: t1@t1_pkey
+        │ equality: (rowid) = (rowid)
+        │ equality cols are key
         │
         └── • vector search
               table: t1@idx_vec_visible

pkg/sql/opt/exec/execbuilder/testdata/vector_search

@@ -32,8 +32,10 @@ vectorized: false
 │
 └── • render
     │
-    └── • index join
+    └── • lookup join
         │ table: t@t_pkey
+        │ equality: (k) = (k)
+        │ equality cols are key
         │
         └── • vector search
               table: t@t_v_idx
@@ -60,11 +62,12 @@ vectorized: false
         │ render k: k
         │ render v: v
         │
-        └── • index join
+        └── • lookup join (inner)
             │ columns: (k, v)
             │ estimated row count: 2 (missing stats)
             │ table: t@t_pkey
-            │ key columns: k
+            │ equality: (k) = (k)
+            │ equality cols are key
             │
             └── • vector search
                   columns: (k)
@@ -94,11 +97,12 @@ vectorized: false
         │ render k: k
         │ render v: v
         │
-        └── • index join
+        └── • lookup join (inner)
             │ columns: (k, v)
             │ estimated row count: 10 (missing stats)
             │ table: t@t_pkey
-            │ key columns: k
+            │ equality: (k) = (k)
+            │ equality cols are key
             │
             └── • vector search
                   columns: (k)
@@ -107,12 +111,11 @@ vectorized: false
                   target count: 5
                   query vector: '[1,2,3]'
 
-# NOTE: the vector-search rule doesn't yet handle prefix columns.
 query T
 EXPLAIN (VERBOSE) SELECT * FROM t_multi WHERE a = 1 AND b = 2 ORDER BY v <-> '[1, 2, 3]' LIMIT 1;
 ----
 distribution: local
-vectorized: true
+vectorized: false
 ·
 • project
 │ columns: (x, y, a, b, c, v)
@@ -133,16 +136,60 @@ vectorized: true
         │ render c: c
         │ render v: v
         │
-        └── • filter
+        └── • lookup join (inner)
             │ columns: (x, y, a, b, c, v)
-            │ estimated row count: 1 (missing stats)
-            │ filter: (a = 1) AND (b = 2)
+            │ estimated row count: 2 (missing stats)
+            │ table: t_multi@t_multi_pkey
+            │ equality: (x, y) = (x, y)
+            │ equality cols are key
             │
-            └── • scan
-                  columns: (x, y, a, b, c, v)
-                  estimated row count: 1,000 (missing stats)
-                  table: t_multi@t_multi_pkey
-                  spans: FULL SCAN
+            └── • vector search
+                  columns: (x, y)
+                  estimated row count: 2 (missing stats)
+                  table: t_multi@t_multi_a_b_v_idx
+                  target count: 1
+                  prefix spans: /1/2-/1/3
+                  query vector: '[1,2,3]'
+
+query T
+EXPLAIN (VERBOSE) SELECT * FROM t_multi WHERE (a, b) IN ((1, 2), (3, 4), (5, 6)) ORDER BY v <-> '[1, 2, 3]' LIMIT 1;
+----
+distribution: local
+vectorized: false
+·
+• project
+│ columns: (x, y, a, b, c, v)
+│
+└── • top-k
+    │ columns: (column11, x, y, a, b, c, v)
+    │ estimated row count: 1 (missing stats)
+    │ order: +column11
+    │ k: 1
+    │
+    └── • render
+        │ columns: (column11, x, y, a, b, c, v)
+        │ render column11: v <-> '[1,2,3]'
+        │ render x: x
+        │ render y: y
+        │ render a: a
+        │ render b: b
+        │ render c: c
+        │ render v: v
+        │
+        └── • lookup join (inner)
+            │ columns: (x, y, a, b, c, v)
+            │ estimated row count: 2 (missing stats)
+            │ table: t_multi@t_multi_pkey
+            │ equality: (x, y) = (x, y)
+            │ equality cols are key
+            │
+            └── • vector search
+                  columns: (x, y)
+                  estimated row count: 2 (missing stats)
+                  table: t_multi@t_multi_a_b_v_idx
+                  target count: 1
+                  prefix spans: /1/2-/1/3 /3/4-/3/5 /5/6-/5/7
+                  query vector: '[1,2,3]'
 
 # ==============================================================================
 # Vector Mutation Search Tests

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

@@ -944,10 +944,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/idxconstraint/BUILD.bazel

@@ -7,6 +7,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/sql/opt",
+        "//pkg/sql/opt/cat",
         "//pkg/sql/opt/constraint",
         "//pkg/sql/opt/memo",
         "//pkg/sql/opt/norm",