what
21b80bdea4
- exec/scanner: 用 *interface{} 替换 **json.RawMessage 扫描目标，兼容 DuckDB 返回 map[string]interface{} 的场景；新增 toJSONRawMessage 转换函数
- exec/scanner: ScanVal 支持结构体指针，通过 JSON 中间层转换（DuckDB STRUCT 列）
- exec/scanner: 将 *sql.RawBytes 和 *[]byte 的处理从 ScanValContext 移入 scanner.ScanVal
- exec/query_executor: 简化 ScanValContext，移除私有 scan 方法
- exec: 补充 scanner 级别 ScanVal 测试用例
- internal/util/reflect: 重写 SafeSetVarValue，修复非指针 src 及 nil 指针字段的 panic
- internal/util/column_map: 恢复非匿名带标签结构体字段的展开逻辑
- schema: 新增 vector 列类型支持
- engine: 补充 DuckDB 相关配置
- dialect/sqlite3/vtab: 完善虚拟表适配器
- 各方言测试改用 sqlmock 虚拟连接
102 lines
3.0 KiB
Go
102 lines
3.0 KiB
Go
// Package vec 在 mattn/go-sqlite3 上集成 sqlite-vec 向量检索扩展。
|
|
//
|
|
// 使用该驱动后，可直接创建 vec0 虚拟表并执行向量相似度检索：
|
|
//
|
|
// db, _ := sql.Open(vec.DriverName, "data.db")
|
|
// db.Exec(`CREATE VIRTUAL TABLE IF NOT EXISTS embeddings USING vec0(vector float[1536])`)
|
|
// db.Exec(`INSERT INTO embeddings(rowid, vector) VALUES (?, ?)`, id, vec.SerializeFloat32(embedding))
|
|
// db.QueryRow(`SELECT rowid, distance FROM embeddings WHERE vector MATCH ? ORDER BY distance LIMIT 10`, vec.SerializeFloat32(query))
|
|
package vec
|
|
|
|
import (
|
|
"database/sql"
|
|
"encoding/binary"
|
|
"encoding/json"
|
|
"math"
|
|
"regexp"
|
|
|
|
"git.fsdpf.net/go/db"
|
|
"git.fsdpf.net/go/db/dialect/sqlite3"
|
|
sqlitevec "github.com/asg017/sqlite-vec-go-bindings/cgo"
|
|
gosqlite3 "github.com/mattn/go-sqlite3"
|
|
)
|
|
|
|
// DriverName 是加载了 sqlite-vec 扩展的 SQLite3 驱动名称。
|
|
const DriverName = "sqlite3_vec"
|
|
|
|
func init() {
|
|
sqlitevec.Auto() // 全局加载 sqlite-vec 扩展到所有后续连接
|
|
|
|
sql.Register(DriverName, &gosqlite3.SQLiteDriver{
|
|
ConnectHook: func(conn *gosqlite3.SQLiteConn) error {
|
|
// sqlitevec.LoadIntoConn 在当前版本库中不存在，改用 init() 中的 Auto() 全局注册
|
|
// if err := sqlitevec.LoadIntoConn(conn); err != nil {
|
|
// return err
|
|
// }
|
|
|
|
if err := conn.RegisterFunc("IF", func(cond int64, trueVal, falseVal interface{}) interface{} {
|
|
if cond != 0 {
|
|
return trueVal
|
|
}
|
|
return falseVal
|
|
}, true); err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := conn.RegisterFunc("REGEXP", func(expr, item string) (bool, error) {
|
|
return regexp.MatchString(expr, item)
|
|
}, true); err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := conn.RegisterFunc("vec_serialize", func(jsonText string) ([]byte, error) {
|
|
var floats []float64
|
|
if err := json.Unmarshal([]byte(jsonText), &floats); err != nil {
|
|
return nil, err
|
|
}
|
|
f32 := make([]float32, len(floats))
|
|
for i, f := range floats {
|
|
f32[i] = float32(f)
|
|
}
|
|
return SerializeFloat32(f32), nil
|
|
}, true); err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := conn.RegisterFunc("vec_deserialize", func(b []byte) (string, error) {
|
|
floats := DeserializeFloat32(b)
|
|
f64 := make([]float64, len(floats))
|
|
for i, f := range floats {
|
|
f64[i] = float64(f)
|
|
}
|
|
out, err := json.Marshal(f64)
|
|
return string(out), err
|
|
}, true); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
},
|
|
})
|
|
|
|
db.RegisterDialect(DriverName, sqlite3.DialectOptions())
|
|
}
|
|
|
|
// SerializeFloat32 将 float32 切片序列化为 sqlite-vec 接受的小端 IEEE 754 字节序列。
|
|
func SerializeFloat32(v []float32) []byte {
|
|
buf := make([]byte, len(v)*4)
|
|
for i, f := range v {
|
|
binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(f))
|
|
}
|
|
return buf
|
|
}
|
|
|
|
// DeserializeFloat32 将 sqlite-vec 返回的字节序列反序列化为 float32 切片。
|
|
func DeserializeFloat32(b []byte) []float32 {
|
|
v := make([]float32, len(b)/4)
|
|
for i := range v {
|
|
v[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[i*4:]))
|
|
}
|
|
return v
|
|
}
|