Hein
105 lines
2.5 KiB
Go
105 lines
2.5 KiB
Go
package pgx
|
|
|
|
import (
|
|
"database/sql/driver"
|
|
"fmt"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
"github.com/pgvector/pgvector-go"
|
|
)
|
|
|
|
type SparseVectorCodec struct{}
|
|
|
|
func (SparseVectorCodec) FormatSupported(format int16) bool {
|
|
return format == pgx.BinaryFormatCode || format == pgx.TextFormatCode
|
|
}
|
|
|
|
func (SparseVectorCodec) PreferredFormat() int16 {
|
|
return pgx.BinaryFormatCode
|
|
}
|
|
|
|
func (SparseVectorCodec) PlanEncode(m *pgtype.Map, oid uint32, format int16, value any) pgtype.EncodePlan {
|
|
_, ok := value.(pgvector.SparseVector)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
switch format {
|
|
case pgx.BinaryFormatCode:
|
|
return encodePlanSparseVectorCodecBinary{}
|
|
case pgx.TextFormatCode:
|
|
return encodePlanSparseVectorCodecText{}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
type encodePlanSparseVectorCodecBinary struct{}
|
|
|
|
func (encodePlanSparseVectorCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) {
|
|
v := value.(pgvector.SparseVector)
|
|
return v.EncodeBinary(buf)
|
|
}
|
|
|
|
type encodePlanSparseVectorCodecText struct{}
|
|
|
|
func (encodePlanSparseVectorCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) {
|
|
v := value.(pgvector.SparseVector)
|
|
// use String() for now to avoid adding another method to SparseVector
|
|
return append(buf, v.String()...), nil
|
|
}
|
|
|
|
func (SparseVectorCodec) PlanScan(m *pgtype.Map, oid uint32, format int16, target any) pgtype.ScanPlan {
|
|
_, ok := target.(*pgvector.SparseVector)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
switch format {
|
|
case pgx.BinaryFormatCode:
|
|
return scanPlanSparseVectorCodecBinary{}
|
|
case pgx.TextFormatCode:
|
|
return scanPlanSparseVectorCodecText{}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
type scanPlanSparseVectorCodecBinary struct{}
|
|
|
|
func (scanPlanSparseVectorCodecBinary) Scan(src []byte, dst any) error {
|
|
v := (dst).(*pgvector.SparseVector)
|
|
return v.DecodeBinary(src)
|
|
}
|
|
|
|
type scanPlanSparseVectorCodecText struct{}
|
|
|
|
func (scanPlanSparseVectorCodecText) Scan(src []byte, dst any) error {
|
|
v := (dst).(*pgvector.SparseVector)
|
|
return v.Scan(src)
|
|
}
|
|
|
|
func (c SparseVectorCodec) DecodeDatabaseSQLValue(m *pgtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) {
|
|
return c.DecodeValue(m, oid, format, src)
|
|
}
|
|
|
|
func (c SparseVectorCodec) DecodeValue(m *pgtype.Map, oid uint32, format int16, src []byte) (any, error) {
|
|
if src == nil {
|
|
return nil, nil
|
|
}
|
|
|
|
var vec pgvector.SparseVector
|
|
scanPlan := c.PlanScan(m, oid, format, &vec)
|
|
if scanPlan == nil {
|
|
return nil, fmt.Errorf("Unable to decode sparsevec type")
|
|
}
|
|
|
|
err := scanPlan.Scan(src, &vec)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return vec, nil
|
|
}
|