Store vectors in the engine as byte array instead of float array. It simplifies the logic a lot.

nicktobey · nicktobey · commit 6ba7124510dd · 2025-08-19T20:27:33.000-07:00
diff --git a/enginetest/memory_engine_test.go b/enginetest/memory_engine_test.go
@@ -1001,9 +1001,6 @@ func TestDatabaseCollationWire(t *testing.T) {
 }
 
 func TestTypesOverWire(t *testing.T) {
-	if _, ok := os.LookupEnv("CI_TEST"); !ok {
-		t.Skip("Skipping test that requires CI_TEST=true")
-	}
 	harness := enginetest.NewDefaultMemoryHarness()
 	enginetest.TestTypesOverWire(t, harness, harness.SessionBuilder())
 }
diff --git a/enginetest/queries/type_wire_queries.go b/enginetest/queries/type_wire_queries.go
@@ -15,8 +15,6 @@
 package queries
 
 import (
-	"encoding/binary"
-
 	"github.com/dolthub/go-mysql-server/sql"
 )
 
@@ -30,9 +28,11 @@ type TypeWireTest struct {
 }
 
 func floatsToString(fs ...float32) string {
-	result := make([]byte, 4*len(fs))
-	binary.Encode(result, binary.LittleEndian, fs)
-	return string(result)
+	return string(sql.EncodeVector(fs))
+}
+
+func floatsToBytes(fs ...float32) []byte {
+	return sql.EncodeVector(fs)
 }
 
 // TypeWireTests are used to ensure that types are properly represented over the wire (vs being directly returned from
diff --git a/enginetest/queries/vector_ddl_queries.go b/enginetest/queries/vector_ddl_queries.go
@@ -43,17 +43,17 @@ var VectorDDLQueries = []ScriptTest{
 			},
 			{
 				Query:    `SELECT id, small_vec, medium_vec FROM test_vectors WHERE id = 2`,
-				Expected: []sql.Row{{2, []float32{3.5, 4.5}, nil}},
+				Expected: []sql.Row{{2, floatsToBytes(3.5, 4.5), nil}},
 			},
 			{
 				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}}},
+				Expected: []sql.Row{{1, floatsToBytes(1.0, 2.0), floatsToBytes(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}},
+					{1, floatsToBytes(1.0, 2.0)},
+					{2, floatsToBytes(3.5, 4.5)},
 				},
 			},
 			{
@@ -62,15 +62,15 @@ var VectorDDLQueries = []ScriptTest{
 			},
 			{
 				Query:    `SELECT small_vec FROM test_vectors WHERE id = 1`,
-				Expected: []sql.Row{{[]float32{10.0, 20.0}}},
+				Expected: []sql.Row{{floatsToBytes(10.0, 20.0)}},
 			},
 			{
 				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{10.0, 10.0}}},
+				Expected: []sql.Row{{floatsToBytes(10.0, 10.0)}},
 			},
 		},
 	},
@@ -138,9 +138,9 @@ var VectorDDLQueries = []ScriptTest{
 			{
 				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}},
+					{1, floatsToBytes(1.0, 2.0)},
+					{2, floatsToBytes(3.0, 4.0)},
+					{3, floatsToBytes(5.5, 6700)},
 				},
 			},
 		},
