Add test cases for vector indexes on vectors.

Author: nicktobey

enginetest/queries/vector_index_queries.go

@@ -21,7 +21,7 @@ import (
 
 var VectorIndexQueries = []ScriptTest{
 	{
-		Name: "basic vector index",
+		Name: "basic JSON vector index",
 		SetUpScript: []string{
 			"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]');`,
@@ -117,19 +117,119 @@ var VectorIndexQueries = []ScriptTest{
 		},
 	},
 	{
-		Name: "vector index errors",
+		Name: "basic VECTOR vector index",
 		SetUpScript: []string{
-			"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 table vectors (id int primary key, v vector(2));",
+			`insert into vectors values (1, STRING_TO_VECTOR('[4.0,3.0]')), (2, STRING_TO_VECTOR('[0.0,0.0]')), (3, STRING_TO_VECTOR('[-1.0,1.0]')), (4, STRING_TO_VECTOR('[0.0,-2.0]'));`,
 			`create vector index v_idx on vectors(v);`,
 			`set @query_vec = '[0.0,0.0]';`,
 		},
 		Assertions: []ScriptTestAssertion{
 			{
-				Query: `create vector index v_idx2 on vectors(id);`,
+				Query: "show create table vectors",
 				Expected: []sql.Row{
-					{"vectors", "CREATE TABLE `vectors` (\n  `id` int NOT NULL,\n  `v` json,\n  PRIMARY KEY (`id`),\n  VECTOR KEY `v_idx` (`v`)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"},
+					{"vectors", "CREATE TABLE `vectors` (\n  `id` int NOT NULL,\n  `v` VECTOR(2),\n  PRIMARY KEY (`id`),\n  VECTOR KEY `v_idx` (`v`)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"},
+				},
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v) limit 4",
+				Expected: []sql.Row{
+					{2, floatsToBytes(0.0, 0.0)},
+					{3, floatsToBytes(-1.0, 1.0)},
+					{4, floatsToBytes(0.0, -2.0)},
+					{1, floatsToBytes(4.0, 3.0)},
+				},
+				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, floatsToBytes(0.0, 0.0)},
+					{3, floatsToBytes(-1.0, 1.0)},
+					{4, floatsToBytes(0.0, -2.0)},
+					{1, floatsToBytes(4.0, 3.0)},
+				},
+				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)",
+				Expected: []sql.Row{
+					{2, floatsToBytes(0.0, 0.0)},
+					{3, floatsToBytes(-1.0, 1.0)},
+					{4, floatsToBytes(0.0, -2.0)},
+					{1, floatsToBytes(4.0, 3.0)},
 				},
+				ExpectedIndexes: nil,
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE_L2_SQUARED('[0.0,-2.0]', v) limit 4",
+				Expected: []sql.Row{
+					{4, floatsToBytes(0.0, -2.0)},
+					{2, floatsToBytes(0.0, 0.0)},
+					{3, floatsToBytes(-1.0, 1.0)},
+					{1, floatsToBytes(4.0, 3.0)},
+				},
+				ExpectedIndexes: []string{"v_idx"},
+			},
+			{
+				// Ensure vector index is not used for range lookups.
+				Query: "select * from vectors order by v limit 4",
+				Expected: []sql.Row{
+					{2, floatsToBytes(0.0, 0.0)},
+					{4, floatsToBytes(0.0, -2.0)},
+					{1, floatsToBytes(4.0, 3.0)},
+					{3, floatsToBytes(-1.0, 1.0)},
+				},
+				ExpectedIndexes: []string{},
+			},
+			{
+				// Modify the index after creation.
+				Query: "insert into vectors values (5, STRING_TO_VECTOR('[1.0,0.0]'))",
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",
+				Expected: []sql.Row{
+					{2, floatsToBytes(0.0, 0.0)},
+					{5, floatsToBytes(1.0, 0.0)},
+					{3, floatsToBytes(-1.0, 1.0)},
+					{4, floatsToBytes(0.0, -2.0)},
+					{1, floatsToBytes(4.0, 3.0)},
+				},
+				ExpectedIndexes: []string{},
+			},
+		},
+	},
+	{
+		Name: "vector index errors",
+		SetUpScript: []string{
+			"create table vectors (id int primary key, j json, v vector(2));",
+			`insert into vectors values
+                        (1, '[4.0,3.0]', STRING_TO_VECTOR('[4.0,3.0]')),
+                        (2, '[0.0,0.0]', STRING_TO_VECTOR('[0.0,0.0]')),
+                        (3, '[-1.0,1.0]', STRING_TO_VECTOR('[-1.0,1.0]')),
+                        (4, '[0.0,-2.0]', STRING_TO_VECTOR('[0.0,-2.0]'));`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:          `create vector index v_idx2 on vectors(j, v);`,
+				ExpectedErrStr: "a vector index must have exactly one column",
+			},
+			{
+				Query:          `create vector index v_idx2 on vectors(id);`,
+				ExpectedErrStr: "a vector index colum must be a vector or JSON",
 			},
 		},
 	},