Allow using vector index when the queried vector is provided in a user variable.

enginetest/queries/vector_index_queries.go

@@ -26,6 +26,7 @@ var VectorIndexQueries = []ScriptTest{
 			"create table vectors (id int primary key, v json);",
 			`insert into vectors values (1, '[4.0,3.0]'), (2, '[0.0,0.0]'), (3, '[-1.0,1.0]'), (4, '[0.0,-2.0]');`,
 			`create vector index v_idx on vectors(v);`,
+			`set @query_vec = '[0.0,0.0]';`,
 		},
 		Assertions: []ScriptTestAssertion{
 			{
@@ -44,6 +45,17 @@ var VectorIndexQueries = []ScriptTest{
 				},
 				ExpectedIndexes: []string{"v_idx"},
 			},
+			{
+				// Queries against a user var can be optimized.
+				Query: "select * from vectors order by VEC_DISTANCE(@query_vec, v) limit 4",
+				Expected: []sql.Row{
+					{2, types.MustJSON(`[0.0, 0.0]`)},
+					{3, types.MustJSON(`[-1.0, 1.0]`)},
+					{4, types.MustJSON(`[0.0, -2.0]`)},
+					{1, types.MustJSON(`[4.0, 3.0]`)},
+				},
+				ExpectedIndexes: []string{"v_idx"},
+			},
 			{
 				// Only queries with a limit can use a vector index.
 				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",

sql/analyzer/replace_order_by_distance.go

@@ -59,15 +59,24 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 		if !isDistance {
 			return n, transform.SameTree, nil
 		}
+
+		// We currently require that the query vector to the distance function is a constant value that does not
+		// depend on the row. Right now that can be a Literal or a UserVar.
+		isLiteral := func(expr sql.Expression) bool {
+			switch expr.(type) {
+			case *expression.Literal, *expression.UserVar:
+				return true
+			}
+			return false
+		}
+
 		var column sql.Expression
 		var literal sql.Expression
-		_, leftIsLiteral := distance.LeftChild.(*expression.Literal)
-		if leftIsLiteral {
+		if isLiteral(distance.LeftChild) {
 			column = distance.RightChild
 			literal = distance.LeftChild
 		} else {
-			_, rightIsLiteral := distance.RightChild.(*expression.Literal)
-			if rightIsLiteral {
+			if isLiteral(distance.RightChild) {
 				column = distance.LeftChild
 				literal = distance.RightChild
 			} else {