fix(go.sum): update ResolveSpec dependency to v1.0.87
This commit is contained in:
Hein
@@ -0,0 +1,4 @@
|
||||
/go.sum
|
||||
/test/ent/*
|
||||
!/test/ent/generate.go
|
||||
!/test/ent/schema
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
## 0.3.0 (2025-03-11)
|
||||
|
||||
- Added distance functions for Ent
|
||||
- Added support for binary format for `HalfVector` and pgx
|
||||
- Dropped support for Go < 1.23
|
||||
|
||||
## 0.2.3 (2025-01-15)
|
||||
|
||||
- Added support for Postgres arrays for pgx
|
||||
|
||||
## 0.2.2 (2024-08-10)
|
||||
|
||||
- Added support for `CopyFrom` with `string` values
|
||||
|
||||
## 0.2.1 (2024-07-23)
|
||||
|
||||
- Added `pgx` package
|
||||
- Added `EncodeBinary` and `DecodeBinary` methods to `Vector`
|
||||
- Added `EncodeBinary` and `DecodeBinary` methods to `SparseVector`
|
||||
|
||||
## 0.2.0 (2024-06-25)
|
||||
|
||||
- Added support for `halfvec` and `sparsevec` types
|
||||
- Added support for JSON serialization
|
||||
- Improved performance of serialization
|
||||
- Dropped support for Go < 1.21
|
||||
|
||||
## 0.1.1 (2023-04-20)
|
||||
|
||||
- Fixed `Scan` for pgx
|
||||
|
||||
## 0.1.0 (2023-03-13)
|
||||
|
||||
- First release
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2021-2025 Andrew Kane
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
+483
@@ -0,0 +1,483 @@
|
||||
# pgvector-go
|
||||
|
||||
[pgvector](https://github.com/pgvector/pgvector) support for Go
|
||||
|
||||
Supports [pgx](https://github.com/jackc/pgx), [pg](https://github.com/go-pg/pg), [Bun](https://github.com/uptrace/bun), [Ent](https://github.com/ent/ent), [GORM](https://github.com/go-gorm/gorm), and [sqlx](https://github.com/jmoiron/sqlx)
|
||||
|
||||
[](https://github.com/pgvector/pgvector-go/actions)
|
||||
|
||||
## Getting Started
|
||||
|
||||
Run:
|
||||
|
||||
```sh
|
||||
go get github.com/pgvector/pgvector-go
|
||||
```
|
||||
|
||||
And follow the instructions for your database library:
|
||||
|
||||
- [pgx](#pgx)
|
||||
- [pg](#pg)
|
||||
- [Bun](#bun)
|
||||
- [Ent](#ent)
|
||||
- [GORM](#gorm)
|
||||
- [sqlx](#sqlx)
|
||||
|
||||
Or check out some examples:
|
||||
|
||||
- [Embeddings](examples/openai/main.go) with OpenAI
|
||||
- [Binary embeddings](examples/cohere/main.go) with Cohere
|
||||
- [Hybrid search](examples/hybrid/main.go) with Ollama (Reciprocal Rank Fusion)
|
||||
- [Sparse search](examples/sparse/main.go) with Text Embeddings Inference
|
||||
- [Recommendations](examples/disco/main.go) with Disco
|
||||
- [Horizontal scaling](examples/citus/main.go) with Citus
|
||||
- [Bulk loading](examples/loading/main.go) with `COPY`
|
||||
|
||||
## pgx
|
||||
|
||||
Import the packages
|
||||
|
||||
```go
|
||||
import (
|
||||
"github.com/pgvector/pgvector-go"
|
||||
pgxvec "github.com/pgvector/pgvector-go/pgx"
|
||||
)
|
||||
```
|
||||
|
||||
Enable the extension
|
||||
|
||||
```go
|
||||
_, err := conn.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector")
|
||||
```
|
||||
|
||||
Register the types with the connection
|
||||
|
||||
```go
|
||||
err := pgxvec.RegisterTypes(ctx, conn)
|
||||
```
|
||||
|
||||
or the pool
|
||||
|
||||
```go
|
||||
config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
|
||||
return pgxvec.RegisterTypes(ctx, conn)
|
||||
}
|
||||
```
|
||||
|
||||
Create a table
|
||||
|
||||
```go
|
||||
_, err := conn.Exec(ctx, "CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))")
|
||||
```
|
||||
|
||||
Insert a vector
|
||||
|
||||
```go
|
||||
_, err := conn.Exec(ctx, "INSERT INTO items (embedding) VALUES ($1)", pgvector.NewVector([]float32{1, 2, 3}))
|
||||
```
|
||||
|
||||
Get the nearest neighbors to a vector
|
||||
|
||||
```go
|
||||
rows, err := conn.Query(ctx, "SELECT id FROM items ORDER BY embedding <-> $1 LIMIT 5", pgvector.NewVector([]float32{1, 2, 3}))
|
||||
```
|
||||
|
||||
Add an approximate index
|
||||
|
||||
```go
|
||||
_, err := conn.Exec(ctx, "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
|
||||
// or
|
||||
_, err := conn.Exec(ctx, "CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")
|
||||
```
|
||||
|
||||
Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
|
||||
|
||||
See a [full example](pgx_test.go)
|
||||
|
||||
## pg
|
||||
|
||||
Import the package
|
||||
|
||||
```go
|
||||
import "github.com/pgvector/pgvector-go"
|
||||
```
|
||||
|
||||
Enable the extension
|
||||
|
||||
```go
|
||||
_, err := db.Exec("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
```
|
||||
|
||||
Add a vector column
|
||||
|
||||
```go
|
||||
type Item struct {
|
||||
Embedding pgvector.Vector `pg:"type:vector(3)"`
|
||||
}
|
||||
```
|
||||
|
||||
Insert a vector
|
||||
|
||||
```go
|
||||
item := Item{
|
||||
Embedding: pgvector.NewVector([]float32{1, 2, 3}),
|
||||
}
|
||||
_, err := db.Model(&item).Insert()
|
||||
```
|
||||
|
||||
Get the nearest neighbors to a vector
|
||||
|
||||
```go
|
||||
var items []Item
|
||||
err := db.Model(&items).
|
||||
OrderExpr("embedding <-> ?", pgvector.NewVector([]float32{1, 2, 3})).
|
||||
Limit(5).
|
||||
Select()
|
||||
```
|
||||
|
||||
Add an approximate index
|
||||
|
||||
```go
|
||||
_, err := conn.Exec(ctx, "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
|
||||
// or
|
||||
_, err := conn.Exec(ctx, "CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")
|
||||
```
|
||||
|
||||
Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
|
||||
|
||||
See a [full example](pg_test.go)
|
||||
|
||||
## Bun
|
||||
|
||||
Import the package
|
||||
|
||||
```go
|
||||
import "github.com/pgvector/pgvector-go"
|
||||
```
|
||||
|
||||
Enable the extension
|
||||
|
||||
```go
|
||||
_, err := db.Exec("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
```
|
||||
|
||||
Add a vector column
|
||||
|
||||
```go
|
||||
type Item struct {
|
||||
Embedding pgvector.Vector `bun:"type:vector(3)"`
|
||||
}
|
||||
```
|
||||
|
||||
Insert a vector
|
||||
|
||||
```go
|
||||
item := Item{
|
||||
Embedding: pgvector.NewVector([]float32{1, 2, 3}),
|
||||
}
|
||||
_, err := db.NewInsert().Model(&item).Exec(ctx)
|
||||
```
|
||||
|
||||
Get the nearest neighbors to a vector
|
||||
|
||||
```go
|
||||
var items []Item
|
||||
err := db.NewSelect().
|
||||
Model(&items).
|
||||
OrderExpr("embedding <-> ?", pgvector.NewVector([]float32{1, 2, 3})).
|
||||
Limit(5).
|
||||
Scan(ctx)
|
||||
```
|
||||
|
||||
Add an approximate index
|
||||
|
||||
```go
|
||||
var _ bun.AfterCreateTableHook = (*Item)(nil)
|
||||
|
||||
func (*Item) AfterCreateTable(ctx context.Context, query *bun.CreateTableQuery) error {
|
||||
_, err := query.DB().NewCreateIndex().
|
||||
Model((*Item)(nil)).
|
||||
Index("items_embedding_idx").
|
||||
ColumnExpr("embedding vector_l2_ops").
|
||||
Using("hnsw").
|
||||
Exec(ctx)
|
||||
return err
|
||||
}
|
||||
```
|
||||
|
||||
Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
|
||||
|
||||
See a [full example](bun_test.go)
|
||||
|
||||
## Ent
|
||||
|
||||
Import the package
|
||||
|
||||
```go
|
||||
import "github.com/pgvector/pgvector-go"
|
||||
import entvec "github.com/pgvector/pgvector-go/ent"
|
||||
```
|
||||
|
||||
Enable the extension (requires the [sql/execquery](https://entgo.io/docs/feature-flags/#sql-raw-api) feature)
|
||||
|
||||
```go
|
||||
_, err := client.ExecContext(ctx, "CREATE EXTENSION IF NOT EXISTS vector")
|
||||
```
|
||||
|
||||
Add a vector column
|
||||
|
||||
```go
|
||||
func (Item) Fields() []ent.Field {
|
||||
return []ent.Field{
|
||||
field.Other("embedding", pgvector.Vector{}).
|
||||
SchemaType(map[string]string{
|
||||
dialect.Postgres: "vector(3)",
|
||||
}),
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Insert a vector
|
||||
|
||||
```go
|
||||
_, err := client.Item.
|
||||
Create().
|
||||
SetEmbedding(pgvector.NewVector([]float32{1, 2, 3})).
|
||||
Save(ctx)
|
||||
```
|
||||
|
||||
Get the nearest neighbors to a vector
|
||||
|
||||
```go
|
||||
items, err := client.Item.
|
||||
Query().
|
||||
Order(func(s *sql.Selector) {
|
||||
s.OrderExpr(entvec.L2Distance("embedding", pgvector.NewVector([]float32{1, 2, 3})))
|
||||
}).
|
||||
Limit(5).
|
||||
All(ctx)
|
||||
```
|
||||
|
||||
Also supports `MaxInnerProduct`, `CosineDistance`, `L1Distance`, `HammingDistance`, and `JaccardDistance`
|
||||
|
||||
Add an approximate index
|
||||
|
||||
```go
|
||||
func (Item) Indexes() []ent.Index {
|
||||
return []ent.Index{
|
||||
index.Fields("embedding").
|
||||
Annotations(
|
||||
entsql.IndexType("hnsw"),
|
||||
entsql.OpClass("vector_l2_ops"),
|
||||
),
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
|
||||
|
||||
See a [full example](ent_test.go)
|
||||
|
||||
## GORM
|
||||
|
||||
Import the package
|
||||
|
||||
```go
|
||||
import "github.com/pgvector/pgvector-go"
|
||||
```
|
||||
|
||||
Enable the extension
|
||||
|
||||
```go
|
||||
db.Exec("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
```
|
||||
|
||||
Add a vector column
|
||||
|
||||
```go
|
||||
type Item struct {
|
||||
Embedding pgvector.Vector `gorm:"type:vector(3)"`
|
||||
}
|
||||
```
|
||||
|
||||
Insert a vector
|
||||
|
||||
```go
|
||||
item := Item{
|
||||
Embedding: pgvector.NewVector([]float32{1, 2, 3}),
|
||||
}
|
||||
result := db.Create(&item)
|
||||
```
|
||||
|
||||
Get the nearest neighbors to a vector
|
||||
|
||||
```go
|
||||
var items []Item
|
||||
db.Clauses(clause.OrderBy{
|
||||
Expression: clause.Expr{SQL: "embedding <-> ?", Vars: []interface{}{pgvector.NewVector([]float32{1, 1, 1})}},
|
||||
}).Limit(5).Find(&items)
|
||||
```
|
||||
|
||||
Add an approximate index
|
||||
|
||||
```go
|
||||
db.Exec("CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
|
||||
// or
|
||||
db.Exec("CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")
|
||||
```
|
||||
|
||||
Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
|
||||
|
||||
See a [full example](gorm_test.go)
|
||||
|
||||
## sqlx
|
||||
|
||||
Import the package
|
||||
|
||||
```go
|
||||
import "github.com/pgvector/pgvector-go"
|
||||
```
|
||||
|
||||
Enable the extension
|
||||
|
||||
```go
|
||||
db.MustExec("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
```
|
||||
|
||||
Add a vector column
|
||||
|
||||
```go
|
||||
type Item struct {
|
||||
Embedding pgvector.Vector
|
||||
}
|
||||
```
|
||||
|
||||
Insert a vector
|
||||
|
||||
```go
|
||||
item := Item{
|
||||
Embedding: pgvector.NewVector([]float32{1, 2, 3}),
|
||||
}
|
||||
_, err := db.NamedExec(`INSERT INTO items (embedding) VALUES (:embedding)`, item)
|
||||
```
|
||||
|
||||
Get the nearest neighbors to a vector
|
||||
|
||||
```go
|
||||
var items []Item
|
||||
db.Select(&items, "SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5", pgvector.NewVector([]float32{1, 1, 1}))
|
||||
```
|
||||
|
||||
Add an approximate index
|
||||
|
||||
```go
|
||||
db.MustExec("CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
|
||||
// or
|
||||
db.MustExec("CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")
|
||||
```
|
||||
|
||||
Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
|
||||
|
||||
See a [full example](sqlx_test.go)
|
||||
|
||||
## Reference
|
||||
|
||||
### Vectors
|
||||
|
||||
Create a vector from a slice
|
||||
|
||||
```go
|
||||
vec := pgvector.NewVector([]float32{1, 2, 3})
|
||||
```
|
||||
|
||||
Get a slice
|
||||
|
||||
```go
|
||||
slice := vec.Slice()
|
||||
```
|
||||
|
||||
### Half Vectors
|
||||
|
||||
Create a half vector from a slice
|
||||
|
||||
```go
|
||||
vec := pgvector.NewHalfVector([]float32{1, 2, 3})
|
||||
```
|
||||
|
||||
Get a slice
|
||||
|
||||
```go
|
||||
slice := vec.Slice()
|
||||
```
|
||||
|
||||
### Sparse Vectors
|
||||
|
||||
Create a sparse vector from a slice
|
||||
|
||||
```go
|
||||
vec := pgvector.NewSparseVector([]float32{1, 0, 2, 0, 3, 0})
|
||||
```
|
||||
|
||||
Or a map of non-zero elements
|
||||
|
||||
```go
|
||||
elements := map[int32]float32{0: 1, 2: 2, 4: 3}
|
||||
vec := pgvector.NewSparseVectorFromMap(elements, 6)
|
||||
```
|
||||
|
||||
Note: Indices start at 0
|
||||
|
||||
Get the number of dimensions
|
||||
|
||||
```go
|
||||
dim := vec.Dimensions()
|
||||
```
|
||||
|
||||
Get the indices of non-zero elements
|
||||
|
||||
```go
|
||||
indices := vec.Indices()
|
||||
```
|
||||
|
||||
Get the values of non-zero elements
|
||||
|
||||
```go
|
||||
values := vec.Values()
|
||||
```
|
||||
|
||||
Get a slice
|
||||
|
||||
```go
|
||||
slice := vec.Slice()
|
||||
```
|
||||
|
||||
## History
|
||||
|
||||
View the [changelog](https://github.com/pgvector/pgvector-go/blob/master/CHANGELOG.md)
|
||||
|
||||
## Contributing
|
||||
|
||||
Everyone is encouraged to help improve this project. Here are a few ways you can help:
|
||||
|
||||
- [Report bugs](https://github.com/pgvector/pgvector-go/issues)
|
||||
- Fix bugs and [submit pull requests](https://github.com/pgvector/pgvector-go/pulls)
|
||||
- Write, clarify, or fix documentation
|
||||
- Suggest or add new features
|
||||
|
||||
To get started with development:
|
||||
|
||||
```sh
|
||||
git clone https://github.com/pgvector/pgvector-go.git
|
||||
cd pgvector-go
|
||||
go mod tidy
|
||||
createdb pgvector_go_test
|
||||
go generate ./test/ent
|
||||
go test -v
|
||||
```
|
||||
|
||||
To run an example:
|
||||
|
||||
```sh
|
||||
createdb pgvector_example
|
||||
go run ./examples/loading
|
||||
```
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
package pgvector
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"database/sql/driver"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// HalfVector is a wrapper for []float32 to implement sql.Scanner and driver.Valuer.
|
||||
type HalfVector struct {
|
||||
vec []float32
|
||||
}
|
||||
|
||||
// NewHalfVector creates a new HalfVector from a slice of float32.
|
||||
func NewHalfVector(vec []float32) HalfVector {
|
||||
return HalfVector{vec: vec}
|
||||
}
|
||||
|
||||
// Slice returns the underlying slice of float32.
|
||||
func (v HalfVector) Slice() []float32 {
|
||||
return v.vec
|
||||
}
|
||||
|
||||
// String returns a string representation of the half vector.
|
||||
func (v HalfVector) String() string {
|
||||
// should never throw an error
|
||||
// but returning an empty string is fine if it does
|
||||
buf, _ := v.EncodeText(nil)
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
// Parse parses a string representation of a half vector.
|
||||
func (v *HalfVector) Parse(s string) error {
|
||||
sp := strings.Split(s[1:len(s)-1], ",")
|
||||
v.vec = make([]float32, 0, len(sp))
|
||||
for i := 0; i < len(sp); i++ {
|
||||
n, err := strconv.ParseFloat(sp[i], 32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
v.vec = append(v.vec, float32(n))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// EncodeText encodes a text representation of the half vector.
|
||||
func (v HalfVector) EncodeText(buf []byte) (newBuf []byte, err error) {
|
||||
buf = slices.Grow(buf, 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 buf, nil
|
||||
}
|
||||
|
||||
// SetSlice sets the underlying slice of float32.
|
||||
func (v *HalfVector) SetSlice(vec []float32) {
|
||||
v.vec = vec
|
||||
}
|
||||
|
||||
// statically assert that HalfVector implements sql.Scanner.
|
||||
var _ sql.Scanner = (*HalfVector)(nil)
|
||||
|
||||
// Scan implements the sql.Scanner interface.
|
||||
func (v *HalfVector) Scan(src interface{}) (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 HalfVector implements driver.Valuer.
|
||||
var _ driver.Valuer = (*HalfVector)(nil)
|
||||
|
||||
// Value implements the driver.Valuer interface.
|
||||
func (v HalfVector) Value() (driver.Value, error) {
|
||||
return v.String(), nil
|
||||
}
|
||||
|
||||
// statically assert that HalfVector implements json.Marshaler.
|
||||
var _ json.Marshaler = (*HalfVector)(nil)
|
||||
|
||||
// MarshalJSON implements the json.Marshaler interface.
|
||||
func (v HalfVector) MarshalJSON() ([]byte, error) {
|
||||
return json.Marshal(v.vec)
|
||||
}
|
||||
|
||||
// statically assert that HalfVector implements json.Unmarshaler.
|
||||
var _ json.Unmarshaler = (*HalfVector)(nil)
|
||||
|
||||
// UnmarshalJSON implements the json.Unmarshaler interface.
|
||||
func (v *HalfVector) UnmarshalJSON(data []byte) error {
|
||||
return json.Unmarshal(data, &v.vec)
|
||||
}
|
||||
+128
@@ -0,0 +1,128 @@
|
||||
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
|
||||
}
|
||||
+49
@@ -0,0 +1,49 @@
|
||||
package pgx
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
)
|
||||
|
||||
func RegisterTypes(ctx context.Context, conn *pgx.Conn) error {
|
||||
var vectorOid *uint32
|
||||
var vectorArrayOid *uint32
|
||||
var halfvecOid *uint32
|
||||
var halfvecArrayOid *uint32
|
||||
var sparsevecOid *uint32
|
||||
var sparsevecArrayOid *uint32
|
||||
err := conn.QueryRow(ctx, "SELECT to_regtype('vector')::oid, to_regtype('_vector')::oid, to_regtype('halfvec')::oid, to_regtype('_halfvec')::oid, to_regtype('sparsevec')::oid, to_regtype('_sparsevec')::oid").Scan(&vectorOid, &vectorArrayOid, &halfvecOid, &halfvecArrayOid, &sparsevecOid, &sparsevecArrayOid)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if vectorOid == nil {
|
||||
return fmt.Errorf("vector type not found in the database")
|
||||
}
|
||||
|
||||
tm := conn.TypeMap()
|
||||
registerType(tm, "vector", vectorOid, vectorArrayOid, &VectorCodec{})
|
||||
|
||||
if halfvecOid != nil {
|
||||
registerType(tm, "halfvec", halfvecOid, halfvecArrayOid, &HalfVectorCodec{})
|
||||
}
|
||||
|
||||
if sparsevecOid != nil {
|
||||
registerType(tm, "sparsevec", sparsevecOid, sparsevecArrayOid, &SparseVectorCodec{})
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func registerType(tm *pgtype.Map, name string, oid *uint32, arrayOid *uint32, codec pgtype.Codec) {
|
||||
t := pgtype.Type{Name: name, OID: *oid, Codec: codec}
|
||||
tm.RegisterType(&t)
|
||||
|
||||
// should never be nil
|
||||
if arrayOid != nil {
|
||||
tm.RegisterType(&pgtype.Type{Name: "_" + name, OID: *arrayOid, Codec: &pgtype.ArrayCodec{ElementType: &t}})
|
||||
}
|
||||
}
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
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
|
||||
}
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
package pgx
|
||||
|
||||
import (
|
||||
"database/sql/driver"
|
||||
"fmt"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
"github.com/pgvector/pgvector-go"
|
||||
)
|
||||
|
||||
type VectorCodec struct{}
|
||||
|
||||
func (VectorCodec) FormatSupported(format int16) bool {
|
||||
return format == pgx.BinaryFormatCode || format == pgx.TextFormatCode
|
||||
}
|
||||
|
||||
func (VectorCodec) PreferredFormat() int16 {
|
||||
return pgx.BinaryFormatCode
|
||||
}
|
||||
|
||||
func (VectorCodec) PlanEncode(m *pgtype.Map, oid uint32, format int16, value any) pgtype.EncodePlan {
|
||||
_, ok := value.(pgvector.Vector)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch format {
|
||||
case pgx.BinaryFormatCode:
|
||||
return encodePlanVectorCodecBinary{}
|
||||
case pgx.TextFormatCode:
|
||||
return encodePlanVectorCodecText{}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type encodePlanVectorCodecBinary struct{}
|
||||
|
||||
func (encodePlanVectorCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) {
|
||||
v := value.(pgvector.Vector)
|
||||
return v.EncodeBinary(buf)
|
||||
}
|
||||
|
||||
type encodePlanVectorCodecText struct{}
|
||||
|
||||
func (encodePlanVectorCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) {
|
||||
v := value.(pgvector.Vector)
|
||||
// use String() for now to avoid adding another method to Vector
|
||||
return append(buf, v.String()...), nil
|
||||
}
|
||||
|
||||
func (VectorCodec) PlanScan(m *pgtype.Map, oid uint32, format int16, target any) pgtype.ScanPlan {
|
||||
_, ok := target.(*pgvector.Vector)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch format {
|
||||
case pgx.BinaryFormatCode:
|
||||
return scanPlanVectorCodecBinary{}
|
||||
case pgx.TextFormatCode:
|
||||
return scanPlanVectorCodecText{}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type scanPlanVectorCodecBinary struct{}
|
||||
|
||||
func (scanPlanVectorCodecBinary) Scan(src []byte, dst any) error {
|
||||
v := (dst).(*pgvector.Vector)
|
||||
return v.DecodeBinary(src)
|
||||
}
|
||||
|
||||
type scanPlanVectorCodecText struct{}
|
||||
|
||||
func (scanPlanVectorCodecText) Scan(src []byte, dst any) error {
|
||||
v := (dst).(*pgvector.Vector)
|
||||
return v.Scan(src)
|
||||
}
|
||||
|
||||
func (c VectorCodec) DecodeDatabaseSQLValue(m *pgtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) {
|
||||
return c.DecodeValue(m, oid, format, src)
|
||||
}
|
||||
|
||||
func (c VectorCodec) DecodeValue(m *pgtype.Map, oid uint32, format int16, src []byte) (any, error) {
|
||||
if src == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var vec pgvector.Vector
|
||||
scanPlan := c.PlanScan(m, oid, format, &vec)
|
||||
if scanPlan == nil {
|
||||
return nil, fmt.Errorf("Unable to decode vector type")
|
||||
}
|
||||
|
||||
err := scanPlan.Scan(src, &vec)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return vec, nil
|
||||
}
|
||||
+192
@@ -0,0 +1,192 @@
|
||||
package pgvector
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"database/sql/driver"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"math"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// SparseVector is a wrapper to implement sql.Scanner and driver.Valuer.
|
||||
type SparseVector struct {
|
||||
dim int32
|
||||
indices []int32
|
||||
values []float32
|
||||
}
|
||||
|
||||
// NewSparseVector creates a new SparseVector from a slice of float32.
|
||||
func NewSparseVector(vec []float32) SparseVector {
|
||||
dim := int32(len(vec))
|
||||
indices := make([]int32, 0)
|
||||
values := make([]float32, 0)
|
||||
for i := 0; i < len(vec); i++ {
|
||||
if vec[i] != 0 {
|
||||
indices = append(indices, int32(i))
|
||||
values = append(values, vec[i])
|
||||
}
|
||||
}
|
||||
return SparseVector{dim: dim, indices: indices, values: values}
|
||||
}
|
||||
|
||||
// NewSparseVectorFromMap creates a new SparseVector from a map of non-zero elements.
|
||||
func NewSparseVectorFromMap(elements map[int32]float32, dim int32) SparseVector {
|
||||
indices := make([]int32, 0, len(elements))
|
||||
values := make([]float32, 0, len(elements))
|
||||
for k, v := range elements {
|
||||
if v != 0 {
|
||||
indices = append(indices, k)
|
||||
}
|
||||
}
|
||||
slices.Sort(indices)
|
||||
for _, k := range indices {
|
||||
values = append(values, elements[k])
|
||||
}
|
||||
return SparseVector{dim: dim, indices: indices, values: values}
|
||||
}
|
||||
|
||||
// Dimensions returns the number of dimensions.
|
||||
func (v SparseVector) Dimensions() int32 {
|
||||
return v.dim
|
||||
}
|
||||
|
||||
// Indices returns the non-zero indices.
|
||||
func (v SparseVector) Indices() []int32 {
|
||||
return v.indices
|
||||
}
|
||||
|
||||
// Values returns the non-zero values.
|
||||
func (v SparseVector) Values() []float32 {
|
||||
return v.values
|
||||
}
|
||||
|
||||
// Slice returns a slice of float32.
|
||||
func (v SparseVector) Slice() []float32 {
|
||||
vec := make([]float32, v.dim)
|
||||
for i := 0; i < len(v.indices); i++ {
|
||||
vec[v.indices[i]] = v.values[i]
|
||||
}
|
||||
return vec
|
||||
}
|
||||
|
||||
// String returns a string representation of the sparse vector.
|
||||
func (v SparseVector) String() string {
|
||||
buf := make([]byte, 0, 13+27*len(v.indices))
|
||||
buf = append(buf, '{')
|
||||
|
||||
for i := 0; i < len(v.indices); i++ {
|
||||
if i > 0 {
|
||||
buf = append(buf, ',')
|
||||
}
|
||||
buf = strconv.AppendInt(buf, int64(v.indices[i])+1, 10)
|
||||
buf = append(buf, ':')
|
||||
buf = strconv.AppendFloat(buf, float64(v.values[i]), 'f', -1, 32)
|
||||
}
|
||||
|
||||
buf = append(buf, '}')
|
||||
buf = append(buf, '/')
|
||||
buf = strconv.AppendInt(buf, int64(v.dim), 10)
|
||||
return string(buf)
|
||||
}
|
||||
|
||||
// Parse parses a string representation of a sparse vector.
|
||||
func (v *SparseVector) Parse(s string) error {
|
||||
sp := strings.SplitN(s, "/", 2)
|
||||
|
||||
dim, err := strconv.ParseInt(sp[1], 10, 32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
elements := strings.Split(sp[0][1:len(sp[0])-1], ",")
|
||||
v.dim = int32(dim)
|
||||
v.indices = make([]int32, 0, len(elements))
|
||||
v.values = make([]float32, 0, len(elements))
|
||||
|
||||
for i := 0; i < len(elements); i++ {
|
||||
ep := strings.SplitN(elements[i], ":", 2)
|
||||
|
||||
n, err := strconv.ParseInt(ep[0], 10, 32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
v.indices = append(v.indices, int32(n-1))
|
||||
|
||||
n2, err := strconv.ParseFloat(ep[1], 32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
v.values = append(v.values, float32(n2))
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// EncodeBinary encodes a binary representation of the sparse vector.
|
||||
func (v SparseVector) EncodeBinary(buf []byte) (newBuf []byte, err error) {
|
||||
nnz := len(v.indices)
|
||||
buf = slices.Grow(buf, 12+8*nnz)
|
||||
buf = binary.BigEndian.AppendUint32(buf, uint32(v.dim))
|
||||
buf = binary.BigEndian.AppendUint32(buf, uint32(nnz))
|
||||
buf = binary.BigEndian.AppendUint32(buf, 0)
|
||||
for _, v := range v.indices {
|
||||
buf = binary.BigEndian.AppendUint32(buf, uint32(v))
|
||||
}
|
||||
for _, v := range v.values {
|
||||
buf = binary.BigEndian.AppendUint32(buf, math.Float32bits(v))
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
// DecodeBinary decodes a binary representation of a sparse vector.
|
||||
func (v *SparseVector) DecodeBinary(buf []byte) error {
|
||||
dim := binary.BigEndian.Uint32(buf[0:4])
|
||||
nnz := int(binary.BigEndian.Uint32(buf[4:8]))
|
||||
unused := binary.BigEndian.Uint32(buf[8:12])
|
||||
if unused != 0 {
|
||||
return fmt.Errorf("expected unused to be 0")
|
||||
}
|
||||
|
||||
v.dim = int32(dim)
|
||||
v.indices = make([]int32, 0, nnz)
|
||||
v.values = make([]float32, 0, nnz)
|
||||
offset := 12
|
||||
|
||||
for i := 0; i < nnz; i++ {
|
||||
v.indices = append(v.indices, int32(binary.BigEndian.Uint32(buf[offset:offset+4])))
|
||||
offset += 4
|
||||
}
|
||||
|
||||
for i := 0; i < nnz; i++ {
|
||||
v.values = append(v.values, math.Float32frombits(binary.BigEndian.Uint32(buf[offset:offset+4])))
|
||||
offset += 4
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// statically assert that SparseVector implements sql.Scanner.
|
||||
var _ sql.Scanner = (*SparseVector)(nil)
|
||||
|
||||
// Scan implements the sql.Scanner interface.
|
||||
func (v *SparseVector) Scan(src interface{}) (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 SparseVector implements driver.Valuer.
|
||||
var _ driver.Valuer = (*SparseVector)(nil)
|
||||
|
||||
// Value implements the driver.Valuer interface.
|
||||
func (v SparseVector) Value() (driver.Value, error) {
|
||||
return v.String(), nil
|
||||
}
|
||||
+126
@@ -0,0 +1,126 @@
|
||||
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 {
|
||||
sp := strings.Split(s[1:len(s)-1], ",")
|
||||
v.vec = make([]float32, 0, len(sp))
|
||||
for i := 0; i < len(sp); i++ {
|
||||
n, err := strconv.ParseFloat(sp[i], 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 {
|
||||
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")
|
||||
}
|
||||
|
||||
v.vec = make([]float32, 0, dim)
|
||||
offset := 4
|
||||
for i := 0; i < dim; i++ {
|
||||
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 interface{}) (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)
|
||||
}
|
||||
Reference in New Issue
Block a user