feat: 完善扫描器、exec 及 schema 相关功能
- 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 虚拟连接
This commit is contained in:
@@ -0,0 +1,101 @@
|
||||
// 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
|
||||
}
|
||||
Reference in New Issue
Block a user