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

Author: DrewKimball

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)

pkg/sql/opt/memo/testdata/stats/vector-search

@@ -107,17 +107,20 @@ top-k
       ├── stats: [rows=2]
       ├── key: (1)
       ├── fd: (1)-->(2), (2)-->(5)
-      ├── index-join t
+      ├── inner-join (lookup t)
       │    ├── columns: x:1(int!null) v:2(vector)
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── stats: [rows=2]
       │    ├── key: (1)
       │    ├── fd: (1)-->(2)
-      │    └── vector-search t@t_v_idx,vector
-      │         ├── columns: x:1(int!null)
-      │         ├── target nearest neighbors: 1
-      │         ├── stats: [rows=2]
-      │         ├── key: (1)
-      │         └── '[1,2,3]' [type=vector]
+      │    ├── vector-search t@t_v_idx,vector
+      │    │    ├── columns: x:1(int!null)
+      │    │    ├── target nearest neighbors: 1
+      │    │    ├── stats: [rows=2]
+      │    │    ├── key: (1)
+      │    │    └── '[1,2,3]' [type=vector]
+      │    └── filters (true)
       └── projections
            └── v:2 <-> '[1,2,3]' [as=column5:5, type=float, outer=(2), immutable]
 
@@ -140,17 +143,20 @@ top-k
       ├── stats: [rows=10]
       ├── key: (1)
       ├── fd: (1)-->(2), (2)-->(5)
-      ├── index-join t
+      ├── inner-join (lookup t)
       │    ├── columns: x:1(int!null) v:2(vector)
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── stats: [rows=10]
       │    ├── key: (1)
       │    ├── fd: (1)-->(2)
-      │    └── vector-search t@t_v_idx,vector
-      │         ├── columns: x:1(int!null)
-      │         ├── target nearest neighbors: 5
-      │         ├── stats: [rows=10]
-      │         ├── key: (1)
-      │         └── '[1,2,3]' [type=vector]
+      │    ├── vector-search t@t_v_idx,vector
+      │    │    ├── columns: x:1(int!null)
+      │    │    ├── target nearest neighbors: 5
+      │    │    ├── stats: [rows=10]
+      │    │    ├── key: (1)
+      │    │    └── '[1,2,3]' [type=vector]
+      │    └── filters (true)
       └── projections
            └── v:2 <-> '[1,2,3]' [as=column5:5, type=float, outer=(2), immutable]
 
@@ -174,16 +180,19 @@ top-k
       ├── stats: [rows=1000]
       ├── key: (1)
       ├── fd: (1)-->(2), (2)-->(5)
-      ├── index-join t
+      ├── inner-join (lookup t)
       │    ├── columns: x:1(int!null) v:2(vector)
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── stats: [rows=1000]
       │    ├── key: (1)
       │    ├── fd: (1)-->(2)
-      │    └── vector-search t@t_v_idx,vector
-      │         ├── columns: x:1(int!null)
-      │         ├── target nearest neighbors: 10000
-      │         ├── stats: [rows=1000]
-      │         ├── key: (1)
-      │         └── '[1,2,3]' [type=vector]
+      │    ├── vector-search t@t_v_idx,vector
+      │    │    ├── columns: x:1(int!null)
+      │    │    ├── target nearest neighbors: 10000
+      │    │    ├── stats: [rows=1000]
+      │    │    ├── key: (1)
+      │    │    └── '[1,2,3]' [type=vector]
+      │    └── filters (true)
       └── projections
            └── v:2 <-> '[1,2,3]' [as=column5:5, type=float, outer=(2), immutable]

pkg/sql/opt/xform/limit_funcs.go

@@ -487,11 +487,24 @@ func (c *CustomFuncs) TryGenerateVectorSearch(
 				TargetNeighborCount: limitInt,
 			},
 		)
