Hein
129 lines
3.0 KiB
Go
129 lines
3.0 KiB
Go
package pgx
|
|
|
|
import (
|
|
"database/sql/driver"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"slices"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
"github.com/pgvector/pgvector-go"
|
|
"github.com/x448/float16"
|
|
)
|
|
|
|
type HalfVectorCodec struct{}
|
|
|
|
func (HalfVectorCodec) FormatSupported(format int16) bool {
|
|
return format == pgx.BinaryFormatCode || format == pgx.TextFormatCode
|
|
}
|
|
|
|
func (HalfVectorCodec) PreferredFormat() int16 {
|
|
return pgx.BinaryFormatCode
|
|
}
|
|
|
|
func (HalfVectorCodec) PlanEncode(m *pgtype.Map, oid uint32, format int16, value any) pgtype.EncodePlan {
|
|
_, ok := value.(pgvector.HalfVector)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
switch format {
|
|
case pgx.BinaryFormatCode:
|
|
return encodePlanHalfVectorCodecBinary{}
|
|
case pgx.TextFormatCode:
|
|
return encodePlanHalfVectorCodecText{}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
type encodePlanHalfVectorCodecBinary struct{}
|
|
|
|
func (encodePlanHalfVectorCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) {
|
|
v := value.(pgvector.HalfVector)
|
|
vec := v.Slice()
|
|
dim := len(vec)
|
|
buf = slices.Grow(buf, 4+2*dim)
|
|
buf = binary.BigEndian.AppendUint16(buf, uint16(dim))
|
|
buf = binary.BigEndian.AppendUint16(buf, 0)
|
|
for _, v := range vec {
|
|
buf = binary.BigEndian.AppendUint16(buf, float16.Fromfloat32(v).Bits())
|
|
}
|
|
return buf, nil
|
|
}
|
|
|
|
type encodePlanHalfVectorCodecText struct{}
|
|
|
|
func (encodePlanHalfVectorCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) {
|
|
v := value.(pgvector.HalfVector)
|
|
return v.EncodeText(buf)
|
|
}
|
|
|
|
func (HalfVectorCodec) PlanScan(m *pgtype.Map, oid uint32, format int16, target any) pgtype.ScanPlan {
|
|
_, ok := target.(*pgvector.HalfVector)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
switch format {
|
|
case pgx.BinaryFormatCode:
|
|
return scanPlanHalfVectorCodecBinary{}
|
|
case pgx.TextFormatCode:
|
|
return scanPlanHalfVectorCodecText{}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
type scanPlanHalfVectorCodecBinary struct{}
|
|
|
|
func (scanPlanHalfVectorCodecBinary) Scan(src []byte, dst any) error {
|
|
v := (dst).(*pgvector.HalfVector)
|
|
buf := src
|
|
dim := int(binary.BigEndian.Uint16(buf[0:2]))
|
|
unused := binary.BigEndian.Uint16(buf[2:4])
|
|
if unused != 0 {
|
|
return fmt.Errorf("expected unused to be 0")
|
|
}
|
|
|
|
vec := make([]float32, 0, dim)
|
|
offset := 4
|
|
for i := 0; i < dim; i++ {
|
|
vec = append(vec, float16.Frombits(binary.BigEndian.Uint16(buf[offset:offset+2])).Float32())
|
|
offset += 2
|
|
}
|
|
v.SetSlice(vec)
|
|
return nil
|
|
}
|
|
|
|
type scanPlanHalfVectorCodecText struct{}
|
|
|
|
func (scanPlanHalfVectorCodecText) Scan(src []byte, dst any) error {
|
|
v := (dst).(*pgvector.HalfVector)
|
|
return v.Scan(src)
|
|
}
|
|
|
|
func (c HalfVectorCodec) DecodeDatabaseSQLValue(m *pgtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) {
|
|
return c.DecodeValue(m, oid, format, src)
|
|
}
|
|
|
|
func (c HalfVectorCodec) DecodeValue(m *pgtype.Map, oid uint32, format int16, src []byte) (any, error) {
|
|
if src == nil {
|
|
return nil, nil
|
|
}
|
|
|
|
var vec pgvector.HalfVector
|
|
scanPlan := c.PlanScan(m, oid, format, &vec)
|
|
if scanPlan == nil {
|
|
return nil, fmt.Errorf("Unable to decode halfvec type")
|
|
}
|
|
|
|
err := scanPlan.Scan(src, &vec)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return vec, nil
|
|
}
|