115 lines
3.0 KiB
Go
115 lines
3.0 KiB
Go
package pgvector
|
|
|
|
import (
|
|
"database/sql/driver"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"math"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
// Vector represents a vector for similarity search.
|
|
// Stub implementation for building without PostgreSQL/pgvector extension.
|
|
type Vector []float32
|
|
|
|
// NewVector creates a new vector from a slice of float32.
|
|
func NewVector(slice []float32) Vector {
|
|
return Vector(slice)
|
|
}
|
|
|
|
// String returns the string representation of the vector.
|
|
func (v Vector) String() string {
|
|
return ""
|
|
}
|
|
|
|
// Dimensions returns the number of dimensions in the vector.
|
|
func (v Vector) Dimensions() int {
|
|
return len(v)
|
|
}
|
|
|
|
// Scan implements sql.Scanner so GORM/database/sql can scan vector columns
|
|
// from PostgreSQL into the Vector type.
|
|
//
|
|
// PostgreSQL pgvector returns the vector as a string like "[0.1,0.2,0.3]"
|
|
// via text protocol, or as raw bytes via binary protocol.
|
|
// We handle both cases.
|
|
func (v *Vector) Scan(src any) error {
|
|
if src == nil {
|
|
*v = nil
|
|
return nil
|
|
}
|
|
|
|
switch val := src.(type) {
|
|
case string:
|
|
return v.parseVectorString(val)
|
|
case []byte:
|
|
// Try text protocol first: "[0.1,0.2,...]"
|
|
s := string(val)
|
|
if strings.HasPrefix(s, "[") {
|
|
return v.parseVectorString(s)
|
|
}
|
|
// Binary protocol: pgvector binary format
|
|
// Header: 2 bytes unused (version), 2 bytes dim count, then dim*4 bytes float32 big-endian
|
|
return v.parseVectorBinary(val)
|
|
default:
|
|
return fmt.Errorf("pgvector: cannot scan %T into Vector", src)
|
|
}
|
|
}
|
|
|
|
// Value implements driver.Valuer so GORM/database/sql can use Vector
|
|
// in parameter binding. Returns the PostgreSQL text representation.
|
|
func (v Vector) Value() (driver.Value, error) {
|
|
if len(v) == 0 {
|
|
return nil, nil
|
|
}
|
|
strs := make([]string, len(v))
|
|
for i, f := range v {
|
|
strs[i] = strconv.FormatFloat(float64(f), 'f', -1, 32)
|
|
}
|
|
return "[" + strings.Join(strs, ",") + "]", nil
|
|
}
|
|
|
|
// parseVectorString parses a PostgreSQL vector text representation "[0.1,0.2,0.3]".
|
|
func (v *Vector) parseVectorString(s string) error {
|
|
s = strings.TrimSpace(s)
|
|
s = strings.TrimPrefix(s, "[")
|
|
s = strings.TrimSuffix(s, "]")
|
|
if s == "" {
|
|
*v = Vector{}
|
|
return nil
|
|
}
|
|
parts := strings.Split(s, ",")
|
|
result := make(Vector, len(parts))
|
|
for i, part := range parts {
|
|
f, err := strconv.ParseFloat(strings.TrimSpace(part), 32)
|
|
if err != nil {
|
|
return fmt.Errorf("pgvector: parse float %q: %w", part, err)
|
|
}
|
|
result[i] = float32(f)
|
|
}
|
|
*v = result
|
|
return nil
|
|
}
|
|
|
|
// parseVectorBinary parses pgvector binary format.
|
|
// Format: 2 bytes (version, currently 0), 2 bytes (dims), then dims * 4 bytes float32.
|
|
func (v *Vector) parseVectorBinary(data []byte) error {
|
|
if len(data) < 4 {
|
|
return fmt.Errorf("pgvector: binary vector too short (%d bytes)", len(data))
|
|
}
|
|
dims := int(binary.BigEndian.Uint16(data[2:4]))
|
|
expected := 4 + dims*4
|
|
if len(data) < expected {
|
|
return fmt.Errorf("pgvector: binary vector expected %d bytes, got %d", expected, len(data))
|
|
}
|
|
result := make(Vector, dims)
|
|
for i := 0; i < dims; i++ {
|
|
offset := 4 + i*4
|
|
bits := binary.BigEndian.Uint32(data[offset : offset+4])
|
|
result[i] = math.Float32frombits(bits)
|
|
}
|
|
*v = result
|
|
return nil
|
|
}
|