-		// Add an index join to get the rest of the columns. The index join is
-		// always necessary because the vector column is not projected by the
-		// VectorSearch operator.
-		indexJoinPrivate := memo.IndexJoinPrivate{Table: sp.Table, Cols: sp.Cols}
-		vectorSearch = c.e.f.ConstructIndexJoin(vectorSearch, &indexJoinPrivate)
+
+		// Add a lookup-join against the primary index to get the rest of the
+		// columns. The primary-key join is always necessary because the vector
+		// column is not projected by the VectorSearch operator.
+		primaryIndex := c.e.mem.Metadata().Table(sp.Table).Index(cat.PrimaryIndex)
+		lookupCols := make(opt.ColList, primaryIndex.KeyColumnCount())
+		for i := 0; i < primaryIndex.KeyColumnCount(); i++ {
+			lookupCols[i] = sp.Table.IndexColumnID(primaryIndex, i)
+		}
+		lookupPrivate := &memo.LookupJoinPrivate{
+			JoinType:              opt.InnerJoinOp,
+			Table:                 sp.Table,
+			Index:                 cat.PrimaryIndex,
+			KeyCols:               lookupCols,
+			Cols:                  sp.Cols,
+			LookupColsAreTableKey: true,
+		}
+		vectorSearch = c.e.f.ConstructLookupJoin(vectorSearch, nil /* on */, lookupPrivate)
 
 		// Project the distance column.
 		projections := memo.ProjectionsExpr{c.e.f.ConstructProjectionsItem(distanceExpr, distanceCol)}

pkg/sql/opt/xform/testdata/rules/limit

@@ -3505,15 +3505,18 @@ top-k
       ├── immutable
       ├── key: (1)
       ├── fd: (1)-->(2-11), (9)-->(15)
-      ├── index-join index_tab
+      ├── inner-join (lookup index_tab)
       │    ├── columns: id:1!null val:2 region:3 latitude:4 longitude:5 data1:6!null data2:7!null geom:8 vec1:9 vec2:10 vec3:11
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── key: (1)
       │    ├── fd: (1)-->(2-11)
-      │    └── vector-search index_tab@index_tab_vec1_idx,vector
-      │         ├── columns: id:1!null
-      │         ├── target nearest neighbors: 1
-      │         ├── key: (1)
-      │         └── '[3,1,2]'
+      │    ├── vector-search index_tab@index_tab_vec1_idx,vector
+      │    │    ├── columns: id:1!null
+      │    │    ├── target nearest neighbors: 1
+      │    │    ├── key: (1)
+      │    │    └── '[3,1,2]'
+      │    └── filters (true)
       └── projections
            └── vec1:9 <-> '[3,1,2]' [as=column15:15, outer=(9), immutable]
 
@@ -3535,15 +3538,18 @@ top-k
       ├── immutable
       ├── key: (1)
       ├── fd: (1)-->(2-11), (9)-->(15)
-      ├── index-join index_tab
+      ├── inner-join (lookup index_tab)
       │    ├── columns: id:1!null val:2 region:3 latitude:4 longitude:5 data1:6!null data2:7!null geom:8 vec1:9 vec2:10 vec3:11
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── key: (1)
       │    ├── fd: (1)-->(2-11)
-      │    └── vector-search index_tab@index_tab_vec1_idx,vector
-      │         ├── columns: id:1!null
-      │         ├── target nearest neighbors: 5
-      │         ├── key: (1)
-      │         └── '[3,1,2]'
+      │    ├── vector-search index_tab@index_tab_vec1_idx,vector
+      │    │    ├── columns: id:1!null
+      │    │    ├── target nearest neighbors: 5
+      │    │    ├── key: (1)
+      │    │    └── '[3,1,2]'
+      │    └── filters (true)
       └── projections
            └── vec1:9 <-> '[3,1,2]' [as=column15:15, outer=(9), immutable]
 
@@ -3565,15 +3571,18 @@ top-k
       ├── immutable
       ├── key: (1)
       ├── fd: (1)-->(2-11), (10)-->(15)
-      ├── index-join index_tab
+      ├── inner-join (lookup index_tab)
       │    ├── columns: id:1!null val:2 region:3 latitude:4 longitude:5 data1:6!null data2:7!null geom:8 vec1:9 vec2:10 vec3:11
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── key: (1)
       │    ├── fd: (1)-->(2-11)
-      │    └── vector-search index_tab@index_tab_vec2_idx,vector
-      │         ├── columns: id:1!null
-      │         ├── target nearest neighbors: 5
-      │         ├── key: (1)
-      │         └── '[3,1,2]'
+      │    ├── vector-search index_tab@index_tab_vec2_idx,vector
+      │    │    ├── columns: id:1!null
+      │    │    ├── target nearest neighbors: 5
+      │    │    ├── key: (1)
+      │    │    └── '[3,1,2]'
+      │    └── filters (true)
       └── projections
            └── vec2:10 <-> '[3,1,2]' [as=column15:15, outer=(10), immutable]
 
@@ -3595,15 +3604,18 @@ top-k
       ├── immutable
       ├── key: (1)
       ├── fd: (1)-->(9), (9)-->(15)
