Implement missing functions for vector types.

Author: nicktobey

enginetest/queries/type_wire_queries.go

@@ -820,7 +820,7 @@ var TypeWireTests = []TypeWireTest{
 		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]'));`,
+			`INSERT INTO test VALUES (1, VEC_FROMTEXT('[1.0, 2.0]'), VEC_FROMTEXT('[1.0, 2.0, 3.0]')), (2, VEC_FROMTEXT('[4.0, 5.0]'), VEC_FROMTEXT('[4.0, 5.0, 6.0]'));`,
 		},
 		Queries: []string{
 			`SELECT * FROM test ORDER BY pk;`,

enginetest/queries/vector_ddl_queries.go

@@ -38,15 +38,15 @@ var VectorDDLQueries = []ScriptTest{
 		},
 		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:    `SHOW CREATE TABLE test_vectors`,
+				Expected: []sql.Row{{"test_vectors", "CREATE TABLE `test_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`,
+				Query:    `SELECT id, small_vec, medium_vec FROM test_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`,
+				Query:    `SELECT id, small_vec, medium_vec FROM test_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}}},
 			},
 			{
@@ -57,21 +57,20 @@ var VectorDDLQueries = []ScriptTest{
 				},
 			},
 			{
-				Query:    `UPDATE test_vectors SET small_vec = '[10.0, 20.0]' WHERE id = 1`,
+				Query:    `UPDATE test_vectors SET small_vec = STRING_TO_VECTOR('[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)`,
+				Query:    `INSERT INTO test_vectors VALUES (3, 0x0000204100002041, NULL, NULL)`, // [10.0, 10.0]
 				Expected: []sql.Row{{types.NewOkResult(1)}},
 			},
 			{
 				Query:    `SELECT small_vec FROM test_vectors WHERE id = 3`,
-				Expected: []sql.Row{{[]float32{1.0, 2.0}}},
+				Expected: []sql.Row{{[]float32{10.0, 10.0}}},
 			},
 		},
 	},
@@ -83,7 +82,7 @@ var VectorDDLQueries = []ScriptTest{
 		Assertions: []ScriptTestAssertion{
 			{
 				Query:          `INSERT INTO error_vectors VALUES (1, '[1.0, 2.0]')`,
-				ExpectedErrStr: "",
+				ExpectedErrStr: "value of type string cannot be converted to 'vector' type",
 			},
 			{
 				Query:          `INSERT INTO error_vectors VALUES (1, STRING_TO_VECTOR('[1.0, 2.0]'))`,
@@ -95,23 +94,23 @@ var VectorDDLQueries = []ScriptTest{
 			},
 			{
 				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",
+				ExpectedErrStr: "can't convert JSON to vector: 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",
+				ExpectedErrStr: "can't convert JSON to vector: invalid character 'i' looking for beginning of value",
 			},
 			{
 				Query:          `INSERT INTO error_vectors VALUES (5, STRING_TO_VECTOR('[1.0, "not an array"]'))`,
-				ExpectedErrStr: "VECTOR must be in JSON array format",
+				ExpectedErrStr: "can't convert JSON to vector; expected array of floats, but array contained string",
 			},
 			{
 				Query:          `INSERT INTO error_vectors VALUES (5, STRING_TO_VECTOR('"not an array"'))`,
-				ExpectedErrStr: "VECTOR must be in JSON array format",
+				ExpectedErrStr: "can't convert JSON to vector; expected array, got string",
 			},
 			{
-				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(-3))`,
+				ExpectedErrStr: "syntax error at position 62 near 'VECTOR'",
 			},
 			{
 				Query:       `CREATE TABLE error_vectors (id INT PRIMARY KEY, vec3 VECTOR(0))`,

enginetest/queries/vector_function_queries.go

@@ -42,7 +42,7 @@ var VectorFunctionTestCases = []VectorFunctionTestCase{
 		hex:            "0000803F",
 		length:         4,
 		base64:         "AACAPw==",
-		md5:            "84b7e6bf8e3cd674011866c606a8011d",
+		md5:            "429d81ed2795e3c586906c6c335aa136",
 		sha1:           "5bb96baed2a67ef718989bf7de91433ca9b9f8cf",
 		sha2:           "e00e5eb9444182f352323374ef4e08ebcb784725fdd4fd612d7730540b3e0c8c",
 	},
@@ -53,7 +53,7 @@ var VectorFunctionTestCases = []VectorFunctionTestCase{
 		charLength:     8,
 		hex:            "0000004000004040",
 		length:         8,
-		base64:         "AAAAQAAAQEA",
+		base64:         "AAAAQAAAQEA=",
 		md5:            "f37b6e459e9e2d49261fe42d3a7bff07",
 		sha1:           "fd3352c0e141970e5b1c45d1755760d018cfe32d",
 		sha2:           "2fd848aa90e817e10e20985de4e8ac6a09b0fe70623d6b952e46800be6b025b9",
@@ -122,15 +122,15 @@ var VectorFunctionQueries = []ScriptTest{
 			},
 			{
 				Query:    `SELECT VECTOR_TO_STRING(STRING_TO_VECTOR("[1.0, 2.0]"));`,
-				Expected: []sql.Row{{"[1.0, 2.0]"}},
+				Expected: []sql.Row{{"[1, 2]"}},
 			},
 			{
 				Query:    `SELECT FROM_VECTOR(TO_VECTOR("[1.0, 2.0]"));`,
-				Expected: []sql.Row{{"[1.0, 2.0]"}},
+				Expected: []sql.Row{{"[1, 2]"}},
 			},
 			{
 				Query:    `SELECT VEC_ToText(VEC_FromText("[1.0, 2.0]"));`,
-				Expected: []sql.Row{{"[1.0, 2.0]"}},
+				Expected: []sql.Row{{"[1, 2]"}},
 			},
 		},
 	},
@@ -151,58 +151,19 @@ var VectorFunctionQueries = []ScriptTest{
 			},
 			{
 				Query:    "select VEC_DISTANCE_EUCLIDEAN('[1.0, 2.0]', '[5.0, 5.0]');",
-				Expected: []sql.Row{{float32(5.0)}},
+				Expected: []sql.Row{{5.0}},
 			},
 			{
 				Query:    `SELECT DISTANCE(STRING_TO_VECTOR("[0.0, 0.0]"), STRING_TO_VECTOR("[3.0, 4.0]"), "EUCLIDEAN");`,
-				Expected: []sql.Row{{float32(5.0)}},
+				Expected: []sql.Row{{5.0}},
 			},
 			{
 				Query:    "select VEC_DISTANCE_COSINE(STRING_TO_VECTOR('[0.0, 3.0]'), '[5.0, 5.0]');",
-				Expected: []sql.Row{{float32(15.0)}},
+				Expected: []sql.Row{{0.29289321881345254}},
 			},
 			{
-				Query:    `SELECT DISTANCE("[0.0, 3.0]", STRING_TO_VECTOR("[5.0, 5.0]"), "COSINE");`,
-				Expected: []sql.Row{{float32(5.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]`)},
-				},
-			},
-		},
-	},
-	{
-		Name: "test that existing functions accept vectors",
-		Assertions: []ScriptTestAssertion{
-			{
-				Query:    "select VEC_DISTANCE_EUCLIDEAN('[1.0, 2.0]', '[5.0, 5.0]');",
-				Expected: []sql.Row{{float32(5.0)}},
-			},
-			{
-				Query:    `SELECT DISTANCE(STRING_TO_VECTOR("[0.0, 0.0]"), STRING_TO_VECTOR("[3.0, 4.0]"), "EUCLIDEAN");`,
-				Expected: []sql.Row{{float32(5.0)}},
-			},
-			{
-				Query:    "select VEC_DISTANCE_COSINE(STRING_TO_VECTOR('[0.0, 3.0]'), '[5.0, 5.0]');",
-				Expected: []sql.Row{{float32(15.0)}},
-			},
-			{
-				Query:    `SELECT DISTANCE("[0.0, 3.0]", STRING_TO_VECTOR("[5.0, 5.0]"), "COSINE");`,
-				Expected: []sql.Row{{float32(5.0)}},
+				Query:    `SELECT DISTANCE("[1.0, 1.0]", STRING_TO_VECTOR("[-1.0, 1.0]"), "COSINE");`,
+				Expected: []sql.Row{{1.0}},
 			},
 			{
 				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",

sql/expression/function/length.go

@@ -123,6 +123,10 @@ func (l *Length) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
 		return int32(wrapper.MaxByteLength()), nil
 	}
 
+	// Getting the length of a vector doesn't require converting it to a string, it just returns the length in bytes
+	if vec, ok := val.([]float32); ok {
+		return int32(4 * len(vec)), nil
+	}
 	content, collation, err := types.ConvertToCollatedString(ctx, val, l.Child.Type())
 	if err != nil {
 		return nil, err

sql/expression/function/registry.go

@@ -327,6 +327,15 @@ var BuiltIns = []sql.Function{
 	sql.Function1{Name: "var_samp", Fn: func(e sql.Expression) sql.Expression { return aggregation.NewVarSamp(e) }},
 	sql.Function2{Name: "vec_distance", Fn: vector.NewL2SquaredDistance},
 	sql.Function2{Name: "vec_distance_l2_squared", Fn: vector.NewL2SquaredDistance},
+	sql.Function2{Name: "vec_distance_euclidean", Fn: vector.NewEuclideanDistance},
+	sql.Function2{Name: "vec_distance_cosine", Fn: vector.NewCosineDistance},
+	sql.FunctionN{Name: "distance", Fn: vector.NewGenericDistance},
+	sql.Function1{Name: "string_to_vector", Fn: vector.NewStringToVector},
+	sql.Function1{Name: "to_vector", Fn: vector.NewStringToVector},
+	sql.Function1{Name: "vec_fromtext", Fn: vector.NewStringToVector},
+	sql.Function1{Name: "vector_to_string", Fn: vector.NewVectorToString},
+	sql.Function1{Name: "from_vector", Fn: vector.NewVectorToString},
+	sql.Function1{Name: "vec_totext", Fn: vector.NewVectorToString},
 	sql.Function1{Name: "weekday", Fn: NewWeekday},
 	sql.Function1{Name: "weekofyear", Fn: NewWeekOfYear},
 	sql.Function1{Name: "year", Fn: NewYear},

sql/expression/function/string.go

@@ -16,6 +16,7 @@ package function
 
 import (
 	"bytes"
+	"encoding/binary"
 	"encoding/hex"
 	"fmt"
 	"math"
@@ -256,6 +257,13 @@ func (h *Hex) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
 		}
 		return hexForString(string(b)), nil
 
+	case []float32:
+		buf := make([]byte, 4*len(val))
+		for i, v := range val {
+			binary.Encode(buf[4*i:], binary.LittleEndian, v)
+		}
+		return hexForString(string(buf)), nil
+
 	case types.GeometryValue:
 		return hexForString(string(val.Serialize())), nil
 
@@ -601,6 +609,8 @@ func (h *Bitlength) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
 		return 64, nil
 	case string:
 		return 8 * len([]byte(val)), nil
+	case []float32:
+		return 32 * len(val), nil
 	case time.Time:
 		return 128, nil
 	}

sql/expression/function/vector/conversion.go

@@ -0,0 +1,126 @@
+// 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 vector
+
+import (
+	"fmt"
+
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/expression"
+	"github.com/dolthub/go-mysql-server/sql/types"
+)
+
+// StringToVector converts a JSON string representation to a vector
+type StringToVector struct {
+	expression.UnaryExpression
+}
+
+var _ sql.Expression = (*StringToVector)(nil)
+var _ sql.FunctionExpression = (*StringToVector)(nil)
+var _ sql.CollationCoercible = (*StringToVector)(nil)
+
+func NewStringToVector(e sql.Expression) sql.Expression {
+	return &StringToVector{UnaryExpression: expression.UnaryExpression{Child: e}}
+}
+
+func (s *StringToVector) FunctionName() string {
+	return "string_to_vector"
+}
+
+func (s *StringToVector) Description() string {
+	return "converts a JSON array string to a vector"
+}
+
+func (s *StringToVector) Type() sql.Type {
+	return types.VectorType{}
+}
+
+func (s *StringToVector) CollationCoercibility(_ *sql.Context) (collation sql.CollationID, coercibility byte) {
+	return sql.Collation_binary, 5
+}
+
+func (s *StringToVector) String() string {
+	return fmt.Sprintf("STRING_TO_VECTOR(%s)", s.Child)
+}
+
+func (s *StringToVector) WithChildren(children ...sql.Expression) (sql.Expression, error) {
+	if len(children) != 1 {
+		return nil, sql.ErrInvalidChildrenNumber.New(s, len(children), 1)
+	}
+	return NewStringToVector(children[0]), nil
+}
+
+func (s *StringToVector) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	val, err := s.Child.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+	if val == nil {
+		return nil, nil
+	}
+
+	return sql.ConvertToVector(ctx, val)
+}
+
+// VectorToString converts a vector to a JSON string representation
+type VectorToString struct {
+	expression.UnaryExpression
+}
+
+var _ sql.Expression = (*VectorToString)(nil)
+var _ sql.FunctionExpression = (*VectorToString)(nil)
+var _ sql.CollationCoercible = (*VectorToString)(nil)
+
+func NewVectorToString(e sql.Expression) sql.Expression {
+	return &VectorToString{UnaryExpression: expression.UnaryExpression{Child: e}}
+}
+
+func (v *VectorToString) FunctionName() string {
+	return "vector_to_string"
+}
+
+func (v *VectorToString) Description() string {
+	return "converts a vector to a JSON array string"
+}
+
+func (v *VectorToString) Type() sql.Type {
+	return types.LongText
+}
+
+func (v *VectorToString) CollationCoercibility(_ *sql.Context) (collation sql.CollationID, coercibility byte) {
+	return sql.Collation_binary, 5
+}
+
+func (v *VectorToString) String() string {
+	return fmt.Sprintf("VECTOR_TO_STRING(%s)", v.Child)
+}
+
+func (v *VectorToString) WithChildren(children ...sql.Expression) (sql.Expression, error) {
+	if len(children) != 1 {
+		return nil, sql.ErrInvalidChildrenNumber.New(v, len(children), 1)
+	}
+	return NewVectorToString(children[0]), nil
+}
+
+func (v *VectorToString) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	val, err := v.Child.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+	if val == nil {
+		return nil, nil
+	}
+	return types.JSONDocument{Val: val}.JSONString()
+}