Add integration tests for vector indexes.

Author: nicktobey

enginetest/enginetests.go

@@ -5870,6 +5870,18 @@ func TestIndexes(t *testing.T, h Harness) {
 	}
 }
 
+func TestVectorIndexes(t *testing.T, h Harness) {
+	for _, tt := range queries.VectorIndexQueries {
+		TestScript(t, h, tt)
+	}
+}
+
+func TestVectorFunctions(t *testing.T, h Harness) {
+	for _, tt := range queries.VectorFunctionQueries {
+		TestScript(t, h, tt)
+	}
+}
+
 func TestIndexPrefix(t *testing.T, h Harness) {
 	for _, tt := range queries.IndexPrefixQueries {
 		TestScript(t, h, tt)

enginetest/memory_engine_test.go

@@ -902,6 +902,14 @@ func TestIndexes(t *testing.T) {
 	enginetest.TestIndexes(t, enginetest.NewDefaultMemoryHarness())
 }
 
+func TestVectorIndexes(t *testing.T) {
+	enginetest.TestVectorIndexes(t, enginetest.NewDefaultMemoryHarness())
+}
+
+func TestVectorFunctions(t *testing.T) {
+	enginetest.TestVectorFunctions(t, enginetest.NewDefaultMemoryHarness())
+}
+
 func TestIndexPrefix(t *testing.T) {
 	enginetest.TestIndexPrefix(t, enginetest.NewDefaultMemoryHarness())
 }

enginetest/queries/vector_function_queries.go

@@ -0,0 +1,58 @@
+// Copyright 2024 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package queries
+
+import (
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/types"
+)
+
+var VectorFunctionQueries = []ScriptTest{
+	{
+		Name: "basic usage of VEC_DISTANCE without index",
+		SetUpScript: []string{
+			"create table vectors (id int primary key, v json);",
+			`insert into vectors values (1, '[3.0,4.0]'), (2, '[0.0,0.0]'), (3, '[1.0,-1.0]'), (4, '[-2.0,0.0]');`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "select VEC_DISTANCE('[10.0]', '[20.0]');",
+				Expected: []sql.Row{{100.0}},
+			},
+			{
+				Query:    "select VEC_DISTANCE_L2_SQUARED('[1.0, 2.0]', '[5.0, 5.0]');",
+				Expected: []sql.Row{{25.0}},
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",
+				Expected: []sql.Row{
+					{2, types.MustJSON(`[0.0, 0.0]`)},
+					{3, types.MustJSON(`[1.0, -1.0]`)},
+					{4, types.MustJSON(`[-2.0, 0.0]`)},
+					{1, types.MustJSON(`[3.0, 4.0]`)},
+				},
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE_L2_SQUARED('[-2.0,0.0]', v)",
+				Expected: []sql.Row{
+					{4, types.MustJSON(`[-2.0, 0.0]`)},
+					{2, types.MustJSON(`[0.0, 0.0]`)},
+					{3, types.MustJSON(`[1.0, -1.0]`)},
+					{1, types.MustJSON(`[3.0, 4.0]`)},
+				},
+			},
+		},
+	},
+}

enginetest/queries/vector_index_queries.go

@@ -0,0 +1,64 @@
+// Copyright 2024 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package queries
+
+import (
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/types"
+)
+
+var VectorIndexQueries = []ScriptTest{
+	{
+		Name: "basic vector index",
+		SetUpScript: []string{
+			"create table vectors (id int primary key, v json);",
+			`insert into vectors values (1, '[3.0,4.0]'), (2, '[0.0,0.0]'), (3, '[1.0,-1.0]'), (4, '[-2.0,0.0]');`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: `create vector index v_idx on vectors(v);`,
+				Expected: []sql.Row{
+					{types.OkResult{RowsAffected: 0}},
+				},
+			},
+			{
+				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"},
+				},
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",
+				Expected: []sql.Row{
+					{2, types.MustJSON(`[0.0, 0.0]`)},
+					{3, types.MustJSON(`[1.0, -1.0]`)},
+					{4, types.MustJSON(`[-2.0, 0.0]`)},
+					{1, types.MustJSON(`[3.0, 4.0]`)},
+				},
+				ExpectedIndexes: []string{"v_idx"},
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE_L2_SQUARED('[-2.0,0.0]', v)",
+				Expected: []sql.Row{
+					{4, types.MustJSON(`[-2.0, 0.0]`)},
+					{2, types.MustJSON(`[0.0, 0.0]`)},
+					{3, types.MustJSON(`[1.0, -1.0]`)},
+					{1, types.MustJSON(`[3.0, 4.0]`)},
+				},
+				ExpectedIndexes: []string{"v_idx"},
+			},
+		},
+	},
+}