Use vector index when the SELECT cause has a projection.

enginetest/queries/vector_index_queries.go

@@ -44,6 +44,17 @@ var VectorIndexQueries = []ScriptTest{
 				},
 				ExpectedIndexes: []string{"v_idx"},
 			},
+			{
+				// Use the index even when there's a projection involved.
+				Query: "select `id`+1 from vectors order by VEC_DISTANCE('[0.0,0.0]', v) limit 4",
+				Expected: []sql.Row{
+					{3},
+					{4},
+					{5},
+					{2},
+				},
+				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)",
@@ -76,6 +87,21 @@ var VectorIndexQueries = []ScriptTest{
 				},
 				ExpectedIndexes: []string{},
 			},
+			{
+				// Modify the index after creation.
+				Query: "insert into vectors values (5, '[1.0,0.0]')",
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",
+				Expected: []sql.Row{
+					{2, types.MustJSON(`[0.0, 0.0]`)},
+					{5, types.MustJSON(`[1.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{},
+			},
 		},
 	},
 }

sql/analyzer/replace_order_by_distance.go

@@ -10,15 +10,20 @@ import (
 
 // replaceIdxSort applies an IndexAccess when there is an `OrderBy` over a prefix of any columns with Indexes
 func replaceIdxOrderByDistance(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope, sel RuleSelector, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
-	return replaceIdxOrderByDistanceHelper(ctx, scope, n, nil)
+	return replaceIdxOrderByDistanceHelper(ctx, scope, n, nil, nil)
 }
 
-func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node sql.Node, sortNode *plan.TopN) (sql.Node, transform.TreeIdentity, error) {
+func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node sql.Node, sortNode plan.Sortable, limit sql.Expression) (sql.Node, transform.TreeIdentity, error) {
 	switch n := node.(type) {
 	case *plan.TopN:
 		sortNode = n // lowest parent sort node
+		limit = n.Limit
+	case plan.Sortable:
+		sortNode = n
+	case *plan.Limit:
+		limit = n.Limit
 	case *plan.ResolvedTable:
-		if sortNode == nil {
+		if sortNode == nil || limit == nil {
 			return n, transform.SameTree, nil
 		}
 
@@ -44,7 +49,7 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 
 		// Column references have not been assigned their final indexes yet, so do that for the ORDER BY expression now.
 		// We can safely do this because an expression that references other tables won't pass `isSortFieldsValidPrefix` below.
-		sortNode = offsetAssignIndexes(sortNode).(*plan.TopN)
+		sortNode = offsetAssignIndexes(sortNode).(plan.Sortable)
 
 		sfExprs := normalizeExpressions(tableAliases, sortNode.GetSortFields().ToExpressions()...)
 		sfAliases := aliasedExpressionsInNode(sortNode)
@@ -91,8 +96,6 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 			return n, transform.SameTree, nil
 		}
 
-		limit := sortNode.Limit
-
 		lookup := sql.IndexLookup{
 			Index:  idx,
 			Ranges: sql.MySQLRangeCollection{},
@@ -116,7 +119,7 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 		same := transform.SameTree
 		switch c := child.(type) {
 		case *plan.Project, *plan.TableAlias, *plan.ResolvedTable, *plan.Filter, *plan.Limit, *plan.TopN, *plan.Offset, *plan.Sort, *plan.IndexedTableAccess:
-			newChildren[i], same, err = replaceIdxOrderByDistanceHelper(ctx, scope, child, sortNode)
+			newChildren[i], same, err = replaceIdxOrderByDistanceHelper(ctx, scope, child, sortNode, limit)
 		default:
 			newChildren[i] = c
 		}

sql/analyzer/vector_index_test.go

@@ -114,7 +114,7 @@ func TestVectorIndex(t *testing.T) {
 
 	for _, testCase := range vectorIndexTestCases(t, db, vectorIndexTable) {
 		t.Run(testCase.name, func(t *testing.T) {
-			res, same, err := replaceIdxOrderByDistanceHelper(nil, nil, testCase.inputPlan, nil)
+			res, same, err := replaceIdxOrderByDistanceHelper(nil, nil, testCase.inputPlan, nil, nil)
 			require.NoError(t, err)
 			require.Equal(t, testCase.usesVectorIndex, !bool(same))
 			res = offsetAssignIndexes(res)
@@ -218,7 +218,6 @@ func (i vectorIndexTable) SkipIndexCosting() bool {
 }
 
 func (i vectorIndexTable) IndexWithPrefix(ctx *sql.Context, expressions []string) (sql.Index, error) {
-	//TODO implement me
 	panic("implement me")
 }