From bbd3a63b3ff4becbd8a202a9a5bcc8f58197273c Mon Sep 17 00:00:00 2001
From: Nick Tobey <nick@dolthub.com>
Date: Wed, 15 Jan 2025 01:35:08 -0800
Subject: [PATCH 1/2] Allow using vector index with user variable.

---
 sql/analyzer/replace_order_by_distance.go | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/sql/analyzer/replace_order_by_distance.go b/sql/analyzer/replace_order_by_distance.go
index e1bfc5b2d1..6ac58fa0c0 100644
--- a/sql/analyzer/replace_order_by_distance.go
+++ b/sql/analyzer/replace_order_by_distance.go
@@ -59,15 +59,22 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 		if !isDistance {
 			return n, transform.SameTree, nil
 		}
+
+		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 {

From 1354b88746f1e1c55b062d97293b08812f8d11bb Mon Sep 17 00:00:00 2001
From: Nick Tobey <nick@dolthub.com>
Date: Wed, 15 Jan 2025 14:20:53 -0800
Subject: [PATCH 2/2] Add comment and test for user var vector index query.

---
 enginetest/queries/vector_index_queries.go | 12 ++++++++++++
 sql/analyzer/replace_order_by_distance.go  |  2 ++
 2 files changed, 14 insertions(+)

diff --git a/enginetest/queries/vector_index_queries.go b/enginetest/queries/vector_index_queries.go
index fd1b17a8d9..2c740b696b 100644
--- a/enginetest/queries/vector_index_queries.go
+++ b/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)",
diff --git a/sql/analyzer/replace_order_by_distance.go b/sql/analyzer/replace_order_by_distance.go
index 6ac58fa0c0..bc80de1768 100644
--- a/sql/analyzer/replace_order_by_distance.go
+++ b/sql/analyzer/replace_order_by_distance.go
@@ -60,6 +60,8 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 			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: