opt: handle locking in the vector search rule

This commit updates `GenerateVectorSearch` so that the lookup-join has the
same locking clause as the originally matched scan. This ensures that any
locking/nowait behavior is preserved, and makes the rule's behavior match
those for other index types more closely.

Epic: CRDB-42943

Release note: None

Author: DrewKimball

pkg/sql/opt/xform/limit_funcs.go

@@ -503,6 +503,7 @@ func (c *CustomFuncs) TryGenerateVectorSearch(
 			KeyCols:               lookupCols,
 			Cols:                  sp.Cols,
 			LookupColsAreTableKey: true,
+			Locking:               sp.Locking,
 		}
 		vectorSearch = c.e.f.ConstructLookupJoin(vectorSearch, nil /* on */, lookupPrivate)
 

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

@@ -3654,6 +3654,118 @@ top-k
       └── projections
            └── vec1:9 <-> $1 [as=column15:15, outer=(9), immutable]
 
+# Locking should transfer from the scan to the lookup-join.
+opt expect=GenerateVectorSearch
+SELECT * FROM index_tab ORDER BY vec1 <-> '[3,1,2]' LIMIT 5 FOR UPDATE;
+----
+top-k
+ ├── 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  [hidden: column15:15]
+ ├── internal-ordering: +15
+ ├── k: 5
+ ├── cardinality: [0 - 5]
+ ├── volatile
+ ├── key: (1)
+ ├── fd: (1)-->(2-11), (9)-->(15)
+ ├── ordering: +15
+ └── project
+      ├── columns: column15:15 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
+      ├── volatile
+      ├── key: (1)
+      ├── fd: (1)-->(2-11), (9)-->(15)
+      ├── 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
+      │    ├── locking: for-update
+      │    ├── volatile
+      │    ├── 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]'
+      │    └── filters (true)
+      └── projections
+           └── vec1:9 <-> '[3,1,2]' [as=column15:15, outer=(9), immutable]
+
+opt expect=GenerateVectorSearch
+SELECT * FROM index_tab ORDER BY vec1 <-> '[3,1,2]' LIMIT 5 FOR SHARE SKIP LOCKED;
+----
+top-k
+ ├── 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  [hidden: column15:15]
+ ├── internal-ordering: +15
+ ├── k: 5
+ ├── cardinality: [0 - 5]
+ ├── volatile
+ ├── key: (1)
+ ├── fd: (1)-->(2-11), (9)-->(15)
+ ├── ordering: +15
+ └── project
+      ├── columns: column15:15 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
+      ├── volatile
+      ├── key: (1)
+      ├── fd: (1)-->(2-11), (9)-->(15)
+      ├── 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
+      │    ├── locking: for-share,skip-locked
+      │    ├── volatile
+      │    ├── 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]'
+      │    └── filters (true)
+      └── projections
+           └── vec1:9 <-> '[3,1,2]' [as=column15:15, outer=(9), immutable]
+
+# The lock operator handles locking, so the lookup-join is non-locking.
+opt expect=GenerateVectorSearch set=optimizer_use_lock_op_for_serializable=true
+SELECT * FROM index_tab ORDER BY vec1 <-> '[3,1,2]' LIMIT 5 FOR SHARE;
+----
+lock 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  [hidden: column15:15]
+ ├── key columns: id:1
+ ├── lock columns: (16-26)
+ ├── locking: for-share
+ ├── cardinality: [0 - 5]
+ ├── volatile
+ ├── key: (1)
+ ├── fd: (1)-->(2-11), (9)-->(15)
+ ├── ordering: +15
+ └── top-k
+      ├── 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 column15:15
+      ├── internal-ordering: +15
+      ├── k: 5
+      ├── cardinality: [0 - 5]
+      ├── immutable
+      ├── key: (1)
+      ├── fd: (1)-->(2-11), (9)-->(15)
+      ├── ordering: +15
+      └── project
+           ├── columns: column15:15 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
+           ├── immutable
+           ├── key: (1)
+           ├── fd: (1)-->(2-11), (9)-->(15)
+           ├── 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]'
+           │    └── filters (true)
+           └── projections
+                └── vec1:9 <-> '[3,1,2]' [as=column15:15, outer=(9), immutable]
+
 # The index can be forced.
 opt expect=GenerateVectorSearch
 SELECT * FROM index_tab@index_tab_vec2_idx ORDER BY vec2 <-> '[3,1,2]' LIMIT 500;