@an5/adapters 0.2.11 → 0.2.14
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +137 -5
- package/dist/an5Adapter.d.ts +42 -1
- package/dist/an5Adapter.js +146 -38
- package/dist/base/types.d.ts +14 -0
- package/dist/index.d.ts +2 -0
- package/dist/index.js +8 -1
- package/dist/sqlite/browserEngine.d.ts +13 -0
- package/dist/sqlite/browserEngine.js +13 -0
- package/dist/sqlite/engine.d.ts +5 -0
- package/dist/sqlite/engine.js +10 -0
- package/dist/sqlite/index.d.ts +1 -0
- package/dist/sqlite/index.js +1 -0
- package/dist/sqlite/vector.d.ts +155 -0
- package/dist/sqlite/vector.js +432 -0
- package/dotnet/An5Adapter.cs +99 -28
- package/dotnet/Base/Core.cs +14 -0
- package/dotnet/Mssql/MssqlEngine.cs +44 -1
- package/dotnet/Postgres/PostgresEngine.cs +2 -0
- package/dotnet/Sqlite/SqliteEngine.cs +82 -2
- package/dotnet/Sqlite/SqliteVectors.cs +401 -0
- package/golang/an5_adapter.go +84 -32
- package/golang/base/vectors.go +263 -0
- package/golang/sqlite/vector.go +248 -0
- package/golang/sqlite_vector_test.go +288 -0
- package/golang/table_client.go +67 -3
- package/java/pom.xml +177 -0
- package/java/src/main/java/an5/adapters/An5Adapter.java +374 -0
- package/java/src/main/java/an5/adapters/An5Aggregate.java +116 -0
- package/java/src/main/java/an5/adapters/An5GroupBy.java +97 -0
- package/java/src/main/java/an5/adapters/An5Query.java +100 -0
- package/java/src/main/java/an5/adapters/An5TableClient.java +1030 -0
- package/java/src/main/java/an5/adapters/An5ViewClient.java +107 -0
- package/java/src/main/java/an5/adapters/base/Dialect.java +103 -0
- package/java/src/main/java/an5/adapters/base/FieldMeta.java +42 -0
- package/java/src/main/java/an5/adapters/base/JdbcUrls.java +201 -0
- package/java/src/main/java/an5/adapters/base/Metadata.java +195 -0
- package/java/src/main/java/an5/adapters/base/RelationDef.java +26 -0
- package/java/src/main/java/an5/adapters/base/SqlBuilder.java +422 -0
- package/java/src/main/java/an5/adapters/base/SqliteVectors.java +375 -0
- package/java/src/main/java/an5/adapters/base/Vectors.java +123 -0
- package/java/src/test/java/An5AdaptersSmoke.java +733 -0
- package/java/src/test/java/Json.java +174 -0
- package/java/src/test/java/QuerySemantics.java +111 -0
- package/kotlin/build.gradle.kts +145 -0
- package/kotlin/settings.gradle.kts +10 -0
- package/kotlin/src/main/kotlin/an5/adapters/Aggregate.kt +89 -0
- package/kotlin/src/main/kotlin/an5/adapters/An5.kt +91 -0
- package/kotlin/src/main/kotlin/an5/adapters/Query.kt +130 -0
- package/kotlin/src/main/kotlin/an5/adapters/Rows.kt +166 -0
- package/kotlin/src/main/kotlin/an5/adapters/TableClient.kt +169 -0
- package/kotlin/src/main/kotlin/an5/adapters/Types.kt +77 -0
- package/kotlin/src/test/kotlin/an5/adapters/An5Smoke.kt +329 -0
- package/native/sqlite/CMakeLists.txt +12 -0
- package/native/sqlite/README.md +117 -0
- package/native/sqlite/an5_vector.c +209 -0
- package/package.json +33 -7
- package/python/an5_adapter.py +37 -6
- package/python/base/__init__.py +15 -0
- package/python/base/vectors.py +431 -0
- package/python/sqlite/__init__.py +2 -2
- package/python/sqlite/provider.py +21 -1
- package/python/table_client.py +129 -59
- package/rust/src/adapter.rs +342 -24
- package/rust/src/base/mod.rs +7 -0
- package/rust/src/base/vectors.rs +293 -0
- package/rust/src/lib.rs +3 -3
- package/scripts/sqlite-native-build.js +32 -0
- package/swift/Package.swift +22 -0
- package/swift/Sources/An5Adapters/An5Adapter.swift +160 -0
- package/swift/Sources/An5Adapters/Metadata.swift +183 -0
- package/swift/Sources/An5Adapters/Query.swift +105 -0
- package/swift/Sources/An5Adapters/Row+Accessors.swift +129 -0
- package/swift/Sources/An5Adapters/SQLBuilder.swift +359 -0
- package/swift/Sources/An5Adapters/SQLiteDriver.swift +301 -0
- package/swift/Sources/An5Adapters/SqliteVectors.swift +243 -0
- package/swift/Sources/An5Adapters/TableClient.swift +671 -0
- package/swift/Sources/An5Adapters/Types.swift +127 -0
- package/swift/Sources/An5Adapters/Vectors.swift +86 -0
- package/swift/Sources/An5Adapters/ViewClient.swift +50 -0
- package/swift/Sources/CSQLite/module.modulemap +8 -0
- package/swift/Sources/CSQLite/shim.h +9 -0
- package/swift/Tests/An5AdaptersTests/QuerySemanticsTests.swift +51 -0
- package/swift/Tests/An5AdaptersTests/RuntimeSmokeTests.swift +488 -0
- package/typescript/src/an5Adapter.ts +155 -31
- package/typescript/src/base/types.ts +15 -0
- package/typescript/src/index.ts +13 -0
- package/typescript/src/sqlite/browserEngine.ts +24 -0
- package/typescript/src/sqlite/engine.ts +13 -0
- package/typescript/src/sqlite/index.ts +1 -1
- package/typescript/src/sqlite/vector.ts +537 -0
- package/golang/sqlite/sqlite.go +0 -14
package/golang/an5_adapter.go
CHANGED
|
@@ -68,6 +68,9 @@ func GetRelationsForModel(modelName string) map[string]RelationDef {
|
|
|
68
68
|
type An5Adapter struct {
|
|
69
69
|
DB *sql.DB
|
|
70
70
|
Dialect Dialect
|
|
71
|
+
// VectorSupport carries the SQLite vector settings. Set it before the first
|
|
72
|
+
// search to load sqlite-vec or to pin a strategy.
|
|
73
|
+
VectorSupport sqlite.VectorSupport
|
|
71
74
|
}
|
|
72
75
|
|
|
73
76
|
// NewAn5Adapter constructs an An5Adapter instance.
|
|
@@ -167,14 +170,47 @@ func (a *An5Adapter) VectorSearch(ctx context.Context, tableName string, targetV
|
|
|
167
170
|
vecBytes, _ := json.Marshal(targetVector)
|
|
168
171
|
vecStr := string(vecBytes)
|
|
169
172
|
|
|
170
|
-
// 1.
|
|
173
|
+
// 1. SQLite ranks through sqlite-vec, a driver function or json_each,
|
|
174
|
+
// whichever this connection supports, so the table never has to leave the
|
|
175
|
+
// database. `database/sql` cannot register a user function, so the UDF
|
|
176
|
+
// strategy needs the driver to install them; see sqlite.RegisterVectorFunctions.
|
|
177
|
+
if a.Dialect == base.DialectSqlite {
|
|
178
|
+
tail := ""
|
|
179
|
+
if strings.TrimSpace(whereClause) != "" {
|
|
180
|
+
tail = " WHERE " + whereClause
|
|
181
|
+
}
|
|
182
|
+
support := a.VectorSupport
|
|
183
|
+
ranked, ran, err := sqlite.RunVectorSearch(
|
|
184
|
+
ctx,
|
|
185
|
+
a.DB,
|
|
186
|
+
a.QueryRaw,
|
|
187
|
+
base.QuoteTable(tableName, a.Dialect),
|
|
188
|
+
sqlite.NormalizeTableName(tableName),
|
|
189
|
+
base.QuoteIdentifier(vectorField, a.Dialect),
|
|
190
|
+
vectorField,
|
|
191
|
+
targetVector,
|
|
192
|
+
distanceMetric,
|
|
193
|
+
take,
|
|
194
|
+
tail,
|
|
195
|
+
args,
|
|
196
|
+
support,
|
|
197
|
+
)
|
|
198
|
+
if err != nil {
|
|
199
|
+
return nil, err
|
|
200
|
+
}
|
|
201
|
+
if ran {
|
|
202
|
+
return DecodeVectorRows(ranked, vectorField), nil
|
|
203
|
+
}
|
|
204
|
+
// No in-database strategy: score the column in Go instead.
|
|
205
|
+
return a.vectorSearchInMemory(ctx, tableName, targetVector, take, whereClause, vectorField, distanceMetric, args)
|
|
206
|
+
}
|
|
207
|
+
|
|
208
|
+
// 2. Primary path: Native database SQL vector query execution via dialect provider (mssql / postgres)
|
|
171
209
|
var sqlQuery string
|
|
172
210
|
if a.Dialect == base.DialectPostgres {
|
|
173
211
|
sqlQuery = postgres.BuildVectorSearchQuery(tableName, vectorField, distanceMetric, dim, take, whereClause)
|
|
174
212
|
} else if a.Dialect == base.DialectMssql {
|
|
175
213
|
sqlQuery = mssql.BuildVectorSearchQuery(tableName, vectorField, distanceMetric, dim, take, whereClause)
|
|
176
|
-
} else if a.Dialect == base.DialectSqlite {
|
|
177
|
-
sqlQuery = sqlite.BuildVectorSearchQuery(tableName, vectorField, distanceMetric, dim, take, whereClause)
|
|
178
214
|
}
|
|
179
215
|
|
|
180
216
|
if sqlQuery != "" {
|
|
@@ -185,7 +221,7 @@ func (a *An5Adapter) VectorSearch(ctx context.Context, tableName string, targetV
|
|
|
185
221
|
}
|
|
186
222
|
}
|
|
187
223
|
|
|
188
|
-
//
|
|
224
|
+
// 3. Secondary fallback: Fetch rows and compute vector distance in-memory if DB lacks native vector extensions
|
|
189
225
|
fallbackQuery := "SELECT * FROM " + base.QuoteTable(tableName, a.Dialect)
|
|
190
226
|
if strings.TrimSpace(whereClause) != "" {
|
|
191
227
|
fallbackQuery += " WHERE " + whereClause
|
|
@@ -199,6 +235,45 @@ func (a *An5Adapter) VectorSearch(ctx context.Context, tableName string, targetV
|
|
|
199
235
|
return VectorSearchFallback(rows, targetVector, take, vectorField, distanceMetric), nil
|
|
200
236
|
}
|
|
201
237
|
|
|
238
|
+
// vectorSearchInMemory loads the table and scores the vector column in Go.
|
|
239
|
+
//
|
|
240
|
+
// The last resort for every provider, and the only option on a SQLite connection
|
|
241
|
+
// that offers neither sqlite-vec, a driver function nor JSON1. The whole table is
|
|
242
|
+
// loaded, so it is only reasonable for a table small enough to fit in memory.
|
|
243
|
+
func (a *An5Adapter) vectorSearchInMemory(ctx context.Context, tableName string, targetVector []float64, take int, whereClause string, vectorField string, distanceMetric string, args []interface{}) ([]map[string]interface{}, error) {
|
|
244
|
+
query := "SELECT * FROM " + base.QuoteTable(tableName, a.Dialect)
|
|
245
|
+
if strings.TrimSpace(whereClause) != "" {
|
|
246
|
+
query += " WHERE " + whereClause
|
|
247
|
+
}
|
|
248
|
+
rows, err := a.QueryRaw(ctx, query, args...)
|
|
249
|
+
if err != nil {
|
|
250
|
+
return nil, err
|
|
251
|
+
}
|
|
252
|
+
return VectorSearchFallback(rows, targetVector, take, vectorField, distanceMetric), nil
|
|
253
|
+
}
|
|
254
|
+
|
|
255
|
+
// DecodeVectorRows replaces every stored `VECTOR(n)` column with a decoded
|
|
256
|
+
// float64 slice, so a caller reads numbers rather than float32 bytes or JSON text.
|
|
257
|
+
func DecodeVectorRows(rows []map[string]interface{}, vectorField string) []map[string]interface{} {
|
|
258
|
+
for _, row := range rows {
|
|
259
|
+
for key, value := range row {
|
|
260
|
+
if key != vectorField {
|
|
261
|
+
continue
|
|
262
|
+
}
|
|
263
|
+
decoded := base.DecodeVector(value)
|
|
264
|
+
if decoded == nil {
|
|
265
|
+
continue
|
|
266
|
+
}
|
|
267
|
+
// A `distance` column is a real number, not a stored vector.
|
|
268
|
+
if key == "distance" {
|
|
269
|
+
continue
|
|
270
|
+
}
|
|
271
|
+
row[key] = decoded
|
|
272
|
+
}
|
|
273
|
+
}
|
|
274
|
+
return rows
|
|
275
|
+
}
|
|
276
|
+
|
|
202
277
|
// VectorSearchFallback ranks rows in-memory.
|
|
203
278
|
func VectorSearchFallback(rows []map[string]interface{}, targetVector []float64, take int, vectorField string, distanceMetric string) []map[string]interface{} {
|
|
204
279
|
type scored struct {
|
|
@@ -213,36 +288,14 @@ func VectorSearchFallback(rows []map[string]interface{}, targetVector []float64,
|
|
|
213
288
|
continue
|
|
214
289
|
}
|
|
215
290
|
|
|
216
|
-
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
|
|
221
|
-
_ = json.Unmarshal(v, &rowVec)
|
|
222
|
-
case []float64:
|
|
223
|
-
rowVec = v
|
|
224
|
-
case []interface{}:
|
|
225
|
-
for _, item := range v {
|
|
226
|
-
if f, ok := item.(float64); ok {
|
|
227
|
-
rowVec = append(rowVec, f)
|
|
228
|
-
}
|
|
229
|
-
}
|
|
230
|
-
}
|
|
231
|
-
|
|
232
|
-
if len(rowVec) != len(targetVector) {
|
|
291
|
+
// A float32 BLOB, a legacy JSON text column and an already-decoded slice
|
|
292
|
+
// all read the same way.
|
|
293
|
+
rowVec := base.DecodeVector(raw)
|
|
294
|
+
dist, ok := base.VectorDistance(targetVector, rowVec, distanceMetric)
|
|
295
|
+
if !ok {
|
|
233
296
|
continue
|
|
234
297
|
}
|
|
235
298
|
|
|
236
|
-
var dist float64
|
|
237
|
-
switch strings.ToLower(distanceMetric) {
|
|
238
|
-
case "euclidean":
|
|
239
|
-
dist = base.EuclideanDistance(targetVector, rowVec)
|
|
240
|
-
case "dot":
|
|
241
|
-
dist = -base.DotProduct(targetVector, rowVec)
|
|
242
|
-
default: // cosine
|
|
243
|
-
dist = 1.0 - base.CosineSimilarity(targetVector, rowVec)
|
|
244
|
-
}
|
|
245
|
-
|
|
246
299
|
rCopy := make(map[string]interface{})
|
|
247
300
|
for k, val := range r {
|
|
248
301
|
rCopy[k] = val
|
|
@@ -311,4 +364,3 @@ func (a *An5Adapter) ExecuteProc(ctx context.Context, procName string, args ...i
|
|
|
311
364
|
}
|
|
312
365
|
return a.ExecuteRaw(ctx, sql, args...)
|
|
313
366
|
}
|
|
314
|
-
|
|
@@ -0,0 +1,263 @@
|
|
|
1
|
+
package base
|
|
2
|
+
|
|
3
|
+
import (
|
|
4
|
+
"encoding/binary"
|
|
5
|
+
"encoding/json"
|
|
6
|
+
"math"
|
|
7
|
+
"strings"
|
|
8
|
+
)
|
|
9
|
+
|
|
10
|
+
// ─── SQLite vector codec ──────────────────────────────────────────────────────────
|
|
11
|
+
//
|
|
12
|
+
// SQLite has no vector type, so a `VECTOR(n)` column stores a BLOB of little-endian
|
|
13
|
+
// float32 values. These helpers are shared by the adapter (which reads and writes
|
|
14
|
+
// the column) and by the `sqlite` package (which builds the ranking SQL).
|
|
15
|
+
//
|
|
16
|
+
// BytesPerVectorFloat is the size of one stored value; the column length in bytes
|
|
17
|
+
// is the dimension times this, which is what the sqlite-vec strategy filters on.
|
|
18
|
+
const BytesPerVectorFloat = 4
|
|
19
|
+
|
|
20
|
+
// Vector field kinds the generated clients can describe.
|
|
21
|
+
const (
|
|
22
|
+
VectorStrategySqliteVec = "sqlite-vec"
|
|
23
|
+
VectorStrategyUdf = "udf"
|
|
24
|
+
VectorStrategySQL = "sql"
|
|
25
|
+
VectorStrategyMemory = "memory"
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
// An5VectorFunctions are the scalar functions the driver may register so SQLite
|
|
29
|
+
// itself ranks the column.
|
|
30
|
+
var An5VectorFunctions = map[string]string{
|
|
31
|
+
"cosine": "an5_vec_cosine",
|
|
32
|
+
"euclidean": "an5_vec_l2",
|
|
33
|
+
"dot": "an5_vec_ip",
|
|
34
|
+
}
|
|
35
|
+
|
|
36
|
+
// SqliteVecFunctions are the scalar functions the sqlite-vec extension provides.
|
|
37
|
+
var SqliteVecFunctions = map[string]string{
|
|
38
|
+
"cosine": "vec_distance_cosine",
|
|
39
|
+
"euclidean": "vec_distance_l2",
|
|
40
|
+
"dot": "vec_distance_ip",
|
|
41
|
+
}
|
|
42
|
+
|
|
43
|
+
// EncodeVector encodes a vector as the little-endian float32 BLOB the column stores.
|
|
44
|
+
func EncodeVector(values []float64) []byte {
|
|
45
|
+
bytes := make([]byte, len(values)*BytesPerVectorFloat)
|
|
46
|
+
for i, v := range values {
|
|
47
|
+
binary.LittleEndian.PutUint32(bytes[i*BytesPerVectorFloat:], math.Float32bits(float32(v)))
|
|
48
|
+
}
|
|
49
|
+
return bytes
|
|
50
|
+
}
|
|
51
|
+
|
|
52
|
+
// ParseVector reads a numeric sequence, or JSON text holding one.
|
|
53
|
+
func ParseVector(value interface{}) []float64 {
|
|
54
|
+
switch v := value.(type) {
|
|
55
|
+
case nil:
|
|
56
|
+
return nil
|
|
57
|
+
case []float64:
|
|
58
|
+
return v
|
|
59
|
+
case float32:
|
|
60
|
+
return []float64{float64(v)}
|
|
61
|
+
case float64:
|
|
62
|
+
return []float64{v}
|
|
63
|
+
case []float32:
|
|
64
|
+
out := make([]float64, len(v))
|
|
65
|
+
for i, f := range v {
|
|
66
|
+
out[i] = float64(f)
|
|
67
|
+
}
|
|
68
|
+
return out
|
|
69
|
+
case string:
|
|
70
|
+
trimmed := strings.TrimSpace(v)
|
|
71
|
+
if !strings.HasPrefix(trimmed, "[") {
|
|
72
|
+
return nil
|
|
73
|
+
}
|
|
74
|
+
var parsed []float64
|
|
75
|
+
if err := json.Unmarshal([]byte(trimmed), &parsed); err != nil {
|
|
76
|
+
return nil
|
|
77
|
+
}
|
|
78
|
+
return ParseVector(parsed)
|
|
79
|
+
case []byte:
|
|
80
|
+
// `database/sql` hands a BLOB column over as bytes; whether that is
|
|
81
|
+
// float32 storage or JSON text is decided by its length and first byte.
|
|
82
|
+
if len(v) > 0 && v[0] == '[' {
|
|
83
|
+
var parsed []float64
|
|
84
|
+
if err := json.Unmarshal(v, &parsed); err != nil {
|
|
85
|
+
return nil
|
|
86
|
+
}
|
|
87
|
+
return ParseVector(parsed)
|
|
88
|
+
}
|
|
89
|
+
return DecodeVector(v)
|
|
90
|
+
case []interface{}:
|
|
91
|
+
out := make([]float64, 0, len(v))
|
|
92
|
+
for _, item := range v {
|
|
93
|
+
f, ok := item.(float64)
|
|
94
|
+
if !ok {
|
|
95
|
+
return nil
|
|
96
|
+
}
|
|
97
|
+
out = append(out, f)
|
|
98
|
+
}
|
|
99
|
+
return out
|
|
100
|
+
default:
|
|
101
|
+
return nil
|
|
102
|
+
}
|
|
103
|
+
}
|
|
104
|
+
|
|
105
|
+
// DecodeVector decodes a stored vector: a float32 BLOB, a legacy JSON text
|
|
106
|
+
// column, or a sequence the driver already decoded.
|
|
107
|
+
func DecodeVector(value interface{}) []float64 {
|
|
108
|
+
if bytes, ok := value.([]byte); ok {
|
|
109
|
+
if len(bytes) == 0 || len(bytes)%BytesPerVectorFloat != 0 {
|
|
110
|
+
// Not float32 storage; it may still be JSON text in the same column.
|
|
111
|
+
return ParseVector(string(bytes))
|
|
112
|
+
}
|
|
113
|
+
if len(bytes) > 0 && bytes[0] == '[' {
|
|
114
|
+
var parsed []float64
|
|
115
|
+
if json.Unmarshal(bytes, &parsed) == nil { return parsed }
|
|
116
|
+
}
|
|
117
|
+
out := make([]float64, len(bytes)/BytesPerVectorFloat)
|
|
118
|
+
for i := range out {
|
|
119
|
+
out[i] = float64(math.Float32frombits(binary.LittleEndian.Uint32(bytes[i*BytesPerVectorFloat:])))
|
|
120
|
+
}
|
|
121
|
+
return out
|
|
122
|
+
}
|
|
123
|
+
return ParseVector(value)
|
|
124
|
+
}
|
|
125
|
+
|
|
126
|
+
// NeedsVectorEncoding reports whether a value still has to become bytes on its
|
|
127
|
+
// way into a `VECTOR(n)` column. `[]float64` and `[]interface{}` are the shapes
|
|
128
|
+
// a caller writes; text and bytes are already stored forms and pass through.
|
|
129
|
+
func NeedsVectorEncoding(value interface{}) bool {
|
|
130
|
+
switch value.(type) {
|
|
131
|
+
case nil, string, []byte:
|
|
132
|
+
return false
|
|
133
|
+
case []float64, []float32:
|
|
134
|
+
return true
|
|
135
|
+
case []interface{}:
|
|
136
|
+
return len(value.([]interface{})) > 0
|
|
137
|
+
default:
|
|
138
|
+
return false
|
|
139
|
+
}
|
|
140
|
+
}
|
|
141
|
+
|
|
142
|
+
// IsVectorField reports whether generated metadata describes a `VECTOR(n)` column.
|
|
143
|
+
// The metadata value is a map of field name to a small record, so both a typed
|
|
144
|
+
// struct and a plain map are accepted.
|
|
145
|
+
func IsVectorField(definition interface{}) bool {
|
|
146
|
+
switch def := definition.(type) {
|
|
147
|
+
case nil:
|
|
148
|
+
return false
|
|
149
|
+
case string:
|
|
150
|
+
return strings.Contains(def, "[]") || strings.HasPrefix(strings.ToUpper(strings.TrimSpace(def)), "VECTOR")
|
|
151
|
+
case map[string]interface{}:
|
|
152
|
+
if kind, _ := def["kind"].(string); kind == "vector" {
|
|
153
|
+
return true
|
|
154
|
+
}
|
|
155
|
+
for _, key := range []string{"sql", "type"} {
|
|
156
|
+
if sql, ok := def[key].(string); ok && sqlLooksVector(sql) {
|
|
157
|
+
return true
|
|
158
|
+
}
|
|
159
|
+
}
|
|
160
|
+
ts, _ := def["ts"].(string)
|
|
161
|
+
return strings.Contains(ts, "[]")
|
|
162
|
+
default:
|
|
163
|
+
return false
|
|
164
|
+
}
|
|
165
|
+
}
|
|
166
|
+
|
|
167
|
+
// sqlLooksVector reports whether a declared SQL type names a vector column.
|
|
168
|
+
func sqlLooksVector(sql string) bool {
|
|
169
|
+
trimmed := strings.ToUpper(strings.TrimSpace(sql))
|
|
170
|
+
if !strings.HasPrefix(trimmed, "VECTOR") {
|
|
171
|
+
return false
|
|
172
|
+
}
|
|
173
|
+
rest := trimmed[len("VECTOR"):]
|
|
174
|
+
return rest == "" || strings.HasPrefix(rest, "(") || strings.HasPrefix(rest, " ")
|
|
175
|
+
}
|
|
176
|
+
|
|
177
|
+
// PlanVectorStrategies returns the ordered strategies to try for a connection.
|
|
178
|
+
//
|
|
179
|
+
// capabilities reports what the connection supports; declaredType is the
|
|
180
|
+
// column's DDL type, because the JSON strategy only reaches rows stored as text
|
|
181
|
+
// and so a BLOB column drops it from the order. preference pins one strategy.
|
|
182
|
+
func PlanVectorStrategies(capabilities map[string]bool, declaredType, preference string) []string {
|
|
183
|
+
if preference != "" && preference != "auto" {
|
|
184
|
+
if preference == VectorStrategyMemory {
|
|
185
|
+
return nil
|
|
186
|
+
}
|
|
187
|
+
return []string{preference}
|
|
188
|
+
}
|
|
189
|
+
|
|
190
|
+
var available []string
|
|
191
|
+
if capabilities["vec"] {
|
|
192
|
+
available = append(available, VectorStrategySqliteVec)
|
|
193
|
+
}
|
|
194
|
+
if capabilities["udf"] {
|
|
195
|
+
available = append(available, VectorStrategyUdf)
|
|
196
|
+
}
|
|
197
|
+
if capabilities["json1"] {
|
|
198
|
+
available = append(available, VectorStrategySQL)
|
|
199
|
+
}
|
|
200
|
+
|
|
201
|
+
declared := strings.ToUpper(strings.TrimSpace(declaredType))
|
|
202
|
+
if declared != "" && !strings.Contains(declared, "TEXT") && !strings.Contains(declared, "CHAR") &&
|
|
203
|
+
!strings.Contains(declared, "CLOB") && !strings.Contains(declared, "STRING") &&
|
|
204
|
+
!strings.Contains(declared, "JSON") {
|
|
205
|
+
return without(available, VectorStrategySQL)
|
|
206
|
+
}
|
|
207
|
+
return available
|
|
208
|
+
}
|
|
209
|
+
|
|
210
|
+
func without(list []string, item string) []string {
|
|
211
|
+
out := make([]string, 0, len(list))
|
|
212
|
+
for _, entry := range list {
|
|
213
|
+
if entry != item {
|
|
214
|
+
out = append(out, entry)
|
|
215
|
+
}
|
|
216
|
+
}
|
|
217
|
+
return out
|
|
218
|
+
}
|
|
219
|
+
|
|
220
|
+
// NormalizeMetric returns a known metric name, defaulting to cosine.
|
|
221
|
+
func NormalizeMetric(metric string) string {
|
|
222
|
+
switch strings.ToLower(strings.TrimSpace(metric)) {
|
|
223
|
+
case "euclidean", "dot":
|
|
224
|
+
return strings.ToLower(strings.TrimSpace(metric))
|
|
225
|
+
default:
|
|
226
|
+
return "cosine"
|
|
227
|
+
}
|
|
228
|
+
}
|
|
229
|
+
|
|
230
|
+
// VectorDistance returns the distance between two equal-length vectors, lower is
|
|
231
|
+
// closer. It mirrors the SQL functions so a search ranks the same way whichever
|
|
232
|
+
// strategy ran. The second result is false when the pair cannot be scored, which
|
|
233
|
+
// is how a stored vector of another dimension is kept out of the ranking.
|
|
234
|
+
func VectorDistance(a, b []float64, metric string) (float64, bool) {
|
|
235
|
+
if len(a) == 0 || len(a) != len(b) {
|
|
236
|
+
return 0, false
|
|
237
|
+
}
|
|
238
|
+
switch NormalizeMetric(metric) {
|
|
239
|
+
case "euclidean":
|
|
240
|
+
return EuclideanDistance(a, b), true
|
|
241
|
+
case "dot":
|
|
242
|
+
return -DotProduct(a, b), true
|
|
243
|
+
default:
|
|
244
|
+
m1 := math.Sqrt(sumOfSquares(a))
|
|
245
|
+
m2 := math.Sqrt(sumOfSquares(b))
|
|
246
|
+
if m1 == 0 || m2 == 0 {
|
|
247
|
+
return 1.0, true
|
|
248
|
+
}
|
|
249
|
+
var dot float64
|
|
250
|
+
for i := range a {
|
|
251
|
+
dot += a[i] * b[i]
|
|
252
|
+
}
|
|
253
|
+
return 1.0 - dot/(m1*m2), true
|
|
254
|
+
}
|
|
255
|
+
}
|
|
256
|
+
|
|
257
|
+
func sumOfSquares(values []float64) float64 {
|
|
258
|
+
var sum float64
|
|
259
|
+
for _, v := range values {
|
|
260
|
+
sum += v * v
|
|
261
|
+
}
|
|
262
|
+
return sum
|
|
263
|
+
}
|
|
@@ -0,0 +1,248 @@
|
|
|
1
|
+
// Package sqlite provides SQLite specific dialect execution and query generation for AN5 Go Adapters.
|
|
2
|
+
//
|
|
3
|
+
// SQLite has no vector type, so a `VECTOR(n)` column stores a BLOB of little-endian
|
|
4
|
+
// float32 values and is ranked in one of four ways, in this order:
|
|
5
|
+
//
|
|
6
|
+
// 1. sqlite-vec native scalar distances, when the driver loaded the extension.
|
|
7
|
+
// 2. udf an5_vec_cosine / an5_vec_l2 / an5_vec_ip registered with the
|
|
8
|
+
// driver. Reads the BLOB and the legacy JSON text.
|
|
9
|
+
// 3. sql json_each brute force in plain SQL. Needs no user function, so it
|
|
10
|
+
// works with any driver, but only reaches rows stored as JSON text.
|
|
11
|
+
// 4. memory the column is loaded and scored in Go.
|
|
12
|
+
//
|
|
13
|
+
// Strategies 1-3 rank inside the database and transfer only the matching rows,
|
|
14
|
+
// which is why they are preferred over the in-memory path.
|
|
15
|
+
package sqlite
|
|
16
|
+
|
|
17
|
+
import (
|
|
18
|
+
"context"
|
|
19
|
+
"database/sql"
|
|
20
|
+
"encoding/json"
|
|
21
|
+
"fmt"
|
|
22
|
+
"strings"
|
|
23
|
+
|
|
24
|
+
"an5adapters/base"
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
// BuildVectorSearchQuery returns empty string as SQLite ranks through its own
|
|
28
|
+
// strategies; see RunVectorSearch for the statements it builds.
|
|
29
|
+
func BuildVectorSearchQuery(table, field, metric string, dim, take int, where string) string {
|
|
30
|
+
_ = base.QuoteTable(table, base.DialectSqlite)
|
|
31
|
+
_ = base.QuoteIdentifier(field, base.DialectSqlite)
|
|
32
|
+
// Native vector extension not present in standard sqlite; in-memory fallback is used.
|
|
33
|
+
return ""
|
|
34
|
+
}
|
|
35
|
+
|
|
36
|
+
// VectorSupport carries what a SQLite connection needs for in-database vector
|
|
37
|
+
// search: the path to the sqlite-vec extension to load and a strategy to pin
|
|
38
|
+
// instead of probing.
|
|
39
|
+
type VectorSupport struct {
|
|
40
|
+
// SqliteVecPath names the sqlite-vec extension binary. It is only used to
|
|
41
|
+
// report what the caller asked for; loading the extension is the driver's
|
|
42
|
+
// job (mattn/go-sqlite3 takes it in the DSN as `_loc`/`_x` style options, and
|
|
43
|
+
// modernc.org/sqlite through its own LoadExtension API).
|
|
44
|
+
SqliteVecPath string
|
|
45
|
+
// VectorStrategy pins "sqlite-vec", "udf", "sql" or "memory". Empty probes.
|
|
46
|
+
VectorStrategy string
|
|
47
|
+
}
|
|
48
|
+
|
|
49
|
+
// ScalarFuncRegistrar installs a scalar function with a driver. `database/sql`
|
|
50
|
+
// has no API for this, so a driver has to hand one over:
|
|
51
|
+
//
|
|
52
|
+
// sql.Register("an5-sqlite3", &sqlite3.SQLiteDriver{
|
|
53
|
+
// ConnectHook: func(c *sqlite3.SQLiteConn) error {
|
|
54
|
+
// return sqlite.RegisterVectorFunctions(c.RegisterFunc)
|
|
55
|
+
// },
|
|
56
|
+
// })
|
|
57
|
+
func RegisterVectorFunctions(register func(name string, impl interface{}, pure bool) error) error {
|
|
58
|
+
for metric, name := range base.An5VectorFunctions {
|
|
59
|
+
captured := metric
|
|
60
|
+
if err := register(name, func(a, b interface{}) interface{} {
|
|
61
|
+
distance, ok := base.VectorDistance(base.DecodeVector(a), base.DecodeVector(b), captured)
|
|
62
|
+
if !ok { return nil }
|
|
63
|
+
return distance
|
|
64
|
+
}, true); err != nil {
|
|
65
|
+
return err
|
|
66
|
+
}
|
|
67
|
+
}
|
|
68
|
+
return nil
|
|
69
|
+
}
|
|
70
|
+
|
|
71
|
+
// Capabilities probes what this connection can do, once per call. Both probes are
|
|
72
|
+
// cheap `SELECT`s; the sqlite-vec one fails on a connection without the extension,
|
|
73
|
+
// and the user-function one on a driver that never registered them.
|
|
74
|
+
func Capabilities(ctx context.Context, db *sql.DB) map[string]bool {
|
|
75
|
+
caps := map[string]bool{"vec": false, "udf": false, "json1": false}
|
|
76
|
+
if db == nil {
|
|
77
|
+
return caps
|
|
78
|
+
}
|
|
79
|
+
|
|
80
|
+
row := db.QueryRowContext(ctx, "SELECT vec_version()")
|
|
81
|
+
var version interface{}
|
|
82
|
+
if err := row.Scan(&version); err == nil {
|
|
83
|
+
caps["vec"] = true
|
|
84
|
+
}
|
|
85
|
+
|
|
86
|
+
// The probe has to reach a real distance: a driver that resolved the name
|
|
87
|
+
// would still fail on the arguments otherwise.
|
|
88
|
+
row = db.QueryRowContext(ctx, "SELECT an5_vec_cosine(zeroblob(4), zeroblob(4))")
|
|
89
|
+
var distance interface{}
|
|
90
|
+
if err := row.Scan(&distance); err == nil {
|
|
91
|
+
caps["udf"] = true
|
|
92
|
+
}
|
|
93
|
+
|
|
94
|
+
row = db.QueryRowContext(ctx, "SELECT json_valid('[1]')")
|
|
95
|
+
var jsonOK interface{}
|
|
96
|
+
if err := row.Scan(&jsonOK); err == nil {
|
|
97
|
+
caps["json1"] = true
|
|
98
|
+
}
|
|
99
|
+
|
|
100
|
+
return caps
|
|
101
|
+
}
|
|
102
|
+
|
|
103
|
+
// ReadColumnType returns a vector column's declared DDL type, or "" when the
|
|
104
|
+
// table or column is not there.
|
|
105
|
+
func ReadColumnType(ctx context.Context, db *sql.DB, table, column string) string {
|
|
106
|
+
row := db.QueryRowContext(ctx, "SELECT type FROM pragma_table_info(?) WHERE name = ?", table, column)
|
|
107
|
+
var declared sql.NullString
|
|
108
|
+
if err := row.Scan(&declared); err != nil || !declared.Valid {
|
|
109
|
+
return ""
|
|
110
|
+
}
|
|
111
|
+
return declared.String
|
|
112
|
+
}
|
|
113
|
+
|
|
114
|
+
// jsonText guards the JSON path: `json_valid` fails on a float32 BLOB.
|
|
115
|
+
func jsonText(column string) string {
|
|
116
|
+
return fmt.Sprintf("CASE WHEN json_valid(%s) THEN %s ELSE '[]' END", column, column)
|
|
117
|
+
}
|
|
118
|
+
|
|
119
|
+
func jsonDistanceExpr(metric, column string) string {
|
|
120
|
+
row := jsonText(column)
|
|
121
|
+
sameLength := fmt.Sprintf("(SELECT COUNT(*) FROM json_each(%s)) = (SELECT COUNT(*) FROM q)", row)
|
|
122
|
+
|
|
123
|
+
var terms string
|
|
124
|
+
switch metric {
|
|
125
|
+
case "euclidean":
|
|
126
|
+
terms = fmt.Sprintf("sqrt((SELECT SUM((je.value - q.v) * (je.value - q.v)) FROM json_each(%s) je JOIN q ON q.k = je.key))", row)
|
|
127
|
+
case "dot":
|
|
128
|
+
terms = fmt.Sprintf("-(SELECT SUM(je.value * q.v) FROM json_each(%s) je JOIN q ON q.k = je.key)", row)
|
|
129
|
+
default:
|
|
130
|
+
terms = fmt.Sprintf("1.0 - (SELECT SUM(je.value * q.v) FROM json_each(%s) je JOIN q ON q.k = je.key)"+
|
|
131
|
+
" / NULLIF(sqrt((SELECT SUM(q.v * q.v) FROM q)) * sqrt((SELECT SUM(je.value * je.value) FROM json_each(%s) je)), 0)", row, row)
|
|
132
|
+
}
|
|
133
|
+
return fmt.Sprintf("CASE WHEN %s THEN %s END", sameLength, terms)
|
|
134
|
+
}
|
|
135
|
+
|
|
136
|
+
// BuildRankingQuery returns the ranking SQL for one strategy, with the arguments
|
|
137
|
+
// the query vector needs appended to `args`.
|
|
138
|
+
//
|
|
139
|
+
// sqlite-vec and udf bind the query vector as a float32 BLOB; sql binds it as
|
|
140
|
+
// JSON text because json_each reads text.
|
|
141
|
+
//
|
|
142
|
+
// A row whose stored vector cannot be scored (a NULL, a text column, another
|
|
143
|
+
// dimension) yields a NULL distance; those rows are dropped rather than reported,
|
|
144
|
+
// so a caller never has to sort them out of the result.
|
|
145
|
+
func BuildRankingQuery(strategy, metric, table, column string, vector []float64, take int, tail string, args []interface{}) (string, []interface{}) {
|
|
146
|
+
rank := func(inner string) string {
|
|
147
|
+
return fmt.Sprintf("SELECT * FROM (%s) AS an5_ranked WHERE distance IS NOT NULL ORDER BY distance ASC LIMIT %d", inner, take)
|
|
148
|
+
}
|
|
149
|
+
|
|
150
|
+
switch metric {
|
|
151
|
+
case "euclidean", "dot":
|
|
152
|
+
default:
|
|
153
|
+
metric = "cosine"
|
|
154
|
+
}
|
|
155
|
+
|
|
156
|
+
if strategy == base.VectorStrategySqliteVec {
|
|
157
|
+
fn, ok := base.SqliteVecFunctions[metric]
|
|
158
|
+
if !ok {
|
|
159
|
+
return "", args
|
|
160
|
+
}
|
|
161
|
+
// sqlite-vec only understands float32 BLOB operands, so rows stored any
|
|
162
|
+
// other way are excluded instead of aborting the query. The guard joins the
|
|
163
|
+
// caller's WHERE with AND, since the ranking query is where both apply.
|
|
164
|
+
guard := fmt.Sprintf("%s IS NOT NULL", column)
|
|
165
|
+
// `tail` carries its own leading " WHERE", so it is appended after AND.
|
|
166
|
+
if tail != "" {
|
|
167
|
+
guard += " AND" + strings.Replace(tail, " WHERE ", " ", 1)
|
|
168
|
+
}
|
|
169
|
+
inner := fmt.Sprintf("SELECT *, CASE WHEN vec_length(vec_f32(%s)) = %d THEN %s(vec_f32(%s), ?) END AS distance FROM %s WHERE %s",
|
|
170
|
+
column, len(vector), fn, column, table, guard)
|
|
171
|
+
return rank(inner), append([]interface{}{base.EncodeVector(vector)}, args...)
|
|
172
|
+
}
|
|
173
|
+
|
|
174
|
+
if strategy == base.VectorStrategyUdf {
|
|
175
|
+
fn, ok := base.An5VectorFunctions[metric]
|
|
176
|
+
if !ok {
|
|
177
|
+
return "", args
|
|
178
|
+
}
|
|
179
|
+
inner := fmt.Sprintf("SELECT *, %s(%s, ?) AS distance FROM %s%s", fn, column, table, tail)
|
|
180
|
+
return rank(inner), append([]interface{}{base.EncodeVector(vector)}, args...)
|
|
181
|
+
}
|
|
182
|
+
|
|
183
|
+
payload, err := json.Marshal(vector)
|
|
184
|
+
if err != nil {
|
|
185
|
+
return "", args
|
|
186
|
+
}
|
|
187
|
+
inner := fmt.Sprintf("SELECT *, %s AS distance FROM %s%s", jsonDistanceExpr(metric, column), table, tail)
|
|
188
|
+
prefix := "WITH q AS (SELECT je.key AS k, CAST(je.value AS REAL) AS v FROM json_each(?) je) "
|
|
189
|
+
return prefix + rank(inner), append([]interface{}{string(payload)}, args...)
|
|
190
|
+
}
|
|
191
|
+
|
|
192
|
+
// RunVectorSearch ranks a SQLite table inside the database, trying each available
|
|
193
|
+
// strategy in order. It reports whether a strategy ran: when none does, the
|
|
194
|
+
// caller falls back to scoring in memory.
|
|
195
|
+
func RunVectorSearch(
|
|
196
|
+
ctx context.Context,
|
|
197
|
+
db *sql.DB,
|
|
198
|
+
exec func(ctx context.Context, query string, args ...interface{}) ([]map[string]interface{}, error),
|
|
199
|
+
table, rawTable, column, rawColumn string,
|
|
200
|
+
vector []float64,
|
|
201
|
+
metric string,
|
|
202
|
+
take int,
|
|
203
|
+
tail string,
|
|
204
|
+
args []interface{},
|
|
205
|
+
support VectorSupport,
|
|
206
|
+
) ([]map[string]interface{}, bool, error) {
|
|
207
|
+
if len(vector) == 0 {
|
|
208
|
+
return nil, false, nil
|
|
209
|
+
}
|
|
210
|
+
|
|
211
|
+
caps := Capabilities(ctx, db)
|
|
212
|
+
if db == nil {
|
|
213
|
+
return nil, false, nil
|
|
214
|
+
}
|
|
215
|
+
declared := ReadColumnType(ctx, db, rawTable, rawColumn)
|
|
216
|
+
|
|
217
|
+
for _, strategy := range base.PlanVectorStrategies(caps, declared, support.VectorStrategy) {
|
|
218
|
+
if strategy == base.VectorStrategySQL {
|
|
219
|
+
clause := " WHERE "
|
|
220
|
+
if tail != "" { clause = tail + " AND " }
|
|
221
|
+
binary, err := exec(ctx, "SELECT 1 FROM " + table + clause + "typeof(" + column + ") = 'blob' LIMIT 1", args...)
|
|
222
|
+
if err != nil || len(binary) > 0 { continue }
|
|
223
|
+
}
|
|
224
|
+
query, queryArgs := BuildRankingQuery(strategy, metric, table, column, vector, take, tail, args)
|
|
225
|
+
if query == "" {
|
|
226
|
+
continue
|
|
227
|
+
}
|
|
228
|
+
rows, err := exec(ctx, query, queryArgs...)
|
|
229
|
+
if err == nil {
|
|
230
|
+
return rows, true, nil
|
|
231
|
+
}
|
|
232
|
+
// The connection advertised the capability but the query failed, e.g. an
|
|
233
|
+
// extension that did not really load. The next strategy is cheaper than
|
|
234
|
+
// reporting the failure, and memory is the last resort.
|
|
235
|
+
}
|
|
236
|
+
return nil, false, nil
|
|
237
|
+
}
|
|
238
|
+
|
|
239
|
+
// NormalizeTableName reduces a mapped table name to the bare name
|
|
240
|
+
// `pragma_table_info` takes as a plain string: the last segment of a qualified
|
|
241
|
+
// name, with the quoting a generated client may have left around it removed.
|
|
242
|
+
func NormalizeTableName(name string) string {
|
|
243
|
+
trimmed := strings.TrimSpace(name)
|
|
244
|
+
if index := strings.LastIndex(trimmed, "."); index >= 0 {
|
|
245
|
+
trimmed = trimmed[index+1:]
|
|
246
|
+
}
|
|
247
|
+
return strings.Trim(trimmed, "[]`\"")
|
|
248
|
+
}
|