pgvector-go/vector.go at master · pgvector/pgvector-go · GitHub

History

148 lines (124 loc) · 3.42 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

package pgvector

import (

"database/sql"

"database/sql/driver"

"encoding/binary"

"encoding/json"

"fmt"

"math"

"slices"

"strconv"

"strings"

)

// Vector is a wrapper for []float32 to implement sql.Scanner and driver.Valuer.

type Vector struct {

vec []float32

}

// NewVector creates a new Vector from a slice of float32.

func NewVector(vec []float32) Vector {

return Vector{vec: vec}

}

// Slice returns the underlying slice of float32.

func (v Vector) Slice() []float32 {

return v.vec

}

// String returns a string representation of the vector.

func (v Vector) String() string {

buf := make([]byte, 0, 2+16*len(v.vec))

buf = append(buf, '[')

for i := 0; i < len(v.vec); i++ {

if i > 0 {

buf = append(buf, ',')

}

buf = strconv.AppendFloat(buf, float64(v.vec[i]), 'f', -1, 32)

}

buf = append(buf, ']')

return string(buf)

}

// Parse parses a string representation of a vector.

func (v *Vector) Parse(s string) error {

// TODO check brackets in 0.5.0

if len(s) < 2 {

return fmt.Errorf("malformed vector literal")

}

if len(s) == 2 {

v.vec = []float32{}

return nil

}

sp := strings.Split(s[1:len(s)-1], ",")

v.vec = make([]float32, 0, len(sp))

for _, si := range sp {

n, err := strconv.ParseFloat(si, 32)

if err != nil {

return err

}

v.vec = append(v.vec, float32(n))

}

return nil

}

// EncodeBinary encodes a binary representation of the vector.

func (v Vector) EncodeBinary(buf []byte) (newBuf []byte, err error) {

dim := len(v.vec)

buf = slices.Grow(buf, 4+4*dim)

buf = binary.BigEndian.AppendUint16(buf, uint16(dim))

buf = binary.BigEndian.AppendUint16(buf, 0)

for _, v := range v.vec {

buf = binary.BigEndian.AppendUint32(buf, math.Float32bits(v))

}

return buf, nil

}

// DecodeBinary decodes a binary representation of a vector.

func (v *Vector) DecodeBinary(buf []byte) error {

if len(buf) < 4 {

return fmt.Errorf("invalid length")

}

dim := int(binary.BigEndian.Uint16(buf[0:2]))

if dim < 0 {

return fmt.Errorf("vector cannot have negative dimensions")

}

unused := binary.BigEndian.Uint16(buf[2:4])

if unused != 0 {

return fmt.Errorf("expected unused to be 0")

}

if (len(buf)-4)/4 != dim || len(buf)%4 != 0 {

return fmt.Errorf("invalid length")

}

v.vec = make([]float32, 0, dim)

offset := 4

for range dim {

v.vec = append(v.vec, math.Float32frombits(binary.BigEndian.Uint32(buf[offset:offset+4])))

offset += 4

}

return nil

}

// statically assert that Vector implements sql.Scanner.

var _ sql.Scanner = (*Vector)(nil)

// Scan implements the sql.Scanner interface.

func (v *Vector) Scan(src any) (err error) {

switch src := src.(type) {

case []byte:

return v.Parse(string(src))

case string:

return v.Parse(src)

default:

return fmt.Errorf("unsupported data type: %T", src)

}

// statically assert that Vector implements driver.Valuer.

var _ driver.Valuer = (*Vector)(nil)

// Value implements the driver.Valuer interface.

func (v Vector) Value() (driver.Value, error) {

return v.String(), nil

}

// statically assert that Vector implements json.Marshaler.

var _ json.Marshaler = (*Vector)(nil)

// MarshalJSON implements the json.Marshaler interface.

func (v Vector) MarshalJSON() ([]byte, error) {

return json.Marshal(v.vec)

}

// statically assert that Vector implements json.Unmarshaler.

var _ json.Unmarshaler = (*Vector)(nil)

// UnmarshalJSON implements the json.Unmarshaler interface.

func (v *Vector) UnmarshalJSON(data []byte) error {

return json.Unmarshal(data, &v.vec)

}

pgvector-go/vector.go at master · pgvector/pgvector-go · GitHub

Uh oh!

Expand file tree

Breadcrumbs

vector.go

Latest commit

History

Breadcrumbs

vector.go

File metadata and controls