Add initial vector type implementation and tests.

Author: nicktobey

enginetest/enginetests.go

@@ -6128,6 +6128,15 @@ func TestVectorFunctions(t *testing.T, h Harness) {
 	for _, tt := range queries.VectorFunctionQueries {
 		TestScript(t, h, tt)
 	}
+	for _, testCase := range queries.VectorFunctionTestCases {
+		TestScript(t, h, queries.MakeVectorFunctionTest(testCase))
+	}
+}
+
+func TestVectorType(t *testing.T, h Harness) {
+	for _, tt := range queries.VectorDDLQueries {
+		TestScript(t, h, tt)
+	}
 }
 
 func TestIndexPrefix(t *testing.T, h Harness) {

enginetest/memory_engine_test.go

@@ -922,6 +922,10 @@ func TestVectorFunctions(t *testing.T) {
 	enginetest.TestVectorFunctions(t, enginetest.NewDefaultMemoryHarness())
 }
 
+func TestVectorType(t *testing.T) {
+	enginetest.TestVectorType(t, enginetest.NewDefaultMemoryHarness())
+}
+
 func TestIndexPrefix(t *testing.T) {
 	enginetest.TestIndexPrefix(t, enginetest.NewDefaultMemoryHarness())
 }

enginetest/queries/create_table_queries.go

@@ -318,6 +318,12 @@ var CreateTableQueries = []WriteQueryTest{
 		SelectQuery:         `SHOW CREATE TABLE t1`,
 		ExpectedSelect:      []sql.Row{{"t1", "CREATE TABLE `t1` (\n  `pk` bit(2) DEFAULT b'10'\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"}},
 	},
+	{
+		WriteQuery:          `CREATE TABLE embeddings (id INT, vector_col VECTOR(128) NOT NULL, small_vec VECTOR(1))`,
+		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
+		SelectQuery:         `SHOW CREATE TABLE embeddings`,
+		ExpectedSelect:      []sql.Row{{"embeddings", "CREATE TABLE `embeddings` (\n  `id` int,\n  `vector_col` VECTOR(128) NOT NULL,\n  `small_vec` VECTOR(1)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"}},
+	},
 }
 
 var CreateTableScriptTests = []ScriptTest{

enginetest/queries/type_wire_queries.go

@@ -816,4 +816,21 @@ var TypeWireTests = []TypeWireTest{
 			{{"1", "[[\"a\",1]]"}, {"2", "[{\"key1\":\"value1\",\"key2\":\"value2\"}]"}},
 		},
 	},
+	{
+		Name: "VECTOR",
+		SetUpScript: []string{
+			`CREATE TABLE test (pk INT PRIMARY KEY, v1 VECTOR(2), v2 VECTOR(3));`,
+			`INSERT INTO test VALUES (1, VEC_FROMTEXT('[1.0, 2.0]', '[1.0, 2.0, 3.0]'), (2, '[4.0, 5.0]', '[4.0, 5.0, 6.0]'));`,
+		},
+		Queries: []string{
+			`SELECT * FROM test ORDER BY pk;`,
+			`SELECT v1, v2 FROM test ORDER BY pk;`,
+			`SELECT pk, v1 FROM test WHERE pk = 1;`,
+		},
+		Results: [][]sql.Row{
+			{{"1", "[1,2]", "[1,2,3]"}, {"2", "[4,5]", "[4,5,6]"}},
+			{{"[1,2]", "[1,2,3]"}, {"[4,5]", "[4,5,6]"}},
+			{{"1", "[1,2]"}},
+		},
+	},
 }

enginetest/queries/vector_ddl_queries.go