-      ├── index-join index_tab
+      ├── inner-join (lookup index_tab)
       │    ├── columns: id:1!null vec1:9
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── key: (1)
       │    ├── fd: (1)-->(9)
-      │    └── vector-search index_tab@index_tab_vec1_idx,vector
-      │         ├── columns: id:1!null
-      │         ├── target nearest neighbors: 5
-      │         ├── key: (1)
-      │         └── '[3,1,2]'
+      │    ├── vector-search index_tab@index_tab_vec1_idx,vector
+      │    │    ├── columns: id:1!null
+      │    │    ├── target nearest neighbors: 5
+      │    │    ├── key: (1)
+      │    │    └── '[3,1,2]'
+      │    └── filters (true)
       └── projections
            └── vec1:9 <-> '[3,1,2]' [as=column15:15, outer=(9), immutable]
 
@@ -3625,17 +3637,20 @@ top-k
       ├── immutable, has-placeholder
       ├── key: (1)
       ├── fd: (1)-->(2-11), (9)-->(15)
-      ├── index-join index_tab
+      ├── inner-join (lookup index_tab)
       │    ├── columns: id:1!null val:2 region:3 latitude:4 longitude:5 data1:6!null data2:7!null geom:8 vec1:9 vec2:10 vec3:11
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── has-placeholder
       │    ├── key: (1)
       │    ├── fd: (1)-->(2-11)
-      │    └── vector-search index_tab@index_tab_vec1_idx,vector
-      │         ├── columns: id:1!null
-      │         ├── target nearest neighbors: 5
-      │         ├── has-placeholder
-      │         ├── key: (1)
-      │         └── $1
+      │    ├── vector-search index_tab@index_tab_vec1_idx,vector
+      │    │    ├── columns: id:1!null
+      │    │    ├── target nearest neighbors: 5
+      │    │    ├── has-placeholder
+      │    │    ├── key: (1)
+      │    │    └── $1
+      │    └── filters (true)
       └── projections
            └── vec1:9 <-> $1 [as=column15:15, outer=(9), immutable]
 
@@ -3657,15 +3672,18 @@ top-k
       ├── immutable
       ├── key: (1)
       ├── fd: (1)-->(2-11), (10)-->(15)
-      ├── index-join index_tab
+      ├── inner-join (lookup index_tab)
       │    ├── columns: id:1!null val:2 region:3 latitude:4 longitude:5 data1:6!null data2:7!null geom:8 vec1:9 vec2:10 vec3:11
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
       │    ├── key: (1)
       │    ├── fd: (1)-->(2-11)
-      │    └── vector-search index_tab@index_tab_vec2_idx,vector
-      │         ├── columns: id:1!null
-      │         ├── target nearest neighbors: 500
-      │         ├── key: (1)
-      │         └── '[3,1,2]'
+      │    ├── vector-search index_tab@index_tab_vec2_idx,vector
+      │    │    ├── columns: id:1!null
+      │    │    ├── target nearest neighbors: 500
+      │    │    ├── key: (1)
+      │    │    └── '[3,1,2]'
+      │    └── filters (true)
       └── projections
            └── vec2:10 <-> '[3,1,2]' [as=column15:15, outer=(10), immutable]
 
@@ -3903,12 +3921,15 @@ top-k
  └── project
       ├── columns: column8:8 photo_id:2!null photo_url:3!null description:4
       ├── immutable
-      ├── index-join image_embeddings
+      ├── inner-join (lookup image_embeddings)
       │    ├── columns: photo_id:2!null photo_url:3!null description:4 embedding:5!null
-      │    └── vector-search image_embeddings@image_embeddings_embedding_idx,vector
-      │         ├── columns: id:1!null
-      │         ├── target nearest neighbors: 5
-      │         ├── key: (1)
-      │         └── '[0.016068,-0.033417,-0.020309,-0.031494,-0.014076,0.03653]'
+      │    ├── key columns: [1] = [1]
+      │    ├── lookup columns are key
+      │    ├── vector-search image_embeddings@image_embeddings_embedding_idx,vector
+      │    │    ├── columns: id:1!null
+      │    │    ├── target nearest neighbors: 5
+      │    │    ├── key: (1)
+      │    │    └── '[0.016068,-0.033417,-0.020309,-0.031494,-0.014076,0.03653]'
+      │    └── filters (true)
       └── projections
            └── embedding:5 <-> '[0.016068,-0.033417,-0.020309,-0.031494,-0.014076,0.03653]' [as=column8:8, outer=(5), immutable]