diff --git a/enginetest/queries/index_query_plans.go b/enginetest/queries/index_query_plans.go index 07e472ac1b..d1240738b8 100644 --- a/enginetest/queries/index_query_plans.go +++ b/enginetest/queries/index_query_plans.go @@ -16287,4 +16287,116 @@ var IndexPlanTests = []QueryPlanTest{ " └─ columns: [pk v1 v2]\n" + "", }, + { + Query: `select * from comp_vector_index_t0 order by vec_distance('[50,50]', v2) limit 5`, + ExpectedPlan: "Limit(5)\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED('[50,50]', comp_vector_index_t0.v2) LIMIT 5 (bigint)\n" + + " ├─ colSet: (1-3)\n" + + " ├─ tableId: 1\n" + + " └─ Table\n" + + " ├─ name: comp_vector_index_t0\n" + + " └─ columns: [pk v1 v2]\n" + + "", + ExpectedEstimates: "Limit(5)\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED('[50,50]', comp_vector_index_t0.v2) LIMIT 5 (bigint)\n" + + " └─ columns: [pk v1 v2]\n" + + "", + ExpectedAnalysis: "Limit(5)\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED('[50,50]', comp_vector_index_t0.v2) LIMIT 5 (bigint)\n" + + " └─ columns: [pk v1 v2]\n" + + "", + }, + { + Query: `select * from comp_vector_index_t0 order by vec_distance(v2, '[50,50]') limit 5`, + ExpectedPlan: "Limit(5)\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED(comp_vector_index_t0.v2, '[50,50]') LIMIT 5 (bigint)\n" + + " ├─ colSet: (1-3)\n" + + " ├─ tableId: 1\n" + + " └─ Table\n" + + " ├─ name: comp_vector_index_t0\n" + + " └─ columns: [pk v1 v2]\n" + + "", + ExpectedEstimates: "Limit(5)\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED(comp_vector_index_t0.v2, '[50,50]') LIMIT 5 (bigint)\n" + + " └─ columns: [pk v1 v2]\n" + + "", + ExpectedAnalysis: "Limit(5)\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED(comp_vector_index_t0.v2, '[50,50]') LIMIT 5 (bigint)\n" + + " └─ columns: [pk v1 v2]\n" + + "", + }, + { + Query: `select pk+1 from comp_vector_index_t0 order by vec_distance('[50,50]', v2) limit 5`, + ExpectedPlan: "Limit(5)\n" + + " └─ Project\n" + + " ├─ columns: [(comp_vector_index_t0.pk:0!null + 1 (tinyint)) as pk+1]\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED('[50,50]', comp_vector_index_t0.v2) LIMIT 5 (bigint)\n" + + " ├─ colSet: (1-3)\n" + + " ├─ tableId: 1\n" + + " └─ Table\n" + + " ├─ name: comp_vector_index_t0\n" + + " └─ columns: [pk v2]\n" + + "", + ExpectedEstimates: "Limit(5)\n" + + " └─ Project\n" + + " ├─ columns: [(comp_vector_index_t0.pk + 1) as pk+1]\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED('[50,50]', comp_vector_index_t0.v2) LIMIT 5 (bigint)\n" + + " └─ columns: [pk v2]\n" + + "", + ExpectedAnalysis: "Limit(5)\n" + + " └─ Project\n" + + " ├─ columns: [(comp_vector_index_t0.pk + 1) as pk+1]\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED('[50,50]', comp_vector_index_t0.v2) LIMIT 5 (bigint)\n" + + " └─ columns: [pk v2]\n" + + "", + }, + { + Query: `select v1+1 from comp_vector_index_t0 order by vec_distance(v2, '[50,50]') limit 5`, + ExpectedPlan: "Limit(5)\n" + + " └─ Project\n" + + " ├─ columns: [(comp_vector_index_t0.v1:0 + 1 (tinyint)) as v1+1]\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED(comp_vector_index_t0.v2, '[50,50]') LIMIT 5 (bigint)\n" + + " ├─ colSet: (1-3)\n" + + " ├─ tableId: 1\n" + + " └─ Table\n" + + " ├─ name: comp_vector_index_t0\n" + + " └─ columns: [v1 v2]\n" + + "", + ExpectedEstimates: "Limit(5)\n" + + " └─ Project\n" + + " ├─ columns: [(comp_vector_index_t0.v1 + 1) as v1+1]\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED(comp_vector_index_t0.v2, '[50,50]') LIMIT 5 (bigint)\n" + + " └─ columns: [v1 v2]\n" + + "", + ExpectedAnalysis: "Limit(5)\n" + + " └─ Project\n" + + " ├─ columns: [(comp_vector_index_t0.v1 + 1) as v1+1]\n" + + " └─ IndexedTableAccess(comp_vector_index_t0)\n" + + " ├─ index: [comp_vector_index_t0.v2]\n" + + " ├─ order: VEC_DISTANCE_L2_SQUARED(comp_vector_index_t0.v2, '[50,50]') LIMIT 5 (bigint)\n" + + " └─ columns: [v1 v2]\n" + + "", + }, } diff --git a/enginetest/queries/vector_index_queries.go b/enginetest/queries/vector_index_queries.go index 2c740b696b..c83fac4303 100644 --- a/enginetest/queries/vector_index_queries.go +++ b/enginetest/queries/vector_index_queries.go @@ -56,6 +56,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)", @@ -88,6 +99,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{}, + }, }, }, } diff --git a/enginetest/scriptgen/setup/scripts/comp_index_tables b/enginetest/scriptgen/setup/scripts/comp_index_tables index 3b32457b4a..e67787e6d7 100644 --- a/enginetest/scriptgen/setup/scripts/comp_index_tables +++ b/enginetest/scriptgen/setup/scripts/comp_index_tables @@ -95,4 +95,28 @@ create table pref_index_t3 (v1 varchar(10), v2 varchar(10), unique index (v1(3), exec create table pref_index_t4 (i int primary key, v1 varchar(10), v2 varchar(10), unique index (v1(3),v2(5))); +---- + +exec +CREATE TABLE comp_vector_index_t0 (pk BIGINT PRIMARY KEY, v1 BIGINT, v2 JSON); +---- + +exec +INSERT INTO comp_vector_index_t0 VALUES (0,0,"[3,16]"),(1,2,"[65,9]"),(2,3,"[38,37]"),(3,3,"[99,99]"),(4,5,"[17,42]"),(5,6,"[6,76]"),(6,6,"[81,33]"), +(7,7,"[33,51]"),(8,7,"[37,42]"),(9,8,"[9,21]"),(10,8,"[37,90]"),(11,9,"[39,20]"),(12,9,"[71,82]"),(13,10,"[16,21]"),(14,10,"[32,46]"),(15,10,"[47,36]"), +(16,12,"[44,84]"),(17,12,"[66,40]"),(18,13,"[47,30]"),(19,13,"[56,41]"),(20,14,"[38,24]"),(21,14,"[91,1]"),(22,15,"[2,69]"),(23,16,"[40,36]"), +(24,20,"[29,93]"),(25,21,"[9,89]"),(26,21,"[42,76]"),(27,23,"[13,53]"),(28,23,"[28,68]"),(29,23,"[28,90]"),(30,23,"[30,44]"),(31,24,"[20,8]"), +(32,25,"[49,88]"),(33,26,"[15,28]"),(34,27,"[35,12]"),(35,28,"[39,84]"),(36,29,"[7,38]"),(37,29,"[21,74]"),(38,29,"[27,48]"),(39,29,"[77,46]"), +(40,31,"[47,21]"),(41,31,"[47,91]"),(42,32,"[40,76]"),(43,33,"[70,50]"),(44,34,"[27,58]"),(45,35,"[32,36]"),(46,36,"[4,36]"),(47,36,"[84,75]"), +(48,37,"[27,32]"),(49,38,"[88,68]"),(50,41,"[17,68]"),(51,41,"[77,26]"),(52,42,"[80,85]"),(53,45,"[1,57]"),(54,46,"[58,8]"),(55,49,"[26,11]"), +(56,50,"[49,20]"),(57,50,"[86,6]"),(58,54,"[13,78]"),(59,54,"[57,83]"),(60,55,"[45,46]"),(61,55,"[81,80]"),(62,56,"[0,97]"),(63,56,"[8,78]"), +(64,56,"[58,4]"),(65,56,"[66,33]"),(66,57,"[7,52]"),(67,59,"[77,53]"),(68,60,"[8,70]"),(69,61,"[11,25]"),(70,63,"[85,23]"),(71,65,"[17,9]"), +(72,66,"[46,46]"),(73,66,"[73,4]"),(74,67,"[55,27]"),(75,70,"[8,54]"),(76,70,"[58,33]"),(77,71,"[39,15]"),(78,72,"[65,64]"),(79,74,"[78,26]"), +(80,75,"[91,35]"),(81,76,"[40,52]"),(82,76,"[44,87]"),(83,81,"[32,4]"),(84,82,"[11,6]"),(85,82,"[46,32]"),(86,84,"[40,8]"),(87,84,"[93,37]"), +(88,85,"[53,50]"),(89,86,"[63,79]"),(90,87,"[22,34]"),(91,87,"[57,62]"),(92,88,"[88,42]"),(93,90,"[30,67]"),(94,91,"[15,15]"),(95,93,"[7,26]"), +(96,94,"[92,38]"),(97,95,"[89,66]"),(98,97,"[63,19]"),(99,98,"[31,21]"),(100,98,"[42,22]") +---- + +exec +create VECTOR INDEX v_idx on comp_vector_index_t0 (v2) ---- \ No newline at end of file diff --git a/enginetest/scriptgen/setup/setup_data.sg.go b/enginetest/scriptgen/setup/setup_data.sg.go index f7110f80bd..b986162c4c 100755 --- a/enginetest/scriptgen/setup/setup_data.sg.go +++ b/enginetest/scriptgen/setup/setup_data.sg.go @@ -103,6 +103,21 @@ var Comp_index_tablesData = []SetupScript{{ `create table pref_index_t1 (i int primary key, v1 text, v2 text, unique index (v1(3),v2(5)));`, `create table pref_index_t3 (v1 varchar(10), v2 varchar(10), unique index (v1(3),v2(5)));`, `create table pref_index_t4 (i int primary key, v1 varchar(10), v2 varchar(10), unique index (v1(3),v2(5)));`, + `CREATE TABLE comp_vector_index_t0 (pk BIGINT PRIMARY KEY, v1 BIGINT, v2 JSON);`, + `INSERT INTO comp_vector_index_t0 VALUES (0,0,"[3,16]"),(1,2,"[65,9]"),(2,3,"[38,37]"),(3,3,"[99,99]"),(4,5,"[17,42]"),(5,6,"[6,76]"),(6,6,"[81,33]"), +(7,7,"[33,51]"),(8,7,"[37,42]"),(9,8,"[9,21]"),(10,8,"[37,90]"),(11,9,"[39,20]"),(12,9,"[71,82]"),(13,10,"[16,21]"),(14,10,"[32,46]"),(15,10,"[47,36]"), +(16,12,"[44,84]"),(17,12,"[66,40]"),(18,13,"[47,30]"),(19,13,"[56,41]"),(20,14,"[38,24]"),(21,14,"[91,1]"),(22,15,"[2,69]"),(23,16,"[40,36]"), +(24,20,"[29,93]"),(25,21,"[9,89]"),(26,21,"[42,76]"),(27,23,"[13,53]"),(28,23,"[28,68]"),(29,23,"[28,90]"),(30,23,"[30,44]"),(31,24,"[20,8]"), +(32,25,"[49,88]"),(33,26,"[15,28]"),(34,27,"[35,12]"),(35,28,"[39,84]"),(36,29,"[7,38]"),(37,29,"[21,74]"),(38,29,"[27,48]"),(39,29,"[77,46]"), +(40,31,"[47,21]"),(41,31,"[47,91]"),(42,32,"[40,76]"),(43,33,"[70,50]"),(44,34,"[27,58]"),(45,35,"[32,36]"),(46,36,"[4,36]"),(47,36,"[84,75]"), +(48,37,"[27,32]"),(49,38,"[88,68]"),(50,41,"[17,68]"),(51,41,"[77,26]"),(52,42,"[80,85]"),(53,45,"[1,57]"),(54,46,"[58,8]"),(55,49,"[26,11]"), +(56,50,"[49,20]"),(57,50,"[86,6]"),(58,54,"[13,78]"),(59,54,"[57,83]"),(60,55,"[45,46]"),(61,55,"[81,80]"),(62,56,"[0,97]"),(63,56,"[8,78]"), +(64,56,"[58,4]"),(65,56,"[66,33]"),(66,57,"[7,52]"),(67,59,"[77,53]"),(68,60,"[8,70]"),(69,61,"[11,25]"),(70,63,"[85,23]"),(71,65,"[17,9]"), +(72,66,"[46,46]"),(73,66,"[73,4]"),(74,67,"[55,27]"),(75,70,"[8,54]"),(76,70,"[58,33]"),(77,71,"[39,15]"),(78,72,"[65,64]"),(79,74,"[78,26]"), +(80,75,"[91,35]"),(81,76,"[40,52]"),(82,76,"[44,87]"),(83,81,"[32,4]"),(84,82,"[11,6]"),(85,82,"[46,32]"),(86,84,"[40,8]"),(87,84,"[93,37]"), +(88,85,"[53,50]"),(89,86,"[63,79]"),(90,87,"[22,34]"),(91,87,"[57,62]"),(92,88,"[88,42]"),(93,90,"[30,67]"),(94,91,"[15,15]"),(95,93,"[7,26]"), +(96,94,"[92,38]"),(97,95,"[89,66]"),(98,97,"[63,19]"),(99,98,"[31,21]"),(100,98,"[42,22]")`, + `create VECTOR INDEX v_idx on comp_vector_index_t0 (v2)`, }} var DatetimetableData = []SetupScript{{ diff --git a/sql/analyzer/replace_order_by_distance.go b/sql/analyzer/replace_order_by_distance.go index bc80de1768..6434d181d7 100644 --- a/sql/analyzer/replace_order_by_distance.go +++ b/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) @@ -100,8 +105,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{}, @@ -125,7 +128,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 } diff --git a/sql/analyzer/vector_index_test.go b/sql/analyzer/vector_index_test.go index 2757e15860..a65d492249 100644 --- a/sql/analyzer/vector_index_test.go +++ b/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") }