@@ -0,0 +1,149 @@
+// Copyright 2022 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/plan"
+	"github.com/dolthub/go-mysql-server/sql/types"
+)
+
+// VectorDDLQueries tests VECTOR type creation, insertion, and querying
+var VectorDDLQueries = []ScriptTest{
+	{
+		Name: "basic VECTOR type creation and manipulation",
+		SetUpScript: []string{
+			`CREATE TABLE test_vectors (
+				id INT PRIMARY KEY,
+				small_vec VECTOR(2),
+				medium_vec VECTOR(10),
+				large_vec VECTOR(1000)
+			)`,
+			`INSERT INTO test_vectors VALUES 
+				(1, STRING_TO_VECTOR('[1.0, 2.0]'), STRING_TO_VECTOR('[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]'), NULL),
+				(2, STRING_TO_VECTOR('[3.5, 4.5]'), NULL, NULL
+			)`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    `SHOW CREATE TABLE vectors`,
+				Expected: []sql.Row{{"vectors", "CREATE TABLE `vectors` (\n  `id` int NOT NULL,\n  `small_vec` VECTOR(2),\n  `medium_vec` VECTOR(10),\n  `large_vec` VECTOR(1000), \n PRIMARY KEY (`id`)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"}},
+			},
+			{
+				Query:    `SELECT id, small_vec, medium_vec FROM vectors WHERE id = 2`,
+				Expected: []sql.Row{{2, []float32{3.5, 4.5}, nil}},
+			},
+			{
+				Query:    `SELECT id, small_vec, medium_vec FROM vectors WHERE id = 1`,
+				Expected: []sql.Row{{1, []float32{1.0, 2.0}, []float32{1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0}}},
+			},
+			{
+				Query: `SELECT id, small_vec FROM test_vectors ORDER BY id`,
+				Expected: []sql.Row{
+					{1, []float32{1.0, 2.0}},
+					{2, []float32{3.5, 4.5}},
+				},
+			},
+			{
+				Query:    `UPDATE test_vectors SET small_vec = '[10.0, 20.0]' WHERE id = 1`,
+				Expected: []sql.Row{{types.OkResult{RowsAffected: 1, Info: plan.UpdateInfo{Matched: 1, Updated: 1}}}},
+			},
+			{
+				Query:    `SELECT small_vec FROM test_vectors WHERE id = 1`,
+				Expected: []sql.Row{{[]float32{10.0, 20.0}}},
+			},
+			{
+				Query: `INSERT INTO test_vectors VALUES 
+					(3, 0x00F0803F00004040, NULL, NULL)`,
+				Expected: []sql.Row{{types.NewOkResult(1)}},
+			},
+			{
+				Query:    `SELECT small_vec FROM test_vectors WHERE id = 3`,
+				Expected: []sql.Row{{[]float32{1.0, 2.0}}},
+			},
+		},
+	},
+	{
+		Name: "VECTOR type error conditions",
+		SetUpScript: []string{
+			`CREATE TABLE error_vectors (id INT PRIMARY KEY, vec3 VECTOR(3))`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:          `INSERT INTO error_vectors VALUES (1, '[1.0, 2.0]')`,
+				ExpectedErrStr: "",
+			},
+			{
+				Query:          `INSERT INTO error_vectors VALUES (1, STRING_TO_VECTOR('[1.0, 2.0]'))`,
+				ExpectedErrStr: "VECTOR dimension mismatch: expected 3, got 2",
+			},
+			{
+				Query:          `INSERT INTO error_vectors VALUES (2, STRING_TO_VECTOR('[1.0, 2.0, 3.0, 4.0]'))`,
+				ExpectedErrStr: "VECTOR dimension mismatch: expected 3, got 4",
+			},
+			{
+				Query:          `INSERT INTO error_vectors VALUES (3, STRING_TO_VECTOR('[1.0, invalid, 3.0]'))`,
+				ExpectedErrStr: "invalid VECTOR JSON format: invalid character 'i' looking for beginning of value",
+			},
+			{
+				Query:          `INSERT INTO error_vectors VALUES (4, STRING_TO_VECTOR('invalid_json'))`,
+				ExpectedErrStr: "VECTOR must be in JSON array format",
+			},
+			{
+				Query:          `INSERT INTO error_vectors VALUES (5, STRING_TO_VECTOR('[1.0, "not an array"]'))`,
+				ExpectedErrStr: "VECTOR must be in JSON array format",
+			},
+			{
+				Query:          `INSERT INTO error_vectors VALUES (5, STRING_TO_VECTOR('"not an array"'))`,
+				ExpectedErrStr: "VECTOR must be in JSON array format",
+			},
+			{
+				Query:       `CREATE TABLE error_vectors (id INT PRIMARY KEY, vec3 VECTOR(-3))`,
+				ExpectedErr: sql.ErrInvalidColTypeDefinition,
+			},
+			{
+				Query:       `CREATE TABLE error_vectors (id INT PRIMARY KEY, vec3 VECTOR(0))`,
+				ExpectedErr: sql.ErrInvalidColTypeDefinition,
+			},
+			{
+				Query:       `CREATE TABLE error_vectors (id INT PRIMARY KEY, vec3 VECTOR(17000))`,
+				ExpectedErr: sql.ErrInvalidColTypeDefinition,
+			},
+		},
+	},
+	{
+		Name: "VECTOR type with different data formats",
+		SetUpScript: []string{
+			`CREATE TABLE format_vectors (id INT PRIMARY KEY, vec2 VECTOR(2))`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: `INSERT INTO format_vectors VALUES 
+					(1, STRING_TO_VECTOR('[1.0, 2.0]')),
+					(2, STRING_TO_VECTOR('[3, 4]')),
+					(3, STRING_TO_VECTOR('[55e-1, 67e2]'))`,
+				Expected: []sql.Row{{types.NewOkResult(3)}},
+			},
+			{
+				Query: `SELECT id, vec2 FROM format_vectors ORDER BY id`,
+				Expected: []sql.Row{
+					{1, []float32{1.0, 2.0}},
+					{2, []float32{3.0, 4.0}},
+					{3, []float32{5.5, 6700}},
+				},
+			},
+		},
+	},
+}