diff --git a/enginetest/queries/vector_function_queries.go b/enginetest/queries/vector_function_queries.go
@@ -111,20 +111,25 @@ var VectorFunctionQueries = []ScriptTest{
 		Assertions: []ScriptTestAssertion{
 			{
 				Query:    `SELECT STRING_TO_VECTOR("[1.0, 2.0]");`,
-				Expected: []sql.Row{{[]float32{1.0, 2.0}}},
+				Expected: []sql.Row{{floatsToBytes(1.0, 2.0)}},
 			},
 			{
 				Query:    `SELECT TO_VECTOR("[1.0, 2.0]");`,
-				Expected: []sql.Row{{[]float32{1.0, 2.0}}},
+				Expected: []sql.Row{{floatsToBytes(1.0, 2.0)}},
 			},
 			{
 				Query:    `SELECT VEC_FromText("[1.0, 2.0]");`,
-				Expected: []sql.Row{{[]float32{1.0, 2.0}}},
+				Expected: []sql.Row{{floatsToBytes(1.0, 2.0)}},
 			},
 			{
 				Query:    `SELECT VECTOR_TO_STRING(STRING_TO_VECTOR("[1.0, 2.0]"));`,
 				Expected: []sql.Row{{"[1, 2]"}},
 			},
+			{
+				Query:    `select VECTOR_TO_STRING(0x0000803F);`,
+				Expected: []sql.Row{{"[1]"}},
+			},
+
 			{
 				Query:    `SELECT FROM_VECTOR(TO_VECTOR("[1.0, 2.0]"));`,
 				Expected: []sql.Row{{"[1, 2]"}},
diff --git a/enginetest/server_engine.go b/enginetest/server_engine.go
@@ -16,7 +16,6 @@ package enginetest
 
 import (
 	gosql "database/sql"
-	"encoding/binary"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -398,13 +397,6 @@ func convertValue(ctx *sql.Context, sch sql.Schema, row sql.Row) sql.Row {
 					row[i] = r
 				}
 			}
-		case query.Type_VECTOR:
-			if row[i] != nil {
-				r := row[i].([]byte)
-				dimensions := len(r) / 4
-				row[i] = make([]float32, dimensions)
-				binary.Decode(r, binary.LittleEndian, row[i])
-			}
 		case query.Type_TIME:
 			if row[i] != nil {
 				r, _, err := types.TimespanType_{}.Convert(ctx, string(row[i].([]byte)))
diff --git a/sql/core.go b/sql/core.go
@@ -18,11 +18,13 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/dolthub/go-mysql-server/sql/values"
 	trace2 "runtime/trace"
 	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"github.com/shopspring/decimal"
 )
@@ -325,10 +327,22 @@ func ConvertToBool(ctx *Context, v interface{}) (bool, error) {
 	}
 }
 
+// DecodeVector decodes a byte slice that represents a vector. This is needed for distance functions.
+func DecodeVector(buf []byte) []float32 {
+	return unsafe.Slice((*float32)(unsafe.Pointer(&buf[0])), len(buf)/int(values.Float32Size))
+}
+
+// EncodeVector encodes a byte slice that represents a vector.
+func EncodeVector(floats []float32) []byte {
+	return unsafe.Slice((*byte)(unsafe.Pointer(&floats[0])), len(floats)*int(values.Float32Size))
+}
+
 func ConvertToVector(ctx context.Context, v interface{}) ([]float32, error) {
 	switch b := v.(type) {
 	case []float32:
 		return b, nil
+	case []byte:
+		return DecodeVector(b), nil
 	case string:
 		var val interface{}
 		err := json.Unmarshal([]byte(b), &val)
diff --git a/sql/expression/function/length.go b/sql/expression/function/length.go
@@ -123,10 +123,6 @@ 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
diff --git a/sql/expression/function/string.go b/sql/expression/function/string.go
@@ -16,7 +16,6 @@ package function
 
 import (
 	"bytes"
-	"encoding/binary"
 	"encoding/hex"
 	"fmt"
 	"math"
@@ -257,11 +256,6 @@ 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))
-		binary.Encode(buf, binary.LittleEndian, val)
-		return hexForString(string(buf)), nil
-
 	case types.GeometryValue:
 		return hexForString(string(val.Serialize())), nil
 
@@ -607,8 +601,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 []byte:
+		return 8 * len(val), nil
 	case time.Time:
 		return 128, nil
 	}
diff --git a/sql/expression/function/vector/conversion.go b/sql/expression/function/vector/conversion.go
@@ -71,7 +71,13 @@ func (s *StringToVector) Eval(ctx *sql.Context, row sql.Row) (interface{}, error
 		return nil, nil
 	}
 
-	return sql.ConvertToVector(ctx, val)
+	// TODO: Instead of using the JSON parser and then encoding, it would be more efficient to parse and encode
+	// in a single step
+	floats, err := sql.ConvertToVector(ctx, val)
+	if err != nil {
+		return nil, err
+	}
+	return sql.EncodeVector(floats), nil
 }
 
 // VectorToString converts a vector to a JSON string representation
@@ -122,5 +128,9 @@ func (v *VectorToString) Eval(ctx *sql.Context, row sql.Row) (interface{}, error
 	if val == nil {
 		return nil, nil
 	}
-	return types.JSONDocument{Val: val}.JSONString()
+	b, ok := val.([]byte)
+	if !ok {
+		return nil, fmt.Errorf("incorrect argument to VECTOR_TO_STRING: expected a vector, got %T", val)
+	}
+	return types.JSONDocument{Val: sql.DecodeVector(b)}.JSONString()
 }
diff --git a/sql/types/strings.go b/sql/types/strings.go
@@ -16,7 +16,6 @@ package types
 
 import (
 	"context"
-	"encoding/binary"
 	"fmt"
 	"reflect"
 	"strconv"
@@ -414,9 +413,6 @@ func ConvertToBytes(ctx context.Context, v interface{}, t sql.StringType, dest [
 		// We'll check for that below, immediately before extending the slice.
 		val = s
 		start = 0
-	case []float32:
-		val = make([]byte, 4*len(s))
-		binary.Encode(val, binary.LittleEndian, s)
 	case time.Time:
 		val = s.AppendFormat(dest, sql.TimestampDatetimeLayout)
 	case decimal.Decimal:
diff --git a/sql/types/vector.go b/sql/types/vector.go
@@ -17,7 +17,6 @@ package types
 import (
 	"bytes"
 	"context"
-	"encoding/binary"
 	"fmt"
 	"reflect"
 	"strconv"
@@ -28,8 +27,6 @@ import (
 	"github.com/dolthub/go-mysql-server/sql"
 )
 
-var vectorValueType = reflect.TypeOf([]float32{})
-
 // VectorType represents the VECTOR(N) type.
 // It stores a fixed-length array of N floating point numbers.
 type VectorType struct {
@@ -68,21 +65,10 @@ func (t VectorType) Compare(ctx context.Context, a interface{}, b interface{}) (
 		return 0, err
 	}
 
-	avec := av.([]float32)
-	bvec := bv.([]float32)
-
-	aBytes := make([]byte, 4)
-	bBytes := make([]byte, 4)
+	avec := av.([]byte)
+	bvec := bv.([]byte)
 
-	for i := 0; i < len(avec); i++ {
-		binary.Encode(aBytes, binary.LittleEndian, avec[i])
-		binary.Encode(bBytes, binary.LittleEndian, bvec[i])
-		cmp := bytes.Compare(aBytes, bBytes)
-		if cmp != 0 {
-			return cmp, nil
-		}
-	}
-	return 0, nil
+	return bytes.Compare(avec, bvec), nil
 }
 
 // Convert implements Type interface.
@@ -93,19 +79,19 @@ func (t VectorType) Convert(ctx context.Context, v interface{}) (interface{}, sq
 
 	switch val := v.(type) {
 	case []byte:
-		if len(val) != 4*t.Dimensions {
-			return nil, sql.OutOfRange, fmt.Errorf("cannot convert BINARY(%d) to VECTOR(%d), need BINARY(%d)", len(val), t.Dimensions, 4*t.Dimensions)
-		}
-		result := make([]float32, t.Dimensions)
-		for i := range result {
-			binary.Decode(val[4*i:4*(i+1)], binary.LittleEndian, &result[i])
-		}
-		return result, sql.InRange, nil
-	case []float32:
-		if t.Dimensions != 0 && len(val) != t.Dimensions {
-			return nil, sql.OutOfRange, fmt.Errorf("VECTOR dimension mismatch: expected %d, got %d", t.Dimensions, len(val))
+		if t.Dimensions != 0 && len(val) != 4*t.Dimensions {
+			if len(val)%4 != 0 {
+				return nil, sql.OutOfRange, fmt.Errorf("cannot convert BINARY(%d) to VECTOR(%d), need BINARY(%d)", len(val), t.Dimensions, 4*t.Dimensions)
+			}
+			return nil, sql.OutOfRange, fmt.Errorf("VECTOR dimension mismatch: expected %d, got %d", t.Dimensions, len(val)/4)
 		}
 		return val, sql.InRange, nil
+	case sql.JSONWrapper:
+		unwrapped, err := val.ToInterface(ctx)
+		if err != nil {
+			return nil, sql.OutOfRange, err
+		}
+		return t.Convert(ctx, unwrapped)
 	case []interface{}:
 		if t.Dimensions != 0 && len(val) != t.Dimensions {
 			return nil, sql.OutOfRange, fmt.Errorf("VECTOR dimension mismatch: expected %d, got %d", t.Dimensions, len(val))
@@ -135,7 +121,7 @@ func (t VectorType) Convert(ctx context.Context, v interface{}) (interface{}, sq
 				}
 			}
 		}
-		return result, sql.InRange, nil
+		return sql.EncodeVector(result), sql.InRange, nil
 	default:
 		return nil, sql.OutOfRange, fmt.Errorf("value of type %T cannot be converted to 'vector' type", v)
 	}
@@ -190,7 +176,7 @@ func (t VectorType) Type() query.Type {
 
 // ValueType implements Type interface.
 func (t VectorType) ValueType() reflect.Type {
-	return vectorValueType
+	return byteValueType
 }
 
 // Zero implements Type interface.
@@ -202,3 +188,17 @@ func (t VectorType) Zero() interface{} {
 func (VectorType) CollationCoercibility(ctx *sql.Context) (collation sql.CollationID, coercibility byte) {
 	return sql.Collation_binary, 5
 }
+
+var _ sql.TypeWithCollation = VectorType{}
+
+func (VectorType) Collation() sql.CollationID {
+	return sql.Collation_binary
+}
+
+func (VectorType) WithNewCollation(sql.CollationID) (sql.Type, error) {
+	return nil, fmt.Errorf("cannot change collation of binary types")
+}
+
+func (t VectorType) StringWithTableCollation(sql.CollationID) string {
+	return fmt.Sprintf("VECTOR(%d)", t.Dimensions)
+}
diff --git a/sql/types/vector_test.go b/sql/types/vector_test.go

Original file line number	Diff line number	Diff line change
`@@ -1001,9 +1001,6 @@ func TestDatabaseCollationWire(t *testing.T) {`
`1001`	`1001`	`}`
`1002`	`1002`
`1003`	`1003`	`func TestTypesOverWire(t *testing.T) {`
`1004`		`- if _, ok := os.LookupEnv("CI_TEST"); !ok {`
`1005`		`- t.Skip("Skipping test that requires CI_TEST=true")`
`1006`		`- }`
`1007`	`1004`	`harness := enginetest.NewDefaultMemoryHarness()`
`1008`	`1005`	`enginetest.TestTypesOverWire(t, harness, harness.SessionBuilder())`
`1009`	`1006`	`}`
Original file line number	Diff line number	Diff line change
`@@ -15,8 +15,6 @@`
`15`	`15`	`package queries`
`16`	`16`
`17`	`17`	`import (`
`18`		`- "encoding/binary"`
`19`		`-`
`20`	`18`	`"github.com/dolthub/go-mysql-server/sql"`
`21`	`19`	`)`
`22`	`20`
`@@ -30,9 +28,11 @@ type TypeWireTest struct {`
`30`	`28`	`}`
`31`	`29`
`32`	`30`	`func floatsToString(fs ...float32) string {`
`33`		`- result := make([]byte, 4*len(fs))`
`34`		`- binary.Encode(result, binary.LittleEndian, fs)`
`35`		`- return string(result)`
	`31`	`+ return string(sql.EncodeVector(fs))`
	`32`	`+}`
	`33`	`+`
	`34`	`+func floatsToBytes(fs ...float32) []byte {`
	`35`	`+ return sql.EncodeVector(fs)`
`36`	`36`	`}`
`37`	`37`
`38`	`38`	`// TypeWireTests are used to ensure that types are properly represented over the wire (vs being directly returned from`
Original file line number	Diff line number	Diff line change
`@@ -111,20 +111,25 @@ var VectorFunctionQueries = []ScriptTest{`
`111`	`111`	`Assertions: []ScriptTestAssertion{`
`112`	`112`	`{`
`113`	`113`	Query: `SELECT STRING_TO_VECTOR("[1.0, 2.0]");`,
`114`		`- Expected: []sql.Row{{[]float32{1.0, 2.0}}},`
	`114`	`+ Expected: []sql.Row{{floatsToBytes(1.0, 2.0)}},`
`115`	`115`	`},`
`116`	`116`	`{`
`117`	`117`	Query: `SELECT TO_VECTOR("[1.0, 2.0]");`,
`118`		`- Expected: []sql.Row{{[]float32{1.0, 2.0}}},`
	`118`	`+ Expected: []sql.Row{{floatsToBytes(1.0, 2.0)}},`
`119`	`119`	`},`
`120`	`120`	`{`
`121`	`121`	Query: `SELECT VEC_FromText("[1.0, 2.0]");`,
`122`		`- Expected: []sql.Row{{[]float32{1.0, 2.0}}},`
	`122`	`+ Expected: []sql.Row{{floatsToBytes(1.0, 2.0)}},`
`123`	`123`	`},`
`124`	`124`	`{`
`125`	`125`	Query: `SELECT VECTOR_TO_STRING(STRING_TO_VECTOR("[1.0, 2.0]"));`,
`126`	`126`	`Expected: []sql.Row{{"[1, 2]"}},`
`127`	`127`	`},`
	`128`	`+ {`
	`129`	+ Query: `select VECTOR_TO_STRING(0x0000803F);`,
	`130`	`+ Expected: []sql.Row{{"[1]"}},`
	`131`	`+ },`
	`132`	`+`
`128`	`133`	`{`
`129`	`134`	Query: `SELECT FROM_VECTOR(TO_VECTOR("[1.0, 2.0]"));`,
`130`	`135`	`Expected: []sql.Row{{"[1, 2]"}},`
Original file line number	Diff line number	Diff line change
`@@ -71,7 +71,13 @@ func (s StringToVector) Eval(ctx sql.Context, row sql.Row) (interface{}, error`
`71`	`71`	`return nil, nil`
`72`	`72`	`}`
`73`	`73`
`74`		`- return sql.ConvertToVector(ctx, val)`
	`74`	`+ // TODO: Instead of using the JSON parser and then encoding, it would be more efficient to parse and encode`
	`75`	`+ // in a single step`
	`76`	`+ floats, err := sql.ConvertToVector(ctx, val)`
	`77`	`+ if err != nil {`
	`78`	`+ return nil, err`
	`79`	`+ }`
	`80`	`+ return sql.EncodeVector(floats), nil`
`75`	`81`	`}`
`76`	`82`
`77`	`83`	`// VectorToString converts a vector to a JSON string representation`
`@@ -122,5 +128,9 @@ func (v VectorToString) Eval(ctx sql.Context, row sql.Row) (interface{}, error`
`122`	`128`	`if val == nil {`
`123`	`129`	`return nil, nil`
`124`	`130`	`}`
`125`		`- return types.JSONDocument{Val: val}.JSONString()`
	`131`	`+ b, ok := val.([]byte)`
	`132`	`+ if !ok {`
	`133`	`+ return nil, fmt.Errorf("incorrect argument to VECTOR_TO_STRING: expected a vector, got %T", val)`
	`134`	`+ }`
	`135`	`+ return types.JSONDocument{Val: sql.DecodeVector(b)}.JSONString()`
`126`	`136`	`}`