rickub
nandi/gleanpublic⑂ Fork 0
⑂ f205f09
Branches
main View all branches → View all tags →
Commits
⬇ Clone ▾
git clone https://git.rickub.com/nandi/glean.git
git clone ssh://git@rickub.com/nandi/glean.git

Host key fingerprint (ed25519): SHA256:iycHnxEyq0Q7uyVpB7JlznP0G7JrTPXLYRcAU5CSLhc — verify it before your first connect.

Spike the two risky parts of a Nim backend port

Before committing to rewriting 14k lines of Go, prove the two things Nim has
no library for. Volume is typing; these were the actual risk.

sqlite-vec and FTS5 both work: sqlite_probe.nim compiles the same
sqlite-vec.c the Go bindings vendor, registers it, and exercises FTS5 MATCH,
a vec0 table, float32 blob binding and a KNN query. The trap is that
-DSQLITE_CORE is load-bearing -- without it the file builds as a loadable
extension whose entry point reads its API pointer from the caller, so
passing nil segfaults rather than erroring. Nim's {.compile: (file, flags).}
tuple form did not reliably apply it, so it lives in nim.cfg.

ES256/DPoP works too, which was the one that could have sunk the port.
bearssl rather than OpenSSL, because JOSE wants a raw 64-byte R||S signature
and br_ecdsa_sign_raw emits exactly that -- no DER to unwrap wrongly.

Self-verification would only show internal consistency, so the proof is
checked by verify/verify.go using Go's stdlib: the same stack the server
runs on today. It also asserts the signature is not DER-wrapped and that a
tampered payload is rejected, since a verifier that accepts everything
proves nothing. 19/19 checks, stable across 30 keypairs -- worth repeating,
as leading zeros in ECDSA values are a classic intermittent R||S bug.

Still unproven and ahead: the OAuth flow itself (PAR, the DPoP nonce retry),
DID resolution, XRPC, Jetstream, and CBOR/CAR.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
nandithebull committed 2026-09-20T20:44:46-07:00 Browse files
f205f09 parent: 1dc6512
added spike/nim/.gitignore +4 -0
new file mode 100644
@@ -0,0 +1,4 @@
1+# nim build output
2+nimcache/
3+sqlite_probe
4+dpop_probe
new file mode 100644
@@ -0,0 +1,4 @@
1+# nim build output
2+nimcache/
3+sqlite_probe
4+dpop_probe
added spike/nim/dpop_probe.nim +125 -0
new file mode 100644
@@ -0,0 +1,125 @@
1+## Spike: can Nim produce ATProto-compatible ES256 / DPoP proofs?
2+##
3+## This is the riskiest part of a Nim port. Glean's Go backend gets OAuth from
4+## bluesky-social/indigo; Nim has no ATProto library, so the port would have to
5+## build the DPoP layer itself. Everything downstream of sign-in depends on it,
6+## so it is worth proving before committing to the rest.
7+##
8+## What has to be exactly right:
9+## * ES256 signatures in JOSE form -- raw R||S, 64 bytes, NOT the DER that
10+## most crypto libraries hand you.
11+## * An RFC 7638 JWK thumbprint, which is a SHA-256 over canonical JSON:
12+## exactly the members crv/kty/x/y, lexicographically ordered, no spaces.
13+## * base64url without padding, everywhere.
14+##
15+## bearssl is used because br_ecdsa_sign_raw already emits R||S, so there is no
16+## DER unwrapping step to get wrong.
17+##
18+## Self-verification would only prove internal consistency, so the signed token
19+## is handed to a Go verifier (verify/verify.go) built on the same stdlib the
20+## real server uses.
21+
22+import std/[base64, json, strutils, times, strformat]
23+import bearssl/[ec, rand]
24+import bearssl/abi/[bearssl_ec, bearssl_rand, bearssl_hash]
25+import nimcrypto/[sha2, hash]
26+
27+type P256Key* = object
28+ priv*: array[32, byte]
29+ pubX*: array[32, byte]
30+ pubY*: array[32, byte]
31+
32+proc toSeq(a: openArray[byte]): seq[byte] =
33+ result = newSeq[byte](a.len)
34+ for i, b in a: result[i] = b
35+
36+proc b64u*(data: openArray[byte]): string =
37+ ## base64url, unpadded (RFC 7515 §2).
38+ base64.encode(data, safe = true).strip(chars = {'='})
39+
40+proc b64u*(s: string): string =
41+ b64u(s.toOpenArrayByte(0, s.high).toSeq)
42+
43+proc generateKey*(rng: var HmacDrbgContext): P256Key =
44+ ## P-256 keypair. bearssl returns the public point uncompressed:
45+ ## 0x04 || X(32) || Y(32).
46+ var
47+ skBuf: array[EC_KBUF_PRIV_MAX_SIZE, byte]
48+ pkBuf: array[EC_KBUF_PUB_MAX_SIZE, byte]
49+ sk: EcPrivateKey
50+ pk: EcPublicKey
51+
52+ let n = ecKeygen(PrngClassPointerConst(addr rng.vtable), addr ecPrimeI31, addr sk,
53+ addr skBuf[0], EC_secp256r1.cint)
54+ doAssert n > 0, "ecKeygen failed"
55+ let m = ecComputePub(addr ecPrimeI31, addr pk, addr pkBuf[0], addr sk)
56+ doAssert m > 0, "ecComputePub failed"
57+
58+ doAssert sk.xlen == 32, &"unexpected private key length {sk.xlen}"
59+ doAssert pk.qlen == 65 and cast[ptr UncheckedArray[byte]](pk.q)[0] == 0x04'u8,
60+ "expected an uncompressed public point"
61+
62+ let
63+ skBytes = cast[ptr UncheckedArray[byte]](sk.x)
64+ pkBytes = cast[ptr UncheckedArray[byte]](pk.q)
65+ for i in 0 ..< 32:
66+ result.priv[i] = skBytes[i]
67+ result.pubX[i] = pkBytes[1 + i]
68+ result.pubY[i] = pkBytes[33 + i]
69+
70+proc publicJwk*(k: P256Key): JsonNode =
71+ %*{"crv": "P-256", "kty": "EC", "x": b64u(k.pubX), "y": b64u(k.pubY)}
72+
73+proc thumbprint*(k: P256Key): string =
74+ ## RFC 7638: SHA-256 over the canonical JWK. The member set and ordering are
75+ ## fixed by the spec, so this is built by hand rather than serialised from a
76+ ## JsonNode whose key order is incidental.
77+ let canonical = &"""{{"crv":"P-256","kty":"EC","x":"{b64u(k.pubX)}","y":"{b64u(k.pubY)}"}}"""
78+ b64u(sha256.digest(canonical).data)
79+
80+proc signEs256*(k: P256Key, rng: var HmacDrbgContext, signingInput: string): seq[byte] =
81+ var
82+ sk: EcPrivateKey
83+ priv = k.priv
84+ sk.curve = EC_secp256r1.cint
85+ sk.x = cast[ptr byte](addr priv[0])
86+ sk.xlen = 32
87+
88+ let digest = sha256.digest(signingInput).data
89+ var sig: array[64, byte]
90+ let n = ecdsaSignRawGetDefault()(
91+ addr ecPrimeI31, addr sha256Vtable, unsafeAddr digest[0], addr sk, addr sig[0])
92+ doAssert n == 64, &"expected a 64-byte raw signature, got {n}"
93+ result = sig.toSeq
94+
95+proc dpopProof*(k: P256Key, rng: var HmacDrbgContext, htm, htu: string,
96+ nonce = ""): string =
97+ ## A DPoP proof JWT (RFC 9449) as an ATProto PDS expects it.
98+ let header = %*{"typ": "dpop+jwt", "alg": "ES256", "jwk": k.publicJwk}
99+ var payload = %*{
100+ "jti": b64u(sha256.digest(&"{htm}{htu}{epochTime()}").data)[0 ..< 16],
101+ "htm": htm,
102+ "htu": htu,
103+ "iat": now().utc.toTime.toUnix,
104+ }
105+ if nonce.len > 0:
106+ payload["nonce"] = %nonce
107+
108+ let signingInput = b64u($header) & "." & b64u($payload)
109+ signingInput & "." & b64u(k.signEs256(rng, signingInput))
110+
111+when isMainModule:
112+ let rngRef = HmacDrbgContext.new()
113+ doAssert rngRef != nil, "no system randomness"
114+ var rng = rngRef[]
115+
116+ let key = generateKey(rng)
117+ let jwt = key.dpopProof(rng, "POST", "https://bsky.social/oauth/token",
118+ nonce = "spike-nonce")
119+
120+ # Emitted for verify/verify.go to check independently.
121+ echo $(%*{
122+ "jwt": jwt,
123+ "jwk": key.publicJwk,
124+ "thumbprint": key.thumbprint,
125+ })
new file mode 100644
@@ -0,0 +1,125 @@
1+## Spike: can Nim produce ATProto-compatible ES256 / DPoP proofs?
2+##
3+## This is the riskiest part of a Nim port. Glean's Go backend gets OAuth from
4+## bluesky-social/indigo; Nim has no ATProto library, so the port would have to
5+## build the DPoP layer itself. Everything downstream of sign-in depends on it,
6+## so it is worth proving before committing to the rest.
7+##
8+## What has to be exactly right:
9+## * ES256 signatures in JOSE form -- raw R||S, 64 bytes, NOT the DER that
10+## most crypto libraries hand you.
11+## * An RFC 7638 JWK thumbprint, which is a SHA-256 over canonical JSON:
12+## exactly the members crv/kty/x/y, lexicographically ordered, no spaces.
13+## * base64url without padding, everywhere.
14+##
15+## bearssl is used because br_ecdsa_sign_raw already emits R||S, so there is no
16+## DER unwrapping step to get wrong.
17+##
18+## Self-verification would only prove internal consistency, so the signed token
19+## is handed to a Go verifier (verify/verify.go) built on the same stdlib the
20+## real server uses.
21+
22+import std/[base64, json, strutils, times, strformat]
23+import bearssl/[ec, rand]
24+import bearssl/abi/[bearssl_ec, bearssl_rand, bearssl_hash]
25+import nimcrypto/[sha2, hash]
26+
27+type P256Key* = object
28+ priv*: array[32, byte]
29+ pubX*: array[32, byte]
30+ pubY*: array[32, byte]
31+
32+proc toSeq(a: openArray[byte]): seq[byte] =
33+ result = newSeq[byte](a.len)
34+ for i, b in a: result[i] = b
35+
36+proc b64u*(data: openArray[byte]): string =
37+ ## base64url, unpadded (RFC 7515 §2).
38+ base64.encode(data, safe = true).strip(chars = {'='})
39+
40+proc b64u*(s: string): string =
41+ b64u(s.toOpenArrayByte(0, s.high).toSeq)
42+
43+proc generateKey*(rng: var HmacDrbgContext): P256Key =
44+ ## P-256 keypair. bearssl returns the public point uncompressed:
45+ ## 0x04 || X(32) || Y(32).
46+ var
47+ skBuf: array[EC_KBUF_PRIV_MAX_SIZE, byte]
48+ pkBuf: array[EC_KBUF_PUB_MAX_SIZE, byte]
49+ sk: EcPrivateKey
50+ pk: EcPublicKey
51+
52+ let n = ecKeygen(PrngClassPointerConst(addr rng.vtable), addr ecPrimeI31, addr sk,
53+ addr skBuf[0], EC_secp256r1.cint)
54+ doAssert n > 0, "ecKeygen failed"
55+ let m = ecComputePub(addr ecPrimeI31, addr pk, addr pkBuf[0], addr sk)
56+ doAssert m > 0, "ecComputePub failed"
57+
58+ doAssert sk.xlen == 32, &"unexpected private key length {sk.xlen}"
59+ doAssert pk.qlen == 65 and cast[ptr UncheckedArray[byte]](pk.q)[0] == 0x04'u8,
60+ "expected an uncompressed public point"
61+
62+ let
63+ skBytes = cast[ptr UncheckedArray[byte]](sk.x)
64+ pkBytes = cast[ptr UncheckedArray[byte]](pk.q)
65+ for i in 0 ..< 32:
66+ result.priv[i] = skBytes[i]
67+ result.pubX[i] = pkBytes[1 + i]
68+ result.pubY[i] = pkBytes[33 + i]
69+
70+proc publicJwk*(k: P256Key): JsonNode =
71+ %*{"crv": "P-256", "kty": "EC", "x": b64u(k.pubX), "y": b64u(k.pubY)}
72+
73+proc thumbprint*(k: P256Key): string =
74+ ## RFC 7638: SHA-256 over the canonical JWK. The member set and ordering are
75+ ## fixed by the spec, so this is built by hand rather than serialised from a
76+ ## JsonNode whose key order is incidental.
77+ let canonical = &"""{{"crv":"P-256","kty":"EC","x":"{b64u(k.pubX)}","y":"{b64u(k.pubY)}"}}"""
78+ b64u(sha256.digest(canonical).data)
79+
80+proc signEs256*(k: P256Key, rng: var HmacDrbgContext, signingInput: string): seq[byte] =
81+ var
82+ sk: EcPrivateKey
83+ priv = k.priv
84+ sk.curve = EC_secp256r1.cint
85+ sk.x = cast[ptr byte](addr priv[0])
86+ sk.xlen = 32
87+
88+ let digest = sha256.digest(signingInput).data
89+ var sig: array[64, byte]
90+ let n = ecdsaSignRawGetDefault()(
91+ addr ecPrimeI31, addr sha256Vtable, unsafeAddr digest[0], addr sk, addr sig[0])
92+ doAssert n == 64, &"expected a 64-byte raw signature, got {n}"
93+ result = sig.toSeq
94+
95+proc dpopProof*(k: P256Key, rng: var HmacDrbgContext, htm, htu: string,
96+ nonce = ""): string =
97+ ## A DPoP proof JWT (RFC 9449) as an ATProto PDS expects it.
98+ let header = %*{"typ": "dpop+jwt", "alg": "ES256", "jwk": k.publicJwk}
99+ var payload = %*{
100+ "jti": b64u(sha256.digest(&"{htm}{htu}{epochTime()}").data)[0 ..< 16],
101+ "htm": htm,
102+ "htu": htu,
103+ "iat": now().utc.toTime.toUnix,
104+ }
105+ if nonce.len > 0:
106+ payload["nonce"] = %nonce
107+
108+ let signingInput = b64u($header) & "." & b64u($payload)
109+ signingInput & "." & b64u(k.signEs256(rng, signingInput))
110+
111+when isMainModule:
112+ let rngRef = HmacDrbgContext.new()
113+ doAssert rngRef != nil, "no system randomness"
114+ var rng = rngRef[]
115+
116+ let key = generateKey(rng)
117+ let jwt = key.dpopProof(rng, "POST", "https://bsky.social/oauth/token",
118+ nonce = "spike-nonce")
119+
120+ # Emitted for verify/verify.go to check independently.
121+ echo $(%*{
122+ "jwt": jwt,
123+ "jwk": key.publicJwk,
124+ "thumbprint": key.thumbprint,
125+ })
added spike/nim/nim.cfg +7 -0
new file mode 100644
@@ -0,0 +1,7 @@
1+# sqlite-vec is compiled into the binary rather than loaded at runtime, the
2+# same way the Go bindings do it. Without SQLITE_CORE the file builds as a
3+# loadable extension whose entry point reads its SQLite API pointer from the
4+# caller -- passing nil there segfaults instead of failing, so this define is
5+# load-bearing, not an optimisation.
6+--passC:"-DSQLITE_CORE"
7+--passL:"-lsqlite3 -lm"
new file mode 100644
@@ -0,0 +1,7 @@
1+# sqlite-vec is compiled into the binary rather than loaded at runtime, the
2+# same way the Go bindings do it. Without SQLITE_CORE the file builds as a
3+# loadable extension whose entry point reads its SQLite API pointer from the
4+# caller -- passing nil there segfaults instead of failing, so this define is
5+# load-bearing, not an optimisation.
6+--passC:"-DSQLITE_CORE"
7+--passL:"-lsqlite3 -lm"
added spike/nim/readme.md +74 -0
new file mode 100644
@@ -0,0 +1,74 @@
1+# Nim port spike
2+
3+Two questions had to be answered before committing to a Nim rewrite of the Go
4+backend. Neither is about volume -- 14k lines is a lot of typing, not a lot of
5+risk. The risk is that Go has libraries Nim does not, and two of them sit under
6+everything.
7+
8+Run both:
9+
10+```
11+cd spike/nim
12+nim c -r sqlite_probe.nim # storage
13+nim c -r dpop_probe.nim && ./dpop_probe | (cd verify && go run .) # auth
14+```
15+
16+## 1. sqlite-vec + FTS5 — works
17+
18+Glean needs FTS5 for article search and sqlite-vec for embedding similarity.
19+Both are C-level concerns, so the question was never "does Nim have a SQLite
20+library" but whether the same extension links and registers the same way.
21+
22+`sqlite_probe.nim` compiles the same `sqlite-vec.c` the Go bindings vendor,
23+calls `sqlite3_vec_init` directly, and exercises `vec_version()`, an FTS5
24+`MATCH`, a `vec0` virtual table, float32 blob binding, and a KNN query.
25+All pass; KNN returns the expected nearest neighbour.
26+
27+**One trap, and it is a nasty one.** sqlite-vec must be compiled with
28+`-DSQLITE_CORE`. Without it the file builds as a *loadable extension* whose
29+entry point takes its SQLite API pointer from the caller — pass `nil` there and
30+it segfaults inside `sqlite3_vec_init` rather than returning an error. Nim's
31+`{.compile: (file, flags).}` tuple form did not reliably apply the define, so
32+it lives in `nim.cfg` where it cannot be forgotten.
33+
34+## 2. ES256 / DPoP — works
35+
36+This was the real risk. Glean gets OAuth from `bluesky-social/indigo`; Nim has
37+no ATProto library, so a port has to build the DPoP layer itself, and
38+everything downstream of sign-in depends on it.
39+
40+`dpop_probe.nim` generates a P-256 key, builds the public JWK, computes an
41+RFC 7638 thumbprint, and signs a DPoP proof JWT (RFC 9449).
42+
43+It uses **bearssl** rather than OpenSSL for one specific reason: JOSE requires
44+the signature as raw `R||S`, 64 bytes, and most crypto libraries hand you DER.
45+`br_ecdsa_sign_raw` emits `R||S` directly, so there is no DER-unwrapping step
46+to get subtly wrong.
47+
48+Self-verification would only prove internal consistency, so the token goes to
49+`verify/verify.go`, which re-implements the verifier from Go's standard
50+library — the same stack the current server runs on. It checks the header
51+claims, that the signature is 64 raw bytes and not DER-wrapped, that the
52+signature verifies against the JWK, that a tampered payload is *rejected*, and
53+that the thumbprint matches when recomputed independently.
54+
55+19/19 checks pass, and the whole cycle passes across 30 independently generated
56+keypairs — worth confirming, because ECDSA values with leading zeros are a
57+classic source of intermittent length bugs in `R||S` encoding.
58+
59+## What this does and does not prove
60+
61+Proven: the cryptography and the storage layer are available to Nim, in the
62+exact formats ATProto and Glean need.
63+
64+Not proven, and still ahead:
65+
66+- The **OAuth protocol flow** — PAR, the authorization request, the DPoP nonce
67+ retry dance (a PDS rejects the first request and returns a nonce to use).
68+ Mechanical, but a lot of it.
69+- **DID/handle resolution** and the PLC directory.
70+- **XRPC** and the Jetstream websocket consumer (`indigo` covers both today).
71+- **CBOR/CAR** parsing for repository records.
72+
73+The honest read: nothing here is blocking, and the two things that could have
74+been blocking are not. The remaining work is large but ordinary.
new file mode 100644
@@ -0,0 +1,74 @@
1+# Nim port spike
2+
3+Two questions had to be answered before committing to a Nim rewrite of the Go
4+backend. Neither is about volume -- 14k lines is a lot of typing, not a lot of
5+risk. The risk is that Go has libraries Nim does not, and two of them sit under
6+everything.
7+
8+Run both:
9+
10+```
11+cd spike/nim
12+nim c -r sqlite_probe.nim # storage
13+nim c -r dpop_probe.nim && ./dpop_probe | (cd verify && go run .) # auth
14+```
15+
16+## 1. sqlite-vec + FTS5 — works
17+
18+Glean needs FTS5 for article search and sqlite-vec for embedding similarity.
19+Both are C-level concerns, so the question was never "does Nim have a SQLite
20+library" but whether the same extension links and registers the same way.
21+
22+`sqlite_probe.nim` compiles the same `sqlite-vec.c` the Go bindings vendor,
23+calls `sqlite3_vec_init` directly, and exercises `vec_version()`, an FTS5
24+`MATCH`, a `vec0` virtual table, float32 blob binding, and a KNN query.
25+All pass; KNN returns the expected nearest neighbour.
26+
27+**One trap, and it is a nasty one.** sqlite-vec must be compiled with
28+`-DSQLITE_CORE`. Without it the file builds as a *loadable extension* whose
29+entry point takes its SQLite API pointer from the caller — pass `nil` there and
30+it segfaults inside `sqlite3_vec_init` rather than returning an error. Nim's
31+`{.compile: (file, flags).}` tuple form did not reliably apply the define, so
32+it lives in `nim.cfg` where it cannot be forgotten.
33+
34+## 2. ES256 / DPoP — works
35+
36+This was the real risk. Glean gets OAuth from `bluesky-social/indigo`; Nim has
37+no ATProto library, so a port has to build the DPoP layer itself, and
38+everything downstream of sign-in depends on it.
39+
40+`dpop_probe.nim` generates a P-256 key, builds the public JWK, computes an
41+RFC 7638 thumbprint, and signs a DPoP proof JWT (RFC 9449).
42+
43+It uses **bearssl** rather than OpenSSL for one specific reason: JOSE requires
44+the signature as raw `R||S`, 64 bytes, and most crypto libraries hand you DER.
45+`br_ecdsa_sign_raw` emits `R||S` directly, so there is no DER-unwrapping step
46+to get subtly wrong.
47+
48+Self-verification would only prove internal consistency, so the token goes to
49+`verify/verify.go`, which re-implements the verifier from Go's standard
50+library — the same stack the current server runs on. It checks the header
51+claims, that the signature is 64 raw bytes and not DER-wrapped, that the
52+signature verifies against the JWK, that a tampered payload is *rejected*, and
53+that the thumbprint matches when recomputed independently.
54+
55+19/19 checks pass, and the whole cycle passes across 30 independently generated
56+keypairs — worth confirming, because ECDSA values with leading zeros are a
57+classic source of intermittent length bugs in `R||S` encoding.
58+
59+## What this does and does not prove
60+
61+Proven: the cryptography and the storage layer are available to Nim, in the
62+exact formats ATProto and Glean need.
63+
64+Not proven, and still ahead:
65+
66+- The **OAuth protocol flow** — PAR, the authorization request, the DPoP nonce
67+ retry dance (a PDS rejects the first request and returns a nonce to use).
68+ Mechanical, but a lot of it.
69+- **DID/handle resolution** and the PLC directory.
70+- **XRPC** and the Jetstream websocket consumer (`indigo` covers both today).
71+- **CBOR/CAR** parsing for repository records.
72+
73+The honest read: nothing here is blocking, and the two things that could have
74+been blocking are not. The remaining work is large but ordinary.
added spike/nim/sqlite_probe.nim +151 -0
new file mode 100644
@@ -0,0 +1,151 @@
1+## Spike: can Nim drive the two SQLite features Glean depends on?
2+##
3+## Glean's Go backend needs FTS5 (full-text search over articles) and
4+## sqlite-vec (embedding similarity for recommendations). Both are C-level
5+## concerns, so the question for a Nim port is not "does Nim have a SQLite
6+## library" but whether the same extension can be linked and registered the
7+## same way. The Go bindings compile sqlite-vec.c with -DSQLITE_CORE and call
8+## sqlite3_vec_init directly; this does exactly that from Nim.
9+
10+{.compile: "vendor/sqlite-vec.c".} # flags live in nim.cfg
11+
12+import std/[strformat, strutils, sequtils]
13+
14+type
15+ Sqlite3 = distinct pointer
16+ Stmt = distinct pointer
17+
18+const
19+ SQLITE_OK = 0
20+ SQLITE_ROW = 100
21+ SQLITE_DONE = 101
22+
23+{.push importc, cdecl.}
24+proc sqlite3_open(filename: cstring, db: var Sqlite3): cint
25+proc sqlite3_close(db: Sqlite3): cint
26+proc sqlite3_errmsg(db: Sqlite3): cstring
27+proc sqlite3_exec(db: Sqlite3, sql: cstring, cb, arg: pointer,
28+ errmsg: ptr cstring): cint
29+proc sqlite3_prepare_v2(db: Sqlite3, sql: cstring, nByte: cint,
30+ stmt: var Stmt, tail: ptr cstring): cint
31+proc sqlite3_step(s: Stmt): cint
32+proc sqlite3_finalize(s: Stmt): cint
33+proc sqlite3_column_text(s: Stmt, col: cint): cstring
34+proc sqlite3_column_double(s: Stmt, col: cint): cdouble
35+proc sqlite3_column_int64(s: Stmt, col: cint): int64
36+proc sqlite3_bind_blob(s: Stmt, col: cint, data: pointer, n: cint,
37+ destructor: pointer): cint
38+proc sqlite3_libversion(): cstring
39+{.pop.}
40+
41+# Provided by the compiled sqlite-vec.c. With SQLITE_CORE it is an ordinary
42+# function rather than a loadable-extension entry point.
43+proc sqlite3_vec_init(db: Sqlite3, errMsg: ptr cstring,
44+ api: pointer): cint {.importc, cdecl.}
45+
46+const SQLITE_TRANSIENT = cast[pointer](-1)
47+
48+var failures = 0
49+
50+proc check(db: Sqlite3, rc: cint, what: string) =
51+ if rc notin [SQLITE_OK, SQLITE_ROW, SQLITE_DONE]:
52+ raise newException(IOError, &"{what}: {sqlite3_errmsg(db)} (rc={rc})")
53+
54+proc exec(db: Sqlite3, sql: string) =
55+ check(db, sqlite3_exec(db, sql.cstring, nil, nil, nil), "exec " & sql[0..min(40, sql.high)])
56+
57+proc report(name: string, ok: bool, detail = "") =
58+ if ok:
59+ echo &" PASS {name}" & (if detail.len > 0: " -- " & detail else: "")
60+ else:
61+ inc failures
62+ echo &" FAIL {name}" & (if detail.len > 0: " -- " & detail else: "")
63+
64+## Embeddings are bound as raw little-endian float32 blobs, the same wire
65+## format sqlite-vec expects from every binding.
66+proc toBlob(v: seq[float32]): string =
67+ result = newString(v.len * 4)
68+ if v.len > 0:
69+ copyMem(result[0].addr, v[0].unsafeAddr, result.len)
70+
71+proc main() =
72+ echo "sqlite-vec / FTS5 spike"
73+ echo " sqlite ", sqlite3_libversion()
74+
75+ var db: Sqlite3
76+ check(db, sqlite3_open(":memory:", db), "open")
77+ defer: discard sqlite3_close(db)
78+
79+ # --- sqlite-vec registration ---
80+ block:
81+ let rc = sqlite3_vec_init(db, nil, nil)
82+ report("sqlite3_vec_init registers", rc == SQLITE_OK, &"rc={rc}")
83+
84+ block:
85+ var s: Stmt
86+ check(db, sqlite3_prepare_v2(db, "SELECT vec_version()", -1, s, nil), "prepare vec_version")
87+ let ok = sqlite3_step(s) == SQLITE_ROW
88+ let ver = if ok: $sqlite3_column_text(s, 0) else: ""
89+ discard sqlite3_finalize(s)
90+ report("vec_version() callable", ok, ver)
91+
92+ # --- FTS5, as used for article search ---
93+ block:
94+ db.exec "CREATE VIRTUAL TABLE docs USING fts5(title, body)"
95+ db.exec "INSERT INTO docs VALUES ('Nim at scale', 'systems programming with a python face')"
96+ db.exec "INSERT INTO docs VALUES ('Go concurrency', 'goroutines and channels')"
97+ var s: Stmt
98+ check(db, sqlite3_prepare_v2(db,
99+ "SELECT title FROM docs WHERE docs MATCH 'systems' ORDER BY rank", -1, s, nil), "prepare fts")
100+ let ok = sqlite3_step(s) == SQLITE_ROW
101+ let title = if ok: $sqlite3_column_text(s, 0) else: ""
102+ discard sqlite3_finalize(s)
103+ report("FTS5 match + rank", ok and title == "Nim at scale", title)
104+
105+ # --- vec0 virtual table: the recommendation path ---
106+ block:
107+ db.exec "CREATE VIRTUAL TABLE embeddings USING vec0(id INTEGER PRIMARY KEY, v float[4])"
108+ report("vec0 table created", true)
109+
110+ # Three vectors; the query is nearest to id 2.
111+ let rows = {
112+ 1'i64: @[1.0'f32, 0.0, 0.0, 0.0],
113+ 2'i64: @[0.0'f32, 1.0, 0.0, 0.0],
114+ 3'i64: @[0.0'f32, 0.0, 1.0, 0.0],
115+ }
116+ for (id, vec) in rows.items:
117+ var s: Stmt
118+ let sql = &"INSERT INTO embeddings(id, v) VALUES ({id}, ?)"
119+ check(db, sqlite3_prepare_v2(db, sql.cstring, -1, s, nil), "prepare insert")
120+ let blob = vec.toBlob
121+ check(db, sqlite3_bind_blob(s, 1, blob[0].unsafeAddr, blob.len.cint,
122+ SQLITE_TRANSIENT), "bind blob")
123+ check(db, sqlite3_step(s), "insert vec")
124+ discard sqlite3_finalize(s)
125+ report("float32 blobs bound", true, &"{rows.len} rows")
126+
127+ block:
128+ let query = @[0.05'f32, 0.95, 0.0, 0.0]
129+ var s: Stmt
130+ check(db, sqlite3_prepare_v2(db,
131+ "SELECT id, distance FROM embeddings WHERE v MATCH ? AND k = 2 ORDER BY distance",
132+ -1, s, nil), "prepare knn")
133+ let blob = query.toBlob
134+ check(db, sqlite3_bind_blob(s, 1, blob[0].unsafeAddr, blob.len.cint,
135+ SQLITE_TRANSIENT), "bind query")
136+ var got: seq[(int64, float)]
137+ while sqlite3_step(s) == SQLITE_ROW:
138+ got.add (sqlite3_column_int64(s, 0), sqlite3_column_double(s, 1).float)
139+ discard sqlite3_finalize(s)
140+ let nearest = if got.len > 0: got[0][0] else: -1
141+ report("KNN returns nearest first", got.len == 2 and nearest == 2,
142+ &"k=2 -> {got.mapIt($it[0]).join(\",\")}")
143+
144+ echo ""
145+ if failures == 0:
146+ echo "RESULT: sqlite-vec + FTS5 work from Nim."
147+ else:
148+ echo &"RESULT: {failures} check(s) failed."
149+ quit 1
150+
151+main()
new file mode 100644
@@ -0,0 +1,151 @@
1+## Spike: can Nim drive the two SQLite features Glean depends on?
2+##
3+## Glean's Go backend needs FTS5 (full-text search over articles) and
4+## sqlite-vec (embedding similarity for recommendations). Both are C-level
5+## concerns, so the question for a Nim port is not "does Nim have a SQLite
6+## library" but whether the same extension can be linked and registered the
7+## same way. The Go bindings compile sqlite-vec.c with -DSQLITE_CORE and call
8+## sqlite3_vec_init directly; this does exactly that from Nim.
9+
10+{.compile: "vendor/sqlite-vec.c".} # flags live in nim.cfg
11+
12+import std/[strformat, strutils, sequtils]
13+
14+type
15+ Sqlite3 = distinct pointer
16+ Stmt = distinct pointer
17+
18+const
19+ SQLITE_OK = 0
20+ SQLITE_ROW = 100
21+ SQLITE_DONE = 101
22+
23+{.push importc, cdecl.}
24+proc sqlite3_open(filename: cstring, db: var Sqlite3): cint
25+proc sqlite3_close(db: Sqlite3): cint
26+proc sqlite3_errmsg(db: Sqlite3): cstring
27+proc sqlite3_exec(db: Sqlite3, sql: cstring, cb, arg: pointer,
28+ errmsg: ptr cstring): cint
29+proc sqlite3_prepare_v2(db: Sqlite3, sql: cstring, nByte: cint,
30+ stmt: var Stmt, tail: ptr cstring): cint
31+proc sqlite3_step(s: Stmt): cint
32+proc sqlite3_finalize(s: Stmt): cint
33+proc sqlite3_column_text(s: Stmt, col: cint): cstring
34+proc sqlite3_column_double(s: Stmt, col: cint): cdouble
35+proc sqlite3_column_int64(s: Stmt, col: cint): int64
36+proc sqlite3_bind_blob(s: Stmt, col: cint, data: pointer, n: cint,
37+ destructor: pointer): cint
38+proc sqlite3_libversion(): cstring
39+{.pop.}
40+
41+# Provided by the compiled sqlite-vec.c. With SQLITE_CORE it is an ordinary
42+# function rather than a loadable-extension entry point.
43+proc sqlite3_vec_init(db: Sqlite3, errMsg: ptr cstring,
44+ api: pointer): cint {.importc, cdecl.}
45+
46+const SQLITE_TRANSIENT = cast[pointer](-1)
47+
48+var failures = 0
49+
50+proc check(db: Sqlite3, rc: cint, what: string) =
51+ if rc notin [SQLITE_OK, SQLITE_ROW, SQLITE_DONE]:
52+ raise newException(IOError, &"{what}: {sqlite3_errmsg(db)} (rc={rc})")
53+
54+proc exec(db: Sqlite3, sql: string) =
55+ check(db, sqlite3_exec(db, sql.cstring, nil, nil, nil), "exec " & sql[0..min(40, sql.high)])
56+
57+proc report(name: string, ok: bool, detail = "") =
58+ if ok:
59+ echo &" PASS {name}" & (if detail.len > 0: " -- " & detail else: "")
60+ else:
61+ inc failures
62+ echo &" FAIL {name}" & (if detail.len > 0: " -- " & detail else: "")
63+
64+## Embeddings are bound as raw little-endian float32 blobs, the same wire
65+## format sqlite-vec expects from every binding.
66+proc toBlob(v: seq[float32]): string =
67+ result = newString(v.len * 4)
68+ if v.len > 0:
69+ copyMem(result[0].addr, v[0].unsafeAddr, result.len)
70+
71+proc main() =
72+ echo "sqlite-vec / FTS5 spike"
73+ echo " sqlite ", sqlite3_libversion()
74+
75+ var db: Sqlite3
76+ check(db, sqlite3_open(":memory:", db), "open")
77+ defer: discard sqlite3_close(db)
78+
79+ # --- sqlite-vec registration ---
80+ block:
81+ let rc = sqlite3_vec_init(db, nil, nil)
82+ report("sqlite3_vec_init registers", rc == SQLITE_OK, &"rc={rc}")
83+
84+ block:
85+ var s: Stmt
86+ check(db, sqlite3_prepare_v2(db, "SELECT vec_version()", -1, s, nil), "prepare vec_version")
87+ let ok = sqlite3_step(s) == SQLITE_ROW
88+ let ver = if ok: $sqlite3_column_text(s, 0) else: ""
89+ discard sqlite3_finalize(s)
90+ report("vec_version() callable", ok, ver)
91+
92+ # --- FTS5, as used for article search ---
93+ block:
94+ db.exec "CREATE VIRTUAL TABLE docs USING fts5(title, body)"
95+ db.exec "INSERT INTO docs VALUES ('Nim at scale', 'systems programming with a python face')"
96+ db.exec "INSERT INTO docs VALUES ('Go concurrency', 'goroutines and channels')"
97+ var s: Stmt
98+ check(db, sqlite3_prepare_v2(db,
99+ "SELECT title FROM docs WHERE docs MATCH 'systems' ORDER BY rank", -1, s, nil), "prepare fts")
100+ let ok = sqlite3_step(s) == SQLITE_ROW
101+ let title = if ok: $sqlite3_column_text(s, 0) else: ""
102+ discard sqlite3_finalize(s)
103+ report("FTS5 match + rank", ok and title == "Nim at scale", title)
104+
105+ # --- vec0 virtual table: the recommendation path ---
106+ block:
107+ db.exec "CREATE VIRTUAL TABLE embeddings USING vec0(id INTEGER PRIMARY KEY, v float[4])"
108+ report("vec0 table created", true)
109+
110+ # Three vectors; the query is nearest to id 2.
111+ let rows = {
112+ 1'i64: @[1.0'f32, 0.0, 0.0, 0.0],
113+ 2'i64: @[0.0'f32, 1.0, 0.0, 0.0],
114+ 3'i64: @[0.0'f32, 0.0, 1.0, 0.0],
115+ }
116+ for (id, vec) in rows.items:
117+ var s: Stmt
118+ let sql = &"INSERT INTO embeddings(id, v) VALUES ({id}, ?)"
119+ check(db, sqlite3_prepare_v2(db, sql.cstring, -1, s, nil), "prepare insert")
120+ let blob = vec.toBlob
121+ check(db, sqlite3_bind_blob(s, 1, blob[0].unsafeAddr, blob.len.cint,
122+ SQLITE_TRANSIENT), "bind blob")
123+ check(db, sqlite3_step(s), "insert vec")
124+ discard sqlite3_finalize(s)
125+ report("float32 blobs bound", true, &"{rows.len} rows")
126+
127+ block:
128+ let query = @[0.05'f32, 0.95, 0.0, 0.0]
129+ var s: Stmt
130+ check(db, sqlite3_prepare_v2(db,
131+ "SELECT id, distance FROM embeddings WHERE v MATCH ? AND k = 2 ORDER BY distance",
132+ -1, s, nil), "prepare knn")
133+ let blob = query.toBlob
134+ check(db, sqlite3_bind_blob(s, 1, blob[0].unsafeAddr, blob.len.cint,
135+ SQLITE_TRANSIENT), "bind query")
136+ var got: seq[(int64, float)]
137+ while sqlite3_step(s) == SQLITE_ROW:
138+ got.add (sqlite3_column_int64(s, 0), sqlite3_column_double(s, 1).float)
139+ discard sqlite3_finalize(s)
140+ let nearest = if got.len > 0: got[0][0] else: -1
141+ report("KNN returns nearest first", got.len == 2 and nearest == 2,
142+ &"k=2 -> {got.mapIt($it[0]).join(\",\")}")
143+
144+ echo ""
145+ if failures == 0:
146+ echo "RESULT: sqlite-vec + FTS5 work from Nim."
147+ else:
148+ echo &"RESULT: {failures} check(s) failed."
149+ quit 1
150+
151+main()
added spike/nim/vendor/sqlite-vec.c +9749 -0
new file mode 100644
@@ -0,0 +1,9749 @@
1+#include "sqlite-vec.h"
2+
3+#include <assert.h>
4+#include <errno.h>
5+#include <float.h>
6+#include <inttypes.h>
7+#include <limits.h>
8+#include <math.h>
9+#include <stdbool.h>
10+#include <stdint.h>
11+#include <stdlib.h>
12+#include <string.h>
13+
14+#ifndef SQLITE_VEC_OMIT_FS
15+#include <stdio.h>
16+#endif
17+
18+#ifndef SQLITE_CORE
19+#include "sqlite3ext.h"
20+SQLITE_EXTENSION_INIT1
21+#else
22+#include "sqlite3.h"
23+#endif
24+
25+#ifndef UINT32_TYPE
26+#ifdef HAVE_UINT32_T
27+#define UINT32_TYPE uint32_t
28+#else
29+#define UINT32_TYPE unsigned int
30+#endif
31+#endif
32+#ifndef UINT16_TYPE
33+#ifdef HAVE_UINT16_T
34+#define UINT16_TYPE uint16_t
35+#else
36+#define UINT16_TYPE unsigned short int
37+#endif
38+#endif
39+#ifndef INT16_TYPE
40+#ifdef HAVE_INT16_T
41+#define INT16_TYPE int16_t
42+#else
43+#define INT16_TYPE short int
44+#endif
45+#endif
46+#ifndef UINT8_TYPE
47+#ifdef HAVE_UINT8_T
48+#define UINT8_TYPE uint8_t
49+#else
50+#define UINT8_TYPE unsigned char
51+#endif
52+#endif
53+#ifndef INT8_TYPE
54+#ifdef HAVE_INT8_T
55+#define INT8_TYPE int8_t
56+#else
57+#define INT8_TYPE signed char
58+#endif
59+#endif
60+#ifndef LONGDOUBLE_TYPE
61+#define LONGDOUBLE_TYPE long double
62+#endif
63+
64+#ifndef _WIN32
65+#ifndef __EMSCRIPTEN__
66+#ifndef __COSMOPOLITAN__
67+#ifndef __wasi__
68+typedef u_int8_t uint8_t;
69+typedef u_int16_t uint16_t;
70+typedef u_int64_t uint64_t;
71+#endif
72+#endif
73+#endif
74+#endif
75+
76+typedef int8_t i8;
77+typedef uint8_t u8;
78+typedef int16_t i16;
79+typedef int32_t i32;
80+typedef sqlite3_int64 i64;
81+typedef uint32_t u32;
82+typedef uint64_t u64;
83+typedef float f32;
84+typedef size_t usize;
85+
86+#ifndef UNUSED_PARAMETER
87+#define UNUSED_PARAMETER(X) (void)(X)
88+#endif
89+
90+// sqlite3_vtab_in() was added in SQLite version 3.38 (2022-02-22)
91+// https://www.sqlite.org/changes.html#version_3_38_0
92+#if SQLITE_VERSION_NUMBER >= 3038000
93+#define COMPILER_SUPPORTS_VTAB_IN 1
94+#endif
95+
96+#ifndef SQLITE_SUBTYPE
97+#define SQLITE_SUBTYPE 0x000100000
98+#endif
99+
100+#ifndef SQLITE_RESULT_SUBTYPE
101+#define SQLITE_RESULT_SUBTYPE 0x001000000
102+#endif
103+
104+#ifndef SQLITE_INDEX_CONSTRAINT_LIMIT
105+#define SQLITE_INDEX_CONSTRAINT_LIMIT 73
106+#endif
107+
108+#ifndef SQLITE_INDEX_CONSTRAINT_OFFSET
109+#define SQLITE_INDEX_CONSTRAINT_OFFSET 74
110+#endif
111+
112+#define countof(x) (sizeof(x) / sizeof((x)[0]))
113+#define min(a, b) (((a) <= (b)) ? (a) : (b))
114+
115+enum VectorElementType {
116+ // clang-format off
117+ SQLITE_VEC_ELEMENT_TYPE_FLOAT32 = 223 + 0,
118+ SQLITE_VEC_ELEMENT_TYPE_BIT = 223 + 1,
119+ SQLITE_VEC_ELEMENT_TYPE_INT8 = 223 + 2,
120+ // clang-format on
121+};
122+
123+#ifdef SQLITE_VEC_ENABLE_AVX
124+#include <immintrin.h>
125+#define PORTABLE_ALIGN32 __attribute__((aligned(32)))
126+#define PORTABLE_ALIGN64 __attribute__((aligned(64)))
127+
128+static f32 l2_sqr_float_avx(const void *pVect1v, const void *pVect2v,
129+ const void *qty_ptr) {
130+ f32 *pVect1 = (f32 *)pVect1v;
131+ f32 *pVect2 = (f32 *)pVect2v;
132+ size_t qty = *((size_t *)qty_ptr);
133+ f32 PORTABLE_ALIGN32 TmpRes[8];
134+ size_t qty16 = qty >> 4;
135+
136+ const f32 *pEnd1 = pVect1 + (qty16 << 4);
137+
138+ __m256 diff, v1, v2;
139+ __m256 sum = _mm256_set1_ps(0);
140+
141+ while (pVect1 < pEnd1) {
142+ v1 = _mm256_loadu_ps(pVect1);
143+ pVect1 += 8;
144+ v2 = _mm256_loadu_ps(pVect2);
145+ pVect2 += 8;
146+ diff = _mm256_sub_ps(v1, v2);
147+ sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff));
148+
149+ v1 = _mm256_loadu_ps(pVect1);
150+ pVect1 += 8;
151+ v2 = _mm256_loadu_ps(pVect2);
152+ pVect2 += 8;
153+ diff = _mm256_sub_ps(v1, v2);
154+ sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff));
155+ }
156+
157+ _mm256_store_ps(TmpRes, sum);
158+ return sqrt(TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] +
159+ TmpRes[5] + TmpRes[6] + TmpRes[7]);
160+}
161+#endif
162+
163+#ifdef SQLITE_VEC_ENABLE_NEON
164+#include <arm_neon.h>
165+
166+#define PORTABLE_ALIGN32 __attribute__((aligned(32)))
167+
168+// thx https://github.com/nmslib/hnswlib/pull/299/files
169+static f32 l2_sqr_float_neon(const void *pVect1v, const void *pVect2v,
170+ const void *qty_ptr) {
171+ f32 *pVect1 = (f32 *)pVect1v;
172+ f32 *pVect2 = (f32 *)pVect2v;
173+ size_t qty = *((size_t *)qty_ptr);
174+ size_t qty16 = qty >> 4;
175+
176+ const f32 *pEnd1 = pVect1 + (qty16 << 4);
177+
178+ float32x4_t diff, v1, v2;
179+ float32x4_t sum0 = vdupq_n_f32(0);
180+ float32x4_t sum1 = vdupq_n_f32(0);
181+ float32x4_t sum2 = vdupq_n_f32(0);
182+ float32x4_t sum3 = vdupq_n_f32(0);
183+
184+ while (pVect1 < pEnd1) {
185+ v1 = vld1q_f32(pVect1);
186+ pVect1 += 4;
187+ v2 = vld1q_f32(pVect2);
188+ pVect2 += 4;
189+ diff = vsubq_f32(v1, v2);
190+ sum0 = vfmaq_f32(sum0, diff, diff);
191+
192+ v1 = vld1q_f32(pVect1);
193+ pVect1 += 4;
194+ v2 = vld1q_f32(pVect2);
195+ pVect2 += 4;
196+ diff = vsubq_f32(v1, v2);
197+ sum1 = vfmaq_f32(sum1, diff, diff);
198+
199+ v1 = vld1q_f32(pVect1);
200+ pVect1 += 4;
201+ v2 = vld1q_f32(pVect2);
202+ pVect2 += 4;
203+ diff = vsubq_f32(v1, v2);
204+ sum2 = vfmaq_f32(sum2, diff, diff);
205+
206+ v1 = vld1q_f32(pVect1);
207+ pVect1 += 4;
208+ v2 = vld1q_f32(pVect2);
209+ pVect2 += 4;
210+ diff = vsubq_f32(v1, v2);
211+ sum3 = vfmaq_f32(sum3, diff, diff);
212+ }
213+
214+ f32 sum_scalar =
215+ vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3)));
216+ const f32 *pEnd2 = pVect1 + (qty - (qty16 << 4));
217+ while (pVect1 < pEnd2) {
218+ f32 diff = *pVect1 - *pVect2;
219+ sum_scalar += diff * diff;
220+ pVect1++;
221+ pVect2++;
222+ }
223+
224+ return sqrt(sum_scalar);
225+}
226+
227+static f32 l2_sqr_int8_neon(const void *pVect1v, const void *pVect2v,
228+ const void *qty_ptr) {
229+ i8 *pVect1 = (i8 *)pVect1v;
230+ i8 *pVect2 = (i8 *)pVect2v;
231+ size_t qty = *((size_t *)qty_ptr);
232+
233+ const i8 *pEnd1 = pVect1 + qty;
234+ i32 sum_scalar = 0;
235+
236+ while (pVect1 < pEnd1 - 7) {
237+ // loading 8 at a time
238+ int8x8_t v1 = vld1_s8(pVect1);
239+ int8x8_t v2 = vld1_s8(pVect2);
240+ pVect1 += 8;
241+ pVect2 += 8;
242+
243+ // widen to protect against overflow
244+ int16x8_t v1_wide = vmovl_s8(v1);
245+ int16x8_t v2_wide = vmovl_s8(v2);
246+
247+ int16x8_t diff = vsubq_s16(v1_wide, v2_wide);
248+ int16x8_t squared_diff = vmulq_s16(diff, diff);
249+ int32x4_t sum = vpaddlq_s16(squared_diff);
250+
251+ sum_scalar += vgetq_lane_s32(sum, 0) + vgetq_lane_s32(sum, 1) +
252+ vgetq_lane_s32(sum, 2) + vgetq_lane_s32(sum, 3);
253+ }
254+
255+ // handle leftovers
256+ while (pVect1 < pEnd1) {
257+ i16 diff = (i16)*pVect1 - (i16)*pVect2;
258+ sum_scalar += diff * diff;
259+ pVect1++;
260+ pVect2++;
261+ }
262+
263+ return sqrtf(sum_scalar);
264+}
265+
266+static i32 l1_int8_neon(const void *pVect1v, const void *pVect2v,
267+ const void *qty_ptr) {
268+ i8 *pVect1 = (i8 *)pVect1v;
269+ i8 *pVect2 = (i8 *)pVect2v;
270+ size_t qty = *((size_t *)qty_ptr);
271+
272+ const int8_t *pEnd1 = pVect1 + qty;
273+
274+ int32x4_t acc1 = vdupq_n_s32(0);
275+ int32x4_t acc2 = vdupq_n_s32(0);
276+ int32x4_t acc3 = vdupq_n_s32(0);
277+ int32x4_t acc4 = vdupq_n_s32(0);
278+
279+ while (pVect1 < pEnd1 - 63) {
280+ int8x16_t v1 = vld1q_s8(pVect1);
281+ int8x16_t v2 = vld1q_s8(pVect2);
282+ int8x16_t diff1 = vabdq_s8(v1, v2);
283+ acc1 = vaddq_s32(acc1, vpaddlq_u16(vpaddlq_u8(diff1)));
284+
285+ v1 = vld1q_s8(pVect1 + 16);
286+ v2 = vld1q_s8(pVect2 + 16);
287+ int8x16_t diff2 = vabdq_s8(v1, v2);
288+ acc2 = vaddq_s32(acc2, vpaddlq_u16(vpaddlq_u8(diff2)));
289+
290+ v1 = vld1q_s8(pVect1 + 32);
291+ v2 = vld1q_s8(pVect2 + 32);
292+ int8x16_t diff3 = vabdq_s8(v1, v2);
293+ acc3 = vaddq_s32(acc3, vpaddlq_u16(vpaddlq_u8(diff3)));
294+
295+ v1 = vld1q_s8(pVect1 + 48);
296+ v2 = vld1q_s8(pVect2 + 48);
297+ int8x16_t diff4 = vabdq_s8(v1, v2);
298+ acc4 = vaddq_s32(acc4, vpaddlq_u16(vpaddlq_u8(diff4)));
299+
300+ pVect1 += 64;
301+ pVect2 += 64;
302+ }
303+
304+ while (pVect1 < pEnd1 - 15) {
305+ int8x16_t v1 = vld1q_s8(pVect1);
306+ int8x16_t v2 = vld1q_s8(pVect2);
307+ int8x16_t diff = vabdq_s8(v1, v2);
308+ acc1 = vaddq_s32(acc1, vpaddlq_u16(vpaddlq_u8(diff)));
309+ pVect1 += 16;
310+ pVect2 += 16;
311+ }
312+
313+ int32x4_t acc = vaddq_s32(vaddq_s32(acc1, acc2), vaddq_s32(acc3, acc4));
314+
315+ int32_t sum = 0;
316+ while (pVect1 < pEnd1) {
317+ int32_t diff = abs((int32_t)*pVect1 - (int32_t)*pVect2);
318+ sum += diff;
319+ pVect1++;
320+ pVect2++;
321+ }
322+
323+ return vaddvq_s32(acc) + sum;
324+}
325+
326+static double l1_f32_neon(const void *pVect1v, const void *pVect2v,
327+ const void *qty_ptr) {
328+ f32 *pVect1 = (f32 *)pVect1v;
329+ f32 *pVect2 = (f32 *)pVect2v;
330+ size_t qty = *((size_t *)qty_ptr);
331+
332+ const f32 *pEnd1 = pVect1 + qty;
333+ float64x2_t acc = vdupq_n_f64(0);
334+
335+ while (pVect1 < pEnd1 - 3) {
336+ float32x4_t v1 = vld1q_f32(pVect1);
337+ float32x4_t v2 = vld1q_f32(pVect2);
338+ pVect1 += 4;
339+ pVect2 += 4;
340+
341+ // f32x4 -> f64x2 pad for overflow
342+ float64x2_t low_diff = vabdq_f64(vcvt_f64_f32(vget_low_f32(v1)),
343+ vcvt_f64_f32(vget_low_f32(v2)));
344+ float64x2_t high_diff =
345+ vabdq_f64(vcvt_high_f64_f32(v1), vcvt_high_f64_f32(v2));
346+
347+ acc = vaddq_f64(acc, vaddq_f64(low_diff, high_diff));
348+ }
349+
350+ double sum = 0;
351+ while (pVect1 < pEnd1) {
352+ sum += fabs((double)*pVect1 - (double)*pVect2);
353+ pVect1++;
354+ pVect2++;
355+ }
356+
357+ return vaddvq_f64(acc) + sum;
358+}
359+#endif
360+
361+static f32 l2_sqr_float(const void *pVect1v, const void *pVect2v,
362+ const void *qty_ptr) {
363+ f32 *pVect1 = (f32 *)pVect1v;
364+ f32 *pVect2 = (f32 *)pVect2v;
365+ size_t qty = *((size_t *)qty_ptr);
366+
367+ f32 res = 0;
368+ for (size_t i = 0; i < qty; i++) {
369+ f32 t = *pVect1 - *pVect2;
370+ pVect1++;
371+ pVect2++;
372+ res += t * t;
373+ }
374+ return sqrt(res);
375+}
376+
377+static f32 l2_sqr_int8(const void *pA, const void *pB, const void *pD) {
378+ i8 *a = (i8 *)pA;
379+ i8 *b = (i8 *)pB;
380+ size_t d = *((size_t *)pD);
381+
382+ f32 res = 0;
383+ for (size_t i = 0; i < d; i++) {
384+ f32 t = *a - *b;
385+ a++;
386+ b++;
387+ res += t * t;
388+ }
389+ return sqrt(res);
390+}
391+
392+static f32 distance_l2_sqr_float(const void *a, const void *b, const void *d) {
393+#ifdef SQLITE_VEC_ENABLE_NEON
394+ if ((*(const size_t *)d) > 16) {
395+ return l2_sqr_float_neon(a, b, d);
396+ }
397+#endif
398+#ifdef SQLITE_VEC_ENABLE_AVX
399+ if (((*(const size_t *)d) % 16 == 0)) {
400+ return l2_sqr_float_avx(a, b, d);
401+ }
402+#endif
403+ return l2_sqr_float(a, b, d);
404+}
405+
406+static f32 distance_l2_sqr_int8(const void *a, const void *b, const void *d) {
407+#ifdef SQLITE_VEC_ENABLE_NEON
408+ if ((*(const size_t *)d) > 7) {
409+ return l2_sqr_int8_neon(a, b, d);
410+ }
411+#endif
412+ return l2_sqr_int8(a, b, d);
413+}
414+
415+static i32 l1_int8(const void *pA, const void *pB, const void *pD) {
416+ i8 *a = (i8 *)pA;
417+ i8 *b = (i8 *)pB;
418+ size_t d = *((size_t *)pD);
419+
420+ i32 res = 0;
421+ for (size_t i = 0; i < d; i++) {
422+ res += abs(*a - *b);
423+ a++;
424+ b++;
425+ }
426+
427+ return res;
428+}
429+
430+static i32 distance_l1_int8(const void *a, const void *b, const void *d) {
431+#ifdef SQLITE_VEC_ENABLE_NEON
432+ if ((*(const size_t *)d) > 15) {
433+ return l1_int8_neon(a, b, d);
434+ }
435+#endif
436+ return l1_int8(a, b, d);
437+}
438+
439+static double l1_f32(const void *pA, const void *pB, const void *pD) {
440+ f32 *a = (f32 *)pA;
441+ f32 *b = (f32 *)pB;
442+ size_t d = *((size_t *)pD);
443+
444+ double res = 0;
445+ for (size_t i = 0; i < d; i++) {
446+ res += fabs((double)*a - (double)*b);
447+ a++;
448+ b++;
449+ }
450+
451+ return res;
452+}
453+
454+static double distance_l1_f32(const void *a, const void *b, const void *d) {
455+#ifdef SQLITE_VEC_ENABLE_NEON
456+ if ((*(const size_t *)d) > 3) {
457+ return l1_f32_neon(a, b, d);
458+ }
459+#endif
460+ return l1_f32(a, b, d);
461+}
462+
463+static f32 distance_cosine_float(const void *pVect1v, const void *pVect2v,
464+ const void *qty_ptr) {
465+ f32 *pVect1 = (f32 *)pVect1v;
466+ f32 *pVect2 = (f32 *)pVect2v;
467+ size_t qty = *((size_t *)qty_ptr);
468+
469+ f32 dot = 0;
470+ f32 aMag = 0;
471+ f32 bMag = 0;
472+ for (size_t i = 0; i < qty; i++) {
473+ dot += *pVect1 * *pVect2;
474+ aMag += *pVect1 * *pVect1;
475+ bMag += *pVect2 * *pVect2;
476+ pVect1++;
477+ pVect2++;
478+ }
479+ return 1 - (dot / (sqrt(aMag) * sqrt(bMag)));
480+}
481+static f32 distance_cosine_int8(const void *pA, const void *pB,
482+ const void *pD) {
483+ i8 *a = (i8 *)pA;
484+ i8 *b = (i8 *)pB;
485+ size_t d = *((size_t *)pD);
486+
487+ f32 dot = 0;
488+ f32 aMag = 0;
489+ f32 bMag = 0;
490+ for (size_t i = 0; i < d; i++) {
491+ dot += *a * *b;
492+ aMag += *a * *a;
493+ bMag += *b * *b;
494+ a++;
495+ b++;
496+ }
497+ return 1 - (dot / (sqrt(aMag) * sqrt(bMag)));
498+}
499+
500+// https://github.com/facebookresearch/faiss/blob/77e2e79cd0a680adc343b9840dd865da724c579e/faiss/utils/hamming_distance/common.h#L34
501+static u8 hamdist_table[256] = {
502+ 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 1, 2, 2, 3, 2, 3, 3, 4,
503+ 2, 3, 3, 4, 3, 4, 4, 5, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
504+ 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 1, 2, 2, 3, 2, 3, 3, 4,
505+ 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
506+ 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6,
507+ 4, 5, 5, 6, 5, 6, 6, 7, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
508+ 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5,
509+ 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
510+ 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6,
511+ 4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
512+ 4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8};
513+
514+static f32 distance_hamming_u8(u8 *a, u8 *b, size_t n) {
515+ int same = 0;
516+ for (unsigned long i = 0; i < n; i++) {
517+ same += hamdist_table[a[i] ^ b[i]];
518+ }
519+ return (f32)same;
520+}
521+
522+#ifdef _MSC_VER
523+#if !defined(__clang__) && (defined(_M_ARM) || defined(_M_ARM64))
524+// From
525+// https://github.com/ngtcp2/ngtcp2/blob/b64f1e77b5e0d880b93d31f474147fae4a1d17cc/lib/ngtcp2_ringbuf.c,
526+// line 34-43
527+static unsigned int __builtin_popcountl(unsigned int x) {
528+ unsigned int c = 0;
529+ for (; x; ++c) {
530+ x &= x - 1;
531+ }
532+ return c;
533+}
534+#else
535+#include <intrin.h>
536+#define __builtin_popcountl __popcnt64
537+#endif
538+#endif
539+
540+static f32 distance_hamming_u64(u64 *a, u64 *b, size_t n) {
541+ int same = 0;
542+ for (unsigned long i = 0; i < n; i++) {
543+ same += __builtin_popcountl(a[i] ^ b[i]);
544+ }
545+ return (f32)same;
546+}
547+
548+/**
549+ * @brief Calculate the hamming distance between two bitvectors.
550+ *
551+ * @param a - first bitvector, MUST have d dimensions
552+ * @param b - second bitvector, MUST have d dimensions
553+ * @param d - pointer to size_t, MUST be divisible by CHAR_BIT
554+ * @return f32
555+ */
556+static f32 distance_hamming(const void *a, const void *b, const void *d) {
557+ size_t dimensions = *((size_t *)d);
558+
559+ if ((dimensions % 64) == 0) {
560+ return distance_hamming_u64((u64 *)a, (u64 *)b, dimensions / 8 / CHAR_BIT);
561+ }
562+ return distance_hamming_u8((u8 *)a, (u8 *)b, dimensions / CHAR_BIT);
563+}
564+
565+// from SQLite source:
566+// https://github.com/sqlite/sqlite/blob/a509a90958ddb234d1785ed7801880ccb18b497e/src/json.c#L153
567+static const char vecJsonIsSpaceX[] = {
568+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0,
569+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
570+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
571+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
572+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
573+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
574+
575+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
576+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
577+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
578+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
579+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
580+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
581+};
582+
583+#define vecJsonIsspace(x) (vecJsonIsSpaceX[(unsigned char)x])
584+
585+typedef void (*vector_cleanup)(void *p);
586+
587+void vector_cleanup_noop(void *_) { UNUSED_PARAMETER(_); }
588+
589+#define JSON_SUBTYPE 74
590+
591+void vtab_set_error(sqlite3_vtab *pVTab, const char *zFormat, ...) {
592+ va_list args;
593+ sqlite3_free(pVTab->zErrMsg);
594+ va_start(args, zFormat);
595+ pVTab->zErrMsg = sqlite3_vmprintf(zFormat, args);
596+ va_end(args);
597+}
598+struct Array {
599+ size_t element_size;
600+ size_t length;
601+ size_t capacity;
602+ void *z;
603+};
604+
605+/**
606+ * @brief Initial an array with the given element size and capacity.
607+ *
608+ * @param array
609+ * @param element_size
610+ * @param init_capacity
611+ * @return SQLITE_OK on success, error code on failure. Only error is
612+ * SQLITE_NOMEM
613+ */
614+int array_init(struct Array *array, size_t element_size, size_t init_capacity) {
615+ int sz = element_size * init_capacity;
616+ void *z = sqlite3_malloc(sz);
617+ if (!z) {
618+ return SQLITE_NOMEM;
619+ }
620+ memset(z, 0, sz);
621+
622+ array->element_size = element_size;
623+ array->length = 0;
624+ array->capacity = init_capacity;
625+ array->z = z;
626+ return SQLITE_OK;
627+}
628+
629+int array_append(struct Array *array, const void *element) {
630+ if (array->length == array->capacity) {
631+ size_t new_capacity = array->capacity * 2 + 100;
632+ void *z = sqlite3_realloc64(array->z, array->element_size * new_capacity);
633+ if (z) {
634+ array->capacity = new_capacity;
635+ array->z = z;
636+ } else {
637+ return SQLITE_NOMEM;
638+ }
639+ }
640+ memcpy(&((unsigned char *)array->z)[array->length * array->element_size],
641+ element, array->element_size);
642+ array->length++;
643+ return SQLITE_OK;
644+}
645+
646+void array_cleanup(struct Array *array) {
647+ if (!array)
648+ return;
649+ array->element_size = 0;
650+ array->length = 0;
651+ array->capacity = 0;
652+ sqlite3_free(array->z);
653+ array->z = NULL;
654+}
655+
656+char *vector_subtype_name(int subtype) {
657+ switch (subtype) {
658+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32:
659+ return "float32";
660+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
661+ return "int8";
662+ case SQLITE_VEC_ELEMENT_TYPE_BIT:
663+ return "bit";
664+ }
665+ return "";
666+}
667+char *type_name(int type) {
668+ switch (type) {
669+ case SQLITE_INTEGER:
670+ return "INTEGER";
671+ case SQLITE_BLOB:
672+ return "BLOB";
673+ case SQLITE_TEXT:
674+ return "TEXT";
675+ case SQLITE_FLOAT:
676+ return "FLOAT";
677+ case SQLITE_NULL:
678+ return "NULL";
679+ }
680+ return "";
681+}
682+
683+typedef void (*fvec_cleanup)(f32 *vector);
684+
685+void fvec_cleanup_noop(f32 *_) { UNUSED_PARAMETER(_); }
686+
687+static int fvec_from_value(sqlite3_value *value, f32 **vector,
688+ size_t *dimensions, fvec_cleanup *cleanup,
689+ char **pzErr) {
690+ int value_type = sqlite3_value_type(value);
691+
692+ if (value_type == SQLITE_BLOB) {
693+ const void *blob = sqlite3_value_blob(value);
694+ int bytes = sqlite3_value_bytes(value);
695+ if (bytes == 0) {
696+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
697+ return SQLITE_ERROR;
698+ }
699+ if ((bytes % sizeof(f32)) != 0) {
700+ *pzErr = sqlite3_mprintf("invalid float32 vector BLOB length. Must be "
701+ "divisible by %d, found %d",
702+ sizeof(f32), bytes);
703+ return SQLITE_ERROR;
704+ }
705+ *vector = (f32 *)blob;
706+ *dimensions = bytes / sizeof(f32);
707+ *cleanup = fvec_cleanup_noop;
708+ return SQLITE_OK;
709+ }
710+
711+ if (value_type == SQLITE_TEXT) {
712+ const char *source = (const char *)sqlite3_value_text(value);
713+ int source_len = sqlite3_value_bytes(value);
714+ if (source_len == 0) {
715+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
716+ return SQLITE_ERROR;
717+ }
718+ int i = 0;
719+
720+ struct Array x;
721+ int rc = array_init(&x, sizeof(f32), ceil(source_len / 2.0));
722+ if (rc != SQLITE_OK) {
723+ return rc;
724+ }
725+
726+ // advance leading whitespace to first '['
727+ while (i < source_len) {
728+ if (vecJsonIsspace(source[i])) {
729+ i++;
730+ continue;
731+ }
732+ if (source[i] == '[') {
733+ break;
734+ }
735+
736+ *pzErr = sqlite3_mprintf(
737+ "JSON array parsing error: Input does not start with '['");
738+ array_cleanup(&x);
739+ return SQLITE_ERROR;
740+ }
741+ if (source[i] != '[') {
742+ *pzErr = sqlite3_mprintf(
743+ "JSON array parsing error: Input does not start with '['");
744+ array_cleanup(&x);
745+ return SQLITE_ERROR;
746+ }
747+ int offset = i + 1;
748+
749+ while (offset < source_len) {
750+ char *ptr = (char *)&source[offset];
751+ char *endptr;
752+
753+ errno = 0;
754+ double result = strtod(ptr, &endptr);
755+ if ((errno != 0 && result == 0) // some interval error?
756+ || (errno == ERANGE &&
757+ (result == HUGE_VAL || result == -HUGE_VAL)) // too big / smalls
758+ ) {
759+ sqlite3_free(x.z);
760+ *pzErr = sqlite3_mprintf("JSON parsing error");
761+ return SQLITE_ERROR;
762+ }
763+
764+ if (endptr == ptr) {
765+ if (*ptr != ']') {
766+ sqlite3_free(x.z);
767+ *pzErr = sqlite3_mprintf("JSON parsing error");
768+ return SQLITE_ERROR;
769+ }
770+ goto done;
771+ }
772+
773+ f32 res = (f32)result;
774+ array_append(&x, (const void *)&res);
775+
776+ offset += (endptr - ptr);
777+ while (offset < source_len) {
778+ if (vecJsonIsspace(source[offset])) {
779+ offset++;
780+ continue;
781+ }
782+ if (source[offset] == ',') {
783+ offset++;
784+ continue;
785+ }
786+ if (source[offset] == ']')
787+ goto done;
788+ break;
789+ }
790+ }
791+
792+ done:
793+
794+ if (x.length > 0) {
795+ *vector = (f32 *)x.z;
796+ *dimensions = x.length;
797+ *cleanup = (fvec_cleanup)sqlite3_free;
798+ return SQLITE_OK;
799+ }
800+ sqlite3_free(x.z);
801+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
802+ return SQLITE_ERROR;
803+ }
804+
805+ *pzErr = sqlite3_mprintf(
806+ "Input must have type BLOB (compact format) or TEXT (JSON), found %s",
807+ type_name(value_type));
808+ return SQLITE_ERROR;
809+}
810+
811+static int bitvec_from_value(sqlite3_value *value, u8 **vector,
812+ size_t *dimensions, vector_cleanup *cleanup,
813+ char **pzErr) {
814+ int value_type = sqlite3_value_type(value);
815+ if (value_type == SQLITE_BLOB) {
816+ const void *blob = sqlite3_value_blob(value);
817+ int bytes = sqlite3_value_bytes(value);
818+ if (bytes == 0) {
819+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
820+ return SQLITE_ERROR;
821+ }
822+ *vector = (u8 *)blob;
823+ *dimensions = bytes * CHAR_BIT;
824+ *cleanup = vector_cleanup_noop;
825+ return SQLITE_OK;
826+ }
827+ *pzErr = sqlite3_mprintf("Unknown type for bitvector.");
828+ return SQLITE_ERROR;
829+}
830+
831+static int int8_vec_from_value(sqlite3_value *value, i8 **vector,
832+ size_t *dimensions, vector_cleanup *cleanup,
833+ char **pzErr) {
834+ int value_type = sqlite3_value_type(value);
835+ if (value_type == SQLITE_BLOB) {
836+ const void *blob = sqlite3_value_blob(value);
837+ int bytes = sqlite3_value_bytes(value);
838+ if (bytes == 0) {
839+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
840+ return SQLITE_ERROR;
841+ }
842+ *vector = (i8 *)blob;
843+ *dimensions = bytes;
844+ *cleanup = vector_cleanup_noop;
845+ return SQLITE_OK;
846+ }
847+
848+ if (value_type == SQLITE_TEXT) {
849+ const char *source = (const char *)sqlite3_value_text(value);
850+ int source_len = sqlite3_value_bytes(value);
851+ int i = 0;
852+
853+ if (source_len == 0) {
854+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
855+ return SQLITE_ERROR;
856+ }
857+
858+ struct Array x;
859+ int rc = array_init(&x, sizeof(i8), ceil(source_len / 2.0));
860+ if (rc != SQLITE_OK) {
861+ return rc;
862+ }
863+
864+ // advance leading whitespace to first '['
865+ while (i < source_len) {
866+ if (vecJsonIsspace(source[i])) {
867+ i++;
868+ continue;
869+ }
870+ if (source[i] == '[') {
871+ break;
872+ }
873+
874+ *pzErr = sqlite3_mprintf(
875+ "JSON array parsing error: Input does not start with '['");
876+ array_cleanup(&x);
877+ return SQLITE_ERROR;
878+ }
879+ if (source[i] != '[') {
880+ *pzErr = sqlite3_mprintf(
881+ "JSON array parsing error: Input does not start with '['");
882+ array_cleanup(&x);
883+ return SQLITE_ERROR;
884+ }
885+ int offset = i + 1;
886+
887+ while (offset < source_len) {
888+ char *ptr = (char *)&source[offset];
889+ char *endptr;
890+
891+ errno = 0;
892+ long result = strtol(ptr, &endptr, 10);
893+ if ((errno != 0 && result == 0) ||
894+ (errno == ERANGE && (result == LONG_MAX || result == LONG_MIN))) {
895+ sqlite3_free(x.z);
896+ *pzErr = sqlite3_mprintf("JSON parsing error");
897+ return SQLITE_ERROR;
898+ }
899+
900+ if (endptr == ptr) {
901+ if (*ptr != ']') {
902+ sqlite3_free(x.z);
903+ *pzErr = sqlite3_mprintf("JSON parsing error");
904+ return SQLITE_ERROR;
905+ }
906+ goto done;
907+ }
908+
909+ if (result < INT8_MIN || result > INT8_MAX) {
910+ sqlite3_free(x.z);
911+ *pzErr =
912+ sqlite3_mprintf("JSON parsing error: value out of range for int8");
913+ return SQLITE_ERROR;
914+ }
915+
916+ i8 res = (i8)result;
917+ array_append(&x, (const void *)&res);
918+
919+ offset += (endptr - ptr);
920+ while (offset < source_len) {
921+ if (vecJsonIsspace(source[offset])) {
922+ offset++;
923+ continue;
924+ }
925+ if (source[offset] == ',') {
926+ offset++;
927+ continue;
928+ }
929+ if (source[offset] == ']')
930+ goto done;
931+ break;
932+ }
933+ }
934+
935+ done:
936+
937+ if (x.length > 0) {
938+ *vector = (i8 *)x.z;
939+ *dimensions = x.length;
940+ *cleanup = (vector_cleanup)sqlite3_free;
941+ return SQLITE_OK;
942+ }
943+ sqlite3_free(x.z);
944+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
945+ return SQLITE_ERROR;
946+ }
947+
948+ *pzErr = sqlite3_mprintf("Unknown type for int8 vector.");
949+ return SQLITE_ERROR;
950+}
951+
952+/**
953+ * @brief Extract a vector from a sqlite3_value. Can be a float32, int8, or bit
954+ * vector.
955+ *
956+ * @param value: the sqlite3_value to read from.
957+ * @param vector: Output pointer to vector data.
958+ * @param dimensions: Output number of dimensions
959+ * @param dimensions: Output vector element type
960+ * @param cleanup
961+ * @param pzErrorMessage
962+ * @return int SQLITE_OK on success, error code otherwise
963+ */
964+int vector_from_value(sqlite3_value *value, void **vector, size_t *dimensions,
965+ enum VectorElementType *element_type,
966+ vector_cleanup *cleanup, char **pzErrorMessage) {
967+ int subtype = sqlite3_value_subtype(value);
968+ if (!subtype || (subtype == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) ||
969+ (subtype == JSON_SUBTYPE)) {
970+ int rc = fvec_from_value(value, (f32 **)vector, dimensions,
971+ (fvec_cleanup *)cleanup, pzErrorMessage);
972+ if (rc == SQLITE_OK) {
973+ *element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
974+ }
975+ return rc;
976+ }
977+
978+ if (subtype == SQLITE_VEC_ELEMENT_TYPE_BIT) {
979+ int rc = bitvec_from_value(value, (u8 **)vector, dimensions, cleanup,
980+ pzErrorMessage);
981+ if (rc == SQLITE_OK) {
982+ *element_type = SQLITE_VEC_ELEMENT_TYPE_BIT;
983+ }
984+ return rc;
985+ }
986+ if (subtype == SQLITE_VEC_ELEMENT_TYPE_INT8) {
987+ int rc = int8_vec_from_value(value, (i8 **)vector, dimensions, cleanup,
988+ pzErrorMessage);
989+ if (rc == SQLITE_OK) {
990+ *element_type = SQLITE_VEC_ELEMENT_TYPE_INT8;
991+ }
992+ return rc;
993+ }
994+ *pzErrorMessage = sqlite3_mprintf("Unknown subtype: %d", subtype);
995+ return SQLITE_ERROR;
996+}
997+
998+int ensure_vector_match(sqlite3_value *aValue, sqlite3_value *bValue, void **a,
999+ void **b, enum VectorElementType *element_type,
1000+ size_t *dimensions, vector_cleanup *outACleanup,
1001+ vector_cleanup *outBCleanup, char **outError) {
1002+ int rc;
1003+ enum VectorElementType aType, bType;
1004+ size_t aDims, bDims;
1005+ char *error = NULL;
1006+ vector_cleanup aCleanup, bCleanup;
1007+
1008+ rc = vector_from_value(aValue, a, &aDims, &aType, &aCleanup, &error);
1009+ if (rc != SQLITE_OK) {
1010+ *outError = sqlite3_mprintf("Error reading 1st vector: %s", error);
1011+ sqlite3_free(error);
1012+ return SQLITE_ERROR;
1013+ }
1014+
1015+ rc = vector_from_value(bValue, b, &bDims, &bType, &bCleanup, &error);
1016+ if (rc != SQLITE_OK) {
1017+ *outError = sqlite3_mprintf("Error reading 2nd vector: %s", error);
1018+ sqlite3_free(error);
1019+ aCleanup(a);
1020+ return SQLITE_ERROR;
1021+ }
1022+
1023+ if (aType != bType) {
1024+ *outError =
1025+ sqlite3_mprintf("Vector type mistmatch. First vector has type %s, "
1026+ "while the second has type %s.",
1027+ vector_subtype_name(aType), vector_subtype_name(bType));
1028+ aCleanup(*a);
1029+ bCleanup(*b);
1030+ return SQLITE_ERROR;
1031+ }
1032+ if (aDims != bDims) {
1033+ *outError = sqlite3_mprintf(
1034+ "Vector dimension mistmatch. First vector has %ld dimensions, "
1035+ "while the second has %ld dimensions.",
1036+ aDims, bDims);
1037+ aCleanup(*a);
1038+ bCleanup(*b);
1039+ return SQLITE_ERROR;
1040+ }
1041+ *element_type = aType;
1042+ *dimensions = aDims;
1043+ *outACleanup = aCleanup;
1044+ *outBCleanup = bCleanup;
1045+ return SQLITE_OK;
1046+}
1047+
1048+int _cmp(const void *a, const void *b) { return (*(i64 *)a - *(i64 *)b); }
1049+
1050+struct VecNpyFile {
1051+ char *path;
1052+ size_t pathLength;
1053+};
1054+#define SQLITE_VEC_NPY_FILE_NAME "vec0-npy-file"
1055+
1056+#ifndef SQLITE_VEC_OMIT_FS
1057+static void vec_npy_file(sqlite3_context *context, int argc,
1058+ sqlite3_value **argv) {
1059+ assert(argc == 1);
1060+ char *path = (char *)sqlite3_value_text(argv[0]);
1061+ size_t pathLength = sqlite3_value_bytes(argv[0]);
1062+ struct VecNpyFile *f;
1063+
1064+ f = sqlite3_malloc(sizeof(*f));
1065+ if (!f) {
1066+ sqlite3_result_error_nomem(context);
1067+ return;
1068+ }
1069+ memset(f, 0, sizeof(*f));
1070+
1071+ f->path = path;
1072+ f->pathLength = pathLength;
1073+ sqlite3_result_pointer(context, f, SQLITE_VEC_NPY_FILE_NAME, sqlite3_free);
1074+}
1075+#endif
1076+
1077+#pragma region scalar functions
1078+static void vec_f32(sqlite3_context *context, int argc, sqlite3_value **argv) {
1079+ assert(argc == 1);
1080+ int rc;
1081+ f32 *vector = NULL;
1082+ size_t dimensions;
1083+ fvec_cleanup cleanup;
1084+ char *errmsg;
1085+ rc = fvec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg);
1086+ if (rc != SQLITE_OK) {
1087+ sqlite3_result_error(context, errmsg, -1);
1088+ sqlite3_free(errmsg);
1089+ return;
1090+ }
1091+ sqlite3_result_blob(context, vector, dimensions * sizeof(f32),
1092+ (void (*)(void *))cleanup);
1093+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1094+}
1095+
1096+static void vec_bit(sqlite3_context *context, int argc, sqlite3_value **argv) {
1097+ assert(argc == 1);
1098+ int rc;
1099+ u8 *vector;
1100+ size_t dimensions;
1101+ vector_cleanup cleanup;
1102+ char *errmsg;
1103+ rc = bitvec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg);
1104+ if (rc != SQLITE_OK) {
1105+ sqlite3_result_error(context, errmsg, -1);
1106+ sqlite3_free(errmsg);
1107+ return;
1108+ }
1109+ sqlite3_result_blob(context, vector, dimensions / CHAR_BIT, SQLITE_TRANSIENT);
1110+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT);
1111+ cleanup(vector);
1112+}
1113+static void vec_int8(sqlite3_context *context, int argc, sqlite3_value **argv) {
1114+ assert(argc == 1);
1115+ int rc;
1116+ i8 *vector;
1117+ size_t dimensions;
1118+ vector_cleanup cleanup;
1119+ char *errmsg;
1120+ rc = int8_vec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg);
1121+ if (rc != SQLITE_OK) {
1122+ sqlite3_result_error(context, errmsg, -1);
1123+ sqlite3_free(errmsg);
1124+ return;
1125+ }
1126+ sqlite3_result_blob(context, vector, dimensions, SQLITE_TRANSIENT);
1127+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1128+ cleanup(vector);
1129+}
1130+
1131+static void vec_length(sqlite3_context *context, int argc,
1132+ sqlite3_value **argv) {
1133+ assert(argc == 1);
1134+ int rc;
1135+ void *vector;
1136+ size_t dimensions;
1137+ vector_cleanup cleanup;
1138+ char *errmsg;
1139+ enum VectorElementType elementType;
1140+ rc = vector_from_value(argv[0], &vector, &dimensions, &elementType, &cleanup,
1141+ &errmsg);
1142+ if (rc != SQLITE_OK) {
1143+ sqlite3_result_error(context, errmsg, -1);
1144+ sqlite3_free(errmsg);
1145+ return;
1146+ }
1147+ sqlite3_result_int64(context, dimensions);
1148+ cleanup(vector);
1149+}
1150+
1151+static void vec_distance_cosine(sqlite3_context *context, int argc,
1152+ sqlite3_value **argv) {
1153+ assert(argc == 2);
1154+ int rc;
1155+ void *a = NULL, *b = NULL;
1156+ size_t dimensions;
1157+ vector_cleanup aCleanup, bCleanup;
1158+ char *error;
1159+ enum VectorElementType elementType;
1160+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1161+ &aCleanup, &bCleanup, &error);
1162+ if (rc != SQLITE_OK) {
1163+ sqlite3_result_error(context, error, -1);
1164+ sqlite3_free(error);
1165+ return;
1166+ }
1167+
1168+ switch (elementType) {
1169+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1170+ sqlite3_result_error(
1171+ context, "Cannot calculate cosine distance between two bitvectors.",
1172+ -1);
1173+ goto finish;
1174+ }
1175+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1176+ f32 result = distance_cosine_float(a, b, &dimensions);
1177+ sqlite3_result_double(context, result);
1178+ goto finish;
1179+ }
1180+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1181+ f32 result = distance_cosine_int8(a, b, &dimensions);
1182+ sqlite3_result_double(context, result);
1183+ goto finish;
1184+ }
1185+ }
1186+
1187+finish:
1188+ aCleanup(a);
1189+ bCleanup(b);
1190+ return;
1191+}
1192+
1193+static void vec_distance_l2(sqlite3_context *context, int argc,
1194+ sqlite3_value **argv) {
1195+ assert(argc == 2);
1196+ int rc;
1197+ void *a = NULL, *b = NULL;
1198+ size_t dimensions;
1199+ vector_cleanup aCleanup, bCleanup;
1200+ char *error;
1201+ enum VectorElementType elementType;
1202+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1203+ &aCleanup, &bCleanup, &error);
1204+ if (rc != SQLITE_OK) {
1205+ sqlite3_result_error(context, error, -1);
1206+ sqlite3_free(error);
1207+ return;
1208+ }
1209+
1210+ switch (elementType) {
1211+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1212+ sqlite3_result_error(
1213+ context, "Cannot calculate L2 distance between two bitvectors.", -1);
1214+ goto finish;
1215+ }
1216+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1217+ f32 result = distance_l2_sqr_float(a, b, &dimensions);
1218+ sqlite3_result_double(context, result);
1219+ goto finish;
1220+ }
1221+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1222+ f32 result = distance_l2_sqr_int8(a, b, &dimensions);
1223+ sqlite3_result_double(context, result);
1224+ goto finish;
1225+ }
1226+ }
1227+
1228+finish:
1229+ aCleanup(a);
1230+ bCleanup(b);
1231+ return;
1232+}
1233+
1234+static void vec_distance_l1(sqlite3_context *context, int argc,
1235+ sqlite3_value **argv) {
1236+ assert(argc == 2);
1237+ int rc;
1238+ void *a, *b;
1239+ size_t dimensions;
1240+ vector_cleanup aCleanup, bCleanup;
1241+ char *error;
1242+ enum VectorElementType elementType;
1243+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1244+ &aCleanup, &bCleanup, &error);
1245+ if (rc != SQLITE_OK) {
1246+ sqlite3_result_error(context, error, -1);
1247+ sqlite3_free(error);
1248+ return;
1249+ }
1250+
1251+ switch (elementType) {
1252+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1253+ sqlite3_result_error(
1254+ context, "Cannot calculate L1 distance between two bitvectors.", -1);
1255+ goto finish;
1256+ }
1257+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1258+ double result = distance_l1_f32(a, b, &dimensions);
1259+ sqlite3_result_double(context, result);
1260+ goto finish;
1261+ }
1262+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1263+ i64 result = distance_l1_int8(a, b, &dimensions);
1264+ sqlite3_result_int(context, result);
1265+ goto finish;
1266+ }
1267+ }
1268+
1269+finish:
1270+ aCleanup(a);
1271+ bCleanup(b);
1272+ return;
1273+}
1274+
1275+static void vec_distance_hamming(sqlite3_context *context, int argc,
1276+ sqlite3_value **argv) {
1277+ assert(argc == 2);
1278+ int rc;
1279+ void *a = NULL, *b = NULL;
1280+ size_t dimensions;
1281+ vector_cleanup aCleanup, bCleanup;
1282+ char *error;
1283+ enum VectorElementType elementType;
1284+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1285+ &aCleanup, &bCleanup, &error);
1286+ if (rc != SQLITE_OK) {
1287+ sqlite3_result_error(context, error, -1);
1288+ sqlite3_free(error);
1289+ return;
1290+ }
1291+
1292+ switch (elementType) {
1293+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1294+ sqlite3_result_double(context, distance_hamming(a, b, &dimensions));
1295+ goto finish;
1296+ }
1297+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1298+ sqlite3_result_error(
1299+ context,
1300+ "Cannot calculate hamming distance between two float32 vectors.", -1);
1301+ goto finish;
1302+ }
1303+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1304+ sqlite3_result_error(
1305+ context, "Cannot calculate hamming distance between two int8 vectors.",
1306+ -1);
1307+ goto finish;
1308+ }
1309+ }
1310+
1311+finish:
1312+ aCleanup(a);
1313+ bCleanup(b);
1314+ return;
1315+}
1316+
1317+char *vec_type_name(enum VectorElementType elementType) {
1318+ switch (elementType) {
1319+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32:
1320+ return "float32";
1321+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
1322+ return "int8";
1323+ case SQLITE_VEC_ELEMENT_TYPE_BIT:
1324+ return "bit";
1325+ }
1326+ return "";
1327+}
1328+
1329+static void vec_type(sqlite3_context *context, int argc, sqlite3_value **argv) {
1330+ assert(argc == 1);
1331+ void *vector;
1332+ size_t dimensions;
1333+ vector_cleanup cleanup;
1334+ char *pzError;
1335+ enum VectorElementType elementType;
1336+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1337+ &cleanup, &pzError);
1338+ if (rc != SQLITE_OK) {
1339+ sqlite3_result_error(context, pzError, -1);
1340+ sqlite3_free(pzError);
1341+ return;
1342+ }
1343+ sqlite3_result_text(context, vec_type_name(elementType), -1, SQLITE_STATIC);
1344+ cleanup(vector);
1345+}
1346+static void vec_quantize_binary(sqlite3_context *context, int argc,
1347+ sqlite3_value **argv) {
1348+ assert(argc == 1);
1349+ void *vector;
1350+ size_t dimensions;
1351+ vector_cleanup vectorCleanup;
1352+ char *pzError;
1353+ enum VectorElementType elementType;
1354+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1355+ &vectorCleanup, &pzError);
1356+ if (rc != SQLITE_OK) {
1357+ sqlite3_result_error(context, pzError, -1);
1358+ sqlite3_free(pzError);
1359+ return;
1360+ }
1361+
1362+ if (dimensions <= 0) {
1363+ sqlite3_result_error(context, "Zero length vectors are not supported.", -1);
1364+ goto cleanup;
1365+ return;
1366+ }
1367+ if ((dimensions % CHAR_BIT) != 0) {
1368+ sqlite3_result_error(
1369+ context,
1370+ "Binary quantization requires vectors with a length divisible by 8",
1371+ -1);
1372+ goto cleanup;
1373+ return;
1374+ }
1375+
1376+ int sz = dimensions / CHAR_BIT;
1377+ u8 *out = sqlite3_malloc(sz);
1378+ if (!out) {
1379+ sqlite3_result_error_code(context, SQLITE_NOMEM);
1380+ goto cleanup;
1381+ return;
1382+ }
1383+ memset(out, 0, sz);
1384+
1385+ switch (elementType) {
1386+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1387+
1388+ for (size_t i = 0; i < dimensions; i++) {
1389+ int res = ((f32 *)vector)[i] > 0.0;
1390+ out[i / 8] |= (res << (i % 8));
1391+ }
1392+ break;
1393+ }
1394+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1395+ for (size_t i = 0; i < dimensions; i++) {
1396+ int res = ((i8 *)vector)[i] > 0;
1397+ out[i / 8] |= (res << (i % 8));
1398+ }
1399+ break;
1400+ }
1401+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1402+ sqlite3_result_error(context,
1403+ "Can only binary quantize float or int8 vectors", -1);
1404+ sqlite3_free(out);
1405+ return;
1406+ }
1407+ }
1408+ sqlite3_result_blob(context, out, sz, sqlite3_free);
1409+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT);
1410+
1411+cleanup:
1412+ vectorCleanup(vector);
1413+}
1414+
1415+static void vec_quantize_int8(sqlite3_context *context, int argc,
1416+ sqlite3_value **argv) {
1417+ assert(argc == 2);
1418+ f32 *srcVector;
1419+ size_t dimensions;
1420+ fvec_cleanup srcCleanup;
1421+ char *err;
1422+ i8 *out = NULL;
1423+ int rc = fvec_from_value(argv[0], &srcVector, &dimensions, &srcCleanup, &err);
1424+ if (rc != SQLITE_OK) {
1425+ sqlite3_result_error(context, err, -1);
1426+ sqlite3_free(err);
1427+ return;
1428+ }
1429+
1430+ int sz = dimensions * sizeof(i8);
1431+ out = sqlite3_malloc(sz);
1432+ if (!out) {
1433+ sqlite3_result_error_nomem(context);
1434+ goto cleanup;
1435+ }
1436+ memset(out, 0, sz);
1437+
1438+ if ((sqlite3_value_type(argv[1]) != SQLITE_TEXT) ||
1439+ (sqlite3_value_bytes(argv[1]) != strlen("unit")) ||
1440+ (sqlite3_stricmp((const char *)sqlite3_value_text(argv[1]), "unit") !=
1441+ 0)) {
1442+ sqlite3_result_error(
1443+ context, "2nd argument to vec_quantize_int8() must be 'unit'.", -1);
1444+ sqlite3_free(out);
1445+ goto cleanup;
1446+ }
1447+ f32 step = (1.0 - (-1.0)) / 255;
1448+ for (size_t i = 0; i < dimensions; i++) {
1449+ out[i] = ((srcVector[i] - (-1.0)) / step) - 128;
1450+ }
1451+
1452+ sqlite3_result_blob(context, out, dimensions * sizeof(i8), sqlite3_free);
1453+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1454+
1455+cleanup:
1456+ srcCleanup(srcVector);
1457+}
1458+
1459+static void vec_add(sqlite3_context *context, int argc, sqlite3_value **argv) {
1460+ assert(argc == 2);
1461+ int rc;
1462+ void *a = NULL, *b = NULL;
1463+ size_t dimensions;
1464+ vector_cleanup aCleanup, bCleanup;
1465+ char *error;
1466+ enum VectorElementType elementType;
1467+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1468+ &aCleanup, &bCleanup, &error);
1469+ if (rc != SQLITE_OK) {
1470+ sqlite3_result_error(context, error, -1);
1471+ sqlite3_free(error);
1472+ return;
1473+ }
1474+
1475+ switch (elementType) {
1476+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1477+ sqlite3_result_error(context, "Cannot add two bitvectors together.", -1);
1478+ goto finish;
1479+ }
1480+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1481+ size_t outSize = dimensions * sizeof(f32);
1482+ f32 *out = sqlite3_malloc(outSize);
1483+ if (!out) {
1484+ sqlite3_result_error_nomem(context);
1485+ goto finish;
1486+ }
1487+ memset(out, 0, outSize);
1488+ for (size_t i = 0; i < dimensions; i++) {
1489+ out[i] = ((f32 *)a)[i] + ((f32 *)b)[i];
1490+ }
1491+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1492+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1493+ goto finish;
1494+ }
1495+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1496+ size_t outSize = dimensions * sizeof(i8);
1497+ i8 *out = sqlite3_malloc(outSize);
1498+ if (!out) {
1499+ sqlite3_result_error_nomem(context);
1500+ goto finish;
1501+ }
1502+ memset(out, 0, outSize);
1503+ for (size_t i = 0; i < dimensions; i++) {
1504+ out[i] = ((i8 *)a)[i] + ((i8 *)b)[i];
1505+ }
1506+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1507+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1508+ goto finish;
1509+ }
1510+ }
1511+finish:
1512+ aCleanup(a);
1513+ bCleanup(b);
1514+ return;
1515+}
1516+static void vec_sub(sqlite3_context *context, int argc, sqlite3_value **argv) {
1517+ assert(argc == 2);
1518+ int rc;
1519+ void *a = NULL, *b = NULL;
1520+ size_t dimensions;
1521+ vector_cleanup aCleanup, bCleanup;
1522+ char *error;
1523+ enum VectorElementType elementType;
1524+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1525+ &aCleanup, &bCleanup, &error);
1526+ if (rc != SQLITE_OK) {
1527+ sqlite3_result_error(context, error, -1);
1528+ sqlite3_free(error);
1529+ return;
1530+ }
1531+
1532+ switch (elementType) {
1533+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1534+ sqlite3_result_error(context, "Cannot subtract two bitvectors together.",
1535+ -1);
1536+ goto finish;
1537+ }
1538+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1539+ size_t outSize = dimensions * sizeof(f32);
1540+ f32 *out = sqlite3_malloc(outSize);
1541+ if (!out) {
1542+ sqlite3_result_error_nomem(context);
1543+ goto finish;
1544+ }
1545+ memset(out, 0, outSize);
1546+ for (size_t i = 0; i < dimensions; i++) {
1547+ out[i] = ((f32 *)a)[i] - ((f32 *)b)[i];
1548+ }
1549+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1550+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1551+ goto finish;
1552+ }
1553+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1554+ size_t outSize = dimensions * sizeof(i8);
1555+ i8 *out = sqlite3_malloc(outSize);
1556+ if (!out) {
1557+ sqlite3_result_error_nomem(context);
1558+ goto finish;
1559+ }
1560+ memset(out, 0, outSize);
1561+ for (size_t i = 0; i < dimensions; i++) {
1562+ out[i] = ((i8 *)a)[i] - ((i8 *)b)[i];
1563+ }
1564+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1565+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1566+ goto finish;
1567+ }
1568+ }
1569+finish:
1570+ aCleanup(a);
1571+ bCleanup(b);
1572+ return;
1573+}
1574+static void vec_slice(sqlite3_context *context, int argc,
1575+ sqlite3_value **argv) {
1576+ assert(argc == 3);
1577+
1578+ void *vector;
1579+ size_t dimensions;
1580+ vector_cleanup cleanup;
1581+ char *err;
1582+ enum VectorElementType elementType;
1583+
1584+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1585+ &cleanup, &err);
1586+ if (rc != SQLITE_OK) {
1587+ sqlite3_result_error(context, err, -1);
1588+ sqlite3_free(err);
1589+ return;
1590+ }
1591+
1592+ int start = sqlite3_value_int(argv[1]);
1593+ int end = sqlite3_value_int(argv[2]);
1594+
1595+ if (start < 0) {
1596+ sqlite3_result_error(context,
1597+ "slice 'start' index must be a postive number.", -1);
1598+ goto done;
1599+ }
1600+ if (end < 0) {
1601+ sqlite3_result_error(context, "slice 'end' index must be a postive number.",
1602+ -1);
1603+ goto done;
1604+ }
1605+ if (((size_t)start) > dimensions) {
1606+ sqlite3_result_error(
1607+ context, "slice 'start' index is greater than the number of dimensions",
1608+ -1);
1609+ goto done;
1610+ }
1611+ if (((size_t)end) > dimensions) {
1612+ sqlite3_result_error(
1613+ context, "slice 'end' index is greater than the number of dimensions",
1614+ -1);
1615+ goto done;
1616+ }
1617+ if (start > end) {
1618+ sqlite3_result_error(context,
1619+ "slice 'start' index is greater than 'end' index", -1);
1620+ goto done;
1621+ }
1622+ if (start == end) {
1623+ sqlite3_result_error(context,
1624+ "slice 'start' index is equal to the 'end' index, "
1625+ "vectors must have non-zero length",
1626+ -1);
1627+ goto done;
1628+ }
1629+ size_t n = end - start;
1630+
1631+ switch (elementType) {
1632+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1633+ int outSize = n * sizeof(f32);
1634+ f32 *out = sqlite3_malloc(outSize);
1635+ if (!out) {
1636+ sqlite3_result_error_nomem(context);
1637+ goto done;
1638+ }
1639+ memset(out, 0, outSize);
1640+ for (size_t i = 0; i < n; i++) {
1641+ out[i] = ((f32 *)vector)[start + i];
1642+ }
1643+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1644+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1645+ goto done;
1646+ }
1647+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1648+ int outSize = n * sizeof(i8);
1649+ i8 *out = sqlite3_malloc(outSize);
1650+ if (!out) {
1651+ sqlite3_result_error_nomem(context);
1652+ return;
1653+ }
1654+ memset(out, 0, outSize);
1655+ for (size_t i = 0; i < n; i++) {
1656+ out[i] = ((i8 *)vector)[start + i];
1657+ }
1658+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1659+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1660+ goto done;
1661+ }
1662+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1663+ if ((start % CHAR_BIT) != 0) {
1664+ sqlite3_result_error(context, "start index must be divisible by 8.", -1);
1665+ goto done;
1666+ }
1667+ if ((end % CHAR_BIT) != 0) {
1668+ sqlite3_result_error(context, "end index must be divisible by 8.", -1);
1669+ goto done;
1670+ }
1671+ int outSize = n / CHAR_BIT;
1672+ u8 *out = sqlite3_malloc(outSize);
1673+ if (!out) {
1674+ sqlite3_result_error_nomem(context);
1675+ return;
1676+ }
1677+ memset(out, 0, outSize);
1678+ for (size_t i = 0; i < n / CHAR_BIT; i++) {
1679+ out[i] = ((u8 *)vector)[(start / CHAR_BIT) + i];
1680+ }
1681+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1682+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT);
1683+ goto done;
1684+ }
1685+ }
1686+done:
1687+ cleanup(vector);
1688+}
1689+
1690+static void vec_to_json(sqlite3_context *context, int argc,
1691+ sqlite3_value **argv) {
1692+ assert(argc == 1);
1693+ void *vector;
1694+ size_t dimensions;
1695+ vector_cleanup cleanup;
1696+ char *err;
1697+ enum VectorElementType elementType;
1698+
1699+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1700+ &cleanup, &err);
1701+ if (rc != SQLITE_OK) {
1702+ sqlite3_result_error(context, err, -1);
1703+ sqlite3_free(err);
1704+ return;
1705+ }
1706+
1707+ sqlite3_str *str = sqlite3_str_new(sqlite3_context_db_handle(context));
1708+ sqlite3_str_appendall(str, "[");
1709+ for (size_t i = 0; i < dimensions; i++) {
1710+ if (i != 0) {
1711+ sqlite3_str_appendall(str, ",");
1712+ }
1713+ if (elementType == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
1714+ f32 value = ((f32 *)vector)[i];
1715+ if (isnan(value)) {
1716+ sqlite3_str_appendall(str, "null");
1717+ } else {
1718+ sqlite3_str_appendf(str, "%f", value);
1719+ }
1720+
1721+ } else if (elementType == SQLITE_VEC_ELEMENT_TYPE_INT8) {
1722+ sqlite3_str_appendf(str, "%d", ((i8 *)vector)[i]);
1723+ } else if (elementType == SQLITE_VEC_ELEMENT_TYPE_BIT) {
1724+ u8 b = (((u8 *)vector)[i / 8] >> (i % CHAR_BIT)) & 1;
1725+ sqlite3_str_appendf(str, "%d", b);
1726+ }
1727+ }
1728+ sqlite3_str_appendall(str, "]");
1729+ int len = sqlite3_str_length(str);
1730+ char *s = sqlite3_str_finish(str);
1731+ if (s) {
1732+ sqlite3_result_text(context, s, len, sqlite3_free);
1733+ sqlite3_result_subtype(context, JSON_SUBTYPE);
1734+ } else {
1735+ sqlite3_result_error_nomem(context);
1736+ }
1737+ cleanup(vector);
1738+}
1739+
1740+static void vec_normalize(sqlite3_context *context, int argc,
1741+ sqlite3_value **argv) {
1742+ assert(argc == 1);
1743+ void *vector;
1744+ size_t dimensions;
1745+ vector_cleanup cleanup;
1746+ char *err;
1747+ enum VectorElementType elementType;
1748+
1749+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1750+ &cleanup, &err);
1751+ if (rc != SQLITE_OK) {
1752+ sqlite3_result_error(context, err, -1);
1753+ sqlite3_free(err);
1754+ return;
1755+ }
1756+
1757+ if (elementType != SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
1758+ sqlite3_result_error(
1759+ context, "only float32 vectors are supported when normalizing", -1);
1760+ cleanup(vector);
1761+ return;
1762+ }
1763+
1764+ int outSize = dimensions * sizeof(f32);
1765+ f32 *out = sqlite3_malloc(outSize);
1766+ if (!out) {
1767+ cleanup(vector);
1768+ sqlite3_result_error_code(context, SQLITE_NOMEM);
1769+ return;
1770+ }
1771+ memset(out, 0, outSize);
1772+
1773+ f32 *v = (f32 *)vector;
1774+
1775+ f32 norm = 0;
1776+ for (size_t i = 0; i < dimensions; i++) {
1777+ norm += v[i] * v[i];
1778+ }
1779+ norm = sqrt(norm);
1780+ for (size_t i = 0; i < dimensions; i++) {
1781+ out[i] = v[i] / norm;
1782+ }
1783+
1784+ sqlite3_result_blob(context, out, dimensions * sizeof(f32), sqlite3_free);
1785+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1786+ cleanup(vector);
1787+}
1788+
1789+static void _static_text_func(sqlite3_context *context, int argc,
1790+ sqlite3_value **argv) {
1791+ UNUSED_PARAMETER(argc);
1792+ UNUSED_PARAMETER(argv);
1793+ sqlite3_result_text(context, sqlite3_user_data(context), -1, SQLITE_STATIC);
1794+}
1795+
1796+#pragma endregion
1797+
1798+enum Vec0TokenType {
1799+ TOKEN_TYPE_IDENTIFIER,
1800+ TOKEN_TYPE_DIGIT,
1801+ TOKEN_TYPE_LBRACKET,
1802+ TOKEN_TYPE_RBRACKET,
1803+ TOKEN_TYPE_PLUS,
1804+ TOKEN_TYPE_EQ,
1805+};
1806+struct Vec0Token {
1807+ enum Vec0TokenType token_type;
1808+ char *start;
1809+ char *end;
1810+};
1811+
1812+int is_alpha(char x) {
1813+ return (x >= 'a' && x <= 'z') || (x >= 'A' && x <= 'Z');
1814+}
1815+int is_digit(char x) { return (x >= '0' && x <= '9'); }
1816+int is_whitespace(char x) {
1817+ return x == ' ' || x == '\t' || x == '\n' || x == '\r';
1818+}
1819+
1820+#define VEC0_TOKEN_RESULT_EOF 1
1821+#define VEC0_TOKEN_RESULT_SOME 2
1822+#define VEC0_TOKEN_RESULT_ERROR 3
1823+
1824+int vec0_token_next(char *start, char *end, struct Vec0Token *out) {
1825+ char *ptr = start;
1826+ while (ptr < end) {
1827+ char curr = *ptr;
1828+ if (is_whitespace(curr)) {
1829+ ptr++;
1830+ continue;
1831+ } else if (curr == '+') {
1832+ ptr++;
1833+ out->start = ptr;
1834+ out->end = ptr;
1835+ out->token_type = TOKEN_TYPE_PLUS;
1836+ return VEC0_TOKEN_RESULT_SOME;
1837+ } else if (curr == '[') {
1838+ ptr++;
1839+ out->start = ptr;
1840+ out->end = ptr;
1841+ out->token_type = TOKEN_TYPE_LBRACKET;
1842+ return VEC0_TOKEN_RESULT_SOME;
1843+ } else if (curr == ']') {
1844+ ptr++;
1845+ out->start = ptr;
1846+ out->end = ptr;
1847+ out->token_type = TOKEN_TYPE_RBRACKET;
1848+ return VEC0_TOKEN_RESULT_SOME;
1849+ } else if (curr == '=') {
1850+ ptr++;
1851+ out->start = ptr;
1852+ out->end = ptr;
1853+ out->token_type = TOKEN_TYPE_EQ;
1854+ return VEC0_TOKEN_RESULT_SOME;
1855+ } else if (is_alpha(curr)) {
1856+ char *start = ptr;
1857+ while (ptr < end && (is_alpha(*ptr) || is_digit(*ptr) || *ptr == '_')) {
1858+ ptr++;
1859+ }
1860+ out->start = start;
1861+ out->end = ptr;
1862+ out->token_type = TOKEN_TYPE_IDENTIFIER;
1863+ return VEC0_TOKEN_RESULT_SOME;
1864+ } else if (is_digit(curr)) {
1865+ char *start = ptr;
1866+ while (ptr < end && (is_digit(*ptr))) {
1867+ ptr++;
1868+ }
1869+ out->start = start;
1870+ out->end = ptr;
1871+ out->token_type = TOKEN_TYPE_DIGIT;
1872+ return VEC0_TOKEN_RESULT_SOME;
1873+ } else {
1874+ return VEC0_TOKEN_RESULT_ERROR;
1875+ }
1876+ }
1877+ return VEC0_TOKEN_RESULT_EOF;
1878+}
1879+
1880+struct Vec0Scanner {
1881+ char *start;
1882+ char *end;
1883+ char *ptr;
1884+};
1885+
1886+void vec0_scanner_init(struct Vec0Scanner *scanner, const char *source,
1887+ int source_length) {
1888+ scanner->start = (char *)source;
1889+ scanner->end = (char *)source + source_length;
1890+ scanner->ptr = (char *)source;
1891+}
1892+int vec0_scanner_next(struct Vec0Scanner *scanner, struct Vec0Token *out) {
1893+ int rc = vec0_token_next(scanner->start, scanner->end, out);
1894+ if (rc == VEC0_TOKEN_RESULT_SOME) {
1895+ scanner->start = out->end;
1896+ }
1897+ return rc;
1898+}
1899+
1900+int vec0_parse_table_option(const char *source, int source_length,
1901+ char **out_key, int *out_key_length,
1902+ char **out_value, int *out_value_length) {
1903+ int rc;
1904+ struct Vec0Scanner scanner;
1905+ struct Vec0Token token;
1906+ char *key;
1907+ char *value;
1908+ int keyLength, valueLength;
1909+
1910+ vec0_scanner_init(&scanner, source, source_length);
1911+
1912+ rc = vec0_scanner_next(&scanner, &token);
1913+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1914+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
1915+ return SQLITE_EMPTY;
1916+ }
1917+ key = token.start;
1918+ keyLength = token.end - token.start;
1919+
1920+ rc = vec0_scanner_next(&scanner, &token);
1921+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_EQ) {
1922+ return SQLITE_EMPTY;
1923+ }
1924+
1925+ rc = vec0_scanner_next(&scanner, &token);
1926+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1927+ !((token.token_type == TOKEN_TYPE_IDENTIFIER) ||
1928+ (token.token_type == TOKEN_TYPE_DIGIT))) {
1929+ return SQLITE_ERROR;
1930+ }
1931+ value = token.start;
1932+ valueLength = token.end - token.start;
1933+
1934+ rc = vec0_scanner_next(&scanner, &token);
1935+ if (rc == VEC0_TOKEN_RESULT_EOF) {
1936+ *out_key = key;
1937+ *out_key_length = keyLength;
1938+ *out_value = value;
1939+ *out_value_length = valueLength;
1940+ return SQLITE_OK;
1941+ }
1942+ return SQLITE_ERROR;
1943+}
1944+/**
1945+ * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if
1946+ * it's a PARTITION KEY definition.
1947+ *
1948+ * @param source: argv[i] source string
1949+ * @param source_length: length of the source string
1950+ * @param out_column_name: If it is a partition key, the output column name. Same lifetime
1951+ * as source, points to specific char *
1952+ * @param out_column_name_length: Length of out_column_name in bytes
1953+ * @param out_column_type: SQLITE_TEXT or SQLITE_INTEGER.
1954+ * @return int: SQLITE_EMPTY if not a PK, SQLITE_OK if it is.
1955+ */
1956+int vec0_parse_partition_key_definition(const char *source, int source_length,
1957+ char **out_column_name,
1958+ int *out_column_name_length,
1959+ int *out_column_type) {
1960+ struct Vec0Scanner scanner;
1961+ struct Vec0Token token;
1962+ char *column_name;
1963+ int column_name_length;
1964+ int column_type;
1965+ vec0_scanner_init(&scanner, source, source_length);
1966+
1967+ // Check first token is identifier, will be the column name
1968+ int rc = vec0_scanner_next(&scanner, &token);
1969+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1970+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
1971+ return SQLITE_EMPTY;
1972+ }
1973+
1974+ column_name = token.start;
1975+ column_name_length = token.end - token.start;
1976+
1977+ // Check the next token matches "text" or "integer", as column type
1978+ rc = vec0_scanner_next(&scanner, &token);
1979+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1980+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
1981+ return SQLITE_EMPTY;
1982+ }
1983+ if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) {
1984+ column_type = SQLITE_TEXT;
1985+ } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) ==
1986+ 0 ||
1987+ sqlite3_strnicmp(token.start, "integer",
1988+ token.end - token.start) == 0) {
1989+ column_type = SQLITE_INTEGER;
1990+ } else {
1991+ return SQLITE_EMPTY;
1992+ }
1993+
1994+ // Check the next token is identifier and matches "partition"
1995+ rc = vec0_scanner_next(&scanner, &token);
1996+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1997+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
1998+ return SQLITE_EMPTY;
1999+ }
2000+ if (sqlite3_strnicmp(token.start, "partition", token.end - token.start) != 0) {
2001+ return SQLITE_EMPTY;
2002+ }
2003+
2004+ // Check the next token is identifier and matches "key"
2005+ rc = vec0_scanner_next(&scanner, &token);
2006+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2007+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2008+ return SQLITE_EMPTY;
2009+ }
2010+ if (sqlite3_strnicmp(token.start, "key", token.end - token.start) != 0) {
2011+ return SQLITE_EMPTY;
2012+ }
2013+
2014+ *out_column_name = column_name;
2015+ *out_column_name_length = column_name_length;
2016+ *out_column_type = column_type;
2017+
2018+ return SQLITE_OK;
2019+}
2020+
2021+/**
2022+ * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if
2023+ * it's an auxiliar column definition, ie `+[name] [type]` like `+contents text`
2024+ *
2025+ * @param source: argv[i] source string
2026+ * @param source_length: length of the source string
2027+ * @param out_column_name: If it is a partition key, the output column name. Same lifetime
2028+ * as source, points to specific char *
2029+ * @param out_column_name_length: Length of out_column_name in bytes
2030+ * @param out_column_type: SQLITE_TEXT, SQLITE_INTEGER, SQLITE_FLOAT, or SQLITE_BLOB.
2031+ * @return int: SQLITE_EMPTY if not an aux column, SQLITE_OK if it is.
2032+ */
2033+int vec0_parse_auxiliary_column_definition(const char *source, int source_length,
2034+ char **out_column_name,
2035+ int *out_column_name_length,
2036+ int *out_column_type) {
2037+ struct Vec0Scanner scanner;
2038+ struct Vec0Token token;
2039+ char *column_name;
2040+ int column_name_length;
2041+ int column_type;
2042+ vec0_scanner_init(&scanner, source, source_length);
2043+
2044+ // Check first token is '+', which denotes aux columns
2045+ int rc = vec0_scanner_next(&scanner, &token);
2046+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2047+ token.token_type != TOKEN_TYPE_PLUS) {
2048+ return SQLITE_EMPTY;
2049+ }
2050+
2051+ rc = vec0_scanner_next(&scanner, &token);
2052+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2053+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2054+ return SQLITE_EMPTY;
2055+ }
2056+
2057+ column_name = token.start;
2058+ column_name_length = token.end - token.start;
2059+
2060+ // Check the next token matches "text" or "integer", as column type
2061+ rc = vec0_scanner_next(&scanner, &token);
2062+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2063+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2064+ return SQLITE_EMPTY;
2065+ }
2066+ if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) {
2067+ column_type = SQLITE_TEXT;
2068+ } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) ==
2069+ 0 ||
2070+ sqlite3_strnicmp(token.start, "integer",
2071+ token.end - token.start) == 0) {
2072+ column_type = SQLITE_INTEGER;
2073+ } else if (sqlite3_strnicmp(token.start, "float", token.end - token.start) ==
2074+ 0 ||
2075+ sqlite3_strnicmp(token.start, "double",
2076+ token.end - token.start) == 0) {
2077+ column_type = SQLITE_FLOAT;
2078+ } else if (sqlite3_strnicmp(token.start, "blob", token.end - token.start) ==0) {
2079+ column_type = SQLITE_BLOB;
2080+ } else {
2081+ return SQLITE_EMPTY;
2082+ }
2083+
2084+ *out_column_name = column_name;
2085+ *out_column_name_length = column_name_length;
2086+ *out_column_type = column_type;
2087+
2088+ return SQLITE_OK;
2089+}
2090+
2091+typedef enum {
2092+ VEC0_METADATA_COLUMN_KIND_BOOLEAN,
2093+ VEC0_METADATA_COLUMN_KIND_INTEGER,
2094+ VEC0_METADATA_COLUMN_KIND_FLOAT,
2095+ VEC0_METADATA_COLUMN_KIND_TEXT,
2096+ // future: blob, date, datetime
2097+} vec0_metadata_column_kind;
2098+
2099+/**
2100+ * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if
2101+ * it's an metadata column definition, ie `[name] [type]` like `is_released boolean`
2102+ *
2103+ * @param source: argv[i] source string
2104+ * @param source_length: length of the source string
2105+ * @param out_column_name: If it is a metadata column, the output column name. Same lifetime
2106+ * as source, points to specific char *
2107+ * @param out_column_name_length: Length of out_column_name in bytes
2108+ * @param out_column_type: one of vec0_metadata_column_kind
2109+ * @return int: SQLITE_EMPTY if not an metadata column, SQLITE_OK if it is.
2110+ */
2111+int vec0_parse_metadata_column_definition(const char *source, int source_length,
2112+ char **out_column_name,
2113+ int *out_column_name_length,
2114+ vec0_metadata_column_kind *out_column_type) {
2115+ struct Vec0Scanner scanner;
2116+ struct Vec0Token token;
2117+ char *column_name;
2118+ int column_name_length;
2119+ vec0_metadata_column_kind column_type;
2120+ int rc;
2121+ vec0_scanner_init(&scanner, source, source_length);
2122+
2123+ rc = vec0_scanner_next(&scanner, &token);
2124+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2125+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2126+ return SQLITE_EMPTY;
2127+ }
2128+
2129+ column_name = token.start;
2130+ column_name_length = token.end - token.start;
2131+
2132+ // Check the next token matches a valid metadata type
2133+ rc = vec0_scanner_next(&scanner, &token);
2134+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2135+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2136+ return SQLITE_EMPTY;
2137+ }
2138+ char * t = token.start;
2139+ int n = token.end - token.start;
2140+ if (sqlite3_strnicmp(t, "boolean", n) == 0 || sqlite3_strnicmp(t, "bool", n) == 0) {
2141+ column_type = VEC0_METADATA_COLUMN_KIND_BOOLEAN;
2142+ }else if (sqlite3_strnicmp(t, "int64", n) == 0 || sqlite3_strnicmp(t, "integer64", n) == 0 || sqlite3_strnicmp(t, "integer", n) == 0 || sqlite3_strnicmp(t, "int", n) == 0) {
2143+ column_type = VEC0_METADATA_COLUMN_KIND_INTEGER;
2144+ }else if (sqlite3_strnicmp(t, "float", n) == 0 || sqlite3_strnicmp(t, "double", n) == 0 || sqlite3_strnicmp(t, "float64", n) == 0 || sqlite3_strnicmp(t, "f64", n) == 0) {
2145+ column_type = VEC0_METADATA_COLUMN_KIND_FLOAT;
2146+ } else if (sqlite3_strnicmp(t, "text", n) == 0) {
2147+ column_type = VEC0_METADATA_COLUMN_KIND_TEXT;
2148+ } else {
2149+ return SQLITE_EMPTY;
2150+ }
2151+
2152+ *out_column_name = column_name;
2153+ *out_column_name_length = column_name_length;
2154+ *out_column_type = column_type;
2155+
2156+ return SQLITE_OK;
2157+}
2158+
2159+/**
2160+ * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if
2161+ * it's a PRIMARY KEY definition.
2162+ *
2163+ * @param source: argv[i] source string
2164+ * @param source_length: length of the source string
2165+ * @param out_column_name: If it is a PK, the output column name. Same lifetime
2166+ * as source, points to specific char *
2167+ * @param out_column_name_length: Length of out_column_name in bytes
2168+ * @param out_column_type: SQLITE_TEXT or SQLITE_INTEGER.
2169+ * @return int: SQLITE_EMPTY if not a PK, SQLITE_OK if it is.
2170+ */
2171+int vec0_parse_primary_key_definition(const char *source, int source_length,
2172+ char **out_column_name,
2173+ int *out_column_name_length,
2174+ int *out_column_type) {
2175+ struct Vec0Scanner scanner;
2176+ struct Vec0Token token;
2177+ char *column_name;
2178+ int column_name_length;
2179+ int column_type;
2180+ vec0_scanner_init(&scanner, source, source_length);
2181+
2182+ // Check first token is identifier, will be the column name
2183+ int rc = vec0_scanner_next(&scanner, &token);
2184+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2185+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2186+ return SQLITE_EMPTY;
2187+ }
2188+
2189+ column_name = token.start;
2190+ column_name_length = token.end - token.start;
2191+
2192+ // Check the next token matches "text" or "integer", as column type
2193+ rc = vec0_scanner_next(&scanner, &token);
2194+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2195+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2196+ return SQLITE_EMPTY;
2197+ }
2198+ if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) {
2199+ column_type = SQLITE_TEXT;
2200+ } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) ==
2201+ 0 ||
2202+ sqlite3_strnicmp(token.start, "integer",
2203+ token.end - token.start) == 0) {
2204+ column_type = SQLITE_INTEGER;
2205+ } else {
2206+ return SQLITE_EMPTY;
2207+ }
2208+
2209+ // Check the next token is identifier and matches "primary"
2210+ rc = vec0_scanner_next(&scanner, &token);
2211+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2212+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2213+ return SQLITE_EMPTY;
2214+ }
2215+ if (sqlite3_strnicmp(token.start, "primary", token.end - token.start) != 0) {
2216+ return SQLITE_EMPTY;
2217+ }
2218+
2219+ // Check the next token is identifier and matches "key"
2220+ rc = vec0_scanner_next(&scanner, &token);
2221+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2222+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2223+ return SQLITE_EMPTY;
2224+ }
2225+ if (sqlite3_strnicmp(token.start, "key", token.end - token.start) != 0) {
2226+ return SQLITE_EMPTY;
2227+ }
2228+
2229+ *out_column_name = column_name;
2230+ *out_column_name_length = column_name_length;
2231+ *out_column_type = column_type;
2232+
2233+ return SQLITE_OK;
2234+}
2235+
2236+enum Vec0DistanceMetrics {
2237+ VEC0_DISTANCE_METRIC_L2 = 1,
2238+ VEC0_DISTANCE_METRIC_COSINE = 2,
2239+ VEC0_DISTANCE_METRIC_L1 = 3,
2240+};
2241+
2242+struct VectorColumnDefinition {
2243+ char *name;
2244+ int name_length;
2245+ size_t dimensions;
2246+ enum VectorElementType element_type;
2247+ enum Vec0DistanceMetrics distance_metric;
2248+};
2249+
2250+struct Vec0PartitionColumnDefinition {
2251+ int type;
2252+ char * name;
2253+ int name_length;
2254+};
2255+
2256+struct Vec0AuxiliaryColumnDefinition {
2257+ int type;
2258+ char * name;
2259+ int name_length;
2260+};
2261+struct Vec0MetadataColumnDefinition {
2262+ vec0_metadata_column_kind kind;
2263+ char * name;
2264+ int name_length;
2265+};
2266+
2267+size_t vector_byte_size(enum VectorElementType element_type,
2268+ size_t dimensions) {
2269+ switch (element_type) {
2270+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32:
2271+ return dimensions * sizeof(f32);
2272+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
2273+ return dimensions * sizeof(i8);
2274+ case SQLITE_VEC_ELEMENT_TYPE_BIT:
2275+ return dimensions / CHAR_BIT;
2276+ }
2277+ return 0;
2278+}
2279+
2280+size_t vector_column_byte_size(struct VectorColumnDefinition column) {
2281+ return vector_byte_size(column.element_type, column.dimensions);
2282+}
2283+
2284+/**
2285+ * @brief Parse an vec0 vtab argv[i] column definition and see if
2286+ * it's a vector column defintion, ex `contents_embedding float[768]`.
2287+ *
2288+ * @param source vec0 argv[i] item
2289+ * @param source_length length of source in bytes
2290+ * @param outColumn Output the parse vector column to this struct, if success
2291+ * @return int SQLITE_OK on success, SQLITE_EMPTY is it's not a vector column
2292+ * definition, SQLITE_ERROR on error.
2293+ */
2294+int vec0_parse_vector_column(const char *source, int source_length,
2295+ struct VectorColumnDefinition *outColumn) {
2296+ // parses a vector column definition like so:
2297+ // "abc float[123]", "abc_123 bit[1234]", eetc.
2298+ // https://github.com/asg017/sqlite-vec/issues/46
2299+ int rc;
2300+ struct Vec0Scanner scanner;
2301+ struct Vec0Token token;
2302+
2303+ char *name;
2304+ int nameLength;
2305+ enum VectorElementType elementType;
2306+ enum Vec0DistanceMetrics distanceMetric = VEC0_DISTANCE_METRIC_L2;
2307+ int dimensions;
2308+
2309+ vec0_scanner_init(&scanner, source, source_length);
2310+
2311+ // starts with an identifier
2312+ rc = vec0_scanner_next(&scanner, &token);
2313+
2314+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2315+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2316+ return SQLITE_EMPTY;
2317+ }
2318+
2319+ name = token.start;
2320+ nameLength = token.end - token.start;
2321+
2322+ // vector column type comes next: float, int, or bit
2323+ rc = vec0_scanner_next(&scanner, &token);
2324+
2325+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2326+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2327+ return SQLITE_EMPTY;
2328+ }
2329+ if (sqlite3_strnicmp(token.start, "float", 5) == 0 ||
2330+ sqlite3_strnicmp(token.start, "f32", 3) == 0) {
2331+ elementType = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
2332+ } else if (sqlite3_strnicmp(token.start, "int8", 4) == 0 ||
2333+ sqlite3_strnicmp(token.start, "i8", 2) == 0) {
2334+ elementType = SQLITE_VEC_ELEMENT_TYPE_INT8;
2335+ } else if (sqlite3_strnicmp(token.start, "bit", 3) == 0) {
2336+ elementType = SQLITE_VEC_ELEMENT_TYPE_BIT;
2337+ } else {
2338+ return SQLITE_EMPTY;
2339+ }
2340+
2341+ // left '[' bracket
2342+ rc = vec0_scanner_next(&scanner, &token);
2343+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_LBRACKET) {
2344+ return SQLITE_EMPTY;
2345+ }
2346+
2347+ // digit, for vector dimension length
2348+ rc = vec0_scanner_next(&scanner, &token);
2349+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_DIGIT) {
2350+ return SQLITE_ERROR;
2351+ }
2352+ dimensions = atoi(token.start);
2353+ if (dimensions <= 0) {
2354+ return SQLITE_ERROR;
2355+ }
2356+
2357+ // // right ']' bracket
2358+ rc = vec0_scanner_next(&scanner, &token);
2359+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_RBRACKET) {
2360+ return SQLITE_ERROR;
2361+ }
2362+
2363+ // any other tokens left should be column-level options , ex `key=value`
2364+ // ex `distance_metric=L2 distance_metric=cosine` should error
2365+ while (1) {
2366+ // should be EOF or identifier (option key)
2367+ rc = vec0_scanner_next(&scanner, &token);
2368+ if (rc == VEC0_TOKEN_RESULT_EOF) {
2369+ break;
2370+ }
2371+
2372+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2373+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2374+ return SQLITE_ERROR;
2375+ }
2376+
2377+ char *key = token.start;
2378+ int keyLength = token.end - token.start;
2379+
2380+ if (sqlite3_strnicmp(key, "distance_metric", keyLength) == 0) {
2381+
2382+ if (elementType == SQLITE_VEC_ELEMENT_TYPE_BIT) {
2383+ return SQLITE_ERROR;
2384+ }
2385+ // ensure equal sign after distance_metric
2386+ rc = vec0_scanner_next(&scanner, &token);
2387+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_EQ) {
2388+ return SQLITE_ERROR;
2389+ }
2390+
2391+ // distance_metric value, an identifier (L2, cosine, etc)
2392+ rc = vec0_scanner_next(&scanner, &token);
2393+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2394+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2395+ return SQLITE_ERROR;
2396+ }
2397+
2398+ char *value = token.start;
2399+ int valueLength = token.end - token.start;
2400+ if (sqlite3_strnicmp(value, "l2", valueLength) == 0) {
2401+ distanceMetric = VEC0_DISTANCE_METRIC_L2;
2402+ } else if (sqlite3_strnicmp(value, "l1", valueLength) == 0) {
2403+ distanceMetric = VEC0_DISTANCE_METRIC_L1;
2404+ } else if (sqlite3_strnicmp(value, "cosine", valueLength) == 0) {
2405+ distanceMetric = VEC0_DISTANCE_METRIC_COSINE;
2406+ } else {
2407+ return SQLITE_ERROR;
2408+ }
2409+ }
2410+ // unknown key
2411+ else {
2412+ return SQLITE_ERROR;
2413+ }
2414+ }
2415+
2416+ outColumn->name = sqlite3_mprintf("%.*s", nameLength, name);
2417+ if (!outColumn->name) {
2418+ return SQLITE_ERROR;
2419+ }
2420+ outColumn->name_length = nameLength;
2421+ outColumn->distance_metric = distanceMetric;
2422+ outColumn->element_type = elementType;
2423+ outColumn->dimensions = dimensions;
2424+ return SQLITE_OK;
2425+}
2426+
2427+#pragma region vec_each table function
2428+
2429+typedef struct vec_each_vtab vec_each_vtab;
2430+struct vec_each_vtab {
2431+ sqlite3_vtab base;
2432+};
2433+
2434+typedef struct vec_each_cursor vec_each_cursor;
2435+struct vec_each_cursor {
2436+ sqlite3_vtab_cursor base;
2437+ i64 iRowid;
2438+ enum VectorElementType vector_type;
2439+ void *vector;
2440+ size_t dimensions;
2441+ vector_cleanup cleanup;
2442+};
2443+
2444+static int vec_eachConnect(sqlite3 *db, void *pAux, int argc,
2445+ const char *const *argv, sqlite3_vtab **ppVtab,
2446+ char **pzErr) {
2447+ UNUSED_PARAMETER(pAux);
2448+ UNUSED_PARAMETER(argc);
2449+ UNUSED_PARAMETER(argv);
2450+ UNUSED_PARAMETER(pzErr);
2451+ vec_each_vtab *pNew;
2452+ int rc;
2453+
2454+ rc = sqlite3_declare_vtab(db, "CREATE TABLE x(value, vector hidden)");
2455+#define VEC_EACH_COLUMN_VALUE 0
2456+#define VEC_EACH_COLUMN_VECTOR 1
2457+ if (rc == SQLITE_OK) {
2458+ pNew = sqlite3_malloc(sizeof(*pNew));
2459+ *ppVtab = (sqlite3_vtab *)pNew;
2460+ if (pNew == 0)
2461+ return SQLITE_NOMEM;
2462+ memset(pNew, 0, sizeof(*pNew));
2463+ }
2464+ return rc;
2465+}
2466+
2467+static int vec_eachDisconnect(sqlite3_vtab *pVtab) {
2468+ vec_each_vtab *p = (vec_each_vtab *)pVtab;
2469+ sqlite3_free(p);
2470+ return SQLITE_OK;
2471+}
2472+
2473+static int vec_eachOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) {
2474+ UNUSED_PARAMETER(p);
2475+ vec_each_cursor *pCur;
2476+ pCur = sqlite3_malloc(sizeof(*pCur));
2477+ if (pCur == 0)
2478+ return SQLITE_NOMEM;
2479+ memset(pCur, 0, sizeof(*pCur));
2480+ *ppCursor = &pCur->base;
2481+ return SQLITE_OK;
2482+}
2483+
2484+static int vec_eachClose(sqlite3_vtab_cursor *cur) {
2485+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2486+ pCur->cleanup(pCur->vector);
2487+ sqlite3_free(pCur);
2488+ return SQLITE_OK;
2489+}
2490+
2491+static int vec_eachBestIndex(sqlite3_vtab *pVTab,
2492+ sqlite3_index_info *pIdxInfo) {
2493+ UNUSED_PARAMETER(pVTab);
2494+ int hasVector = 0;
2495+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
2496+ const struct sqlite3_index_constraint *pCons = &pIdxInfo->aConstraint[i];
2497+ // printf("i=%d iColumn=%d, op=%d, usable=%d\n", i, pCons->iColumn,
2498+ // pCons->op, pCons->usable);
2499+ switch (pCons->iColumn) {
2500+ case VEC_EACH_COLUMN_VECTOR: {
2501+ if (pCons->op == SQLITE_INDEX_CONSTRAINT_EQ && pCons->usable) {
2502+ hasVector = 1;
2503+ pIdxInfo->aConstraintUsage[i].argvIndex = 1;
2504+ pIdxInfo->aConstraintUsage[i].omit = 1;
2505+ }
2506+ break;
2507+ }
2508+ }
2509+ }
2510+ if (!hasVector) {
2511+ return SQLITE_CONSTRAINT;
2512+ }
2513+
2514+ pIdxInfo->estimatedCost = (double)100000;
2515+ pIdxInfo->estimatedRows = 100000;
2516+
2517+ return SQLITE_OK;
2518+}
2519+
2520+static int vec_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
2521+ const char *idxStr, int argc, sqlite3_value **argv) {
2522+ UNUSED_PARAMETER(idxNum);
2523+ UNUSED_PARAMETER(idxStr);
2524+ assert(argc == 1);
2525+ vec_each_cursor *pCur = (vec_each_cursor *)pVtabCursor;
2526+
2527+ if (pCur->vector) {
2528+ pCur->cleanup(pCur->vector);
2529+ pCur->vector = NULL;
2530+ }
2531+
2532+ char *pzErrMsg;
2533+ int rc = vector_from_value(argv[0], &pCur->vector, &pCur->dimensions,
2534+ &pCur->vector_type, &pCur->cleanup, &pzErrMsg);
2535+ if (rc != SQLITE_OK) {
2536+ return SQLITE_ERROR;
2537+ }
2538+ pCur->iRowid = 0;
2539+ return SQLITE_OK;
2540+}
2541+
2542+static int vec_eachRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) {
2543+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2544+ *pRowid = pCur->iRowid;
2545+ return SQLITE_OK;
2546+}
2547+
2548+static int vec_eachEof(sqlite3_vtab_cursor *cur) {
2549+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2550+ return pCur->iRowid >= (i64)pCur->dimensions;
2551+}
2552+
2553+static int vec_eachNext(sqlite3_vtab_cursor *cur) {
2554+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2555+ pCur->iRowid++;
2556+ return SQLITE_OK;
2557+}
2558+
2559+static int vec_eachColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context,
2560+ int i) {
2561+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2562+ switch (i) {
2563+ case VEC_EACH_COLUMN_VALUE:
2564+ switch (pCur->vector_type) {
2565+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
2566+ sqlite3_result_double(context, ((f32 *)pCur->vector)[pCur->iRowid]);
2567+ break;
2568+ }
2569+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
2570+ u8 x = ((u8 *)pCur->vector)[pCur->iRowid / CHAR_BIT];
2571+ sqlite3_result_int(context,
2572+ (x & (0b10000000 >> ((pCur->iRowid % CHAR_BIT)))) > 0);
2573+ break;
2574+ }
2575+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
2576+ sqlite3_result_int(context, ((i8 *)pCur->vector)[pCur->iRowid]);
2577+ break;
2578+ }
2579+ }
2580+
2581+ break;
2582+ }
2583+ return SQLITE_OK;
2584+}
2585+
2586+static sqlite3_module vec_eachModule = {
2587+ /* iVersion */ 0,
2588+ /* xCreate */ 0,
2589+ /* xConnect */ vec_eachConnect,
2590+ /* xBestIndex */ vec_eachBestIndex,
2591+ /* xDisconnect */ vec_eachDisconnect,
2592+ /* xDestroy */ 0,
2593+ /* xOpen */ vec_eachOpen,
2594+ /* xClose */ vec_eachClose,
2595+ /* xFilter */ vec_eachFilter,
2596+ /* xNext */ vec_eachNext,
2597+ /* xEof */ vec_eachEof,
2598+ /* xColumn */ vec_eachColumn,
2599+ /* xRowid */ vec_eachRowid,
2600+ /* xUpdate */ 0,
2601+ /* xBegin */ 0,
2602+ /* xSync */ 0,
2603+ /* xCommit */ 0,
2604+ /* xRollback */ 0,
2605+ /* xFindMethod */ 0,
2606+ /* xRename */ 0,
2607+ /* xSavepoint */ 0,
2608+ /* xRelease */ 0,
2609+ /* xRollbackTo */ 0,
2610+ /* xShadowName */ 0,
2611+#if SQLITE_VERSION_NUMBER >= 3044000
2612+ /* xIntegrity */ 0
2613+#endif
2614+};
2615+
2616+#pragma endregion
2617+
2618+#pragma region vec_npy_each table function
2619+
2620+enum NpyTokenType {
2621+ NPY_TOKEN_TYPE_IDENTIFIER,
2622+ NPY_TOKEN_TYPE_NUMBER,
2623+ NPY_TOKEN_TYPE_LPAREN,
2624+ NPY_TOKEN_TYPE_RPAREN,
2625+ NPY_TOKEN_TYPE_LBRACE,
2626+ NPY_TOKEN_TYPE_RBRACE,
2627+ NPY_TOKEN_TYPE_COLON,
2628+ NPY_TOKEN_TYPE_COMMA,
2629+ NPY_TOKEN_TYPE_STRING,
2630+ NPY_TOKEN_TYPE_FALSE,
2631+};
2632+
2633+struct NpyToken {
2634+ enum NpyTokenType token_type;
2635+ unsigned char *start;
2636+ unsigned char *end;
2637+};
2638+
2639+int npy_token_next(unsigned char *start, unsigned char *end,
2640+ struct NpyToken *out) {
2641+ unsigned char *ptr = start;
2642+ while (ptr < end) {
2643+ unsigned char curr = *ptr;
2644+ if (is_whitespace(curr)) {
2645+ ptr++;
2646+ continue;
2647+ } else if (curr == '(') {
2648+ out->start = ptr++;
2649+ out->end = ptr;
2650+ out->token_type = NPY_TOKEN_TYPE_LPAREN;
2651+ return VEC0_TOKEN_RESULT_SOME;
2652+ } else if (curr == ')') {
2653+ out->start = ptr++;
2654+ out->end = ptr;
2655+ out->token_type = NPY_TOKEN_TYPE_RPAREN;
2656+ return VEC0_TOKEN_RESULT_SOME;
2657+ } else if (curr == '{') {
2658+ out->start = ptr++;
2659+ out->end = ptr;
2660+ out->token_type = NPY_TOKEN_TYPE_LBRACE;
2661+ return VEC0_TOKEN_RESULT_SOME;
2662+ } else if (curr == '}') {
2663+ out->start = ptr++;
2664+ out->end = ptr;
2665+ out->token_type = NPY_TOKEN_TYPE_RBRACE;
2666+ return VEC0_TOKEN_RESULT_SOME;
2667+ } else if (curr == ':') {
2668+ out->start = ptr++;
2669+ out->end = ptr;
2670+ out->token_type = NPY_TOKEN_TYPE_COLON;
2671+ return VEC0_TOKEN_RESULT_SOME;
2672+ } else if (curr == ',') {
2673+ out->start = ptr++;
2674+ out->end = ptr;
2675+ out->token_type = NPY_TOKEN_TYPE_COMMA;
2676+ return VEC0_TOKEN_RESULT_SOME;
2677+ } else if (curr == '\'') {
2678+ unsigned char *start = ptr;
2679+ ptr++;
2680+ while (ptr < end) {
2681+ if ((*ptr) == '\'') {
2682+ break;
2683+ }
2684+ ptr++;
2685+ }
2686+ if ((*ptr) != '\'') {
2687+ return VEC0_TOKEN_RESULT_ERROR;
2688+ }
2689+ out->start = start;
2690+ out->end = ++ptr;
2691+ out->token_type = NPY_TOKEN_TYPE_STRING;
2692+ return VEC0_TOKEN_RESULT_SOME;
2693+ } else if (curr == 'F' &&
2694+ strncmp((char *)ptr, "False", strlen("False")) == 0) {
2695+ out->start = ptr;
2696+ out->end = (ptr + (int)strlen("False"));
2697+ ptr = out->end;
2698+ out->token_type = NPY_TOKEN_TYPE_FALSE;
2699+ return VEC0_TOKEN_RESULT_SOME;
2700+ } else if (is_digit(curr)) {
2701+ unsigned char *start = ptr;
2702+ while (ptr < end && (is_digit(*ptr))) {
2703+ ptr++;
2704+ }
2705+ out->start = start;
2706+ out->end = ptr;
2707+ out->token_type = NPY_TOKEN_TYPE_NUMBER;
2708+ return VEC0_TOKEN_RESULT_SOME;
2709+ } else {
2710+ return VEC0_TOKEN_RESULT_ERROR;
2711+ }
2712+ }
2713+ return VEC0_TOKEN_RESULT_ERROR;
2714+}
2715+
2716+struct NpyScanner {
2717+ unsigned char *start;
2718+ unsigned char *end;
2719+ unsigned char *ptr;
2720+};
2721+
2722+void npy_scanner_init(struct NpyScanner *scanner, const unsigned char *source,
2723+ int source_length) {
2724+ scanner->start = (unsigned char *)source;
2725+ scanner->end = (unsigned char *)source + source_length;
2726+ scanner->ptr = (unsigned char *)source;
2727+}
2728+
2729+int npy_scanner_next(struct NpyScanner *scanner, struct NpyToken *out) {
2730+ int rc = npy_token_next(scanner->start, scanner->end, out);
2731+ if (rc == VEC0_TOKEN_RESULT_SOME) {
2732+ scanner->start = out->end;
2733+ }
2734+ return rc;
2735+}
2736+
2737+#define NPY_PARSE_ERROR "Error parsing numpy array: "
2738+int parse_npy_header(sqlite3_vtab *pVTab, const unsigned char *header,
2739+ size_t headerLength,
2740+ enum VectorElementType *out_element_type,
2741+ int *fortran_order, size_t *numElements,
2742+ size_t *numDimensions) {
2743+
2744+ struct NpyScanner scanner;
2745+ struct NpyToken token;
2746+ int rc;
2747+ npy_scanner_init(&scanner, header, headerLength);
2748+
2749+ if (npy_scanner_next(&scanner, &token) != VEC0_TOKEN_RESULT_SOME &&
2750+ token.token_type != NPY_TOKEN_TYPE_LBRACE) {
2751+ vtab_set_error(pVTab,
2752+ NPY_PARSE_ERROR "numpy header did not start with '{'");
2753+ return SQLITE_ERROR;
2754+ }
2755+ while (1) {
2756+ rc = npy_scanner_next(&scanner, &token);
2757+ if (rc != VEC0_TOKEN_RESULT_SOME) {
2758+ vtab_set_error(pVTab, NPY_PARSE_ERROR "expected key in numpy header");
2759+ return SQLITE_ERROR;
2760+ }
2761+
2762+ if (token.token_type == NPY_TOKEN_TYPE_RBRACE) {
2763+ break;
2764+ }
2765+ if (token.token_type != NPY_TOKEN_TYPE_STRING) {
2766+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2767+ "expected a string as key in numpy header");
2768+ return SQLITE_ERROR;
2769+ }
2770+ unsigned char *key = token.start;
2771+
2772+ rc = npy_scanner_next(&scanner, &token);
2773+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2774+ (token.token_type != NPY_TOKEN_TYPE_COLON)) {
2775+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2776+ "expected a ':' after key in numpy header");
2777+ return SQLITE_ERROR;
2778+ }
2779+
2780+ if (strncmp((char *)key, "'descr'", strlen("'descr'")) == 0) {
2781+ rc = npy_scanner_next(&scanner, &token);
2782+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2783+ (token.token_type != NPY_TOKEN_TYPE_STRING)) {
2784+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2785+ "expected a string value after 'descr' key");
2786+ return SQLITE_ERROR;
2787+ }
2788+ if (strncmp((char *)token.start, "'<f4'", strlen("'<f4'")) != 0) {
2789+ vtab_set_error(
2790+ pVTab, NPY_PARSE_ERROR
2791+ "Only '<f4' values are supported in sqlite-vec numpy functions");
2792+ return SQLITE_ERROR;
2793+ }
2794+ *out_element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
2795+ } else if (strncmp((char *)key, "'fortran_order'",
2796+ strlen("'fortran_order'")) == 0) {
2797+ rc = npy_scanner_next(&scanner, &token);
2798+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2799+ token.token_type != NPY_TOKEN_TYPE_FALSE) {
2800+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2801+ "Only fortran_order = False is supported in sqlite-vec "
2802+ "numpy functions");
2803+ return SQLITE_ERROR;
2804+ }
2805+ *fortran_order = 0;
2806+ } else if (strncmp((char *)key, "'shape'", strlen("'shape'")) == 0) {
2807+ // "(xxx, xxx)" OR (xxx,)
2808+ size_t first;
2809+ rc = npy_scanner_next(&scanner, &token);
2810+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2811+ (token.token_type != NPY_TOKEN_TYPE_LPAREN)) {
2812+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2813+ "Expected left parenthesis '(' after shape key");
2814+ return SQLITE_ERROR;
2815+ }
2816+
2817+ rc = npy_scanner_next(&scanner, &token);
2818+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2819+ (token.token_type != NPY_TOKEN_TYPE_NUMBER)) {
2820+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2821+ "Expected an initial number in shape value");
2822+ return SQLITE_ERROR;
2823+ }
2824+ first = strtol((char *)token.start, NULL, 10);
2825+
2826+ rc = npy_scanner_next(&scanner, &token);
2827+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2828+ (token.token_type != NPY_TOKEN_TYPE_COMMA)) {
2829+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2830+ "Expected comma after first shape value");
2831+ return SQLITE_ERROR;
2832+ }
2833+
2834+ rc = npy_scanner_next(&scanner, &token);
2835+ if (rc != VEC0_TOKEN_RESULT_SOME) {
2836+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2837+ "unexpected header EOF while parsing shape");
2838+ return SQLITE_ERROR;
2839+ }
2840+ if (token.token_type == NPY_TOKEN_TYPE_NUMBER) {
2841+ *numElements = first;
2842+ *numDimensions = strtol((char *)token.start, NULL, 10);
2843+ rc = npy_scanner_next(&scanner, &token);
2844+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2845+ (token.token_type != NPY_TOKEN_TYPE_RPAREN)) {
2846+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2847+ "expected right parenthesis after shape value");
2848+ return SQLITE_ERROR;
2849+ }
2850+ } else if (token.token_type == NPY_TOKEN_TYPE_RPAREN) {
2851+ // '(0,)' means an empty array!
2852+ *numElements = first ? 1 : 0;
2853+ *numDimensions = first;
2854+ } else {
2855+ vtab_set_error(pVTab, NPY_PARSE_ERROR "unknown type in shape value");
2856+ return SQLITE_ERROR;
2857+ }
2858+ } else {
2859+ vtab_set_error(pVTab, NPY_PARSE_ERROR "unknown key in numpy header");
2860+ return SQLITE_ERROR;
2861+ }
2862+
2863+ rc = npy_scanner_next(&scanner, &token);
2864+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2865+ (token.token_type != NPY_TOKEN_TYPE_COMMA)) {
2866+ vtab_set_error(pVTab, NPY_PARSE_ERROR "unknown extra token after value");
2867+ return SQLITE_ERROR;
2868+ }
2869+ }
2870+
2871+ return SQLITE_OK;
2872+}
2873+
2874+typedef struct vec_npy_each_vtab vec_npy_each_vtab;
2875+struct vec_npy_each_vtab {
2876+ sqlite3_vtab base;
2877+};
2878+
2879+typedef enum {
2880+ VEC_NPY_EACH_INPUT_BUFFER,
2881+ VEC_NPY_EACH_INPUT_FILE,
2882+} vec_npy_each_input_type;
2883+
2884+typedef struct vec_npy_each_cursor vec_npy_each_cursor;
2885+struct vec_npy_each_cursor {
2886+ sqlite3_vtab_cursor base;
2887+ i64 iRowid;
2888+ // sqlite-vec compatible type of vector
2889+ enum VectorElementType elementType;
2890+ // number of vectors in the npy array
2891+ size_t nElements;
2892+ // number of dimensions each vector has
2893+ size_t nDimensions;
2894+
2895+ vec_npy_each_input_type input_type;
2896+
2897+ // when input_type == VEC_NPY_EACH_INPUT_BUFFER
2898+
2899+ // Buffer containing the vector data, when reading from an in-memory buffer.
2900+ // Size: nElements * nDimensions * element_size
2901+ // Clean up with sqlite3_free() once complete
2902+ void *vector;
2903+
2904+ // when input_type == VEC_NPY_EACH_INPUT_FILE
2905+
2906+ // Opened npy file, when reading from a file.
2907+ // fclose() when complete.
2908+#ifndef SQLITE_VEC_OMIT_FS
2909+ FILE *file;
2910+#endif
2911+
2912+ // an in-memory buffer containing a portion of the npy array.
2913+ // Used for faster reading, instead of calling fread a lot.
2914+ // Will have a byte-size of fileBufferSize
2915+ void *chunksBuffer;
2916+ // size of allocated fileBuffer in bytes
2917+ size_t chunksBufferSize;
2918+ //// Maximum length of the buffer, in terms of number of vectors.
2919+ size_t maxChunks;
2920+
2921+ // Counter index of the current vector into of fileBuffer to yield.
2922+ // Starts at 0 once fileBuffer is read, and iterates to bufferLength.
2923+ // Resets to 0 once that "buffer" is yielded and a new one is read.
2924+ size_t currentChunkIndex;
2925+ size_t currentChunkSize;
2926+
2927+ // 0 when there are still more elements to read/yield, 1 when complete.
2928+ int eof;
2929+};
2930+
2931+static unsigned char NPY_MAGIC[6] = "\x93NUMPY";
2932+
2933+#ifndef SQLITE_VEC_OMIT_FS
2934+int parse_npy_file(sqlite3_vtab *pVTab, FILE *file, vec_npy_each_cursor *pCur) {
2935+ int n;
2936+ fseek(file, 0, SEEK_END);
2937+ long fileSize = ftell(file);
2938+
2939+ fseek(file, 0L, SEEK_SET);
2940+
2941+ unsigned char header[10];
2942+ n = fread(&header, sizeof(unsigned char), 10, file);
2943+ if (n != 10) {
2944+ vtab_set_error(pVTab, "numpy array file too short");
2945+ return SQLITE_ERROR;
2946+ }
2947+
2948+ if (memcmp(NPY_MAGIC, header, sizeof(NPY_MAGIC)) != 0) {
2949+ vtab_set_error(pVTab,
2950+ "numpy array file does not contain the 'magic' header");
2951+ return SQLITE_ERROR;
2952+ }
2953+
2954+ u8 major = header[6];
2955+ u8 minor = header[7];
2956+ uint16_t headerLength = 0;
2957+ memcpy(&headerLength, &header[8], sizeof(uint16_t));
2958+
2959+ size_t totalHeaderLength = sizeof(NPY_MAGIC) + sizeof(major) + sizeof(minor) +
2960+ sizeof(headerLength) + headerLength;
2961+ i32 dataSize = fileSize - totalHeaderLength;
2962+ if (dataSize < 0) {
2963+ vtab_set_error(pVTab, "numpy array file header length is invalid");
2964+ return SQLITE_ERROR;
2965+ }
2966+
2967+ unsigned char *headerX = sqlite3_malloc(headerLength);
2968+ if (headerLength && !headerX) {
2969+ return SQLITE_NOMEM;
2970+ }
2971+
2972+ n = fread(headerX, sizeof(char), headerLength, file);
2973+ if (n != headerLength) {
2974+ sqlite3_free(headerX);
2975+ vtab_set_error(pVTab, "numpy array file header length is invalid");
2976+ return SQLITE_ERROR;
2977+ }
2978+
2979+ int fortran_order;
2980+ enum VectorElementType element_type;
2981+ size_t numElements;
2982+ size_t numDimensions;
2983+ int rc = parse_npy_header(pVTab, headerX, headerLength, &element_type,
2984+ &fortran_order, &numElements, &numDimensions);
2985+ sqlite3_free(headerX);
2986+ if (rc != SQLITE_OK) {
2987+ // parse_npy_header already attackes an error emssage
2988+ return rc;
2989+ }
2990+
2991+ i32 expectedDataSize =
2992+ numElements * vector_byte_size(element_type, numDimensions);
2993+ if (expectedDataSize != dataSize) {
2994+ vtab_set_error(
2995+ pVTab, "numpy array file error: Expected a data size of %d, found %d",
2996+ expectedDataSize, dataSize);
2997+ return SQLITE_ERROR;
2998+ }
2999+
3000+ pCur->maxChunks = 1024;
3001+ pCur->chunksBufferSize =
3002+ (vector_byte_size(element_type, numDimensions)) * pCur->maxChunks;
3003+ pCur->chunksBuffer = sqlite3_malloc(pCur->chunksBufferSize);
3004+ if (pCur->chunksBufferSize && !pCur->chunksBuffer) {
3005+ return SQLITE_NOMEM;
3006+ }
3007+
3008+ pCur->currentChunkSize =
3009+ fread(pCur->chunksBuffer, vector_byte_size(element_type, numDimensions),
3010+ pCur->maxChunks, file);
3011+
3012+ pCur->currentChunkIndex = 0;
3013+ pCur->elementType = element_type;
3014+ pCur->nElements = numElements;
3015+ pCur->nDimensions = numDimensions;
3016+ pCur->input_type = VEC_NPY_EACH_INPUT_FILE;
3017+
3018+ pCur->eof = pCur->currentChunkSize == 0;
3019+ pCur->file = file;
3020+ return SQLITE_OK;
3021+}
3022+#endif
3023+
3024+int parse_npy_buffer(sqlite3_vtab *pVTab, const unsigned char *buffer,
3025+ int bufferLength, void **data, size_t *numElements,
3026+ size_t *numDimensions,
3027+ enum VectorElementType *element_type) {
3028+
3029+ if (bufferLength < 10) {
3030+ // IMP: V03312_20150
3031+ vtab_set_error(pVTab, "numpy array too short");
3032+ return SQLITE_ERROR;
3033+ }
3034+ if (memcmp(NPY_MAGIC, buffer, sizeof(NPY_MAGIC)) != 0) {
3035+ // V11954_28792
3036+ vtab_set_error(pVTab, "numpy array does not contain the 'magic' header");
3037+ return SQLITE_ERROR;
3038+ }
3039+
3040+ u8 major = buffer[6];
3041+ u8 minor = buffer[7];
3042+ uint16_t headerLength = 0;
3043+ memcpy(&headerLength, &buffer[8], sizeof(uint16_t));
3044+
3045+ i32 totalHeaderLength = sizeof(NPY_MAGIC) + sizeof(major) + sizeof(minor) +
3046+ sizeof(headerLength) + headerLength;
3047+ i32 dataSize = bufferLength - totalHeaderLength;
3048+
3049+ if (dataSize < 0) {
3050+ vtab_set_error(pVTab, "numpy array header length is invalid");
3051+ return SQLITE_ERROR;
3052+ }
3053+
3054+ const unsigned char *header = &buffer[10];
3055+ int fortran_order;
3056+
3057+ int rc = parse_npy_header(pVTab, header, headerLength, element_type,
3058+ &fortran_order, numElements, numDimensions);
3059+ if (rc != SQLITE_OK) {
3060+ return rc;
3061+ }
3062+
3063+ i32 expectedDataSize =
3064+ (*numElements * vector_byte_size(*element_type, *numDimensions));
3065+ if (expectedDataSize != dataSize) {
3066+ vtab_set_error(pVTab,
3067+ "numpy array error: Expected a data size of %d, found %d",
3068+ expectedDataSize, dataSize);
3069+ return SQLITE_ERROR;
3070+ }
3071+
3072+ *data = (void *)&buffer[totalHeaderLength];
3073+ return SQLITE_OK;
3074+}
3075+
3076+static int vec_npy_eachConnect(sqlite3 *db, void *pAux, int argc,
3077+ const char *const *argv, sqlite3_vtab **ppVtab,
3078+ char **pzErr) {
3079+ UNUSED_PARAMETER(pAux);
3080+ UNUSED_PARAMETER(argc);
3081+ UNUSED_PARAMETER(argv);
3082+ UNUSED_PARAMETER(pzErr);
3083+ vec_npy_each_vtab *pNew;
3084+ int rc;
3085+
3086+ rc = sqlite3_declare_vtab(db, "CREATE TABLE x(vector, input hidden)");
3087+#define VEC_NPY_EACH_COLUMN_VECTOR 0
3088+#define VEC_NPY_EACH_COLUMN_INPUT 1
3089+ if (rc == SQLITE_OK) {
3090+ pNew = sqlite3_malloc(sizeof(*pNew));
3091+ *ppVtab = (sqlite3_vtab *)pNew;
3092+ if (pNew == 0)
3093+ return SQLITE_NOMEM;
3094+ memset(pNew, 0, sizeof(*pNew));
3095+ }
3096+ return rc;
3097+}
3098+
3099+static int vec_npy_eachDisconnect(sqlite3_vtab *pVtab) {
3100+ vec_npy_each_vtab *p = (vec_npy_each_vtab *)pVtab;
3101+ sqlite3_free(p);
3102+ return SQLITE_OK;
3103+}
3104+
3105+static int vec_npy_eachOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) {
3106+ UNUSED_PARAMETER(p);
3107+ vec_npy_each_cursor *pCur;
3108+ pCur = sqlite3_malloc(sizeof(*pCur));
3109+ if (pCur == 0)
3110+ return SQLITE_NOMEM;
3111+ memset(pCur, 0, sizeof(*pCur));
3112+ *ppCursor = &pCur->base;
3113+ return SQLITE_OK;
3114+}
3115+
3116+static int vec_npy_eachClose(sqlite3_vtab_cursor *cur) {
3117+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3118+#ifndef SQLITE_VEC_OMIT_FS
3119+ if (pCur->file) {
3120+ fclose(pCur->file);
3121+ pCur->file = NULL;
3122+ }
3123+#endif
3124+ if (pCur->chunksBuffer) {
3125+ sqlite3_free(pCur->chunksBuffer);
3126+ pCur->chunksBuffer = NULL;
3127+ }
3128+ if (pCur->vector) {
3129+ pCur->vector = NULL;
3130+ }
3131+ sqlite3_free(pCur);
3132+ return SQLITE_OK;
3133+}
3134+
3135+static int vec_npy_eachBestIndex(sqlite3_vtab *pVTab,
3136+ sqlite3_index_info *pIdxInfo) {
3137+ int hasInput;
3138+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
3139+ const struct sqlite3_index_constraint *pCons = &pIdxInfo->aConstraint[i];
3140+ // printf("i=%d iColumn=%d, op=%d, usable=%d\n", i, pCons->iColumn,
3141+ // pCons->op, pCons->usable);
3142+ switch (pCons->iColumn) {
3143+ case VEC_NPY_EACH_COLUMN_INPUT: {
3144+ if (pCons->op == SQLITE_INDEX_CONSTRAINT_EQ && pCons->usable) {
3145+ hasInput = 1;
3146+ pIdxInfo->aConstraintUsage[i].argvIndex = 1;
3147+ pIdxInfo->aConstraintUsage[i].omit = 1;
3148+ }
3149+ break;
3150+ }
3151+ }
3152+ }
3153+ if (!hasInput) {
3154+ pVTab->zErrMsg = sqlite3_mprintf("input argument is required");
3155+ return SQLITE_ERROR;
3156+ }
3157+
3158+ pIdxInfo->estimatedCost = (double)100000;
3159+ pIdxInfo->estimatedRows = 100000;
3160+
3161+ return SQLITE_OK;
3162+}
3163+
3164+static int vec_npy_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
3165+ const char *idxStr, int argc,
3166+ sqlite3_value **argv) {
3167+ UNUSED_PARAMETER(idxNum);
3168+ UNUSED_PARAMETER(idxStr);
3169+ assert(argc == 1);
3170+ int rc;
3171+
3172+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)pVtabCursor;
3173+
3174+#ifndef SQLITE_VEC_OMIT_FS
3175+ if (pCur->file) {
3176+ fclose(pCur->file);
3177+ pCur->file = NULL;
3178+ }
3179+#endif
3180+ if (pCur->chunksBuffer) {
3181+ sqlite3_free(pCur->chunksBuffer);
3182+ pCur->chunksBuffer = NULL;
3183+ }
3184+ if (pCur->vector) {
3185+ pCur->vector = NULL;
3186+ }
3187+
3188+#ifndef SQLITE_VEC_OMIT_FS
3189+ struct VecNpyFile *f = NULL;
3190+ if ((f = sqlite3_value_pointer(argv[0], SQLITE_VEC_NPY_FILE_NAME))) {
3191+ FILE *file = fopen(f->path, "r");
3192+ if (!file) {
3193+ vtab_set_error(pVtabCursor->pVtab, "Could not open numpy file");
3194+ return SQLITE_ERROR;
3195+ }
3196+
3197+ rc = parse_npy_file(pVtabCursor->pVtab, file, pCur);
3198+ if (rc != SQLITE_OK) {
3199+#ifndef SQLITE_VEC_OMIT_FS
3200+ fclose(file);
3201+#endif
3202+ return rc;
3203+ }
3204+
3205+ } else
3206+#endif
3207+ {
3208+
3209+ const unsigned char *input = sqlite3_value_blob(argv[0]);
3210+ int inputLength = sqlite3_value_bytes(argv[0]);
3211+ void *data;
3212+ size_t numElements;
3213+ size_t numDimensions;
3214+ enum VectorElementType element_type;
3215+
3216+ rc = parse_npy_buffer(pVtabCursor->pVtab, input, inputLength, &data,
3217+ &numElements, &numDimensions, &element_type);
3218+ if (rc != SQLITE_OK) {
3219+ return rc;
3220+ }
3221+
3222+ pCur->vector = data;
3223+ pCur->elementType = element_type;
3224+ pCur->nElements = numElements;
3225+ pCur->nDimensions = numDimensions;
3226+ pCur->input_type = VEC_NPY_EACH_INPUT_BUFFER;
3227+ }
3228+
3229+ pCur->iRowid = 0;
3230+ return SQLITE_OK;
3231+}
3232+
3233+static int vec_npy_eachRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) {
3234+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3235+ *pRowid = pCur->iRowid;
3236+ return SQLITE_OK;
3237+}
3238+
3239+static int vec_npy_eachEof(sqlite3_vtab_cursor *cur) {
3240+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3241+ if (pCur->input_type == VEC_NPY_EACH_INPUT_BUFFER) {
3242+ return (!pCur->nElements) || (size_t)pCur->iRowid >= pCur->nElements;
3243+ }
3244+ return pCur->eof;
3245+}
3246+
3247+static int vec_npy_eachNext(sqlite3_vtab_cursor *cur) {
3248+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3249+ pCur->iRowid++;
3250+ if (pCur->input_type == VEC_NPY_EACH_INPUT_BUFFER) {
3251+ return SQLITE_OK;
3252+ }
3253+
3254+#ifndef SQLITE_VEC_OMIT_FS
3255+ // else: input is a file
3256+ pCur->currentChunkIndex++;
3257+ if (pCur->currentChunkIndex >= pCur->currentChunkSize) {
3258+ pCur->currentChunkSize =
3259+ fread(pCur->chunksBuffer,
3260+ vector_byte_size(pCur->elementType, pCur->nDimensions),
3261+ pCur->maxChunks, pCur->file);
3262+ if (!pCur->currentChunkSize) {
3263+ pCur->eof = 1;
3264+ }
3265+ pCur->currentChunkIndex = 0;
3266+ }
3267+#endif
3268+ return SQLITE_OK;
3269+}
3270+
3271+static int vec_npy_eachColumnBuffer(vec_npy_each_cursor *pCur,
3272+ sqlite3_context *context, int i) {
3273+ switch (i) {
3274+ case VEC_NPY_EACH_COLUMN_VECTOR: {
3275+ sqlite3_result_subtype(context, pCur->elementType);
3276+ switch (pCur->elementType) {
3277+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
3278+ sqlite3_result_blob(
3279+ context,
3280+ &((unsigned char *)
3281+ pCur->vector)[pCur->iRowid * pCur->nDimensions * sizeof(f32)],
3282+ pCur->nDimensions * sizeof(f32), SQLITE_TRANSIENT);
3283+
3284+ break;
3285+ }
3286+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
3287+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
3288+ // https://github.com/asg017/sqlite-vec/issues/42
3289+ sqlite3_result_error(context,
3290+ "vec_npy_each only supports float32 vectors", -1);
3291+ break;
3292+ }
3293+ }
3294+
3295+ break;
3296+ }
3297+ }
3298+ return SQLITE_OK;
3299+}
3300+static int vec_npy_eachColumnFile(vec_npy_each_cursor *pCur,
3301+ sqlite3_context *context, int i) {
3302+ switch (i) {
3303+ case VEC_NPY_EACH_COLUMN_VECTOR: {
3304+ switch (pCur->elementType) {
3305+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
3306+ sqlite3_result_blob(
3307+ context,
3308+ &((unsigned char *)
3309+ pCur->chunksBuffer)[pCur->currentChunkIndex *
3310+ pCur->nDimensions * sizeof(f32)],
3311+ pCur->nDimensions * sizeof(f32), SQLITE_TRANSIENT);
3312+ break;
3313+ }
3314+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
3315+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
3316+ // https://github.com/asg017/sqlite-vec/issues/42
3317+ sqlite3_result_error(context,
3318+ "vec_npy_each only supports float32 vectors", -1);
3319+ break;
3320+ }
3321+ }
3322+ break;
3323+ }
3324+ }
3325+ return SQLITE_OK;
3326+}
3327+static int vec_npy_eachColumn(sqlite3_vtab_cursor *cur,
3328+ sqlite3_context *context, int i) {
3329+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3330+ switch (pCur->input_type) {
3331+ case VEC_NPY_EACH_INPUT_BUFFER:
3332+ return vec_npy_eachColumnBuffer(pCur, context, i);
3333+ case VEC_NPY_EACH_INPUT_FILE:
3334+ return vec_npy_eachColumnFile(pCur, context, i);
3335+ }
3336+ return SQLITE_ERROR;
3337+}
3338+
3339+static sqlite3_module vec_npy_eachModule = {
3340+ /* iVersion */ 0,
3341+ /* xCreate */ 0,
3342+ /* xConnect */ vec_npy_eachConnect,
3343+ /* xBestIndex */ vec_npy_eachBestIndex,
3344+ /* xDisconnect */ vec_npy_eachDisconnect,
3345+ /* xDestroy */ 0,
3346+ /* xOpen */ vec_npy_eachOpen,
3347+ /* xClose */ vec_npy_eachClose,
3348+ /* xFilter */ vec_npy_eachFilter,
3349+ /* xNext */ vec_npy_eachNext,
3350+ /* xEof */ vec_npy_eachEof,
3351+ /* xColumn */ vec_npy_eachColumn,
3352+ /* xRowid */ vec_npy_eachRowid,
3353+ /* xUpdate */ 0,
3354+ /* xBegin */ 0,
3355+ /* xSync */ 0,
3356+ /* xCommit */ 0,
3357+ /* xRollback */ 0,
3358+ /* xFindMethod */ 0,
3359+ /* xRename */ 0,
3360+ /* xSavepoint */ 0,
3361+ /* xRelease */ 0,
3362+ /* xRollbackTo */ 0,
3363+ /* xShadowName */ 0,
3364+#if SQLITE_VERSION_NUMBER >= 3044000
3365+ /* xIntegrity */ 0,
3366+#endif
3367+};
3368+
3369+#pragma endregion
3370+
3371+#pragma region vec0 virtual table
3372+
3373+#define VEC0_COLUMN_ID 0
3374+#define VEC0_COLUMN_USERN_START 1
3375+#define VEC0_COLUMN_OFFSET_DISTANCE 1
3376+#define VEC0_COLUMN_OFFSET_K 2
3377+
3378+#define VEC0_SHADOW_INFO_NAME "\"%w\".\"%w_info\""
3379+
3380+#define VEC0_SHADOW_CHUNKS_NAME "\"%w\".\"%w_chunks\""
3381+/// 1) schema, 2) original vtab table name
3382+#define VEC0_SHADOW_CHUNKS_CREATE \
3383+ "CREATE TABLE " VEC0_SHADOW_CHUNKS_NAME "(" \
3384+ "chunk_id INTEGER PRIMARY KEY AUTOINCREMENT," \
3385+ "size INTEGER NOT NULL," \
3386+ "validity BLOB NOT NULL," \
3387+ "rowids BLOB NOT NULL" \
3388+ ");"
3389+
3390+#define VEC0_SHADOW_ROWIDS_NAME "\"%w\".\"%w_rowids\""
3391+/// 1) schema, 2) original vtab table name
3392+#define VEC0_SHADOW_ROWIDS_CREATE_BASIC \
3393+ "CREATE TABLE " VEC0_SHADOW_ROWIDS_NAME "(" \
3394+ "rowid INTEGER PRIMARY KEY AUTOINCREMENT," \
3395+ "id," \
3396+ "chunk_id INTEGER," \
3397+ "chunk_offset INTEGER" \
3398+ ");"
3399+
3400+// vec0 tables with a text primary keys are still backed by int64 primary keys,
3401+// since a fixed-length rowid is required for vec0 chunks. But we add a new 'id
3402+// text unique' column to emulate a text primary key interface.
3403+#define VEC0_SHADOW_ROWIDS_CREATE_PK_TEXT \
3404+ "CREATE TABLE " VEC0_SHADOW_ROWIDS_NAME "(" \
3405+ "rowid INTEGER PRIMARY KEY AUTOINCREMENT," \
3406+ "id TEXT UNIQUE NOT NULL," \
3407+ "chunk_id INTEGER," \
3408+ "chunk_offset INTEGER" \
3409+ ");"
3410+
3411+/// 1) schema, 2) original vtab table name
3412+#define VEC0_SHADOW_VECTOR_N_NAME "\"%w\".\"%w_vector_chunks%02d\""
3413+
3414+/// 1) schema, 2) original vtab table name
3415+#define VEC0_SHADOW_VECTOR_N_CREATE \
3416+ "CREATE TABLE " VEC0_SHADOW_VECTOR_N_NAME "(" \
3417+ "rowid PRIMARY KEY," \
3418+ "vectors BLOB NOT NULL" \
3419+ ");"
3420+
3421+#define VEC0_SHADOW_AUXILIARY_NAME "\"%w\".\"%w_auxiliary\""
3422+
3423+#define VEC0_SHADOW_METADATA_N_NAME "\"%w\".\"%w_metadatachunks%02d\""
3424+#define VEC0_SHADOW_METADATA_TEXT_DATA_NAME "\"%w\".\"%w_metadatatext%02d\""
3425+
3426+#define VEC_INTERAL_ERROR "Internal sqlite-vec error: "
3427+#define REPORT_URL "https://github.com/asg017/sqlite-vec/issues/new"
3428+
3429+typedef struct vec0_vtab vec0_vtab;
3430+
3431+#define VEC0_MAX_VECTOR_COLUMNS 16
3432+#define VEC0_MAX_PARTITION_COLUMNS 4
3433+#define VEC0_MAX_AUXILIARY_COLUMNS 16
3434+#define VEC0_MAX_METADATA_COLUMNS 16
3435+
3436+#define SQLITE_VEC_VEC0_MAX_DIMENSIONS 8192
3437+#define VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH 16
3438+#define VEC0_METADATA_TEXT_VIEW_DATA_LENGTH 12
3439+
3440+typedef enum {
3441+ // vector column, ie "contents_embedding float[1024]"
3442+ SQLITE_VEC0_USER_COLUMN_KIND_VECTOR = 1,
3443+
3444+ // partition key column, ie "user_id integer partition key"
3445+ SQLITE_VEC0_USER_COLUMN_KIND_PARTITION = 2,
3446+
3447+ //
3448+ SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY = 3,
3449+
3450+ // metadata column that can be filtered, ie "genre text"
3451+ SQLITE_VEC0_USER_COLUMN_KIND_METADATA = 4,
3452+} vec0_user_column_kind;
3453+
3454+struct vec0_vtab {
3455+ sqlite3_vtab base;
3456+
3457+ // the SQLite connection of the host database
3458+ sqlite3 *db;
3459+
3460+ // True if the primary key of the vec0 table has a column type TEXT.
3461+ // Will change the schema of the _rowids table, and insert/query logic.
3462+ int pkIsText;
3463+
3464+ // number of defined vector columns.
3465+ int numVectorColumns;
3466+
3467+ // number of defined PARTITION KEY columns.
3468+ int numPartitionColumns;
3469+
3470+ // number of defined auxiliary columns
3471+ int numAuxiliaryColumns;
3472+
3473+ // number of defined metadata columns
3474+ int numMetadataColumns;
3475+
3476+
3477+ // Name of the schema the table exists on.
3478+ // Must be freed with sqlite3_free()
3479+ char *schemaName;
3480+
3481+ // Name of the table the table exists on.
3482+ // Must be freed with sqlite3_free()
3483+ char *tableName;
3484+
3485+ // Name of the _rowids shadow table.
3486+ // Must be freed with sqlite3_free()
3487+ char *shadowRowidsName;
3488+
3489+ // Name of the _chunks shadow table.
3490+ // Must be freed with sqlite3_free()
3491+ char *shadowChunksName;
3492+
3493+ // contains enum vec0_user_column_kind values for up to
3494+ // numVectorColumns + numPartitionColumns entries
3495+ vec0_user_column_kind user_column_kinds[VEC0_MAX_VECTOR_COLUMNS + VEC0_MAX_PARTITION_COLUMNS + VEC0_MAX_AUXILIARY_COLUMNS + VEC0_MAX_METADATA_COLUMNS];
3496+
3497+ uint8_t user_column_idxs[VEC0_MAX_VECTOR_COLUMNS + VEC0_MAX_PARTITION_COLUMNS + VEC0_MAX_AUXILIARY_COLUMNS + VEC0_MAX_METADATA_COLUMNS];
3498+
3499+
3500+ // Name of all the vector chunk shadow tables.
3501+ // Ex '_vector_chunks00'
3502+ // Only the first numVectorColumns entries will be available.
3503+ // The first numVectorColumns entries must be freed with sqlite3_free()
3504+ char *shadowVectorChunksNames[VEC0_MAX_VECTOR_COLUMNS];
3505+
3506+ // Name of all metadata chunk shadow tables, ie `_metadatachunks00`
3507+ // Only the first numMetadataColumns entries will be available.
3508+ // The first numMetadataColumns entries must be freed with sqlite3_free()
3509+ char *shadowMetadataChunksNames[VEC0_MAX_METADATA_COLUMNS];
3510+
3511+ struct VectorColumnDefinition vector_columns[VEC0_MAX_VECTOR_COLUMNS];
3512+ struct Vec0PartitionColumnDefinition paritition_columns[VEC0_MAX_PARTITION_COLUMNS];
3513+ struct Vec0AuxiliaryColumnDefinition auxiliary_columns[VEC0_MAX_AUXILIARY_COLUMNS];
3514+ struct Vec0MetadataColumnDefinition metadata_columns[VEC0_MAX_METADATA_COLUMNS];
3515+
3516+ int chunk_size;
3517+
3518+ // select latest chunk from _chunks, getting chunk_id
3519+ sqlite3_stmt *stmtLatestChunk;
3520+
3521+ /**
3522+ * Statement to insert a row into the _rowids table, with a rowid.
3523+ * Parameters:
3524+ * 1: int64, rowid to insert
3525+ * Result columns: none
3526+ * SQL: "INSERT INTO _rowids(rowid) VALUES (?)"
3527+ *
3528+ * Must be cleaned up with sqlite3_finalize().
3529+ */
3530+ sqlite3_stmt *stmtRowidsInsertRowid;
3531+
3532+ /**
3533+ * Statement to insert a row into the _rowids table, with an id.
3534+ * The id column isn't a tradition primary key, but instead a unique
3535+ * column to handle "text primary key" vec0 tables. The true int64 rowid
3536+ * can be retrieved after inserting with sqlite3_last_rowid().
3537+ *
3538+ * Parameters:
3539+ * 1: text or null, id to insert
3540+ * Result columns: none
3541+ *
3542+ * Must be cleaned up with sqlite3_finalize().
3543+ */
3544+ sqlite3_stmt *stmtRowidsInsertId;
3545+
3546+ /**
3547+ * Statement to update the "position" columns chunk_id and chunk_offset for
3548+ * a given _rowids row. Used when the "next available" chunk position is found
3549+ * for a vector.
3550+ *
3551+ * Parameters:
3552+ * 1: int64, chunk_id value
3553+ * 2: int64, chunk_offset value
3554+ * 3: int64, rowid value
3555+ * Result columns: none
3556+ *
3557+ * Must be cleaned up with sqlite3_finalize().
3558+ */
3559+ sqlite3_stmt *stmtRowidsUpdatePosition;
3560+
3561+ /**
3562+ * Statement to quickly find the chunk_id + chunk_offset of a given row.
3563+ * Parameters:
3564+ * 1: rowid of the row/vector to lookup
3565+ * Result columns:
3566+ * 0: chunk_id (i64)
3567+ * 1: chunk_offset (i64)
3568+ * SQL: "SELECT id, chunk_id, chunk_offset FROM _rowids WHERE rowid = ?""
3569+ *
3570+ * Must be cleaned up with sqlite3_finalize().
3571+ */
3572+ sqlite3_stmt *stmtRowidsGetChunkPosition;
3573+};
3574+
3575+/**
3576+ * @brief Finalize all the sqlite3_stmt members in a vec0_vtab.
3577+ *
3578+ * @param p vec0_vtab pointer
3579+ */
3580+void vec0_free_resources(vec0_vtab *p) {
3581+ sqlite3_finalize(p->stmtLatestChunk);
3582+ p->stmtLatestChunk = NULL;
3583+ sqlite3_finalize(p->stmtRowidsInsertRowid);
3584+ p->stmtRowidsInsertRowid = NULL;
3585+ sqlite3_finalize(p->stmtRowidsInsertId);
3586+ p->stmtRowidsInsertId = NULL;
3587+ sqlite3_finalize(p->stmtRowidsUpdatePosition);
3588+ p->stmtRowidsUpdatePosition = NULL;
3589+ sqlite3_finalize(p->stmtRowidsGetChunkPosition);
3590+ p->stmtRowidsGetChunkPosition = NULL;
3591+}
3592+
3593+/**
3594+ * @brief Free all memory and sqlite3_stmt members of a vec0_vtab
3595+ *
3596+ * @param p vec0_vtab pointer
3597+ */
3598+void vec0_free(vec0_vtab *p) {
3599+ vec0_free_resources(p);
3600+
3601+ sqlite3_free(p->schemaName);
3602+ p->schemaName = NULL;
3603+ sqlite3_free(p->tableName);
3604+ p->tableName = NULL;
3605+ sqlite3_free(p->shadowChunksName);
3606+ p->shadowChunksName = NULL;
3607+ sqlite3_free(p->shadowRowidsName);
3608+ p->shadowRowidsName = NULL;
3609+
3610+ for (int i = 0; i < p->numVectorColumns; i++) {
3611+ sqlite3_free(p->shadowVectorChunksNames[i]);
3612+ p->shadowVectorChunksNames[i] = NULL;
3613+
3614+ sqlite3_free(p->vector_columns[i].name);
3615+ p->vector_columns[i].name = NULL;
3616+ }
3617+}
3618+
3619+int vec0_num_defined_user_columns(vec0_vtab *p) {
3620+ return p->numVectorColumns + p->numPartitionColumns + p->numAuxiliaryColumns + p->numMetadataColumns;
3621+}
3622+
3623+/**
3624+ * @brief Returns the index of the distance hidden column for the given vec0
3625+ * table.
3626+ *
3627+ * @param p vec0 table
3628+ * @return int
3629+ */
3630+int vec0_column_distance_idx(vec0_vtab *p) {
3631+ return VEC0_COLUMN_USERN_START + (vec0_num_defined_user_columns(p) - 1) +
3632+ VEC0_COLUMN_OFFSET_DISTANCE;
3633+}
3634+
3635+/**
3636+ * @brief Returns the index of the k hidden column for the given vec0 table.
3637+ *
3638+ * @param p vec0 table
3639+ * @return int k column index
3640+ */
3641+int vec0_column_k_idx(vec0_vtab *p) {
3642+ return VEC0_COLUMN_USERN_START + (vec0_num_defined_user_columns(p) - 1) +
3643+ VEC0_COLUMN_OFFSET_K;
3644+}
3645+
3646+/**
3647+ * Returns 1 if the given column-based index is a valid vector column,
3648+ * 0 otherwise.
3649+ */
3650+int vec0_column_idx_is_vector(vec0_vtab *pVtab, int column_idx) {
3651+ return column_idx >= VEC0_COLUMN_USERN_START &&
3652+ column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) &&
3653+ pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_VECTOR;
3654+}
3655+
3656+/**
3657+ * Returns the vector index of the given user column index.
3658+ * ONLY call if validated with vec0_column_idx_is_vector before
3659+ */
3660+int vec0_column_idx_to_vector_idx(vec0_vtab *pVtab, int column_idx) {
3661+ UNUSED_PARAMETER(pVtab);
3662+ return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START];
3663+}
3664+/**
3665+ * Returns 1 if the given column-based index is a "partition key" column,
3666+ * 0 otherwise.
3667+ */
3668+int vec0_column_idx_is_partition(vec0_vtab *pVtab, int column_idx) {
3669+ return column_idx >= VEC0_COLUMN_USERN_START &&
3670+ column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) &&
3671+ pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_PARTITION;
3672+}
3673+
3674+/**
3675+ * Returns the partition column index of the given user column index.
3676+ * ONLY call if validated with vec0_column_idx_is_vector before
3677+ */
3678+int vec0_column_idx_to_partition_idx(vec0_vtab *pVtab, int column_idx) {
3679+ UNUSED_PARAMETER(pVtab);
3680+ return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START];
3681+}
3682+
3683+/**
3684+ * Returns 1 if the given column-based index is a auxiliary column,
3685+ * 0 otherwise.
3686+ */
3687+int vec0_column_idx_is_auxiliary(vec0_vtab *pVtab, int column_idx) {
3688+ return column_idx >= VEC0_COLUMN_USERN_START &&
3689+ column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) &&
3690+ pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY;
3691+}
3692+
3693+/**
3694+ * Returns the auxiliary column index of the given user column index.
3695+ * ONLY call if validated with vec0_column_idx_to_partition_idx before
3696+ */
3697+int vec0_column_idx_to_auxiliary_idx(vec0_vtab *pVtab, int column_idx) {
3698+ UNUSED_PARAMETER(pVtab);
3699+ return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START];
3700+}
3701+
3702+/**
3703+ * Returns 1 if the given column-based index is a metadata column,
3704+ * 0 otherwise.
3705+ */
3706+int vec0_column_idx_is_metadata(vec0_vtab *pVtab, int column_idx) {
3707+ return column_idx >= VEC0_COLUMN_USERN_START &&
3708+ column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) &&
3709+ pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_METADATA;
3710+}
3711+
3712+/**
3713+ * Returns the metadata column index of the given user column index.
3714+ * ONLY call if validated with vec0_column_idx_is_metadata before
3715+ */
3716+int vec0_column_idx_to_metadata_idx(vec0_vtab *pVtab, int column_idx) {
3717+ UNUSED_PARAMETER(pVtab);
3718+ return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START];
3719+}
3720+
3721+/**
3722+ * @brief Retrieve the chunk_id, chunk_offset, and possible "id" value
3723+ * of a vec0_vtab row with the provided rowid
3724+ *
3725+ * @param p vec0_vtab
3726+ * @param rowid the rowid of the row to query
3727+ * @param id output, optional sqlite3_value to provide the id.
3728+ * Useful for text PK rows. Must be freed with sqlite3_value_free()
3729+ * @param chunk_id output, the chunk_id the row belongs to
3730+ * @param chunk_offset output, the offset within the chunk the row belongs to
3731+ * @return SQLITE_ROW on success, error code otherwise. SQLITE_EMPTY if row DNE
3732+ */
3733+int vec0_get_chunk_position(vec0_vtab *p, i64 rowid, sqlite3_value **id,
3734+ i64 *chunk_id, i64 *chunk_offset) {
3735+ int rc;
3736+
3737+ if (!p->stmtRowidsGetChunkPosition) {
3738+ const char *zSql =
3739+ sqlite3_mprintf("SELECT id, chunk_id, chunk_offset "
3740+ "FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE rowid = ?",
3741+ p->schemaName, p->tableName);
3742+ if (!zSql) {
3743+ rc = SQLITE_NOMEM;
3744+ goto cleanup;
3745+ }
3746+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsGetChunkPosition, 0);
3747+ sqlite3_free((void *)zSql);
3748+ if (rc != SQLITE_OK) {
3749+ vtab_set_error(
3750+ &p->base, VEC_INTERAL_ERROR
3751+ "could not initialize 'rowids get chunk position' statement");
3752+ goto cleanup;
3753+ }
3754+ }
3755+
3756+ sqlite3_bind_int64(p->stmtRowidsGetChunkPosition, 1, rowid);
3757+ rc = sqlite3_step(p->stmtRowidsGetChunkPosition);
3758+ // special case: when no results, return SQLITE_EMPTY to convey "that chunk
3759+ // position doesnt exist"
3760+ if (rc == SQLITE_DONE) {
3761+ rc = SQLITE_EMPTY;
3762+ goto cleanup;
3763+ }
3764+ if (rc != SQLITE_ROW) {
3765+ goto cleanup;
3766+ }
3767+
3768+ if (id) {
3769+ sqlite3_value *value =
3770+ sqlite3_column_value(p->stmtRowidsGetChunkPosition, 0);
3771+ *id = sqlite3_value_dup(value);
3772+ if (!*id) {
3773+ rc = SQLITE_NOMEM;
3774+ goto cleanup;
3775+ }
3776+ }
3777+
3778+ if (chunk_id) {
3779+ *chunk_id = sqlite3_column_int64(p->stmtRowidsGetChunkPosition, 1);
3780+ }
3781+ if (chunk_offset) {
3782+ *chunk_offset = sqlite3_column_int64(p->stmtRowidsGetChunkPosition, 2);
3783+ }
3784+
3785+ rc = SQLITE_OK;
3786+
3787+cleanup:
3788+ sqlite3_reset(p->stmtRowidsGetChunkPosition);
3789+ sqlite3_clear_bindings(p->stmtRowidsGetChunkPosition);
3790+ return rc;
3791+}
3792+
3793+/**
3794+ * @brief Return the id value from the _rowids table where _rowids.rowid =
3795+ * rowid.
3796+ *
3797+ * @param pVtab: vec0 table to query
3798+ * @param rowid: rowid of the row to query.
3799+ * @param out: A dup'ed sqlite3_value of the id column. Might be null.
3800+ * Must be cleaned up with sqlite3_value_free().
3801+ * @returns SQLITE_OK on success, error code on failure
3802+ */
3803+int vec0_get_id_value_from_rowid(vec0_vtab *pVtab, i64 rowid,
3804+ sqlite3_value **out) {
3805+ // PERF: different strategy than get_chunk_position?
3806+ return vec0_get_chunk_position((vec0_vtab *)pVtab, rowid, out, NULL, NULL);
3807+}
3808+
3809+int vec0_rowid_from_id(vec0_vtab *p, sqlite3_value *valueId, i64 *rowid) {
3810+ sqlite3_stmt *stmt = NULL;
3811+ int rc;
3812+ char *zSql;
3813+ zSql = sqlite3_mprintf("SELECT rowid"
3814+ " FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE id = ?",
3815+ p->schemaName, p->tableName);
3816+ if (!zSql) {
3817+ rc = SQLITE_NOMEM;
3818+ goto cleanup;
3819+ }
3820+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
3821+ sqlite3_free(zSql);
3822+ if (rc != SQLITE_OK) {
3823+ goto cleanup;
3824+ }
3825+ sqlite3_bind_value(stmt, 1, valueId);
3826+ rc = sqlite3_step(stmt);
3827+ if (rc == SQLITE_DONE) {
3828+ rc = SQLITE_EMPTY;
3829+ goto cleanup;
3830+ }
3831+ if (rc != SQLITE_ROW) {
3832+ goto cleanup;
3833+ }
3834+ *rowid = sqlite3_column_int64(stmt, 0);
3835+ rc = sqlite3_step(stmt);
3836+ if (rc != SQLITE_DONE) {
3837+ goto cleanup;
3838+ }
3839+
3840+ rc = SQLITE_OK;
3841+
3842+cleanup:
3843+ sqlite3_finalize(stmt);
3844+ return rc;
3845+}
3846+
3847+int vec0_result_id(vec0_vtab *p, sqlite3_context *context, i64 rowid) {
3848+ if (!p->pkIsText) {
3849+ sqlite3_result_int64(context, rowid);
3850+ return SQLITE_OK;
3851+ }
3852+ sqlite3_value *valueId;
3853+ int rc = vec0_get_id_value_from_rowid(p, rowid, &valueId);
3854+ if (rc != SQLITE_OK) {
3855+ return rc;
3856+ }
3857+ if (!valueId) {
3858+ sqlite3_result_error_nomem(context);
3859+ } else {
3860+ sqlite3_result_value(context, valueId);
3861+ sqlite3_value_free(valueId);
3862+ }
3863+ return SQLITE_OK;
3864+}
3865+
3866+/**
3867+ * @brief
3868+ *
3869+ * @param pVtab: virtual table to query
3870+ * @param rowid: row to lookup
3871+ * @param vector_column_idx: which vector column to query
3872+ * @param outVector: Output pointer to the vector buffer.
3873+ * Must be sqlite3_free()'ed.
3874+ * @param outVectorSize: Pointer to a int where the size of outVector
3875+ * will be stored.
3876+ * @return int SQLITE_OK on success.
3877+ */
3878+int vec0_get_vector_data(vec0_vtab *pVtab, i64 rowid, int vector_column_idx,
3879+ void **outVector, int *outVectorSize) {
3880+ vec0_vtab *p = pVtab;
3881+ int rc, brc;
3882+ i64 chunk_id;
3883+ i64 chunk_offset;
3884+ size_t size;
3885+ void *buf = NULL;
3886+ int blobOffset;
3887+ sqlite3_blob *vectorBlob = NULL;
3888+ assert((vector_column_idx >= 0) &&
3889+ (vector_column_idx < pVtab->numVectorColumns));
3890+
3891+ rc = vec0_get_chunk_position(pVtab, rowid, NULL, &chunk_id, &chunk_offset);
3892+ if (rc == SQLITE_EMPTY) {
3893+ vtab_set_error(&pVtab->base, "Could not find a row with rowid %lld", rowid);
3894+ goto cleanup;
3895+ }
3896+ if (rc != SQLITE_OK) {
3897+ goto cleanup;
3898+ }
3899+
3900+ rc = sqlite3_blob_open(p->db, p->schemaName,
3901+ p->shadowVectorChunksNames[vector_column_idx],
3902+ "vectors", chunk_id, 0, &vectorBlob);
3903+
3904+ if (rc != SQLITE_OK) {
3905+ vtab_set_error(&pVtab->base,
3906+ "Could not fetch vector data for %lld, opening blob failed",
3907+ rowid);
3908+ rc = SQLITE_ERROR;
3909+ goto cleanup;
3910+ }
3911+
3912+ size = vector_column_byte_size(pVtab->vector_columns[vector_column_idx]);
3913+ blobOffset = chunk_offset * size;
3914+
3915+ buf = sqlite3_malloc(size);
3916+ if (!buf) {
3917+ rc = SQLITE_NOMEM;
3918+ goto cleanup;
3919+ }
3920+
3921+ rc = sqlite3_blob_read(vectorBlob, buf, size, blobOffset);
3922+ if (rc != SQLITE_OK) {
3923+ sqlite3_free(buf);
3924+ buf = NULL;
3925+ vtab_set_error(
3926+ &pVtab->base,
3927+ "Could not fetch vector data for %lld, reading from blob failed",
3928+ rowid);
3929+ rc = SQLITE_ERROR;
3930+ goto cleanup;
3931+ }
3932+
3933+ *outVector = buf;
3934+ if (outVectorSize) {
3935+ *outVectorSize = size;
3936+ }
3937+ rc = SQLITE_OK;
3938+
3939+cleanup:
3940+ brc = sqlite3_blob_close(vectorBlob);
3941+ if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) {
3942+ vtab_set_error(
3943+ &p->base, VEC_INTERAL_ERROR
3944+ "unknown error, could not close vector blob, please file an issue");
3945+ return brc;
3946+ }
3947+
3948+ return rc;
3949+}
3950+
3951+/**
3952+ * @brief Retrieve the sqlite3_value of the i'th partition value for the given row.
3953+ *
3954+ * @param pVtab - the vec0_vtab in questions
3955+ * @param rowid - rowid of target row
3956+ * @param partition_idx - which partition column to retrieve
3957+ * @param outValue - output sqlite3_value
3958+ * @return int - SQLITE_OK on success, otherwise error code
3959+ */
3960+int vec0_get_partition_value_for_rowid(vec0_vtab *pVtab, i64 rowid, int partition_idx, sqlite3_value ** outValue) {
3961+ int rc;
3962+ i64 chunk_id;
3963+ i64 chunk_offset;
3964+ rc = vec0_get_chunk_position(pVtab, rowid, NULL, &chunk_id, &chunk_offset);
3965+ if(rc != SQLITE_OK) {
3966+ return rc;
3967+ }
3968+ sqlite3_stmt * stmt = NULL;
3969+ char * zSql = sqlite3_mprintf("SELECT partition%02d FROM " VEC0_SHADOW_CHUNKS_NAME " WHERE chunk_id = ?", partition_idx, pVtab->schemaName, pVtab->tableName);
3970+ if(!zSql) {
3971+ return SQLITE_NOMEM;
3972+ }
3973+ rc = sqlite3_prepare_v2(pVtab->db, zSql, -1, &stmt, NULL);
3974+ sqlite3_free(zSql);
3975+ if(rc != SQLITE_OK) {
3976+ return rc;
3977+ }
3978+ sqlite3_bind_int64(stmt, 1, chunk_id);
3979+ rc = sqlite3_step(stmt);
3980+ if(rc != SQLITE_ROW) {
3981+ rc = SQLITE_ERROR;
3982+ goto done;
3983+ }
3984+ *outValue = sqlite3_value_dup(sqlite3_column_value(stmt, 0));
3985+ if(!*outValue) {
3986+ rc = SQLITE_NOMEM;
3987+ goto done;
3988+ }
3989+ rc = SQLITE_OK;
3990+
3991+ done:
3992+ sqlite3_finalize(stmt);
3993+ return rc;
3994+
3995+}
3996+
3997+/**
3998+ * @brief Get the value of an auxiliary column for the given rowid
3999+ *
4000+ * @param pVtab vec0_vtab
4001+ * @param rowid the rowid of the row to lookup
4002+ * @param auxiliary_idx aux index of the column we care about
4003+ * @param outValue Output sqlite3_value to store
4004+ * @return int SQLITE_OK on success, error code otherwise
4005+ */
4006+int vec0_get_auxiliary_value_for_rowid(vec0_vtab *pVtab, i64 rowid, int auxiliary_idx, sqlite3_value ** outValue) {
4007+ int rc;
4008+ sqlite3_stmt * stmt = NULL;
4009+ char * zSql = sqlite3_mprintf("SELECT value%02d FROM " VEC0_SHADOW_AUXILIARY_NAME " WHERE rowid = ?", auxiliary_idx, pVtab->schemaName, pVtab->tableName);
4010+ if(!zSql) {
4011+ return SQLITE_NOMEM;
4012+ }
4013+ rc = sqlite3_prepare_v2(pVtab->db, zSql, -1, &stmt, NULL);
4014+ sqlite3_free(zSql);
4015+ if(rc != SQLITE_OK) {
4016+ return rc;
4017+ }
4018+ sqlite3_bind_int64(stmt, 1, rowid);
4019+ rc = sqlite3_step(stmt);
4020+ if(rc != SQLITE_ROW) {
4021+ rc = SQLITE_ERROR;
4022+ goto done;
4023+ }
4024+ *outValue = sqlite3_value_dup(sqlite3_column_value(stmt, 0));
4025+ if(!*outValue) {
4026+ rc = SQLITE_NOMEM;
4027+ goto done;
4028+ }
4029+ rc = SQLITE_OK;
4030+
4031+ done:
4032+ sqlite3_finalize(stmt);
4033+ return rc;
4034+}
4035+
4036+/**
4037+ * @brief Result the given metadata value for the given row and metadata column index.
4038+ * Will traverse the metadatachunksNN table with BLOB I/0 for the given rowid.
4039+ *
4040+ * @param p
4041+ * @param rowid
4042+ * @param metadata_idx
4043+ * @param context
4044+ * @return int
4045+ */
4046+int vec0_result_metadata_value_for_rowid(vec0_vtab *p, i64 rowid, int metadata_idx, sqlite3_context * context) {
4047+ int rc;
4048+ i64 chunk_id;
4049+ i64 chunk_offset;
4050+ rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset);
4051+ if(rc != SQLITE_OK) {
4052+ return rc;
4053+ }
4054+ sqlite3_blob * blobValue;
4055+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 0, &blobValue);
4056+ if(rc != SQLITE_OK) {
4057+ return rc;
4058+ }
4059+
4060+ switch(p->metadata_columns[metadata_idx].kind) {
4061+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
4062+ u8 block;
4063+ rc = sqlite3_blob_read(blobValue, &block, sizeof(block), chunk_offset / CHAR_BIT);
4064+ if(rc != SQLITE_OK) {
4065+ goto done;
4066+ }
4067+ int value = block >> ((chunk_offset % CHAR_BIT)) & 1;
4068+ sqlite3_result_int(context, value);
4069+ break;
4070+ }
4071+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
4072+ i64 value;
4073+ rc = sqlite3_blob_read(blobValue, &value, sizeof(value), chunk_offset * sizeof(i64));
4074+ if(rc != SQLITE_OK) {
4075+ goto done;
4076+ }
4077+ sqlite3_result_int64(context, value);
4078+ break;
4079+ }
4080+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
4081+ double value;
4082+ rc = sqlite3_blob_read(blobValue, &value, sizeof(value), chunk_offset * sizeof(double));
4083+ if(rc != SQLITE_OK) {
4084+ goto done;
4085+ }
4086+ sqlite3_result_double(context, value);
4087+ break;
4088+ }
4089+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
4090+ u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
4091+ rc = sqlite3_blob_read(blobValue, &view, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH, chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
4092+ if(rc != SQLITE_OK) {
4093+ goto done;
4094+ }
4095+ int length = ((int *)view)[0];
4096+ if(length <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
4097+ sqlite3_result_text(context, (const char*) (view + 4), length, SQLITE_TRANSIENT);
4098+ }
4099+ else {
4100+ sqlite3_stmt * stmt;
4101+ const char * zSql = sqlite3_mprintf("SELECT data FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_idx);
4102+ if(!zSql) {
4103+ rc = SQLITE_ERROR;
4104+ goto done;
4105+ }
4106+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
4107+ sqlite3_free((void *) zSql);
4108+ if(rc != SQLITE_OK) {
4109+ goto done;
4110+ }
4111+ sqlite3_bind_int64(stmt, 1, rowid);
4112+ rc = sqlite3_step(stmt);
4113+ if(rc != SQLITE_ROW) {
4114+ sqlite3_finalize(stmt);
4115+ rc = SQLITE_ERROR;
4116+ goto done;
4117+ }
4118+ sqlite3_result_value(context, sqlite3_column_value(stmt, 0));
4119+ sqlite3_finalize(stmt);
4120+ rc = SQLITE_OK;
4121+ }
4122+ break;
4123+ }
4124+ }
4125+ done:
4126+ // blobValue is read-only, will not fail on close
4127+ sqlite3_blob_close(blobValue);
4128+ return rc;
4129+
4130+}
4131+
4132+int vec0_get_latest_chunk_rowid(vec0_vtab *p, i64 *chunk_rowid, sqlite3_value ** partitionKeyValues) {
4133+ int rc;
4134+ const char *zSql;
4135+ // lazy initialize stmtLatestChunk when needed. May be cleared during xSync()
4136+ if (!p->stmtLatestChunk) {
4137+ if(p->numPartitionColumns > 0) {
4138+ sqlite3_str * s = sqlite3_str_new(NULL);
4139+ sqlite3_str_appendf(s, "SELECT max(rowid) FROM " VEC0_SHADOW_CHUNKS_NAME " WHERE ",
4140+ p->schemaName, p->tableName);
4141+
4142+ for(int i = 0; i < p->numPartitionColumns; i++) {
4143+ if(i != 0) {
4144+ sqlite3_str_appendall(s, " AND ");
4145+ }
4146+ sqlite3_str_appendf(s, " partition%02d = ? ", i);
4147+ }
4148+ zSql = sqlite3_str_finish(s);
4149+ }else {
4150+ zSql = sqlite3_mprintf("SELECT max(rowid) FROM " VEC0_SHADOW_CHUNKS_NAME,
4151+ p->schemaName, p->tableName);
4152+ }
4153+
4154+ if (!zSql) {
4155+ rc = SQLITE_NOMEM;
4156+ goto cleanup;
4157+ }
4158+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtLatestChunk, 0);
4159+ sqlite3_free((void *)zSql);
4160+ if (rc != SQLITE_OK) {
4161+ // IMP: V21406_05476
4162+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
4163+ "could not initialize 'latest chunk' statement");
4164+ goto cleanup;
4165+ }
4166+ }
4167+
4168+ for(int i = 0; i < p->numPartitionColumns; i++) {
4169+ sqlite3_bind_value(p->stmtLatestChunk, i+1, (partitionKeyValues[i]));
4170+ }
4171+
4172+ rc = sqlite3_step(p->stmtLatestChunk);
4173+ if (rc != SQLITE_ROW) {
4174+ // IMP: V31559_15629
4175+ vtab_set_error(&p->base, VEC_INTERAL_ERROR "Could not find latest chunk");
4176+ rc = SQLITE_ERROR;
4177+ goto cleanup;
4178+ }
4179+ if(sqlite3_column_type(p->stmtLatestChunk, 0) == SQLITE_NULL){
4180+ rc = SQLITE_EMPTY;
4181+ goto cleanup;
4182+ }
4183+ *chunk_rowid = sqlite3_column_int64(p->stmtLatestChunk, 0);
4184+ rc = sqlite3_step(p->stmtLatestChunk);
4185+ if (rc != SQLITE_DONE) {
4186+ vtab_set_error(&p->base,
4187+ VEC_INTERAL_ERROR
4188+ "unknown result code when closing out stmtLatestChunk. "
4189+ "Please file an issue: " REPORT_URL,
4190+ p->schemaName, p->shadowChunksName);
4191+ goto cleanup;
4192+ }
4193+ rc = SQLITE_OK;
4194+
4195+cleanup:
4196+ if (p->stmtLatestChunk) {
4197+ sqlite3_reset(p->stmtLatestChunk);
4198+ sqlite3_clear_bindings(p->stmtLatestChunk);
4199+ }
4200+ return rc;
4201+}
4202+
4203+int vec0_rowids_insert_rowid(vec0_vtab *p, i64 rowid) {
4204+ int rc = SQLITE_OK;
4205+ int entered = 0;
4206+ UNUSED_PARAMETER(entered); // temporary
4207+ if (!p->stmtRowidsInsertRowid) {
4208+ const char *zSql =
4209+ sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_ROWIDS_NAME "(rowid)"
4210+ "VALUES (?);",
4211+ p->schemaName, p->tableName);
4212+ if (!zSql) {
4213+ rc = SQLITE_NOMEM;
4214+ goto cleanup;
4215+ }
4216+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsInsertRowid, 0);
4217+ sqlite3_free((void *)zSql);
4218+ if (rc != SQLITE_OK) {
4219+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
4220+ "could not initialize 'insert rowids' statement");
4221+ goto cleanup;
4222+ }
4223+ }
4224+
4225+#if SQLITE_THREADSAFE
4226+ if (sqlite3_mutex_enter) {
4227+ sqlite3_mutex_enter(sqlite3_db_mutex(p->db));
4228+ entered = 1;
4229+ }
4230+#endif
4231+ sqlite3_bind_int64(p->stmtRowidsInsertRowid, 1, rowid);
4232+ rc = sqlite3_step(p->stmtRowidsInsertRowid);
4233+
4234+ if (rc != SQLITE_DONE) {
4235+ if (sqlite3_extended_errcode(p->db) == SQLITE_CONSTRAINT_PRIMARYKEY) {
4236+ // IMP: V17090_01160
4237+ vtab_set_error(&p->base, "UNIQUE constraint failed on %s primary key",
4238+ p->tableName);
4239+ } else {
4240+ // IMP: V04679_21517
4241+ vtab_set_error(&p->base,
4242+ "Error inserting rowid into rowids shadow table: %s",
4243+ sqlite3_errmsg(sqlite3_db_handle(p->stmtRowidsInsertId)));
4244+ }
4245+ rc = SQLITE_ERROR;
4246+ goto cleanup;
4247+ }
4248+
4249+ rc = SQLITE_OK;
4250+
4251+cleanup:
4252+ if (p->stmtRowidsInsertRowid) {
4253+ sqlite3_reset(p->stmtRowidsInsertRowid);
4254+ sqlite3_clear_bindings(p->stmtRowidsInsertRowid);
4255+ }
4256+
4257+#if SQLITE_THREADSAFE
4258+ if (sqlite3_mutex_leave && entered) {
4259+ sqlite3_mutex_leave(sqlite3_db_mutex(p->db));
4260+ }
4261+#endif
4262+ return rc;
4263+}
4264+
4265+int vec0_rowids_insert_id(vec0_vtab *p, sqlite3_value *idValue, i64 *rowid) {
4266+ int rc = SQLITE_OK;
4267+ int entered = 0;
4268+ UNUSED_PARAMETER(entered); // temporary
4269+ if (!p->stmtRowidsInsertId) {
4270+ const char *zSql =
4271+ sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_ROWIDS_NAME "(id)"
4272+ "VALUES (?);",
4273+ p->schemaName, p->tableName);
4274+ if (!zSql) {
4275+ rc = SQLITE_NOMEM;
4276+ goto complete;
4277+ }
4278+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsInsertId, 0);
4279+ sqlite3_free((void *)zSql);
4280+ if (rc != SQLITE_OK) {
4281+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
4282+ "could not initialize 'insert rowids id' statement");
4283+ goto complete;
4284+ }
4285+ }
4286+
4287+#if SQLITE_THREADSAFE
4288+ if (sqlite3_mutex_enter) {
4289+ sqlite3_mutex_enter(sqlite3_db_mutex(p->db));
4290+ entered = 1;
4291+ }
4292+#endif
4293+
4294+ if (idValue) {
4295+ sqlite3_bind_value(p->stmtRowidsInsertId, 1, idValue);
4296+ }
4297+ rc = sqlite3_step(p->stmtRowidsInsertId);
4298+
4299+ if (rc != SQLITE_DONE) {
4300+ if (sqlite3_extended_errcode(p->db) == SQLITE_CONSTRAINT_UNIQUE) {
4301+ // IMP: V20497_04568
4302+ vtab_set_error(&p->base, "UNIQUE constraint failed on %s primary key",
4303+ p->tableName);
4304+ } else {
4305+ // IMP: V24016_08086
4306+ // IMP: V15177_32015
4307+ vtab_set_error(&p->base,
4308+ "Error inserting id into rowids shadow table: %s",
4309+ sqlite3_errmsg(sqlite3_db_handle(p->stmtRowidsInsertId)));
4310+ }
4311+ rc = SQLITE_ERROR;
4312+ goto complete;
4313+ }
4314+
4315+ *rowid = sqlite3_last_insert_rowid(p->db);
4316+ rc = SQLITE_OK;
4317+
4318+complete:
4319+ if (p->stmtRowidsInsertId) {
4320+ sqlite3_reset(p->stmtRowidsInsertId);
4321+ sqlite3_clear_bindings(p->stmtRowidsInsertId);
4322+ }
4323+
4324+#if SQLITE_THREADSAFE
4325+ if (sqlite3_mutex_leave && entered) {
4326+ sqlite3_mutex_leave(sqlite3_db_mutex(p->db));
4327+ }
4328+#endif
4329+ return rc;
4330+}
4331+
4332+int vec0_metadata_chunk_size(vec0_metadata_column_kind kind, int chunk_size) {
4333+ switch(kind) {
4334+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN:
4335+ return chunk_size / 8;
4336+ case VEC0_METADATA_COLUMN_KIND_INTEGER:
4337+ return chunk_size * sizeof(i64);
4338+ case VEC0_METADATA_COLUMN_KIND_FLOAT:
4339+ return chunk_size * sizeof(double);
4340+ case VEC0_METADATA_COLUMN_KIND_TEXT:
4341+ return chunk_size * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH;
4342+ }
4343+ return 0;
4344+}
4345+
4346+int vec0_rowids_update_position(vec0_vtab *p, i64 rowid, i64 chunk_rowid,
4347+ i64 chunk_offset) {
4348+ int rc = SQLITE_OK;
4349+
4350+ if (!p->stmtRowidsUpdatePosition) {
4351+ const char *zSql = sqlite3_mprintf(" UPDATE " VEC0_SHADOW_ROWIDS_NAME
4352+ " SET chunk_id = ?, chunk_offset = ?"
4353+ " WHERE rowid = ?",
4354+ p->schemaName, p->tableName);
4355+ if (!zSql) {
4356+ rc = SQLITE_NOMEM;
4357+ goto cleanup;
4358+ }
4359+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsUpdatePosition, 0);
4360+ sqlite3_free((void *)zSql);
4361+ if (rc != SQLITE_OK) {
4362+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
4363+ "could not initialize 'update rowids position' statement");
4364+ goto cleanup;
4365+ }
4366+ }
4367+
4368+ sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 1, chunk_rowid);
4369+ sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 2, chunk_offset);
4370+ sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 3, rowid);
4371+
4372+ rc = sqlite3_step(p->stmtRowidsUpdatePosition);
4373+ if (rc != SQLITE_DONE) {
4374+ // IMP: V21925_05995
4375+ vtab_set_error(&p->base,
4376+ VEC_INTERAL_ERROR
4377+ "could not update rowids position for rowid=%lld, "
4378+ "chunk_rowid=%lld, chunk_offset=%lld",
4379+ rowid, chunk_rowid, chunk_offset);
4380+ rc = SQLITE_ERROR;
4381+ goto cleanup;
4382+ }
4383+ rc = SQLITE_OK;
4384+
4385+cleanup:
4386+ if (p->stmtRowidsUpdatePosition) {
4387+ sqlite3_reset(p->stmtRowidsUpdatePosition);
4388+ sqlite3_clear_bindings(p->stmtRowidsUpdatePosition);
4389+ }
4390+
4391+ return rc;
4392+}
4393+
4394+/**
4395+ * @brief Adds a new chunk for the vec0 table, and the corresponding vector
4396+ * chunks.
4397+ *
4398+ * Inserts a new row into the _chunks table, with blank data, and uses that new
4399+ * rowid to insert new blank rows into _vector_chunksXX tables.
4400+ *
4401+ * @param p: vec0 table to add new chunk
4402+ * @param paritionKeyValues: Array of partition key valeus for the new chunk, if available
4403+ * @param chunk_rowid: Output pointer, if not NULL, then will be filled with the
4404+ * new chunk rowid.
4405+ * @return int SQLITE_OK on success, error code otherwise.
4406+ */
4407+int vec0_new_chunk(vec0_vtab *p, sqlite3_value ** partitionKeyValues, i64 *chunk_rowid) {
4408+ int rc;
4409+ char *zSql;
4410+ sqlite3_stmt *stmt;
4411+ i64 rowid;
4412+
4413+ // Step 1: Insert a new row in _chunks, capture that new rowid
4414+ if(p->numPartitionColumns > 0) {
4415+ sqlite3_str * s = sqlite3_str_new(NULL);
4416+ sqlite3_str_appendf(s, "INSERT INTO " VEC0_SHADOW_CHUNKS_NAME, p->schemaName, p->tableName);
4417+ sqlite3_str_appendall(s, "(size, validity, rowids");
4418+ for(int i = 0; i < p->numPartitionColumns; i++) {
4419+ sqlite3_str_appendf(s, ", partition%02d", i);
4420+ }
4421+ sqlite3_str_appendall(s, ") VALUES (?, ?, ?");
4422+ for(int i = 0; i < p->numPartitionColumns; i++) {
4423+ sqlite3_str_appendall(s, ", ?");
4424+ }
4425+ sqlite3_str_appendall(s, ")");
4426+
4427+ zSql = sqlite3_str_finish(s);
4428+ }else {
4429+ zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_CHUNKS_NAME
4430+ "(size, validity, rowids) "
4431+ "VALUES (?, ?, ?);",
4432+ p->schemaName, p->tableName);
4433+ }
4434+
4435+ if (!zSql) {
4436+ return SQLITE_NOMEM;
4437+ }
4438+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
4439+ sqlite3_free(zSql);
4440+ if (rc != SQLITE_OK) {
4441+ sqlite3_finalize(stmt);
4442+ return rc;
4443+ }
4444+
4445+#if SQLITE_THREADSAFE
4446+ if (sqlite3_mutex_enter) {
4447+ sqlite3_mutex_enter(sqlite3_db_mutex(p->db));
4448+ }
4449+#endif
4450+
4451+ sqlite3_bind_int64(stmt, 1, p->chunk_size); // size
4452+ sqlite3_bind_zeroblob(stmt, 2, p->chunk_size / CHAR_BIT); // validity bitmap
4453+ sqlite3_bind_zeroblob(stmt, 3, p->chunk_size * sizeof(i64)); // rowids
4454+
4455+ for(int i = 0; i < p->numPartitionColumns; i++) {
4456+ sqlite3_bind_value(stmt, 4 + i, partitionKeyValues[i]);
4457+ }
4458+
4459+ rc = sqlite3_step(stmt);
4460+ int failed = rc != SQLITE_DONE;
4461+ rowid = sqlite3_last_insert_rowid(p->db);
4462+#if SQLITE_THREADSAFE
4463+ if (sqlite3_mutex_leave) {
4464+ sqlite3_mutex_leave(sqlite3_db_mutex(p->db));
4465+ }
4466+#endif
4467+ sqlite3_finalize(stmt);
4468+ if (failed) {
4469+ return SQLITE_ERROR;
4470+ }
4471+
4472+ // Step 2: Create new vector chunks for each vector column, with
4473+ // that new chunk_rowid.
4474+
4475+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
4476+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) {
4477+ continue;
4478+ }
4479+ int vector_column_idx = p->user_column_idxs[i];
4480+ i64 vectorsSize =
4481+ p->chunk_size * vector_column_byte_size(p->vector_columns[vector_column_idx]);
4482+
4483+ zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_VECTOR_N_NAME
4484+ "(rowid, vectors)"
4485+ "VALUES (?, ?)",
4486+ p->schemaName, p->tableName, vector_column_idx);
4487+ if (!zSql) {
4488+ return SQLITE_NOMEM;
4489+ }
4490+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
4491+ sqlite3_free(zSql);
4492+
4493+ if (rc != SQLITE_OK) {
4494+ sqlite3_finalize(stmt);
4495+ return rc;
4496+ }
4497+
4498+ sqlite3_bind_int64(stmt, 1, rowid);
4499+ sqlite3_bind_zeroblob64(stmt, 2, vectorsSize);
4500+
4501+ rc = sqlite3_step(stmt);
4502+ sqlite3_finalize(stmt);
4503+ if (rc != SQLITE_DONE) {
4504+ return rc;
4505+ }
4506+ }
4507+
4508+ // Step 3: Create new metadata chunks for each metadata column
4509+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
4510+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) {
4511+ continue;
4512+ }
4513+ int metadata_column_idx = p->user_column_idxs[i];
4514+ zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_METADATA_N_NAME
4515+ "(rowid, data)"
4516+ "VALUES (?, ?)",
4517+ p->schemaName, p->tableName, metadata_column_idx);
4518+ if (!zSql) {
4519+ return SQLITE_NOMEM;
4520+ }
4521+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
4522+ sqlite3_free(zSql);
4523+
4524+ if (rc != SQLITE_OK) {
4525+ sqlite3_finalize(stmt);
4526+ return rc;
4527+ }
4528+
4529+ sqlite3_bind_int64(stmt, 1, rowid);
4530+ sqlite3_bind_zeroblob64(stmt, 2, vec0_metadata_chunk_size(p->metadata_columns[metadata_column_idx].kind, p->chunk_size));
4531+
4532+ rc = sqlite3_step(stmt);
4533+ sqlite3_finalize(stmt);
4534+ if (rc != SQLITE_DONE) {
4535+ return rc;
4536+ }
4537+ }
4538+
4539+
4540+ if (chunk_rowid) {
4541+ *chunk_rowid = rowid;
4542+ }
4543+
4544+ return SQLITE_OK;
4545+}
4546+
4547+struct vec0_query_fullscan_data {
4548+ sqlite3_stmt *rowids_stmt;
4549+ i8 done;
4550+};
4551+void vec0_query_fullscan_data_clear(
4552+ struct vec0_query_fullscan_data *fullscan_data) {
4553+ if (!fullscan_data)
4554+ return;
4555+
4556+ if (fullscan_data->rowids_stmt) {
4557+ sqlite3_finalize(fullscan_data->rowids_stmt);
4558+ fullscan_data->rowids_stmt = NULL;
4559+ }
4560+}
4561+
4562+struct vec0_query_knn_data {
4563+ i64 k;
4564+ i64 k_used;
4565+ // Array of rowids of size k. Must be freed with sqlite3_free().
4566+ i64 *rowids;
4567+ // Array of distances of size k. Must be freed with sqlite3_free().
4568+ f32 *distances;
4569+ i64 current_idx;
4570+};
4571+void vec0_query_knn_data_clear(struct vec0_query_knn_data *knn_data) {
4572+ if (!knn_data)
4573+ return;
4574+
4575+ if (knn_data->rowids) {
4576+ sqlite3_free(knn_data->rowids);
4577+ knn_data->rowids = NULL;
4578+ }
4579+ if (knn_data->distances) {
4580+ sqlite3_free(knn_data->distances);
4581+ knn_data->distances = NULL;
4582+ }
4583+}
4584+
4585+struct vec0_query_point_data {
4586+ i64 rowid;
4587+ void *vectors[VEC0_MAX_VECTOR_COLUMNS];
4588+ int done;
4589+};
4590+void vec0_query_point_data_clear(struct vec0_query_point_data *point_data) {
4591+ if (!point_data)
4592+ return;
4593+ for (int i = 0; i < VEC0_MAX_VECTOR_COLUMNS; i++) {
4594+ sqlite3_free(point_data->vectors[i]);
4595+ point_data->vectors[i] = NULL;
4596+ }
4597+}
4598+
4599+typedef enum {
4600+ // If any values are updated, please update the ARCHITECTURE.md docs accordingly!
4601+
4602+ VEC0_QUERY_PLAN_FULLSCAN = '1',
4603+ VEC0_QUERY_PLAN_POINT = '2',
4604+ VEC0_QUERY_PLAN_KNN = '3',
4605+} vec0_query_plan;
4606+
4607+typedef struct vec0_cursor vec0_cursor;
4608+struct vec0_cursor {
4609+ sqlite3_vtab_cursor base;
4610+
4611+ vec0_query_plan query_plan;
4612+ struct vec0_query_fullscan_data *fullscan_data;
4613+ struct vec0_query_knn_data *knn_data;
4614+ struct vec0_query_point_data *point_data;
4615+};
4616+
4617+void vec0_cursor_clear(vec0_cursor *pCur) {
4618+ if (pCur->fullscan_data) {
4619+ vec0_query_fullscan_data_clear(pCur->fullscan_data);
4620+ sqlite3_free(pCur->fullscan_data);
4621+ pCur->fullscan_data = NULL;
4622+ }
4623+ if (pCur->knn_data) {
4624+ vec0_query_knn_data_clear(pCur->knn_data);
4625+ sqlite3_free(pCur->knn_data);
4626+ pCur->knn_data = NULL;
4627+ }
4628+ if (pCur->point_data) {
4629+ vec0_query_point_data_clear(pCur->point_data);
4630+ sqlite3_free(pCur->point_data);
4631+ pCur->point_data = NULL;
4632+ }
4633+}
4634+
4635+#define VEC_CONSTRUCTOR_ERROR "vec0 constructor error: "
4636+static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv,
4637+ sqlite3_vtab **ppVtab, char **pzErr, bool isCreate) {
4638+ UNUSED_PARAMETER(pAux);
4639+ vec0_vtab *pNew;
4640+ int rc;
4641+ const char *zSql;
4642+
4643+ pNew = sqlite3_malloc(sizeof(*pNew));
4644+ if (pNew == 0)
4645+ return SQLITE_NOMEM;
4646+ memset(pNew, 0, sizeof(*pNew));
4647+
4648+ // Declared chunk_size=N for entire table.
4649+ // -1 to use the defualt, otherwise will get re-assigned on `chunk_size=N`
4650+ // option
4651+ int chunk_size = -1;
4652+ int numVectorColumns = 0;
4653+ int numPartitionColumns = 0;
4654+ int numAuxiliaryColumns = 0;
4655+ int numMetadataColumns = 0;
4656+ int user_column_idx = 0;
4657+
4658+ // track if a "primary key" column is defined
4659+ char *pkColumnName = NULL;
4660+ int pkColumnNameLength;
4661+ int pkColumnType = SQLITE_INTEGER;
4662+
4663+ for (int i = 3; i < argc; i++) {
4664+ struct VectorColumnDefinition vecColumn;
4665+ struct Vec0PartitionColumnDefinition partitionColumn;
4666+ struct Vec0AuxiliaryColumnDefinition auxColumn;
4667+ struct Vec0MetadataColumnDefinition metadataColumn;
4668+ char *cName = NULL;
4669+ int cNameLength;
4670+ int cType;
4671+
4672+ // Scenario #1: Constructor argument is a vector column definition, ie `foo float[1024]`
4673+ rc = vec0_parse_vector_column(argv[i], strlen(argv[i]), &vecColumn);
4674+ if (rc == SQLITE_ERROR) {
4675+ *pzErr = sqlite3_mprintf(
4676+ VEC_CONSTRUCTOR_ERROR "could not parse vector column '%s'", argv[i]);
4677+ goto error;
4678+ }
4679+ if (rc == SQLITE_OK) {
4680+ if (numVectorColumns >= VEC0_MAX_VECTOR_COLUMNS) {
4681+ sqlite3_free(vecColumn.name);
4682+ *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4683+ "Too many provided vector columns, maximum %d",
4684+ VEC0_MAX_VECTOR_COLUMNS);
4685+ goto error;
4686+ }
4687+
4688+ if (vecColumn.dimensions > SQLITE_VEC_VEC0_MAX_DIMENSIONS) {
4689+ sqlite3_free(vecColumn.name);
4690+ *pzErr = sqlite3_mprintf(
4691+ VEC_CONSTRUCTOR_ERROR
4692+ "Dimension on vector column too large, provided %lld, maximum %lld",
4693+ (i64)vecColumn.dimensions, SQLITE_VEC_VEC0_MAX_DIMENSIONS);
4694+ goto error;
4695+ }
4696+ pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_VECTOR;
4697+ pNew->user_column_idxs[user_column_idx] = numVectorColumns;
4698+ memcpy(&pNew->vector_columns[numVectorColumns], &vecColumn, sizeof(vecColumn));
4699+ numVectorColumns++;
4700+ user_column_idx++;
4701+
4702+ continue;
4703+ }
4704+
4705+ // Scenario #2: Constructor argument is a partition key column definition, ie `user_id text partition key`
4706+ rc = vec0_parse_partition_key_definition(argv[i], strlen(argv[i]), &cName,
4707+ &cNameLength, &cType);
4708+ if (rc == SQLITE_OK) {
4709+ if (numPartitionColumns >= VEC0_MAX_PARTITION_COLUMNS) {
4710+ *pzErr = sqlite3_mprintf(
4711+ VEC_CONSTRUCTOR_ERROR
4712+ "More than %d partition key columns were provided",
4713+ VEC0_MAX_PARTITION_COLUMNS);
4714+ goto error;
4715+ }
4716+ partitionColumn.type = cType;
4717+ partitionColumn.name_length = cNameLength;
4718+ partitionColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName);
4719+ if(!partitionColumn.name) {
4720+ rc = SQLITE_NOMEM;
4721+ goto error;
4722+ }
4723+
4724+ pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_PARTITION;
4725+ pNew->user_column_idxs[user_column_idx] = numPartitionColumns;
4726+ memcpy(&pNew->paritition_columns[numPartitionColumns], &partitionColumn, sizeof(partitionColumn));
4727+ numPartitionColumns++;
4728+ user_column_idx++;
4729+ continue;
4730+ }
4731+
4732+ // Scenario #3: Constructor argument is a primary key column definition, ie `article_id text primary key`
4733+ rc = vec0_parse_primary_key_definition(argv[i], strlen(argv[i]), &cName,
4734+ &cNameLength, &cType);
4735+ if (rc == SQLITE_OK) {
4736+ if (pkColumnName) {
4737+ *pzErr = sqlite3_mprintf(
4738+ VEC_CONSTRUCTOR_ERROR
4739+ "More than one primary key definition was provided, vec0 only "
4740+ "suports a single primary key column",
4741+ argv[i]);
4742+ goto error;
4743+ }
4744+ pkColumnName = cName;
4745+ pkColumnNameLength = cNameLength;
4746+ pkColumnType = cType;
4747+ continue;
4748+ }
4749+
4750+ // Scenario #4: Constructor argument is a auxiliary column definition, ie `+contents text`
4751+ rc = vec0_parse_auxiliary_column_definition(argv[i], strlen(argv[i]), &cName,
4752+ &cNameLength, &cType);
4753+ if(rc == SQLITE_OK) {
4754+ if (numAuxiliaryColumns >= VEC0_MAX_AUXILIARY_COLUMNS) {
4755+ *pzErr = sqlite3_mprintf(
4756+ VEC_CONSTRUCTOR_ERROR
4757+ "More than %d auxiliary columns were provided",
4758+ VEC0_MAX_AUXILIARY_COLUMNS);
4759+ goto error;
4760+ }
4761+ auxColumn.type = cType;
4762+ auxColumn.name_length = cNameLength;
4763+ auxColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName);
4764+ if(!auxColumn.name) {
4765+ rc = SQLITE_NOMEM;
4766+ goto error;
4767+ }
4768+
4769+ pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY;
4770+ pNew->user_column_idxs[user_column_idx] = numAuxiliaryColumns;
4771+ memcpy(&pNew->auxiliary_columns[numAuxiliaryColumns], &auxColumn, sizeof(auxColumn));
4772+ numAuxiliaryColumns++;
4773+ user_column_idx++;
4774+ continue;
4775+ }
4776+
4777+ vec0_metadata_column_kind kind;
4778+ rc = vec0_parse_metadata_column_definition(argv[i], strlen(argv[i]), &cName,
4779+ &cNameLength, &kind);
4780+ if(rc == SQLITE_OK) {
4781+ if (numMetadataColumns >= VEC0_MAX_METADATA_COLUMNS) {
4782+ *pzErr = sqlite3_mprintf(
4783+ VEC_CONSTRUCTOR_ERROR
4784+ "More than %d metadata columns were provided",
4785+ VEC0_MAX_METADATA_COLUMNS);
4786+ goto error;
4787+ }
4788+ metadataColumn.kind = kind;
4789+ metadataColumn.name_length = cNameLength;
4790+ metadataColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName);
4791+ if(!metadataColumn.name) {
4792+ rc = SQLITE_NOMEM;
4793+ goto error;
4794+ }
4795+
4796+ pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_METADATA;
4797+ pNew->user_column_idxs[user_column_idx] = numMetadataColumns;
4798+ memcpy(&pNew->metadata_columns[numMetadataColumns], &metadataColumn, sizeof(metadataColumn));
4799+ numMetadataColumns++;
4800+ user_column_idx++;
4801+ continue;
4802+ }
4803+
4804+ // Scenario #4: Constructor argument is a table-level option, ie `chunk_size`
4805+
4806+ char *key;
4807+ char *value;
4808+ int keyLength, valueLength;
4809+ rc = vec0_parse_table_option(argv[i], strlen(argv[i]), &key, &keyLength,
4810+ &value, &valueLength);
4811+ if (rc == SQLITE_ERROR) {
4812+ *pzErr = sqlite3_mprintf(
4813+ VEC_CONSTRUCTOR_ERROR "could not parse table option '%s'", argv[i]);
4814+ goto error;
4815+ }
4816+ if (rc == SQLITE_OK) {
4817+ if (sqlite3_strnicmp(key, "chunk_size", keyLength) == 0) {
4818+ chunk_size = atoi(value);
4819+ if (chunk_size <= 0) {
4820+ // IMP: V01931_18769
4821+ *pzErr =
4822+ sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4823+ "chunk_size must be a non-zero positive integer");
4824+ goto error;
4825+ }
4826+ if ((chunk_size % 8) != 0) {
4827+ // IMP: V14110_30948
4828+ *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4829+ "chunk_size must be divisible by 8");
4830+ goto error;
4831+ }
4832+#define SQLITE_VEC_CHUNK_SIZE_MAX 4096
4833+ if (chunk_size > SQLITE_VEC_CHUNK_SIZE_MAX) {
4834+ *pzErr =
4835+ sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR "chunk_size too large");
4836+ goto error;
4837+ }
4838+ } else {
4839+ // IMP: V27642_11712
4840+ *pzErr = sqlite3_mprintf(
4841+ VEC_CONSTRUCTOR_ERROR "Unknown table option: %.*s", keyLength, key);
4842+ goto error;
4843+ }
4844+ continue;
4845+ }
4846+
4847+ // Scenario #5: Unknown constructor argument
4848+ *pzErr =
4849+ sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR "Could not parse '%s'", argv[i]);
4850+ goto error;
4851+ }
4852+
4853+ if (chunk_size < 0) {
4854+ chunk_size = 1024;
4855+ }
4856+
4857+ if (numVectorColumns <= 0) {
4858+ *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4859+ "At least one vector column is required");
4860+ goto error;
4861+ }
4862+
4863+ sqlite3_str *createStr = sqlite3_str_new(NULL);
4864+ sqlite3_str_appendall(createStr, "CREATE TABLE x(");
4865+ if (pkColumnName) {
4866+ sqlite3_str_appendf(createStr, "\"%.*w\" primary key, ", pkColumnNameLength,
4867+ pkColumnName);
4868+ } else {
4869+ sqlite3_str_appendall(createStr, "rowid, ");
4870+ }
4871+ for (int i = 0; i < numVectorColumns + numPartitionColumns + numAuxiliaryColumns + numMetadataColumns; i++) {
4872+ switch(pNew->user_column_kinds[i]) {
4873+ case SQLITE_VEC0_USER_COLUMN_KIND_VECTOR: {
4874+ int vector_idx = pNew->user_column_idxs[i];
4875+ sqlite3_str_appendf(createStr, "\"%.*w\", ",
4876+ pNew->vector_columns[vector_idx].name_length,
4877+ pNew->vector_columns[vector_idx].name);
4878+ break;
4879+ }
4880+ case SQLITE_VEC0_USER_COLUMN_KIND_PARTITION: {
4881+ int partition_idx = pNew->user_column_idxs[i];
4882+ sqlite3_str_appendf(createStr, "\"%.*w\", ",
4883+ pNew->paritition_columns[partition_idx].name_length,
4884+ pNew->paritition_columns[partition_idx].name);
4885+ break;
4886+ }
4887+ case SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY: {
4888+ int auxiliary_idx = pNew->user_column_idxs[i];
4889+ sqlite3_str_appendf(createStr, "\"%.*w\", ",
4890+ pNew->auxiliary_columns[auxiliary_idx].name_length,
4891+ pNew->auxiliary_columns[auxiliary_idx].name);
4892+ break;
4893+ }
4894+ case SQLITE_VEC0_USER_COLUMN_KIND_METADATA: {
4895+ int metadata_idx = pNew->user_column_idxs[i];
4896+ sqlite3_str_appendf(createStr, "\"%.*w\", ",
4897+ pNew->metadata_columns[metadata_idx].name_length,
4898+ pNew->metadata_columns[metadata_idx].name);
4899+ break;
4900+ }
4901+ }
4902+
4903+ }
4904+ sqlite3_str_appendall(createStr, " distance hidden, k hidden) ");
4905+ if (pkColumnName) {
4906+ sqlite3_str_appendall(createStr, "without rowid ");
4907+ }
4908+ zSql = sqlite3_str_finish(createStr);
4909+ if (!zSql) {
4910+ goto error;
4911+ }
4912+ rc = sqlite3_declare_vtab(db, zSql);
4913+ sqlite3_free((void *)zSql);
4914+ if (rc != SQLITE_OK) {
4915+ *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4916+ "could not declare virtual table, '%s'",
4917+ sqlite3_errmsg(db));
4918+ goto error;
4919+ }
4920+
4921+ const char *schemaName = argv[1];
4922+ const char *tableName = argv[2];
4923+
4924+ pNew->db = db;
4925+ pNew->pkIsText = pkColumnType == SQLITE_TEXT;
4926+ pNew->schemaName = sqlite3_mprintf("%s", schemaName);
4927+ if (!pNew->schemaName) {
4928+ goto error;
4929+ }
4930+ pNew->tableName = sqlite3_mprintf("%s", tableName);
4931+ if (!pNew->tableName) {
4932+ goto error;
4933+ }
4934+ pNew->shadowRowidsName = sqlite3_mprintf("%s_rowids", tableName);
4935+ if (!pNew->shadowRowidsName) {
4936+ goto error;
4937+ }
4938+ pNew->shadowChunksName = sqlite3_mprintf("%s_chunks", tableName);
4939+ if (!pNew->shadowChunksName) {
4940+ goto error;
4941+ }
4942+ pNew->numVectorColumns = numVectorColumns;
4943+ pNew->numPartitionColumns = numPartitionColumns;
4944+ pNew->numAuxiliaryColumns = numAuxiliaryColumns;
4945+ pNew->numMetadataColumns = numMetadataColumns;
4946+
4947+ for (int i = 0; i < pNew->numVectorColumns; i++) {
4948+ pNew->shadowVectorChunksNames[i] =
4949+ sqlite3_mprintf("%s_vector_chunks%02d", tableName, i);
4950+ if (!pNew->shadowVectorChunksNames[i]) {
4951+ goto error;
4952+ }
4953+ }
4954+ for (int i = 0; i < pNew->numMetadataColumns; i++) {
4955+ pNew->shadowMetadataChunksNames[i] =
4956+ sqlite3_mprintf("%s_metadatachunks%02d", tableName, i);
4957+ if (!pNew->shadowMetadataChunksNames[i]) {
4958+ goto error;
4959+ }
4960+ }
4961+ pNew->chunk_size = chunk_size;
4962+
4963+ // if xCreate, then create the necessary shadow tables
4964+ if (isCreate) {
4965+ sqlite3_stmt *stmt;
4966+ int rc;
4967+
4968+ char * zCreateInfo = sqlite3_mprintf("CREATE TABLE "VEC0_SHADOW_INFO_NAME " (key text primary key, value any)", pNew->schemaName, pNew->tableName);
4969+ if(!zCreateInfo) {
4970+ goto error;
4971+ }
4972+ rc = sqlite3_prepare_v2(db, zCreateInfo, -1, &stmt, NULL);
4973+
4974+ sqlite3_free((void *) zCreateInfo);
4975+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
4976+ // TODO(IMP)
4977+ sqlite3_finalize(stmt);
4978+ *pzErr = sqlite3_mprintf("Could not create '_info' shadow table: %s",
4979+ sqlite3_errmsg(db));
4980+ goto error;
4981+ }
4982+ sqlite3_finalize(stmt);
4983+
4984+ char * zSeedInfo = sqlite3_mprintf(
4985+ "INSERT INTO "VEC0_SHADOW_INFO_NAME "(key, value) VALUES "
4986+ "(?1, ?2), (?3, ?4), (?5, ?6), (?7, ?8) ",
4987+ pNew->schemaName, pNew->tableName
4988+ );
4989+ if(!zSeedInfo) {
4990+ goto error;
4991+ }
4992+ rc = sqlite3_prepare_v2(db, zSeedInfo, -1, &stmt, NULL);
4993+ sqlite3_free((void *) zSeedInfo);
4994+ if (rc != SQLITE_OK) {
4995+ // TODO(IMP)
4996+ sqlite3_finalize(stmt);
4997+ *pzErr = sqlite3_mprintf("Could not seed '_info' shadow table: %s",
4998+ sqlite3_errmsg(db));
4999+ goto error;
5000+ }
5001+ sqlite3_bind_text(stmt, 1, "CREATE_VERSION", -1, SQLITE_STATIC);
5002+ sqlite3_bind_text(stmt, 2, SQLITE_VEC_VERSION, -1, SQLITE_STATIC);
5003+ sqlite3_bind_text(stmt, 3, "CREATE_VERSION_MAJOR", -1, SQLITE_STATIC);
5004+ sqlite3_bind_int(stmt, 4, SQLITE_VEC_VERSION_MAJOR);
5005+ sqlite3_bind_text(stmt, 5, "CREATE_VERSION_MINOR", -1, SQLITE_STATIC);
5006+ sqlite3_bind_int(stmt, 6, SQLITE_VEC_VERSION_MINOR);
5007+ sqlite3_bind_text(stmt, 7, "CREATE_VERSION_PATCH", -1, SQLITE_STATIC);
5008+ sqlite3_bind_int(stmt, 8, SQLITE_VEC_VERSION_PATCH);
5009+
5010+ if(sqlite3_step(stmt) != SQLITE_DONE) {
5011+ // TODO(IMP)
5012+ sqlite3_finalize(stmt);
5013+ *pzErr = sqlite3_mprintf("Could not seed '_info' shadow table: %s",
5014+ sqlite3_errmsg(db));
5015+ goto error;
5016+ }
5017+ sqlite3_finalize(stmt);
5018+
5019+
5020+
5021+ // create the _chunks shadow table
5022+ char *zCreateShadowChunks = NULL;
5023+ if(pNew->numPartitionColumns) {
5024+ sqlite3_str * s = sqlite3_str_new(NULL);
5025+ sqlite3_str_appendf(s, "CREATE TABLE " VEC0_SHADOW_CHUNKS_NAME "(", pNew->schemaName, pNew->tableName);
5026+ sqlite3_str_appendall(s, "chunk_id INTEGER PRIMARY KEY AUTOINCREMENT," "size INTEGER NOT NULL,");
5027+ sqlite3_str_appendall(s, "sequence_id integer,");
5028+ for(int i = 0; i < pNew->numPartitionColumns;i++) {
5029+ sqlite3_str_appendf(s, "partition%02d,", i);
5030+ }
5031+ sqlite3_str_appendall(s, "validity BLOB NOT NULL, rowids BLOB NOT NULL);");
5032+ zCreateShadowChunks = sqlite3_str_finish(s);
5033+ }else {
5034+ zCreateShadowChunks = sqlite3_mprintf(VEC0_SHADOW_CHUNKS_CREATE,
5035+ pNew->schemaName, pNew->tableName);
5036+ }
5037+ if (!zCreateShadowChunks) {
5038+ goto error;
5039+ }
5040+ rc = sqlite3_prepare_v2(db, zCreateShadowChunks, -1, &stmt, 0);
5041+ sqlite3_free((void *)zCreateShadowChunks);
5042+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5043+ // IMP: V17740_01811
5044+ sqlite3_finalize(stmt);
5045+ *pzErr = sqlite3_mprintf("Could not create '_chunks' shadow table: %s",
5046+ sqlite3_errmsg(db));
5047+ goto error;
5048+ }
5049+ sqlite3_finalize(stmt);
5050+
5051+ // create the _rowids shadow table
5052+ char *zCreateShadowRowids;
5053+ if (pNew->pkIsText) {
5054+ // adds a "text unique not null" constraint to the id column
5055+ zCreateShadowRowids = sqlite3_mprintf(VEC0_SHADOW_ROWIDS_CREATE_PK_TEXT,
5056+ pNew->schemaName, pNew->tableName);
5057+ } else {
5058+ zCreateShadowRowids = sqlite3_mprintf(VEC0_SHADOW_ROWIDS_CREATE_BASIC,
5059+ pNew->schemaName, pNew->tableName);
5060+ }
5061+ if (!zCreateShadowRowids) {
5062+ goto error;
5063+ }
5064+ rc = sqlite3_prepare_v2(db, zCreateShadowRowids, -1, &stmt, 0);
5065+ sqlite3_free((void *)zCreateShadowRowids);
5066+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5067+ // IMP: V11631_28470
5068+ sqlite3_finalize(stmt);
5069+ *pzErr = sqlite3_mprintf("Could not create '_rowids' shadow table: %s",
5070+ sqlite3_errmsg(db));
5071+ goto error;
5072+ }
5073+ sqlite3_finalize(stmt);
5074+
5075+ for (int i = 0; i < pNew->numVectorColumns; i++) {
5076+ char *zSql = sqlite3_mprintf(VEC0_SHADOW_VECTOR_N_CREATE,
5077+ pNew->schemaName, pNew->tableName, i);
5078+ if (!zSql) {
5079+ goto error;
5080+ }
5081+ rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0);
5082+ sqlite3_free((void *)zSql);
5083+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5084+ // IMP: V25919_09989
5085+ sqlite3_finalize(stmt);
5086+ *pzErr = sqlite3_mprintf(
5087+ "Could not create '_vector_chunks%02d' shadow table: %s", i,
5088+ sqlite3_errmsg(db));
5089+ goto error;
5090+ }
5091+ sqlite3_finalize(stmt);
5092+ }
5093+
5094+ for (int i = 0; i < pNew->numMetadataColumns; i++) {
5095+ char *zSql = sqlite3_mprintf("CREATE TABLE " VEC0_SHADOW_METADATA_N_NAME "(rowid PRIMARY KEY, data BLOB NOT NULL);",
5096+ pNew->schemaName, pNew->tableName, i);
5097+ if (!zSql) {
5098+ goto error;
5099+ }
5100+ rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0);
5101+ sqlite3_free((void *)zSql);
5102+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5103+ sqlite3_finalize(stmt);
5104+ *pzErr = sqlite3_mprintf(
5105+ "Could not create '_metata_chunks%02d' shadow table: %s", i,
5106+ sqlite3_errmsg(db));
5107+ goto error;
5108+ }
5109+ sqlite3_finalize(stmt);
5110+
5111+ if(pNew->metadata_columns[i].kind == VEC0_METADATA_COLUMN_KIND_TEXT) {
5112+ char *zSql = sqlite3_mprintf("CREATE TABLE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME "(rowid PRIMARY KEY, data TEXT);",
5113+ pNew->schemaName, pNew->tableName, i);
5114+ if (!zSql) {
5115+ goto error;
5116+ }
5117+ rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0);
5118+ sqlite3_free((void *)zSql);
5119+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5120+ sqlite3_finalize(stmt);
5121+ *pzErr = sqlite3_mprintf(
5122+ "Could not create '_metadatatext%02d' shadow table: %s", i,
5123+ sqlite3_errmsg(db));
5124+ goto error;
5125+ }
5126+ sqlite3_finalize(stmt);
5127+
5128+ }
5129+ }
5130+
5131+ if(pNew->numAuxiliaryColumns > 0) {
5132+ sqlite3_stmt * stmt;
5133+ sqlite3_str * s = sqlite3_str_new(NULL);
5134+ sqlite3_str_appendf(s, "CREATE TABLE " VEC0_SHADOW_AUXILIARY_NAME "( rowid integer PRIMARY KEY ", pNew->schemaName, pNew->tableName);
5135+ for(int i = 0; i < pNew->numAuxiliaryColumns; i++) {
5136+ sqlite3_str_appendf(s, ", value%02d", i);
5137+ }
5138+ sqlite3_str_appendall(s, ")");
5139+ char *zSql = sqlite3_str_finish(s);
5140+ if(!zSql) {
5141+ goto error;
5142+ }
5143+ rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, NULL);
5144+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5145+ sqlite3_finalize(stmt);
5146+ *pzErr = sqlite3_mprintf(
5147+ "Could not create auxiliary shadow table: %s",
5148+ sqlite3_errmsg(db));
5149+
5150+ goto error;
5151+ }
5152+ sqlite3_finalize(stmt);
5153+ }
5154+ }
5155+
5156+ *ppVtab = (sqlite3_vtab *)pNew;
5157+ return SQLITE_OK;
5158+
5159+error:
5160+ vec0_free(pNew);
5161+ return SQLITE_ERROR;
5162+}
5163+
5164+static int vec0Create(sqlite3 *db, void *pAux, int argc,
5165+ const char *const *argv, sqlite3_vtab **ppVtab,
5166+ char **pzErr) {
5167+ return vec0_init(db, pAux, argc, argv, ppVtab, pzErr, true);
5168+}
5169+static int vec0Connect(sqlite3 *db, void *pAux, int argc,
5170+ const char *const *argv, sqlite3_vtab **ppVtab,
5171+ char **pzErr) {
5172+ return vec0_init(db, pAux, argc, argv, ppVtab, pzErr, false);
5173+}
5174+
5175+static int vec0Disconnect(sqlite3_vtab *pVtab) {
5176+ vec0_vtab *p = (vec0_vtab *)pVtab;
5177+ vec0_free(p);
5178+ sqlite3_free(p);
5179+ return SQLITE_OK;
5180+}
5181+static int vec0Destroy(sqlite3_vtab *pVtab) {
5182+ vec0_vtab *p = (vec0_vtab *)pVtab;
5183+ sqlite3_stmt *stmt;
5184+ int rc;
5185+ const char *zSql;
5186+
5187+ // Free up any sqlite3_stmt, otherwise DROPs on those tables will fail
5188+ vec0_free_resources(p);
5189+
5190+ // TODO(test) later: can't evidence-of here, bc always gives "SQL logic error" instead of
5191+ // provided error
5192+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_CHUNKS_NAME, p->schemaName,
5193+ p->tableName);
5194+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5195+ sqlite3_free((void *)zSql);
5196+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5197+ rc = SQLITE_ERROR;
5198+ vtab_set_error(pVtab, "could not drop chunks shadow table");
5199+ goto done;
5200+ }
5201+ sqlite3_finalize(stmt);
5202+
5203+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_INFO_NAME, p->schemaName,
5204+ p->tableName);
5205+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5206+ sqlite3_free((void *)zSql);
5207+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5208+ rc = SQLITE_ERROR;
5209+ vtab_set_error(pVtab, "could not drop info shadow table");
5210+ goto done;
5211+ }
5212+ sqlite3_finalize(stmt);
5213+
5214+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_ROWIDS_NAME, p->schemaName,
5215+ p->tableName);
5216+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5217+ sqlite3_free((void *)zSql);
5218+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5219+ rc = SQLITE_ERROR;
5220+ goto done;
5221+ }
5222+ sqlite3_finalize(stmt);
5223+
5224+ for (int i = 0; i < p->numVectorColumns; i++) {
5225+ zSql = sqlite3_mprintf("DROP TABLE \"%w\".\"%w\"", p->schemaName,
5226+ p->shadowVectorChunksNames[i]);
5227+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5228+ sqlite3_free((void *)zSql);
5229+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5230+ rc = SQLITE_ERROR;
5231+ goto done;
5232+ }
5233+ sqlite3_finalize(stmt);
5234+ }
5235+
5236+ if(p->numAuxiliaryColumns > 0) {
5237+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_AUXILIARY_NAME, p->schemaName, p->tableName);
5238+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5239+ sqlite3_free((void *)zSql);
5240+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5241+ rc = SQLITE_ERROR;
5242+ goto done;
5243+ }
5244+ sqlite3_finalize(stmt);
5245+ }
5246+
5247+
5248+ for (int i = 0; i < p->numMetadataColumns; i++) {
5249+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_METADATA_N_NAME, p->schemaName,p->tableName, i);
5250+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5251+ sqlite3_free((void *)zSql);
5252+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5253+ rc = SQLITE_ERROR;
5254+ goto done;
5255+ }
5256+ sqlite3_finalize(stmt);
5257+
5258+ if(p->metadata_columns[i].kind == VEC0_METADATA_COLUMN_KIND_TEXT) {
5259+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME, p->schemaName,p->tableName, i);
5260+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5261+ sqlite3_free((void *)zSql);
5262+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5263+ rc = SQLITE_ERROR;
5264+ goto done;
5265+ }
5266+ sqlite3_finalize(stmt);
5267+ }
5268+ }
5269+
5270+ stmt = NULL;
5271+ rc = SQLITE_OK;
5272+
5273+done:
5274+ sqlite3_finalize(stmt);
5275+ vec0_free(p);
5276+ // If there was an error
5277+ if (rc == SQLITE_OK) {
5278+ sqlite3_free(p);
5279+ }
5280+ return rc;
5281+}
5282+
5283+static int vec0Open(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) {
5284+ UNUSED_PARAMETER(p);
5285+ vec0_cursor *pCur;
5286+ pCur = sqlite3_malloc(sizeof(*pCur));
5287+ if (pCur == 0)
5288+ return SQLITE_NOMEM;
5289+ memset(pCur, 0, sizeof(*pCur));
5290+ *ppCursor = &pCur->base;
5291+ return SQLITE_OK;
5292+}
5293+
5294+static int vec0Close(sqlite3_vtab_cursor *cur) {
5295+ vec0_cursor *pCur = (vec0_cursor *)cur;
5296+ vec0_cursor_clear(pCur);
5297+ sqlite3_free(pCur);
5298+ return SQLITE_OK;
5299+}
5300+
5301+// All the different type of "values" provided to argv/argc in vec0Filter.
5302+// These enums denote the use and purpose of all of them.
5303+typedef enum {
5304+ // If any values are updated, please update the ARCHITECTURE.md docs accordingly!
5305+
5306+ VEC0_IDXSTR_KIND_KNN_MATCH = '{',
5307+ VEC0_IDXSTR_KIND_KNN_K = '}',
5308+ VEC0_IDXSTR_KIND_KNN_ROWID_IN = '[',
5309+ VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT = ']',
5310+ VEC0_IDXSTR_KIND_POINT_ID = '!',
5311+ VEC0_IDXSTR_KIND_METADATA_CONSTRAINT = '&',
5312+} vec0_idxstr_kind;
5313+
5314+// The different SQLITE_INDEX_CONSTRAINT values that vec0 partition key columns
5315+// support, but as characters that fit nicely in idxstr.
5316+typedef enum {
5317+ // If any values are updated, please update the ARCHITECTURE.md docs accordingly!
5318+
5319+ VEC0_PARTITION_OPERATOR_EQ = 'a',
5320+ VEC0_PARTITION_OPERATOR_GT = 'b',
5321+ VEC0_PARTITION_OPERATOR_LE = 'c',
5322+ VEC0_PARTITION_OPERATOR_LT = 'd',
5323+ VEC0_PARTITION_OPERATOR_GE = 'e',
5324+ VEC0_PARTITION_OPERATOR_NE = 'f',
5325+} vec0_partition_operator;
5326+typedef enum {
5327+ VEC0_METADATA_OPERATOR_EQ = 'a',
5328+ VEC0_METADATA_OPERATOR_GT = 'b',
5329+ VEC0_METADATA_OPERATOR_LE = 'c',
5330+ VEC0_METADATA_OPERATOR_LT = 'd',
5331+ VEC0_METADATA_OPERATOR_GE = 'e',
5332+ VEC0_METADATA_OPERATOR_NE = 'f',
5333+ VEC0_METADATA_OPERATOR_IN = 'g',
5334+} vec0_metadata_operator;
5335+
5336+static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
5337+ vec0_vtab *p = (vec0_vtab *)pVTab;
5338+ /**
5339+ * Possible query plans are:
5340+ * 1. KNN when:
5341+ * a) An `MATCH` op on vector column
5342+ * b) ORDER BY on distance column
5343+ * c) LIMIT
5344+ * d) rowid in (...) OPTIONAL
5345+ * 2. Point when:
5346+ * a) An `EQ` op on rowid column
5347+ * 3. else: fullscan
5348+ *
5349+ */
5350+ int iMatchTerm = -1;
5351+ int iMatchVectorTerm = -1;
5352+ int iLimitTerm = -1;
5353+ int iRowidTerm = -1;
5354+ int iKTerm = -1;
5355+ int iRowidInTerm = -1;
5356+ int hasAuxConstraint = 0;
5357+
5358+#ifdef SQLITE_VEC_DEBUG
5359+ printf("pIdxInfo->nOrderBy=%d, pIdxInfo->nConstraint=%d\n", pIdxInfo->nOrderBy, pIdxInfo->nConstraint);
5360+#endif
5361+
5362+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
5363+ u8 vtabIn = 0;
5364+
5365+#if COMPILER_SUPPORTS_VTAB_IN
5366+ if (sqlite3_libversion_number() >= 3038000) {
5367+ vtabIn = sqlite3_vtab_in(pIdxInfo, i, -1);
5368+ }
5369+#endif
5370+
5371+#ifdef SQLITE_VEC_DEBUG
5372+ printf("xBestIndex [%d] usable=%d iColumn=%d op=%d vtabin=%d\n", i,
5373+ pIdxInfo->aConstraint[i].usable, pIdxInfo->aConstraint[i].iColumn,
5374+ pIdxInfo->aConstraint[i].op, vtabIn);
5375+#endif
5376+ if (!pIdxInfo->aConstraint[i].usable)
5377+ continue;
5378+
5379+ int iColumn = pIdxInfo->aConstraint[i].iColumn;
5380+ int op = pIdxInfo->aConstraint[i].op;
5381+
5382+ if (op == SQLITE_INDEX_CONSTRAINT_LIMIT) {
5383+ iLimitTerm = i;
5384+ }
5385+ if (op == SQLITE_INDEX_CONSTRAINT_MATCH &&
5386+ vec0_column_idx_is_vector(p, iColumn)) {
5387+ if (iMatchTerm > -1) {
5388+ vtab_set_error(
5389+ pVTab, "only 1 MATCH operator is allowed in a single vec0 query");
5390+ return SQLITE_ERROR;
5391+ }
5392+ iMatchTerm = i;
5393+ iMatchVectorTerm = vec0_column_idx_to_vector_idx(p, iColumn);
5394+ }
5395+ if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == VEC0_COLUMN_ID) {
5396+ if (vtabIn) {
5397+ if (iRowidInTerm != -1) {
5398+ vtab_set_error(pVTab, "only 1 'rowid in (..)' operator is allowed in "
5399+ "a single vec0 query");
5400+ return SQLITE_ERROR;
5401+ }
5402+ iRowidInTerm = i;
5403+
5404+ } else {
5405+ iRowidTerm = i;
5406+ }
5407+ }
5408+ if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == vec0_column_k_idx(p)) {
5409+ iKTerm = i;
5410+ }
5411+ if(
5412+ (op != SQLITE_INDEX_CONSTRAINT_LIMIT && op != SQLITE_INDEX_CONSTRAINT_OFFSET)
5413+ && vec0_column_idx_is_auxiliary(p, iColumn)) {
5414+ hasAuxConstraint = 1;
5415+ }
5416+ }
5417+
5418+ sqlite3_str *idxStr = sqlite3_str_new(NULL);
5419+ int rc;
5420+
5421+ if (iMatchTerm >= 0) {
5422+ if (iLimitTerm < 0 && iKTerm < 0) {
5423+ vtab_set_error(
5424+ pVTab,
5425+ "A LIMIT or 'k = ?' constraint is required on vec0 knn queries.");
5426+ rc = SQLITE_ERROR;
5427+ goto done;
5428+ }
5429+ if (iLimitTerm >= 0 && iKTerm >= 0) {
5430+ vtab_set_error(pVTab, "Only LIMIT or 'k =?' can be provided, not both");
5431+ rc = SQLITE_ERROR;
5432+ goto done;
5433+ }
5434+
5435+ if (pIdxInfo->nOrderBy) {
5436+ if (pIdxInfo->nOrderBy > 1) {
5437+ vtab_set_error(pVTab, "Only a single 'ORDER BY distance' clause is "
5438+ "allowed on vec0 KNN queries");
5439+ rc = SQLITE_ERROR;
5440+ goto done;
5441+ }
5442+ if (pIdxInfo->aOrderBy[0].iColumn != vec0_column_distance_idx(p)) {
5443+ vtab_set_error(pVTab,
5444+ "Only a single 'ORDER BY distance' clause is allowed on "
5445+ "vec0 KNN queries, not on other columns");
5446+ rc = SQLITE_ERROR;
5447+ goto done;
5448+ }
5449+ if (pIdxInfo->aOrderBy[0].desc) {
5450+ vtab_set_error(
5451+ pVTab, "Only ascending in ORDER BY distance clause is supported, "
5452+ "DESC is not supported yet.");
5453+ rc = SQLITE_ERROR;
5454+ goto done;
5455+ }
5456+ }
5457+
5458+ if(hasAuxConstraint) {
5459+ // IMP: V25623_09693
5460+ vtab_set_error(pVTab, "An illegal WHERE constraint was provided on a vec0 auxiliary column in a KNN query.");
5461+ rc = SQLITE_ERROR;
5462+ goto done;
5463+ }
5464+
5465+ sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_KNN);
5466+
5467+ int argvIndex = 1;
5468+ pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = argvIndex++;
5469+ pIdxInfo->aConstraintUsage[iMatchTerm].omit = 1;
5470+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_MATCH);
5471+ sqlite3_str_appendchar(idxStr, 3, '_');
5472+
5473+ if (iLimitTerm >= 0) {
5474+ pIdxInfo->aConstraintUsage[iLimitTerm].argvIndex = argvIndex++;
5475+ pIdxInfo->aConstraintUsage[iLimitTerm].omit = 1;
5476+ } else {
5477+ pIdxInfo->aConstraintUsage[iKTerm].argvIndex = argvIndex++;
5478+ pIdxInfo->aConstraintUsage[iKTerm].omit = 1;
5479+ }
5480+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_K);
5481+ sqlite3_str_appendchar(idxStr, 3, '_');
5482+
5483+#if COMPILER_SUPPORTS_VTAB_IN
5484+ if (iRowidInTerm >= 0) {
5485+ // already validated as >= SQLite 3.38 bc iRowidInTerm is only >= 0 when
5486+ // vtabIn == 1
5487+ sqlite3_vtab_in(pIdxInfo, iRowidInTerm, 1);
5488+ pIdxInfo->aConstraintUsage[iRowidInTerm].argvIndex = argvIndex++;
5489+ pIdxInfo->aConstraintUsage[iRowidInTerm].omit = 1;
5490+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_ROWID_IN);
5491+ sqlite3_str_appendchar(idxStr, 3, '_');
5492+ }
5493+#endif
5494+
5495+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
5496+ if (!pIdxInfo->aConstraint[i].usable)
5497+ continue;
5498+
5499+ int iColumn = pIdxInfo->aConstraint[i].iColumn;
5500+ int op = pIdxInfo->aConstraint[i].op;
5501+ if(op == SQLITE_INDEX_CONSTRAINT_LIMIT || op == SQLITE_INDEX_CONSTRAINT_OFFSET) {
5502+ continue;
5503+ }
5504+ if(!vec0_column_idx_is_partition(p, iColumn)) {
5505+ continue;
5506+ }
5507+
5508+ int partition_idx = vec0_column_idx_to_partition_idx(p, iColumn);
5509+ char value = 0;
5510+
5511+ switch(op) {
5512+ case SQLITE_INDEX_CONSTRAINT_EQ: {
5513+ value = VEC0_PARTITION_OPERATOR_EQ;
5514+ break;
5515+ }
5516+ case SQLITE_INDEX_CONSTRAINT_GT: {
5517+ value = VEC0_PARTITION_OPERATOR_GT;
5518+ break;
5519+ }
5520+ case SQLITE_INDEX_CONSTRAINT_LE: {
5521+ value = VEC0_PARTITION_OPERATOR_LE;
5522+ break;
5523+ }
5524+ case SQLITE_INDEX_CONSTRAINT_LT: {
5525+ value = VEC0_PARTITION_OPERATOR_LT;
5526+ break;
5527+ }
5528+ case SQLITE_INDEX_CONSTRAINT_GE: {
5529+ value = VEC0_PARTITION_OPERATOR_GE;
5530+ break;
5531+ }
5532+ case SQLITE_INDEX_CONSTRAINT_NE: {
5533+ value = VEC0_PARTITION_OPERATOR_NE;
5534+ break;
5535+ }
5536+ }
5537+
5538+ if(value) {
5539+ pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++;
5540+ pIdxInfo->aConstraintUsage[i].omit = 1;
5541+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT);
5542+ sqlite3_str_appendchar(idxStr, 1, 'A' + partition_idx);
5543+ sqlite3_str_appendchar(idxStr, 1, value);
5544+ sqlite3_str_appendchar(idxStr, 1, '_');
5545+ }
5546+
5547+ }
5548+
5549+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
5550+ if (!pIdxInfo->aConstraint[i].usable)
5551+ continue;
5552+
5553+ int iColumn = pIdxInfo->aConstraint[i].iColumn;
5554+ int op = pIdxInfo->aConstraint[i].op;
5555+ if(op == SQLITE_INDEX_CONSTRAINT_LIMIT || op == SQLITE_INDEX_CONSTRAINT_OFFSET) {
5556+ continue;
5557+ }
5558+ if(!vec0_column_idx_is_metadata(p, iColumn)) {
5559+ continue;
5560+ }
5561+
5562+ int metadata_idx = vec0_column_idx_to_metadata_idx(p, iColumn);
5563+ char value = 0;
5564+
5565+ switch(op) {
5566+ case SQLITE_INDEX_CONSTRAINT_EQ: {
5567+ int vtabIn = 0;
5568+ #if COMPILER_SUPPORTS_VTAB_IN
5569+ if (sqlite3_libversion_number() >= 3038000) {
5570+ vtabIn = sqlite3_vtab_in(pIdxInfo, i, -1);
5571+ }
5572+ if(vtabIn) {
5573+ switch(p->metadata_columns[metadata_idx].kind) {
5574+ case VEC0_METADATA_COLUMN_KIND_FLOAT:
5575+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
5576+ // IMP: V15248_32086
5577+ rc = SQLITE_ERROR;
5578+ vtab_set_error(pVTab, "'xxx in (...)' is only available on INTEGER or TEXT metadata columns.");
5579+ goto done;
5580+ break;
5581+ }
5582+ case VEC0_METADATA_COLUMN_KIND_INTEGER:
5583+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
5584+ break;
5585+ }
5586+ }
5587+ value = VEC0_METADATA_OPERATOR_IN;
5588+ sqlite3_vtab_in(pIdxInfo, i, 1);
5589+ }else
5590+ #endif
5591+ {
5592+ value = VEC0_PARTITION_OPERATOR_EQ;
5593+ }
5594+ break;
5595+ }
5596+ case SQLITE_INDEX_CONSTRAINT_GT: {
5597+ value = VEC0_METADATA_OPERATOR_GT;
5598+ break;
5599+ }
5600+ case SQLITE_INDEX_CONSTRAINT_LE: {
5601+ value = VEC0_METADATA_OPERATOR_LE;
5602+ break;
5603+ }
5604+ case SQLITE_INDEX_CONSTRAINT_LT: {
5605+ value = VEC0_METADATA_OPERATOR_LT;
5606+ break;
5607+ }
5608+ case SQLITE_INDEX_CONSTRAINT_GE: {
5609+ value = VEC0_METADATA_OPERATOR_GE;
5610+ break;
5611+ }
5612+ case SQLITE_INDEX_CONSTRAINT_NE: {
5613+ value = VEC0_METADATA_OPERATOR_NE;
5614+ break;
5615+ }
5616+ default: {
5617+ // IMP: V16511_00582
5618+ rc = SQLITE_ERROR;
5619+ vtab_set_error(pVTab,
5620+ "An illegal WHERE constraint was provided on a vec0 metadata column in a KNN query. "
5621+ "Only one of EQUALS, GREATER_THAN, LESS_THAN_OR_EQUAL, LESS_THAN, GREATER_THAN_OR_EQUAL, NOT_EQUALS is allowed."
5622+ );
5623+ goto done;
5624+ }
5625+ }
5626+
5627+ if(p->metadata_columns[metadata_idx].kind == VEC0_METADATA_COLUMN_KIND_BOOLEAN) {
5628+ if(!(value == VEC0_METADATA_OPERATOR_EQ || value == VEC0_METADATA_OPERATOR_NE)) {
5629+ // IMP: V10145_26984
5630+ rc = SQLITE_ERROR;
5631+ vtab_set_error(pVTab, "ONLY EQUALS (=) or NOT_EQUALS (!=) operators are allowed on boolean metadata columns.");
5632+ goto done;
5633+ }
5634+ }
5635+
5636+ pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++;
5637+ pIdxInfo->aConstraintUsage[i].omit = 1;
5638+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_METADATA_CONSTRAINT);
5639+ sqlite3_str_appendchar(idxStr, 1, 'A' + metadata_idx);
5640+ sqlite3_str_appendchar(idxStr, 1, value);
5641+ sqlite3_str_appendchar(idxStr, 1, '_');
5642+
5643+ }
5644+
5645+
5646+
5647+ pIdxInfo->idxNum = iMatchVectorTerm;
5648+ pIdxInfo->estimatedCost = 30.0;
5649+ pIdxInfo->estimatedRows = 10;
5650+
5651+ } else if (iRowidTerm >= 0) {
5652+ sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_POINT);
5653+ pIdxInfo->aConstraintUsage[iRowidTerm].argvIndex = 1;
5654+ pIdxInfo->aConstraintUsage[iRowidTerm].omit = 1;
5655+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_POINT_ID);
5656+ sqlite3_str_appendchar(idxStr, 3, '_');
5657+ pIdxInfo->idxNum = pIdxInfo->colUsed;
5658+ pIdxInfo->estimatedCost = 10.0;
5659+ pIdxInfo->estimatedRows = 1;
5660+ } else {
5661+ sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_FULLSCAN);
5662+ pIdxInfo->estimatedCost = 3000000.0;
5663+ pIdxInfo->estimatedRows = 100000;
5664+ }
5665+ pIdxInfo->idxStr = sqlite3_str_finish(idxStr);
5666+ idxStr = NULL;
5667+ if (!pIdxInfo->idxStr) {
5668+ rc = SQLITE_OK;
5669+ goto done;
5670+ }
5671+ pIdxInfo->needToFreeIdxStr = 1;
5672+
5673+
5674+ rc = SQLITE_OK;
5675+
5676+ done:
5677+ if(idxStr) {
5678+ sqlite3_str_finish(idxStr);
5679+ }
5680+ return rc;
5681+}
5682+
5683+// forward delcaration bc vec0Filter uses it
5684+static int vec0Next(sqlite3_vtab_cursor *cur);
5685+
5686+void merge_sorted_lists(f32 *a, i64 *a_rowids, i64 a_length, f32 *b,
5687+ i64 *b_rowids, i32 *b_top_idxs, i64 b_length, f32 *out,
5688+ i64 *out_rowids, i64 out_length, i64 *out_used) {
5689+ // assert((a_length >= out_length) || (b_length >= out_length));
5690+ i64 ptrA = 0;
5691+ i64 ptrB = 0;
5692+ for (int i = 0; i < out_length; i++) {
5693+ if ((ptrA >= a_length) && (ptrB >= b_length)) {
5694+ *out_used = i;
5695+ return;
5696+ }
5697+ if (ptrA >= a_length) {
5698+ out[i] = b[b_top_idxs[ptrB]];
5699+ out_rowids[i] = b_rowids[b_top_idxs[ptrB]];
5700+ ptrB++;
5701+ } else if (ptrB >= b_length) {
5702+ out[i] = a[ptrA];
5703+ out_rowids[i] = a_rowids[ptrA];
5704+ ptrA++;
5705+ } else {
5706+ if (a[ptrA] <= b[b_top_idxs[ptrB]]) {
5707+ out[i] = a[ptrA];
5708+ out_rowids[i] = a_rowids[ptrA];
5709+ ptrA++;
5710+ } else {
5711+ out[i] = b[b_top_idxs[ptrB]];
5712+ out_rowids[i] = b_rowids[b_top_idxs[ptrB]];
5713+ ptrB++;
5714+ }
5715+ }
5716+ }
5717+
5718+ *out_used = out_length;
5719+}
5720+
5721+u8 *bitmap_new(i32 n) {
5722+ assert(n % 8 == 0);
5723+ u8 *p = sqlite3_malloc(n * sizeof(u8) / CHAR_BIT);
5724+ if (p) {
5725+ memset(p, 0, n * sizeof(u8) / CHAR_BIT);
5726+ }
5727+ return p;
5728+}
5729+u8 *bitmap_new_from(i32 n, u8 *from) {
5730+ assert(n % 8 == 0);
5731+ u8 *p = sqlite3_malloc(n * sizeof(u8) / CHAR_BIT);
5732+ if (p) {
5733+ memcpy(p, from, n / CHAR_BIT);
5734+ }
5735+ return p;
5736+}
5737+
5738+void bitmap_copy(u8 *base, u8 *from, i32 n) {
5739+ assert(n % 8 == 0);
5740+ memcpy(base, from, n / CHAR_BIT);
5741+}
5742+
5743+void bitmap_and_inplace(u8 *base, u8 *other, i32 n) {
5744+ assert((n % 8) == 0);
5745+ for (int i = 0; i < n / CHAR_BIT; i++) {
5746+ base[i] = base[i] & other[i];
5747+ }
5748+}
5749+
5750+void bitmap_set(u8 *bitmap, i32 position, int value) {
5751+ if (value) {
5752+ bitmap[position / CHAR_BIT] |= 1 << (position % CHAR_BIT);
5753+ } else {
5754+ bitmap[position / CHAR_BIT] &= ~(1 << (position % CHAR_BIT));
5755+ }
5756+}
5757+
5758+int bitmap_get(u8 *bitmap, i32 position) {
5759+ return (((bitmap[position / CHAR_BIT]) >> (position % CHAR_BIT)) & 1);
5760+}
5761+
5762+void bitmap_clear(u8 *bitmap, i32 n) {
5763+ assert((n % 8) == 0);
5764+ memset(bitmap, 0, n / CHAR_BIT);
5765+}
5766+
5767+void bitmap_fill(u8 *bitmap, i32 n) {
5768+ assert((n % 8) == 0);
5769+ memset(bitmap, 0xFF, n / CHAR_BIT);
5770+}
5771+
5772+/**
5773+ * @brief Finds the minimum k items in distances, and writes the indicies to
5774+ * out.
5775+ *
5776+ * @param distances input f32 array of size n, the items to consider.
5777+ * @param n: size of distances array.
5778+ * @param out: Output array of size k, will contain at most k element indicies
5779+ * @param k: Size of output array
5780+ * @return int
5781+ */
5782+int min_idx(const f32 *distances, i32 n, u8 *candidates, i32 *out, i32 k,
5783+ u8 *bTaken, i32 *k_used) {
5784+ assert(k > 0);
5785+ assert(k <= n);
5786+
5787+ bitmap_clear(bTaken, n);
5788+
5789+ for (int ik = 0; ik < k; ik++) {
5790+ int min_idx = 0;
5791+ while (min_idx < n &&
5792+ (bitmap_get(bTaken, min_idx) || !bitmap_get(candidates, min_idx))) {
5793+ min_idx++;
5794+ }
5795+ if (min_idx >= n) {
5796+ *k_used = ik;
5797+ return SQLITE_OK;
5798+ }
5799+
5800+ for (int i = 0; i < n; i++) {
5801+ if (distances[i] <= distances[min_idx] && !bitmap_get(bTaken, i) &&
5802+ (bitmap_get(candidates, i))) {
5803+ min_idx = i;
5804+ }
5805+ }
5806+
5807+ out[ik] = min_idx;
5808+ bitmap_set(bTaken, min_idx, 1);
5809+ }
5810+ *k_used = k;
5811+ return SQLITE_OK;
5812+}
5813+
5814+int vec0_get_metadata_text_long_value(
5815+ vec0_vtab * p,
5816+ sqlite3_stmt ** stmt,
5817+ int metadata_idx,
5818+ i64 rowid,
5819+ int *n,
5820+ char ** s) {
5821+ int rc;
5822+ if(!(*stmt)) {
5823+ const char * zSql = sqlite3_mprintf("select data from " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " where rowid = ?", p->schemaName, p->tableName, metadata_idx);
5824+ if(!zSql) {
5825+ rc = SQLITE_NOMEM;
5826+ goto done;
5827+ }
5828+ rc = sqlite3_prepare_v2(p->db, zSql, -1, stmt, NULL);
5829+ sqlite3_free( (void *) zSql);
5830+ if(rc != SQLITE_OK) {
5831+ goto done;
5832+ }
5833+ }
5834+
5835+ sqlite3_reset(*stmt);
5836+ sqlite3_bind_int64(*stmt, 1, rowid);
5837+ rc = sqlite3_step(*stmt);
5838+ if(rc != SQLITE_ROW) {
5839+ rc = SQLITE_ERROR;
5840+ goto done;
5841+ }
5842+ *s = (char *) sqlite3_column_text(*stmt, 0);
5843+ *n = sqlite3_column_bytes(*stmt, 0);
5844+ rc = SQLITE_OK;
5845+ done:
5846+ return rc;
5847+}
5848+
5849+/**
5850+ * @brief Crete at "iterator" (sqlite3_stmt) of chunks with the given constraints
5851+ *
5852+ * Any VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT values in idxStr/argv will be applied
5853+ * as WHERE constraints in the underlying stmt SQL, and any consumer of the stmt
5854+ * can freely step through the stmt with all constraints satisfied.
5855+ *
5856+ * @param p - vec0_vtab
5857+ * @param idxStr - the xBestIndex/xFilter idxstr containing VEC0_IDXSTR values
5858+ * @param argc - number of argv values from xFilter
5859+ * @param argv - array of sqlite3_value from xFilter
5860+ * @param outStmt - output sqlite3_stmt of chunks with all filters applied
5861+ * @return int SQLITE_OK on success, error code otherwise
5862+ */
5863+int vec0_chunks_iter(vec0_vtab * p, const char * idxStr, int argc, sqlite3_value ** argv, sqlite3_stmt** outStmt) {
5864+ // always null terminated, enforced by SQLite
5865+ int idxStrLength = strlen(idxStr);
5866+ // "1" refers to the initial vec0_query_plan char, 4 is the number of chars per "element"
5867+ int numValueEntries = (idxStrLength-1) / 4;
5868+ assert(argc == numValueEntries);
5869+
5870+ int rc;
5871+ sqlite3_str * s = sqlite3_str_new(NULL);
5872+ sqlite3_str_appendf(s, "select chunk_id, validity, rowids "
5873+ " from " VEC0_SHADOW_CHUNKS_NAME,
5874+ p->schemaName, p->tableName);
5875+
5876+ int appendedWhere = 0;
5877+ for(int i = 0; i < numValueEntries; i++) {
5878+ int idx = 1 + (i * 4);
5879+ char kind = idxStr[idx + 0];
5880+ if(kind != VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT) {
5881+ continue;
5882+ }
5883+
5884+ int partition_idx = idxStr[idx + 1] - 'A';
5885+ int operator = idxStr[idx + 2];
5886+ // idxStr[idx + 3] is just null, a '_' placeholder
5887+
5888+ if(!appendedWhere) {
5889+ sqlite3_str_appendall(s, " WHERE ");
5890+ appendedWhere = 1;
5891+ }else {
5892+ sqlite3_str_appendall(s, " AND ");
5893+ }
5894+ switch(operator) {
5895+ case VEC0_PARTITION_OPERATOR_EQ:
5896+ sqlite3_str_appendf(s, " partition%02d = ? ", partition_idx);
5897+ break;
5898+ case VEC0_PARTITION_OPERATOR_GT:
5899+ sqlite3_str_appendf(s, " partition%02d > ? ", partition_idx);
5900+ break;
5901+ case VEC0_PARTITION_OPERATOR_LE:
5902+ sqlite3_str_appendf(s, " partition%02d <= ? ", partition_idx);
5903+ break;
5904+ case VEC0_PARTITION_OPERATOR_LT:
5905+ sqlite3_str_appendf(s, " partition%02d < ? ", partition_idx);
5906+ break;
5907+ case VEC0_PARTITION_OPERATOR_GE:
5908+ sqlite3_str_appendf(s, " partition%02d >= ? ", partition_idx);
5909+ break;
5910+ case VEC0_PARTITION_OPERATOR_NE:
5911+ sqlite3_str_appendf(s, " partition%02d != ? ", partition_idx);
5912+ break;
5913+ default: {
5914+ char * zSql = sqlite3_str_finish(s);
5915+ sqlite3_free(zSql);
5916+ return SQLITE_ERROR;
5917+ }
5918+
5919+ }
5920+
5921+ }
5922+
5923+ char *zSql = sqlite3_str_finish(s);
5924+ if (!zSql) {
5925+ return SQLITE_NOMEM;
5926+ }
5927+
5928+ rc = sqlite3_prepare_v2(p->db, zSql, -1, outStmt, NULL);
5929+ sqlite3_free(zSql);
5930+ if(rc != SQLITE_OK) {
5931+ return rc;
5932+ }
5933+
5934+ int n = 1;
5935+ for(int i = 0; i < numValueEntries; i++) {
5936+ int idx = 1 + (i * 4);
5937+ char kind = idxStr[idx + 0];
5938+ if(kind != VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT) {
5939+ continue;
5940+ }
5941+ sqlite3_bind_value(*outStmt, n++, argv[i]);
5942+ }
5943+
5944+ return rc;
5945+}
5946+
5947+// a single `xxx in (...)` constraint on a metadata column. TEXT or INTEGER only for now.
5948+struct Vec0MetadataIn{
5949+ // index of argv[i]` the constraint is on
5950+ int argv_idx;
5951+ // metadata column index of the constraint, derived from idxStr + argv_idx
5952+ int metadata_idx;
5953+ // array of the copied `(...)` values from sqlite3_vtab_in_first()/sqlite3_vtab_in_next()
5954+ struct Array array;
5955+};
5956+
5957+// Array elements for `xxx in (...)` values for a text column. basically just a string
5958+struct Vec0MetadataInTextEntry {
5959+ int n;
5960+ char * zString;
5961+};
5962+
5963+
5964+int vec0_metadata_filter_text(vec0_vtab * p, sqlite3_value * value, const void * buffer, int size, vec0_metadata_operator op, u8* b, int metadata_idx, int chunk_rowid, struct Array * aMetadataIn, int argv_idx) {
5965+ int rc;
5966+ sqlite3_stmt * stmt = NULL;
5967+ i64 * rowids = NULL;
5968+ sqlite3_blob * rowidsBlob;
5969+ const char * sTarget = (const char *) sqlite3_value_text(value);
5970+ int nTarget = sqlite3_value_bytes(value);
5971+
5972+
5973+ // TODO(perf): only text metadata news the rowids BLOB. Make it so that
5974+ // rowids BLOB is re-used when multiple fitlers on text columns,
5975+ // ex "name BETWEEN 'a' and 'b'""
5976+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "rowids", chunk_rowid, 0, &rowidsBlob);
5977+ if(rc != SQLITE_OK) {
5978+ return rc;
5979+ }
5980+ assert(sqlite3_blob_bytes(rowidsBlob) % sizeof(i64) == 0);
5981+ assert((sqlite3_blob_bytes(rowidsBlob) / sizeof(i64)) == size);
5982+
5983+ rowids = sqlite3_malloc(sqlite3_blob_bytes(rowidsBlob));
5984+ if(!rowids) {
5985+ sqlite3_blob_close(rowidsBlob);
5986+ return SQLITE_NOMEM;
5987+ }
5988+
5989+ rc = sqlite3_blob_read(rowidsBlob, rowids, sqlite3_blob_bytes(rowidsBlob), 0);
5990+ if(rc != SQLITE_OK) {
5991+ sqlite3_blob_close(rowidsBlob);
5992+ return rc;
5993+ }
5994+ sqlite3_blob_close(rowidsBlob);
5995+
5996+ switch(op) {
5997+ int nPrefix;
5998+ char * sPrefix;
5999+ char *sFull;
6000+ int nFull;
6001+ u8 * view;
6002+ case VEC0_METADATA_OPERATOR_EQ: {
6003+ for(int i = 0; i < size; i++) {
6004+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6005+ nPrefix = ((int*) view)[0];
6006+ sPrefix = (char *) &view[4];
6007+
6008+ // for EQ the text lengths must match
6009+ if(nPrefix != nTarget) {
6010+ bitmap_set(b, i, 0);
6011+ continue;
6012+ }
6013+ int cmpPrefix = strncmp(sPrefix, sTarget, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH));
6014+
6015+ // for short strings, use the prefix comparison direclty
6016+ if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6017+ bitmap_set(b, i, cmpPrefix == 0);
6018+ continue;
6019+ }
6020+ // for EQ on longs strings, the prefix must match
6021+ if(cmpPrefix) {
6022+ bitmap_set(b, i, 0);
6023+ continue;
6024+ }
6025+ // consult the full string
6026+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6027+ if(rc != SQLITE_OK) {
6028+ goto done;
6029+ }
6030+ if(nPrefix != nFull) {
6031+ rc = SQLITE_ERROR;
6032+ goto done;
6033+ }
6034+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) == 0);
6035+ }
6036+ break;
6037+ }
6038+ case VEC0_METADATA_OPERATOR_NE: {
6039+ for(int i = 0; i < size; i++) {
6040+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6041+ nPrefix = ((int*) view)[0];
6042+ sPrefix = (char *) &view[4];
6043+
6044+ // for NE if text lengths dont match, it never will
6045+ if(nPrefix != nTarget) {
6046+ bitmap_set(b, i, 1);
6047+ continue;
6048+ }
6049+
6050+ int cmpPrefix = strncmp(sPrefix, sTarget, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH));
6051+
6052+ // for short strings, use the prefix comparison direclty
6053+ if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6054+ bitmap_set(b, i, cmpPrefix != 0);
6055+ continue;
6056+ }
6057+ // for NE on longs strings, if prefixes dont match, then long string wont
6058+ if(cmpPrefix) {
6059+ bitmap_set(b, i, 1);
6060+ continue;
6061+ }
6062+ // consult the full string
6063+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6064+ if(rc != SQLITE_OK) {
6065+ goto done;
6066+ }
6067+ if(nPrefix != nFull) {
6068+ rc = SQLITE_ERROR;
6069+ goto done;
6070+ }
6071+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) != 0);
6072+ }
6073+ break;
6074+ }
6075+ case VEC0_METADATA_OPERATOR_GT: {
6076+ for(int i = 0; i < size; i++) {
6077+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6078+ nPrefix = ((int*) view)[0];
6079+ sPrefix = (char *) &view[4];
6080+ int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget));
6081+
6082+ if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6083+ // if prefix match, check which is longer
6084+ if(cmpPrefix == 0) {
6085+ bitmap_set(b, i, nPrefix > nTarget);
6086+ }
6087+ else {
6088+ bitmap_set(b, i, cmpPrefix > 0);
6089+ }
6090+ continue;
6091+ }
6092+ // TODO(perf): may not need to compare full text in some cases
6093+
6094+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6095+ if(rc != SQLITE_OK) {
6096+ goto done;
6097+ }
6098+ if(nPrefix != nFull) {
6099+ rc = SQLITE_ERROR;
6100+ goto done;
6101+ }
6102+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) > 0);
6103+ }
6104+ break;
6105+ }
6106+ case VEC0_METADATA_OPERATOR_GE: {
6107+ for(int i = 0; i < size; i++) {
6108+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6109+ nPrefix = ((int*) view)[0];
6110+ sPrefix = (char *) &view[4];
6111+ int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget));
6112+
6113+ if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6114+ // if prefix match, check which is longer
6115+ if(cmpPrefix == 0) {
6116+ bitmap_set(b, i, nPrefix >= nTarget);
6117+ }
6118+ else {
6119+ bitmap_set(b, i, cmpPrefix >= 0);
6120+ }
6121+ continue;
6122+ }
6123+ // TODO(perf): may not need to compare full text in some cases
6124+
6125+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6126+ if(rc != SQLITE_OK) {
6127+ goto done;
6128+ }
6129+ if(nPrefix != nFull) {
6130+ rc = SQLITE_ERROR;
6131+ goto done;
6132+ }
6133+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) >= 0);
6134+ }
6135+ break;
6136+ }
6137+ case VEC0_METADATA_OPERATOR_LE: {
6138+ for(int i = 0; i < size; i++) {
6139+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6140+ nPrefix = ((int*) view)[0];
6141+ sPrefix = (char *) &view[4];
6142+ int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget));
6143+
6144+ if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6145+ // if prefix match, check which is longer
6146+ if(cmpPrefix == 0) {
6147+ bitmap_set(b, i, nPrefix <= nTarget);
6148+ }
6149+ else {
6150+ bitmap_set(b, i, cmpPrefix <= 0);
6151+ }
6152+ continue;
6153+ }
6154+ // TODO(perf): may not need to compare full text in some cases
6155+
6156+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6157+ if(rc != SQLITE_OK) {
6158+ goto done;
6159+ }
6160+ if(nPrefix != nFull) {
6161+ rc = SQLITE_ERROR;
6162+ goto done;
6163+ }
6164+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) <= 0);
6165+ }
6166+ break;
6167+ }
6168+ case VEC0_METADATA_OPERATOR_LT: {
6169+ for(int i = 0; i < size; i++) {
6170+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6171+ nPrefix = ((int*) view)[0];
6172+ sPrefix = (char *) &view[4];
6173+ int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget));
6174+
6175+ if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6176+ // if prefix match, check which is longer
6177+ if(cmpPrefix == 0) {
6178+ bitmap_set(b, i, nPrefix < nTarget);
6179+ }
6180+ else {
6181+ bitmap_set(b, i, cmpPrefix < 0);
6182+ }
6183+ continue;
6184+ }
6185+ // TODO(perf): may not need to compare full text in some cases
6186+
6187+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6188+ if(rc != SQLITE_OK) {
6189+ goto done;
6190+ }
6191+ if(nPrefix != nFull) {
6192+ rc = SQLITE_ERROR;
6193+ goto done;
6194+ }
6195+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) < 0);
6196+ }
6197+ break;
6198+ }
6199+
6200+ case VEC0_METADATA_OPERATOR_IN: {
6201+ size_t metadataInIdx = -1;
6202+ for(size_t i = 0; i < aMetadataIn->length; i++) {
6203+ struct Vec0MetadataIn * metadataIn = &(((struct Vec0MetadataIn *) aMetadataIn->z)[i]);
6204+ if(metadataIn->argv_idx == argv_idx) {
6205+ metadataInIdx = i;
6206+ break;
6207+ }
6208+ }
6209+ if(metadataInIdx < 0) {
6210+ rc = SQLITE_ERROR;
6211+ goto done;
6212+ }
6213+
6214+ struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[metadataInIdx];
6215+ struct Array * aTarget = &(metadataIn->array);
6216+
6217+
6218+ int nPrefix;
6219+ char * sPrefix;
6220+ char *sFull;
6221+ int nFull;
6222+ u8 * view;
6223+ for(int i = 0; i < size; i++) {
6224+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6225+ nPrefix = ((int*) view)[0];
6226+ sPrefix = (char *) &view[4];
6227+ for(size_t target_idx = 0; target_idx < aTarget->length; target_idx++) {
6228+ struct Vec0MetadataInTextEntry * entry = &(((struct Vec0MetadataInTextEntry*)aTarget->z)[target_idx]);
6229+ if(entry->n != nPrefix) {
6230+ continue;
6231+ }
6232+ int cmpPrefix = strncmp(sPrefix, entry->zString, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH));
6233+ if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6234+ if(cmpPrefix == 0) {
6235+ bitmap_set(b, i, 1);
6236+ break;
6237+ }
6238+ continue;
6239+ }
6240+ if(cmpPrefix) {
6241+ continue;
6242+ }
6243+
6244+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6245+ if(rc != SQLITE_OK) {
6246+ goto done;
6247+ }
6248+ if(nPrefix != nFull) {
6249+ rc = SQLITE_ERROR;
6250+ goto done;
6251+ }
6252+ if(strncmp(sFull, entry->zString, nFull) == 0) {
6253+ bitmap_set(b, i, 1);
6254+ break;
6255+ }
6256+ }
6257+ }
6258+ break;
6259+ }
6260+
6261+ }
6262+ rc = SQLITE_OK;
6263+
6264+ done:
6265+ sqlite3_finalize(stmt);
6266+ sqlite3_free(rowids);
6267+ return rc;
6268+
6269+}
6270+
6271+/**
6272+ * @brief Fill in bitmap of chunk values, whether or not the values match a metadata constraint
6273+ *
6274+ * @param p vec0_vtab
6275+ * @param metadata_idx index of the metatadata column to perfrom constraints on
6276+ * @param value sqlite3_value of the constraints value
6277+ * @param blob sqlite3_blob that is already opened on the metdata column's shadow chunk table
6278+ * @param chunk_rowid rowid of the chunk to calculate on
6279+ * @param b pre-allocated and zero'd out bitmap to write results to
6280+ * @param size size of the chunk
6281+ * @return int SQLITE_OK on success, error code otherwise
6282+ */
6283+int vec0_set_metadata_filter_bitmap(
6284+ vec0_vtab *p,
6285+ int metadata_idx,
6286+ vec0_metadata_operator op,
6287+ sqlite3_value * value,
6288+ sqlite3_blob * blob,
6289+ i64 chunk_rowid,
6290+ u8* b,
6291+ int size,
6292+ struct Array * aMetadataIn, int argv_idx) {
6293+ // TODO: shouldn't this skip in-valid entries from the chunk's validity bitmap?
6294+
6295+ int rc;
6296+ rc = sqlite3_blob_reopen(blob, chunk_rowid);
6297+ if(rc != SQLITE_OK) {
6298+ return rc;
6299+ }
6300+
6301+ vec0_metadata_column_kind kind = p->metadata_columns[metadata_idx].kind;
6302+ int szMatch = 0;
6303+ int blobSize = sqlite3_blob_bytes(blob);
6304+ switch(kind) {
6305+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
6306+ szMatch = blobSize == size / CHAR_BIT;
6307+ break;
6308+ }
6309+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
6310+ szMatch = blobSize == size * sizeof(i64);
6311+ break;
6312+ }
6313+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
6314+ szMatch = blobSize == size * sizeof(double);
6315+ break;
6316+ }
6317+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
6318+ szMatch = blobSize == size * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH;
6319+ break;
6320+ }
6321+ }
6322+ if(!szMatch) {
6323+ return SQLITE_ERROR;
6324+ }
6325+ void * buffer = sqlite3_malloc(blobSize);
6326+ if(!buffer) {
6327+ return SQLITE_NOMEM;
6328+ }
6329+ rc = sqlite3_blob_read(blob, buffer, blobSize, 0);
6330+ if(rc != SQLITE_OK) {
6331+ goto done;
6332+ }
6333+ switch(kind) {
6334+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
6335+ int target = sqlite3_value_int(value);
6336+ if( (target && op == VEC0_METADATA_OPERATOR_EQ) || (!target && op == VEC0_METADATA_OPERATOR_NE)) {
6337+ for(int i = 0; i < size; i++) { bitmap_set(b, i, bitmap_get((u8*) buffer, i)); }
6338+ }
6339+ else {
6340+ for(int i = 0; i < size; i++) { bitmap_set(b, i, !bitmap_get((u8*) buffer, i)); }
6341+ }
6342+ break;
6343+ }
6344+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
6345+ i64 * array = (i64*) buffer;
6346+ i64 target = sqlite3_value_int64(value);
6347+ switch(op) {
6348+ case VEC0_METADATA_OPERATOR_EQ: {
6349+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] == target); }
6350+ break;
6351+ }
6352+ case VEC0_METADATA_OPERATOR_GT: {
6353+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] > target); }
6354+ break;
6355+ }
6356+ case VEC0_METADATA_OPERATOR_LE: {
6357+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] <= target); }
6358+ break;
6359+ }
6360+ case VEC0_METADATA_OPERATOR_LT: {
6361+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] < target); }
6362+ break;
6363+ }
6364+ case VEC0_METADATA_OPERATOR_GE: {
6365+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] >= target); }
6366+ break;
6367+ }
6368+ case VEC0_METADATA_OPERATOR_NE: {
6369+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] != target); }
6370+ break;
6371+ }
6372+ case VEC0_METADATA_OPERATOR_IN: {
6373+ int metadataInIdx = -1;
6374+ for(size_t i = 0; i < aMetadataIn->length; i++) {
6375+ struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[i];
6376+ if(metadataIn->argv_idx == argv_idx) {
6377+ metadataInIdx = i;
6378+ break;
6379+ }
6380+ }
6381+ if(metadataInIdx < 0) {
6382+ rc = SQLITE_ERROR;
6383+ goto done;
6384+ }
6385+ struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[metadataInIdx];
6386+ struct Array * aTarget = &(metadataIn->array);
6387+
6388+ for(int i = 0; i < size; i++) {
6389+ for(size_t target_idx = 0; target_idx < aTarget->length; target_idx++) {
6390+ if( ((i64*)aTarget->z)[target_idx] == array[i]) {
6391+ bitmap_set(b, i, 1);
6392+ break;
6393+ }
6394+ }
6395+ }
6396+ break;
6397+ }
6398+ }
6399+ break;
6400+ }
6401+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
6402+ double * array = (double*) buffer;
6403+ double target = sqlite3_value_double(value);
6404+ switch(op) {
6405+ case VEC0_METADATA_OPERATOR_EQ: {
6406+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] == target); }
6407+ break;
6408+ }
6409+ case VEC0_METADATA_OPERATOR_GT: {
6410+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] > target); }
6411+ break;
6412+ }
6413+ case VEC0_METADATA_OPERATOR_LE: {
6414+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] <= target); }
6415+ break;
6416+ }
6417+ case VEC0_METADATA_OPERATOR_LT: {
6418+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] < target); }
6419+ break;
6420+ }
6421+ case VEC0_METADATA_OPERATOR_GE: {
6422+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] >= target); }
6423+ break;
6424+ }
6425+ case VEC0_METADATA_OPERATOR_NE: {
6426+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] != target); }
6427+ break;
6428+ }
6429+ case VEC0_METADATA_OPERATOR_IN: {
6430+ // should never be reached
6431+ break;
6432+ }
6433+ }
6434+ break;
6435+ }
6436+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
6437+ rc = vec0_metadata_filter_text(p, value, buffer, size, op, b, metadata_idx, chunk_rowid, aMetadataIn, argv_idx);
6438+ if(rc != SQLITE_OK) {
6439+ goto done;
6440+ }
6441+ break;
6442+ }
6443+ }
6444+ done:
6445+ sqlite3_free(buffer);
6446+ return rc;
6447+}
6448+
6449+int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
6450+ struct VectorColumnDefinition *vector_column,
6451+ int vectorColumnIdx, struct Array *arrayRowidsIn,
6452+ struct Array * aMetadataIn,
6453+ const char * idxStr, int argc, sqlite3_value ** argv,
6454+ void *queryVector, i64 k, i64 **out_topk_rowids,
6455+ f32 **out_topk_distances, i64 *out_used) {
6456+ // for each chunk, get top min(k, chunk_size) rowid + distances to query vec.
6457+ // then reconcile all topk_chunks for a true top k.
6458+ // output only rowids + distances for now
6459+
6460+ int rc = SQLITE_OK;
6461+ sqlite3_blob *blobVectors = NULL;
6462+
6463+ void *baseVectors = NULL; // memory: chunk_size * dimensions * element_size
6464+
6465+ // OWNED BY CALLER ON SUCCESS
6466+ i64 *topk_rowids = NULL; // memory: k * 4
6467+ // OWNED BY CALLER ON SUCCESS
6468+ f32 *topk_distances = NULL; // memory: k * 4
6469+
6470+ i64 *tmp_topk_rowids = NULL; // memory: k * 4
6471+ f32 *tmp_topk_distances = NULL; // memory: k * 4
6472+ f32 *chunk_distances = NULL; // memory: chunk_size * 4
6473+ u8 *b = NULL; // memory: chunk_size / 8
6474+ u8 *bTaken = NULL; // memory: chunk_size / 8
6475+ i32 *chunk_topk_idxs = NULL; // memory: k * 4
6476+ u8 *bmRowids = NULL; // memory: chunk_size / 8
6477+ u8 *bmMetadata = NULL; // memory: chunk_size / 8
6478+ // // total: a lot???
6479+
6480+ // 6 * (k * 4) + (k * 2) + (chunk_size / 8) + (chunk_size * dimensions * 4)
6481+
6482+ topk_rowids = sqlite3_malloc(k * sizeof(i64));
6483+ if (!topk_rowids) {
6484+ rc = SQLITE_NOMEM;
6485+ goto cleanup;
6486+ }
6487+ memset(topk_rowids, 0, k * sizeof(i64));
6488+
6489+ topk_distances = sqlite3_malloc(k * sizeof(f32));
6490+ if (!topk_distances) {
6491+ rc = SQLITE_NOMEM;
6492+ goto cleanup;
6493+ }
6494+ memset(topk_distances, 0, k * sizeof(f32));
6495+
6496+ tmp_topk_rowids = sqlite3_malloc(k * sizeof(i64));
6497+ if (!tmp_topk_rowids) {
6498+ rc = SQLITE_NOMEM;
6499+ goto cleanup;
6500+ }
6501+ memset(tmp_topk_rowids, 0, k * sizeof(i64));
6502+
6503+ tmp_topk_distances = sqlite3_malloc(k * sizeof(f32));
6504+ if (!tmp_topk_distances) {
6505+ rc = SQLITE_NOMEM;
6506+ goto cleanup;
6507+ }
6508+ memset(tmp_topk_distances, 0, k * sizeof(f32));
6509+
6510+ i64 k_used = 0;
6511+ i64 baseVectorsSize = p->chunk_size * vector_column_byte_size(*vector_column);
6512+ baseVectors = sqlite3_malloc(baseVectorsSize);
6513+ if (!baseVectors) {
6514+ rc = SQLITE_NOMEM;
6515+ goto cleanup;
6516+ }
6517+
6518+ chunk_distances = sqlite3_malloc(p->chunk_size * sizeof(f32));
6519+ if (!chunk_distances) {
6520+ rc = SQLITE_NOMEM;
6521+ goto cleanup;
6522+ }
6523+
6524+ b = bitmap_new(p->chunk_size);
6525+ if (!b) {
6526+ rc = SQLITE_NOMEM;
6527+ goto cleanup;
6528+ }
6529+
6530+ bTaken = bitmap_new(p->chunk_size);
6531+ if (!bTaken) {
6532+ rc = SQLITE_NOMEM;
6533+ goto cleanup;
6534+ }
6535+
6536+ chunk_topk_idxs = sqlite3_malloc(k * sizeof(i32));
6537+ if (!chunk_topk_idxs) {
6538+ rc = SQLITE_NOMEM;
6539+ goto cleanup;
6540+ }
6541+
6542+ bmRowids = arrayRowidsIn ? bitmap_new(p->chunk_size) : NULL;
6543+ if (arrayRowidsIn && !bmRowids) {
6544+ rc = SQLITE_NOMEM;
6545+ goto cleanup;
6546+ }
6547+
6548+ sqlite3_blob * metadataBlobs[VEC0_MAX_METADATA_COLUMNS];
6549+ memset(metadataBlobs, 0, sizeof(sqlite3_blob*) * VEC0_MAX_METADATA_COLUMNS);
6550+
6551+ bmMetadata = bitmap_new(p->chunk_size);
6552+ if(!bmMetadata) {
6553+ rc = SQLITE_NOMEM;
6554+ goto cleanup;
6555+ }
6556+
6557+ int idxStrLength = strlen(idxStr);
6558+ int numValueEntries = (idxStrLength-1) / 4;
6559+ assert(numValueEntries == argc);
6560+ int hasMetadataFilters = 0;
6561+ for(int i = 0; i < argc; i++) {
6562+ int idx = 1 + (i * 4);
6563+ char kind = idxStr[idx + 0];
6564+ if(kind == VEC0_IDXSTR_KIND_METADATA_CONSTRAINT) {
6565+ hasMetadataFilters = 1;
6566+ break;
6567+ }
6568+ }
6569+
6570+ while (true) {
6571+ rc = sqlite3_step(stmtChunks);
6572+ if (rc == SQLITE_DONE) {
6573+ break;
6574+ }
6575+ if (rc != SQLITE_ROW) {
6576+ vtab_set_error(&p->base, "chunks iter error");
6577+ rc = SQLITE_ERROR;
6578+ goto cleanup;
6579+ }
6580+ memset(chunk_distances, 0, p->chunk_size * sizeof(f32));
6581+ memset(chunk_topk_idxs, 0, k * sizeof(i32));
6582+ bitmap_clear(b, p->chunk_size);
6583+
6584+ i64 chunk_id = sqlite3_column_int64(stmtChunks, 0);
6585+ unsigned char *chunkValidity =
6586+ (unsigned char *)sqlite3_column_blob(stmtChunks, 1);
6587+ i64 validitySize = sqlite3_column_bytes(stmtChunks, 1);
6588+ if (validitySize != p->chunk_size / CHAR_BIT) {
6589+ // IMP: V05271_22109
6590+ vtab_set_error(
6591+ &p->base,
6592+ "chunk validity size doesn't match - expected %lld, found %lld",
6593+ p->chunk_size / CHAR_BIT, validitySize);
6594+ rc = SQLITE_ERROR;
6595+ goto cleanup;
6596+ }
6597+
6598+ i64 *chunkRowids = (i64 *)sqlite3_column_blob(stmtChunks, 2);
6599+ i64 rowidsSize = sqlite3_column_bytes(stmtChunks, 2);
6600+ if (rowidsSize != p->chunk_size * sizeof(i64)) {
6601+ // IMP: V02796_19635
6602+ vtab_set_error(&p->base, "rowids size doesn't match");
6603+ vtab_set_error(
6604+ &p->base,
6605+ "chunk rowids size doesn't match - expected %lld, found %lld",
6606+ p->chunk_size * sizeof(i64), rowidsSize);
6607+ rc = SQLITE_ERROR;
6608+ goto cleanup;
6609+ }
6610+
6611+ // open the vector chunk blob for the current chunk
6612+ rc = sqlite3_blob_open(p->db, p->schemaName,
6613+ p->shadowVectorChunksNames[vectorColumnIdx],
6614+ "vectors", chunk_id, 0, &blobVectors);
6615+ if (rc != SQLITE_OK) {
6616+ vtab_set_error(&p->base, "could not open vectors blob for chunk %lld",
6617+ chunk_id);
6618+ rc = SQLITE_ERROR;
6619+ goto cleanup;
6620+ }
6621+
6622+ i64 currentBaseVectorsSize = sqlite3_blob_bytes(blobVectors);
6623+ i64 expectedBaseVectorsSize =
6624+ p->chunk_size * vector_column_byte_size(*vector_column);
6625+ if (currentBaseVectorsSize != expectedBaseVectorsSize) {
6626+ // IMP: V16465_00535
6627+ vtab_set_error(
6628+ &p->base,
6629+ "vectors blob size doesn't match - expected %lld, found %lld",
6630+ expectedBaseVectorsSize, currentBaseVectorsSize);
6631+ rc = SQLITE_ERROR;
6632+ goto cleanup;
6633+ }
6634+ rc = sqlite3_blob_read(blobVectors, baseVectors, currentBaseVectorsSize, 0);
6635+
6636+ if (rc != SQLITE_OK) {
6637+ vtab_set_error(&p->base, "vectors blob read error for %lld", chunk_id);
6638+ rc = SQLITE_ERROR;
6639+ goto cleanup;
6640+ }
6641+
6642+ bitmap_copy(b, chunkValidity, p->chunk_size);
6643+ if (arrayRowidsIn) {
6644+ bitmap_clear(bmRowids, p->chunk_size);
6645+
6646+ for (int i = 0; i < p->chunk_size; i++) {
6647+ if (!bitmap_get(chunkValidity, i)) {
6648+ continue;
6649+ }
6650+ i64 rowid = chunkRowids[i];
6651+ void *in = bsearch(&rowid, arrayRowidsIn->z, arrayRowidsIn->length,
6652+ sizeof(i64), _cmp);
6653+ bitmap_set(bmRowids, i, in ? 1 : 0);
6654+ }
6655+ bitmap_and_inplace(b, bmRowids, p->chunk_size);
6656+ }
6657+
6658+ if(hasMetadataFilters) {
6659+ for(int i = 0; i < argc; i++) {
6660+ int idx = 1 + (i * 4);
6661+ char kind = idxStr[idx + 0];
6662+ if(kind != VEC0_IDXSTR_KIND_METADATA_CONSTRAINT) {
6663+ continue;
6664+ }
6665+ int metadata_idx = idxStr[idx + 1] - 'A';
6666+ int operator = idxStr[idx + 2];
6667+
6668+ if(!metadataBlobs[metadata_idx]) {
6669+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 0, &metadataBlobs[metadata_idx]);
6670+ vtab_set_error(&p->base, "Could not open metadata blob");
6671+ if(rc != SQLITE_OK) {
6672+ goto cleanup;
6673+ }
6674+ }
6675+
6676+ bitmap_clear(bmMetadata, p->chunk_size);
6677+ rc = vec0_set_metadata_filter_bitmap(p, metadata_idx, operator, argv[i], metadataBlobs[metadata_idx], chunk_id, bmMetadata, p->chunk_size, aMetadataIn, i);
6678+ if(rc != SQLITE_OK) {
6679+ vtab_set_error(&p->base, "Could not filter metadata fields");
6680+ if(rc != SQLITE_OK) {
6681+ goto cleanup;
6682+ }
6683+ }
6684+ bitmap_and_inplace(b, bmMetadata, p->chunk_size);
6685+ }
6686+ }
6687+
6688+
6689+ for (int i = 0; i < p->chunk_size; i++) {
6690+ if (!bitmap_get(b, i)) {
6691+ continue;
6692+ };
6693+
6694+ f32 result;
6695+ switch (vector_column->element_type) {
6696+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
6697+ const f32 *base_i =
6698+ ((f32 *)baseVectors) + (i * vector_column->dimensions);
6699+ switch (vector_column->distance_metric) {
6700+ case VEC0_DISTANCE_METRIC_L2: {
6701+ result = distance_l2_sqr_float(base_i, (f32 *)queryVector,
6702+ &vector_column->dimensions);
6703+ break;
6704+ }
6705+ case VEC0_DISTANCE_METRIC_L1: {
6706+ result = distance_l1_f32(base_i, (f32 *)queryVector,
6707+ &vector_column->dimensions);
6708+ break;
6709+ }
6710+ case VEC0_DISTANCE_METRIC_COSINE: {
6711+ result = distance_cosine_float(base_i, (f32 *)queryVector,
6712+ &vector_column->dimensions);
6713+ break;
6714+ }
6715+ }
6716+ break;
6717+ }
6718+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
6719+ const i8 *base_i =
6720+ ((i8 *)baseVectors) + (i * vector_column->dimensions);
6721+ switch (vector_column->distance_metric) {
6722+ case VEC0_DISTANCE_METRIC_L2: {
6723+ result = distance_l2_sqr_int8(base_i, (i8 *)queryVector,
6724+ &vector_column->dimensions);
6725+ break;
6726+ }
6727+ case VEC0_DISTANCE_METRIC_L1: {
6728+ result = distance_l1_int8(base_i, (i8 *)queryVector,
6729+ &vector_column->dimensions);
6730+ break;
6731+ }
6732+ case VEC0_DISTANCE_METRIC_COSINE: {
6733+ result = distance_cosine_int8(base_i, (i8 *)queryVector,
6734+ &vector_column->dimensions);
6735+ break;
6736+ }
6737+ }
6738+
6739+ break;
6740+ }
6741+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
6742+ const u8 *base_i =
6743+ ((u8 *)baseVectors) + (i * (vector_column->dimensions / CHAR_BIT));
6744+ result = distance_hamming(base_i, (u8 *)queryVector,
6745+ &vector_column->dimensions);
6746+ break;
6747+ }
6748+ }
6749+
6750+ chunk_distances[i] = result;
6751+ }
6752+
6753+ int used1;
6754+ min_idx(chunk_distances, p->chunk_size, b, chunk_topk_idxs,
6755+ min(k, p->chunk_size), bTaken, &used1);
6756+
6757+ i64 used;
6758+ merge_sorted_lists(topk_distances, topk_rowids, k_used, chunk_distances,
6759+ chunkRowids, chunk_topk_idxs,
6760+ min(min(k, p->chunk_size), used1), tmp_topk_distances,
6761+ tmp_topk_rowids, k, &used);
6762+
6763+ for (int i = 0; i < used; i++) {
6764+ topk_rowids[i] = tmp_topk_rowids[i];
6765+ topk_distances[i] = tmp_topk_distances[i];
6766+ }
6767+ k_used = used;
6768+ // blobVectors is always opened with read-only permissions, so this never
6769+ // fails.
6770+ sqlite3_blob_close(blobVectors);
6771+ blobVectors = NULL;
6772+ }
6773+
6774+ *out_topk_rowids = topk_rowids;
6775+ *out_topk_distances = topk_distances;
6776+ *out_used = k_used;
6777+ rc = SQLITE_OK;
6778+
6779+cleanup:
6780+ if (rc != SQLITE_OK) {
6781+ sqlite3_free(topk_rowids);
6782+ sqlite3_free(topk_distances);
6783+ }
6784+ sqlite3_free(chunk_topk_idxs);
6785+ sqlite3_free(tmp_topk_rowids);
6786+ sqlite3_free(tmp_topk_distances);
6787+ sqlite3_free(b);
6788+ sqlite3_free(bTaken);
6789+ sqlite3_free(bmRowids);
6790+ sqlite3_free(baseVectors);
6791+ sqlite3_free(chunk_distances);
6792+ sqlite3_free(bmMetadata);
6793+ for(int i = 0; i < VEC0_MAX_METADATA_COLUMNS; i++) {
6794+ sqlite3_blob_close(metadataBlobs[i]);
6795+ }
6796+ // blobVectors is always opened with read-only permissions, so this never
6797+ // fails.
6798+ sqlite3_blob_close(blobVectors);
6799+ return rc;
6800+}
6801+
6802+int vec0Filter_knn(vec0_cursor *pCur, vec0_vtab *p, int idxNum,
6803+ const char *idxStr, int argc, sqlite3_value **argv) {
6804+ assert(argc == (strlen(idxStr)-1) / 4);
6805+ int rc;
6806+ struct vec0_query_knn_data *knn_data;
6807+
6808+ int vectorColumnIdx = idxNum;
6809+ struct VectorColumnDefinition *vector_column =
6810+ &p->vector_columns[vectorColumnIdx];
6811+
6812+ struct Array *arrayRowidsIn = NULL;
6813+ sqlite3_stmt *stmtChunks = NULL;
6814+ void *queryVector;
6815+ size_t dimensions;
6816+ enum VectorElementType elementType;
6817+ vector_cleanup queryVectorCleanup = vector_cleanup_noop;
6818+ char *pzError;
6819+ knn_data = sqlite3_malloc(sizeof(*knn_data));
6820+ if (!knn_data) {
6821+ return SQLITE_NOMEM;
6822+ }
6823+ memset(knn_data, 0, sizeof(*knn_data));
6824+ // array of `struct Vec0MetadataIn`, IF there are any `xxx in (...)` metadata constraints
6825+ struct Array * aMetadataIn = NULL;
6826+
6827+ int query_idx =-1;
6828+ int k_idx = -1;
6829+ int rowid_in_idx = -1;
6830+ for(int i = 0; i < argc; i++) {
6831+ if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_MATCH) {
6832+ query_idx = i;
6833+ }
6834+ if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_K) {
6835+ k_idx = i;
6836+ }
6837+ if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_ROWID_IN) {
6838+ rowid_in_idx = i;
6839+ }
6840+ }
6841+ assert(query_idx >= 0);
6842+ assert(k_idx >= 0);
6843+
6844+ // make sure the query vector matches the vector column (type dimensions etc.)
6845+ rc = vector_from_value(argv[query_idx], &queryVector, &dimensions, &elementType,
6846+ &queryVectorCleanup, &pzError);
6847+
6848+ if (rc != SQLITE_OK) {
6849+ vtab_set_error(&p->base,
6850+ "Query vector on the \"%.*s\" column is invalid: %z",
6851+ vector_column->name_length, vector_column->name, pzError);
6852+ rc = SQLITE_ERROR;
6853+ goto cleanup;
6854+ }
6855+ if (elementType != vector_column->element_type) {
6856+ vtab_set_error(
6857+ &p->base,
6858+ "Query vector for the \"%.*s\" column is expected to be of type "
6859+ "%s, but a %s vector was provided.",
6860+ vector_column->name_length, vector_column->name,
6861+ vector_subtype_name(vector_column->element_type),
6862+ vector_subtype_name(elementType));
6863+ rc = SQLITE_ERROR;
6864+ goto cleanup;
6865+ }
6866+ if (dimensions != vector_column->dimensions) {
6867+ vtab_set_error(
6868+ &p->base,
6869+ "Dimension mismatch for query vector for the \"%.*s\" column. "
6870+ "Expected %d dimensions but received %d.",
6871+ vector_column->name_length, vector_column->name,
6872+ vector_column->dimensions, dimensions);
6873+ rc = SQLITE_ERROR;
6874+ goto cleanup;
6875+ }
6876+
6877+ i64 k = sqlite3_value_int64(argv[k_idx]);
6878+ if (k < 0) {
6879+ vtab_set_error(
6880+ &p->base, "k value in knn queries must be greater than or equal to 0.");
6881+ rc = SQLITE_ERROR;
6882+ goto cleanup;
6883+ }
6884+#define SQLITE_VEC_VEC0_K_MAX 4096
6885+ if (k > SQLITE_VEC_VEC0_K_MAX) {
6886+ vtab_set_error(
6887+ &p->base,
6888+ "k value in knn query too large, provided %lld and the limit is %lld",
6889+ k, SQLITE_VEC_VEC0_K_MAX);
6890+ rc = SQLITE_ERROR;
6891+ goto cleanup;
6892+ }
6893+
6894+ if (k == 0) {
6895+ knn_data->k = 0;
6896+ pCur->knn_data = knn_data;
6897+ pCur->query_plan = VEC0_QUERY_PLAN_KNN;
6898+ rc = SQLITE_OK;
6899+ goto cleanup;
6900+ }
6901+
6902+// handle when a `rowid in (...)` operation was provided
6903+// Array of all the rowids that appear in any `rowid in (...)` constraint.
6904+// NULL if none were provided, which means a "full" scan.
6905+#if COMPILER_SUPPORTS_VTAB_IN
6906+ if (rowid_in_idx >= 0) {
6907+ sqlite3_value *item;
6908+ int rc;
6909+ arrayRowidsIn = sqlite3_malloc(sizeof(*arrayRowidsIn));
6910+ if (!arrayRowidsIn) {
6911+ rc = SQLITE_NOMEM;
6912+ goto cleanup;
6913+ }
6914+ memset(arrayRowidsIn, 0, sizeof(*arrayRowidsIn));
6915+
6916+ rc = array_init(arrayRowidsIn, sizeof(i64), 32);
6917+ if (rc != SQLITE_OK) {
6918+ goto cleanup;
6919+ }
6920+ for (rc = sqlite3_vtab_in_first(argv[rowid_in_idx], &item); rc == SQLITE_OK && item;
6921+ rc = sqlite3_vtab_in_next(argv[rowid_in_idx], &item)) {
6922+ i64 rowid;
6923+ if (p->pkIsText) {
6924+ rc = vec0_rowid_from_id(p, item, &rowid);
6925+ if (rc != SQLITE_OK) {
6926+ goto cleanup;
6927+ }
6928+ } else {
6929+ rowid = sqlite3_value_int64(item);
6930+ }
6931+ rc = array_append(arrayRowidsIn, &rowid);
6932+ if (rc != SQLITE_OK) {
6933+ goto cleanup;
6934+ }
6935+ }
6936+ if (rc != SQLITE_DONE) {
6937+ vtab_set_error(&p->base, "error processing rowid in (...) array");
6938+ goto cleanup;
6939+ }
6940+ qsort(arrayRowidsIn->z, arrayRowidsIn->length, arrayRowidsIn->element_size,
6941+ _cmp);
6942+ }
6943+#endif
6944+
6945+ #if COMPILER_SUPPORTS_VTAB_IN
6946+ for(int i = 0; i < argc; i++) {
6947+ if(!(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_METADATA_CONSTRAINT && idxStr[1 + (i*4) + 2] == VEC0_METADATA_OPERATOR_IN)) {
6948+ continue;
6949+ }
6950+ int metadata_idx = idxStr[1 + (i*4) + 1] - 'A';
6951+ if(!aMetadataIn) {
6952+ aMetadataIn = sqlite3_malloc(sizeof(*aMetadataIn));
6953+ if(!aMetadataIn) {
6954+ rc = SQLITE_NOMEM;
6955+ goto cleanup;
6956+ }
6957+ memset(aMetadataIn, 0, sizeof(*aMetadataIn));
6958+ rc = array_init(aMetadataIn, sizeof(struct Vec0MetadataIn), 8);
6959+ if(rc != SQLITE_OK) {
6960+ goto cleanup;
6961+ }
6962+ }
6963+
6964+ struct Vec0MetadataIn item;
6965+ memset(&item, 0, sizeof(item));
6966+ item.metadata_idx=metadata_idx;
6967+ item.argv_idx = i;
6968+
6969+ switch(p->metadata_columns[metadata_idx].kind) {
6970+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
6971+ rc = array_init(&item.array, sizeof(i64), 16);
6972+ if(rc != SQLITE_OK) {
6973+ goto cleanup;
6974+ }
6975+ sqlite3_value *entry;
6976+ for (rc = sqlite3_vtab_in_first(argv[i], &entry); rc == SQLITE_OK && entry; rc = sqlite3_vtab_in_next(argv[i], &entry)) {
6977+ i64 v = sqlite3_value_int64(entry);
6978+ rc = array_append(&item.array, &v);
6979+ if (rc != SQLITE_OK) {
6980+ goto cleanup;
6981+ }
6982+ }
6983+
6984+ if (rc != SQLITE_DONE) {
6985+ vtab_set_error(&p->base, "Error fetching next value in `x in (...)` integer expression");
6986+ goto cleanup;
6987+ }
6988+
6989+ break;
6990+ }
6991+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
6992+ rc = array_init(&item.array, sizeof(struct Vec0MetadataInTextEntry), 16);
6993+ if(rc != SQLITE_OK) {
6994+ goto cleanup;
6995+ }
6996+ sqlite3_value *entry;
6997+ for (rc = sqlite3_vtab_in_first(argv[i], &entry); rc == SQLITE_OK && entry; rc = sqlite3_vtab_in_next(argv[i], &entry)) {
6998+ const char * s = (const char *) sqlite3_value_text(entry);
6999+ int n = sqlite3_value_bytes(entry);
7000+
7001+ struct Vec0MetadataInTextEntry entry;
7002+ entry.zString = sqlite3_mprintf("%.*s", n, s);
7003+ if(!entry.zString) {
7004+ rc = SQLITE_NOMEM;
7005+ goto cleanup;
7006+ }
7007+ entry.n = n;
7008+ rc = array_append(&item.array, &entry);
7009+ if (rc != SQLITE_OK) {
7010+ goto cleanup;
7011+ }
7012+ }
7013+
7014+ if (rc != SQLITE_DONE) {
7015+ vtab_set_error(&p->base, "Error fetching next value in `x in (...)` text expression");
7016+ goto cleanup;
7017+ }
7018+
7019+ break;
7020+ }
7021+ default: {
7022+ vtab_set_error(&p->base, "Internal sqlite-vec error");
7023+ goto cleanup;
7024+ }
7025+ }
7026+
7027+ rc = array_append(aMetadataIn, &item);
7028+ if(rc != SQLITE_OK) {
7029+ goto cleanup;
7030+ }
7031+ }
7032+ #endif
7033+
7034+ rc = vec0_chunks_iter(p, idxStr, argc, argv, &stmtChunks);
7035+ if (rc != SQLITE_OK) {
7036+ // IMP: V06942_23781
7037+ vtab_set_error(&p->base, "Error preparing stmtChunk: %s",
7038+ sqlite3_errmsg(p->db));
7039+ goto cleanup;
7040+ }
7041+
7042+ i64 *topk_rowids = NULL;
7043+ f32 *topk_distances = NULL;
7044+ i64 k_used = 0;
7045+ rc = vec0Filter_knn_chunks_iter(p, stmtChunks, vector_column, vectorColumnIdx,
7046+ arrayRowidsIn, aMetadataIn, idxStr, argc, argv, queryVector, k, &topk_rowids,
7047+ &topk_distances, &k_used);
7048+ if (rc != SQLITE_OK) {
7049+ goto cleanup;
7050+ }
7051+
7052+ knn_data->current_idx = 0;
7053+ knn_data->k = k;
7054+ knn_data->rowids = topk_rowids;
7055+ knn_data->distances = topk_distances;
7056+ knn_data->k_used = k_used;
7057+
7058+ pCur->knn_data = knn_data;
7059+ pCur->query_plan = VEC0_QUERY_PLAN_KNN;
7060+ rc = SQLITE_OK;
7061+
7062+cleanup:
7063+ sqlite3_finalize(stmtChunks);
7064+ array_cleanup(arrayRowidsIn);
7065+ sqlite3_free(arrayRowidsIn);
7066+ queryVectorCleanup(queryVector);
7067+ if(aMetadataIn) {
7068+ for(size_t i = 0; i < aMetadataIn->length; i++) {
7069+ struct Vec0MetadataIn* item = &((struct Vec0MetadataIn *) aMetadataIn->z)[i];
7070+ for(size_t j = 0; j < item->array.length; j++) {
7071+ if(p->metadata_columns[item->metadata_idx].kind == VEC0_METADATA_COLUMN_KIND_TEXT) {
7072+ struct Vec0MetadataInTextEntry entry = ((struct Vec0MetadataInTextEntry*)item->array.z)[j];
7073+ sqlite3_free(entry.zString);
7074+ }
7075+ }
7076+ array_cleanup(&item->array);
7077+ }
7078+ array_cleanup(aMetadataIn);
7079+ }
7080+
7081+ sqlite3_free(aMetadataIn);
7082+
7083+ return rc;
7084+}
7085+
7086+int vec0Filter_fullscan(vec0_vtab *p, vec0_cursor *pCur) {
7087+ int rc;
7088+ char *zSql;
7089+ struct vec0_query_fullscan_data *fullscan_data;
7090+
7091+ fullscan_data = sqlite3_malloc(sizeof(*fullscan_data));
7092+ if (!fullscan_data) {
7093+ return SQLITE_NOMEM;
7094+ }
7095+ memset(fullscan_data, 0, sizeof(*fullscan_data));
7096+
7097+ zSql = sqlite3_mprintf(" SELECT rowid "
7098+ " FROM " VEC0_SHADOW_ROWIDS_NAME
7099+ " ORDER by chunk_id, chunk_offset ",
7100+ p->schemaName, p->tableName);
7101+ if (!zSql) {
7102+ rc = SQLITE_NOMEM;
7103+ goto error;
7104+ }
7105+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &fullscan_data->rowids_stmt, NULL);
7106+ sqlite3_free(zSql);
7107+ if (rc != SQLITE_OK) {
7108+ // IMP: V09901_26739
7109+ vtab_set_error(&p->base, "Error preparing rowid scan: %s",
7110+ sqlite3_errmsg(p->db));
7111+ goto error;
7112+ }
7113+
7114+ rc = sqlite3_step(fullscan_data->rowids_stmt);
7115+
7116+ // DONE when there's no rowids, ROW when there are, both "success"
7117+ if (!(rc == SQLITE_ROW || rc == SQLITE_DONE)) {
7118+ goto error;
7119+ }
7120+
7121+ fullscan_data->done = rc == SQLITE_DONE;
7122+ pCur->query_plan = VEC0_QUERY_PLAN_FULLSCAN;
7123+ pCur->fullscan_data = fullscan_data;
7124+ return SQLITE_OK;
7125+
7126+error:
7127+ vec0_query_fullscan_data_clear(fullscan_data);
7128+ sqlite3_free(fullscan_data);
7129+ return rc;
7130+}
7131+
7132+int vec0Filter_point(vec0_cursor *pCur, vec0_vtab *p, int argc,
7133+ sqlite3_value **argv) {
7134+ int rc;
7135+ assert(argc == 1);
7136+ i64 rowid;
7137+ struct vec0_query_point_data *point_data = NULL;
7138+
7139+ point_data = sqlite3_malloc(sizeof(*point_data));
7140+ if (!point_data) {
7141+ rc = SQLITE_NOMEM;
7142+ goto error;
7143+ }
7144+ memset(point_data, 0, sizeof(*point_data));
7145+
7146+ if (p->pkIsText) {
7147+ rc = vec0_rowid_from_id(p, argv[0], &rowid);
7148+ if (rc == SQLITE_EMPTY) {
7149+ goto eof;
7150+ }
7151+ if (rc != SQLITE_OK) {
7152+ goto error;
7153+ }
7154+ } else {
7155+ rowid = sqlite3_value_int64(argv[0]);
7156+ }
7157+
7158+ for (int i = 0; i < p->numVectorColumns; i++) {
7159+ rc = vec0_get_vector_data(p, rowid, i, &point_data->vectors[i], NULL);
7160+ if (rc == SQLITE_EMPTY) {
7161+ goto eof;
7162+ }
7163+ if (rc != SQLITE_OK) {
7164+ goto error;
7165+ }
7166+ }
7167+
7168+ point_data->rowid = rowid;
7169+ point_data->done = 0;
7170+ pCur->point_data = point_data;
7171+ pCur->query_plan = VEC0_QUERY_PLAN_POINT;
7172+ return SQLITE_OK;
7173+
7174+eof:
7175+ point_data->rowid = rowid;
7176+ point_data->done = 1;
7177+ pCur->point_data = point_data;
7178+ pCur->query_plan = VEC0_QUERY_PLAN_POINT;
7179+ return SQLITE_OK;
7180+
7181+error:
7182+ vec0_query_point_data_clear(point_data);
7183+ sqlite3_free(point_data);
7184+ return rc;
7185+}
7186+
7187+static int vec0Filter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
7188+ const char *idxStr, int argc, sqlite3_value **argv) {
7189+ vec0_vtab *p = (vec0_vtab *)pVtabCursor->pVtab;
7190+ vec0_cursor *pCur = (vec0_cursor *)pVtabCursor;
7191+ vec0_cursor_clear(pCur);
7192+
7193+ int idxStrLength = strlen(idxStr);
7194+ if(idxStrLength <= 0) {
7195+ return SQLITE_ERROR;
7196+ }
7197+ if((idxStrLength-1) % 4 != 0) {
7198+ return SQLITE_ERROR;
7199+ }
7200+ int numValueEntries = (idxStrLength-1) / 4;
7201+ if(numValueEntries != argc) {
7202+ return SQLITE_ERROR;
7203+ }
7204+
7205+ char query_plan = idxStr[0];
7206+ switch(query_plan) {
7207+ case VEC0_QUERY_PLAN_FULLSCAN:
7208+ return vec0Filter_fullscan(p, pCur);
7209+ case VEC0_QUERY_PLAN_KNN:
7210+ return vec0Filter_knn(pCur, p, idxNum, idxStr, argc, argv);
7211+ case VEC0_QUERY_PLAN_POINT:
7212+ return vec0Filter_point(pCur, p, argc, argv);
7213+ default:
7214+ vtab_set_error(pVtabCursor->pVtab, "unknown idxStr '%s'", idxStr);
7215+ return SQLITE_ERROR;
7216+ }
7217+}
7218+
7219+static int vec0Rowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) {
7220+ vec0_cursor *pCur = (vec0_cursor *)cur;
7221+ switch (pCur->query_plan) {
7222+ case VEC0_QUERY_PLAN_FULLSCAN: {
7223+ *pRowid = sqlite3_column_int64(pCur->fullscan_data->rowids_stmt, 0);
7224+ return SQLITE_OK;
7225+ }
7226+ case VEC0_QUERY_PLAN_POINT: {
7227+ *pRowid = pCur->point_data->rowid;
7228+ return SQLITE_OK;
7229+ }
7230+ case VEC0_QUERY_PLAN_KNN: {
7231+ vtab_set_error(cur->pVtab,
7232+ "Internal sqlite-vec error: expected point query plan in "
7233+ "vec0Rowid, found %d",
7234+ pCur->query_plan);
7235+ return SQLITE_ERROR;
7236+ }
7237+ }
7238+ return SQLITE_ERROR;
7239+}
7240+
7241+static int vec0Next(sqlite3_vtab_cursor *cur) {
7242+ vec0_cursor *pCur = (vec0_cursor *)cur;
7243+ switch (pCur->query_plan) {
7244+ case VEC0_QUERY_PLAN_FULLSCAN: {
7245+ if (!pCur->fullscan_data) {
7246+ return SQLITE_ERROR;
7247+ }
7248+ int rc = sqlite3_step(pCur->fullscan_data->rowids_stmt);
7249+ if (rc == SQLITE_DONE) {
7250+ pCur->fullscan_data->done = 1;
7251+ return SQLITE_OK;
7252+ }
7253+ if (rc == SQLITE_ROW) {
7254+ return SQLITE_OK;
7255+ }
7256+ return SQLITE_ERROR;
7257+ }
7258+ case VEC0_QUERY_PLAN_KNN: {
7259+ if (!pCur->knn_data) {
7260+ return SQLITE_ERROR;
7261+ }
7262+
7263+ pCur->knn_data->current_idx++;
7264+ return SQLITE_OK;
7265+ }
7266+ case VEC0_QUERY_PLAN_POINT: {
7267+ if (!pCur->point_data) {
7268+ return SQLITE_ERROR;
7269+ }
7270+ pCur->point_data->done = 1;
7271+ return SQLITE_OK;
7272+ }
7273+ }
7274+ return SQLITE_ERROR;
7275+}
7276+
7277+static int vec0Eof(sqlite3_vtab_cursor *cur) {
7278+ vec0_cursor *pCur = (vec0_cursor *)cur;
7279+ switch (pCur->query_plan) {
7280+ case VEC0_QUERY_PLAN_FULLSCAN: {
7281+ if (!pCur->fullscan_data) {
7282+ return 1;
7283+ }
7284+ return pCur->fullscan_data->done;
7285+ }
7286+ case VEC0_QUERY_PLAN_KNN: {
7287+ if (!pCur->knn_data) {
7288+ return 1;
7289+ }
7290+ // return (pCur->knn_data->current_idx >= pCur->knn_data->k) ||
7291+ // (pCur->knn_data->distances[pCur->knn_data->current_idx] == FLT_MAX);
7292+ return (pCur->knn_data->current_idx >= pCur->knn_data->k_used);
7293+ }
7294+ case VEC0_QUERY_PLAN_POINT: {
7295+ if (!pCur->point_data) {
7296+ return 1;
7297+ }
7298+ return pCur->point_data->done;
7299+ }
7300+ }
7301+ return 1;
7302+}
7303+
7304+static int vec0Column_fullscan(vec0_vtab *pVtab, vec0_cursor *pCur,
7305+ sqlite3_context *context, int i) {
7306+ if (!pCur->fullscan_data) {
7307+ sqlite3_result_error(
7308+ context, "Internal sqlite-vec error: fullscan_data is NULL.", -1);
7309+ return SQLITE_ERROR;
7310+ }
7311+ i64 rowid = sqlite3_column_int64(pCur->fullscan_data->rowids_stmt, 0);
7312+ if (i == VEC0_COLUMN_ID) {
7313+ return vec0_result_id(pVtab, context, rowid);
7314+ }
7315+ else if (vec0_column_idx_is_vector(pVtab, i)) {
7316+ void *v;
7317+ int sz;
7318+ int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i);
7319+ int rc = vec0_get_vector_data(pVtab, rowid, vector_idx, &v, &sz);
7320+ if (rc != SQLITE_OK) {
7321+ return rc;
7322+ }
7323+ sqlite3_result_blob(context, v, sz, sqlite3_free);
7324+ sqlite3_result_subtype(context,
7325+ pVtab->vector_columns[vector_idx].element_type);
7326+
7327+ }
7328+ else if (i == vec0_column_distance_idx(pVtab)) {
7329+ sqlite3_result_null(context);
7330+ }
7331+ else if(vec0_column_idx_is_partition(pVtab, i)) {
7332+ int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i);
7333+ sqlite3_value * v;
7334+ int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v);
7335+ if(rc == SQLITE_OK) {
7336+ sqlite3_result_value(context, v);
7337+ sqlite3_value_free(v);
7338+ }else {
7339+ sqlite3_result_error_code(context, rc);
7340+ }
7341+ }
7342+ else if(vec0_column_idx_is_auxiliary(pVtab, i)) {
7343+ int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i);
7344+ sqlite3_value * v;
7345+ int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v);
7346+ if(rc == SQLITE_OK) {
7347+ sqlite3_result_value(context, v);
7348+ sqlite3_value_free(v);
7349+ }else {
7350+ sqlite3_result_error_code(context, rc);
7351+ }
7352+ }
7353+
7354+ else if(vec0_column_idx_is_metadata(pVtab, i)) {
7355+ if(sqlite3_vtab_nochange(context)) {
7356+ return SQLITE_OK;
7357+ }
7358+ int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i);
7359+ int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context);
7360+ if(rc != SQLITE_OK) {
7361+ // IMP: V15466_32305
7362+ const char * zErr = sqlite3_mprintf(
7363+ "Could not extract metadata value for column %.*s at rowid %lld",
7364+ pVtab->metadata_columns[metadata_idx].name_length,
7365+ pVtab->metadata_columns[metadata_idx].name, rowid
7366+ );
7367+ if(zErr) {
7368+ sqlite3_result_error(context, zErr, -1);
7369+ sqlite3_free((void *) zErr);
7370+ }else {
7371+ sqlite3_result_error_nomem(context);
7372+ }
7373+ }
7374+ }
7375+
7376+ return SQLITE_OK;
7377+}
7378+
7379+static int vec0Column_point(vec0_vtab *pVtab, vec0_cursor *pCur,
7380+ sqlite3_context *context, int i) {
7381+ if (!pCur->point_data) {
7382+ sqlite3_result_error(context,
7383+ "Internal sqlite-vec error: point_data is NULL.", -1);
7384+ return SQLITE_ERROR;
7385+ }
7386+ if (i == VEC0_COLUMN_ID) {
7387+ return vec0_result_id(pVtab, context, pCur->point_data->rowid);
7388+ }
7389+ else if (i == vec0_column_distance_idx(pVtab)) {
7390+ sqlite3_result_null(context);
7391+ return SQLITE_OK;
7392+ }
7393+ else if (vec0_column_idx_is_vector(pVtab, i)) {
7394+ if (sqlite3_vtab_nochange(context)) {
7395+ sqlite3_result_null(context);
7396+ return SQLITE_OK;
7397+ }
7398+ int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i);
7399+ sqlite3_result_blob(
7400+ context, pCur->point_data->vectors[vector_idx],
7401+ vector_column_byte_size(pVtab->vector_columns[vector_idx]),
7402+ SQLITE_TRANSIENT);
7403+ sqlite3_result_subtype(context,
7404+ pVtab->vector_columns[vector_idx].element_type);
7405+ return SQLITE_OK;
7406+ }
7407+ else if(vec0_column_idx_is_partition(pVtab, i)) {
7408+ if(sqlite3_vtab_nochange(context)) {
7409+ return SQLITE_OK;
7410+ }
7411+ int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i);
7412+ i64 rowid = pCur->point_data->rowid;
7413+ sqlite3_value * v;
7414+ int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v);
7415+ if(rc == SQLITE_OK) {
7416+ sqlite3_result_value(context, v);
7417+ sqlite3_value_free(v);
7418+ }else {
7419+ sqlite3_result_error_code(context, rc);
7420+ }
7421+ }
7422+ else if(vec0_column_idx_is_auxiliary(pVtab, i)) {
7423+ if(sqlite3_vtab_nochange(context)) {
7424+ return SQLITE_OK;
7425+ }
7426+ i64 rowid = pCur->point_data->rowid;
7427+ int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i);
7428+ sqlite3_value * v;
7429+ int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v);
7430+ if(rc == SQLITE_OK) {
7431+ sqlite3_result_value(context, v);
7432+ sqlite3_value_free(v);
7433+ }else {
7434+ sqlite3_result_error_code(context, rc);
7435+ }
7436+ }
7437+
7438+ else if(vec0_column_idx_is_metadata(pVtab, i)) {
7439+ if(sqlite3_vtab_nochange(context)) {
7440+ return SQLITE_OK;
7441+ }
7442+ i64 rowid = pCur->point_data->rowid;
7443+ int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i);
7444+ int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context);
7445+ if(rc != SQLITE_OK) {
7446+ const char * zErr = sqlite3_mprintf(
7447+ "Could not extract metadata value for column %.*s at rowid %lld",
7448+ pVtab->metadata_columns[metadata_idx].name_length,
7449+ pVtab->metadata_columns[metadata_idx].name, rowid
7450+ );
7451+ if(zErr) {
7452+ sqlite3_result_error(context, zErr, -1);
7453+ sqlite3_free((void *) zErr);
7454+ }else {
7455+ sqlite3_result_error_nomem(context);
7456+ }
7457+ }
7458+ }
7459+
7460+ return SQLITE_OK;
7461+}
7462+
7463+static int vec0Column_knn(vec0_vtab *pVtab, vec0_cursor *pCur,
7464+ sqlite3_context *context, int i) {
7465+ if (!pCur->knn_data) {
7466+ sqlite3_result_error(context,
7467+ "Internal sqlite-vec error: knn_data is NULL.", -1);
7468+ return SQLITE_ERROR;
7469+ }
7470+ if (i == VEC0_COLUMN_ID) {
7471+ i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx];
7472+ return vec0_result_id(pVtab, context, rowid);
7473+ }
7474+ else if (i == vec0_column_distance_idx(pVtab)) {
7475+ sqlite3_result_double(
7476+ context, pCur->knn_data->distances[pCur->knn_data->current_idx]);
7477+ return SQLITE_OK;
7478+ }
7479+ else if (vec0_column_idx_is_vector(pVtab, i)) {
7480+ void *out;
7481+ int sz;
7482+ int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i);
7483+ int rc = vec0_get_vector_data(
7484+ pVtab, pCur->knn_data->rowids[pCur->knn_data->current_idx], vector_idx,
7485+ &out, &sz);
7486+ if (rc != SQLITE_OK) {
7487+ return rc;
7488+ }
7489+ sqlite3_result_blob(context, out, sz, sqlite3_free);
7490+ sqlite3_result_subtype(context,
7491+ pVtab->vector_columns[vector_idx].element_type);
7492+ return SQLITE_OK;
7493+ }
7494+ else if(vec0_column_idx_is_partition(pVtab, i)) {
7495+ int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i);
7496+ i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx];
7497+ sqlite3_value * v;
7498+ int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v);
7499+ if(rc == SQLITE_OK) {
7500+ sqlite3_result_value(context, v);
7501+ sqlite3_value_free(v);
7502+ }else {
7503+ sqlite3_result_error_code(context, rc);
7504+ }
7505+ }
7506+ else if(vec0_column_idx_is_auxiliary(pVtab, i)) {
7507+ int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i);
7508+ i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx];
7509+ sqlite3_value * v;
7510+ int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v);
7511+ if(rc == SQLITE_OK) {
7512+ sqlite3_result_value(context, v);
7513+ sqlite3_value_free(v);
7514+ }else {
7515+ sqlite3_result_error_code(context, rc);
7516+ }
7517+ }
7518+
7519+ else if(vec0_column_idx_is_metadata(pVtab, i)) {
7520+ int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i);
7521+ i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx];
7522+ int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context);
7523+ if(rc != SQLITE_OK) {
7524+ const char * zErr = sqlite3_mprintf(
7525+ "Could not extract metadata value for column %.*s at rowid %lld",
7526+ pVtab->metadata_columns[metadata_idx].name_length,
7527+ pVtab->metadata_columns[metadata_idx].name, rowid
7528+ );
7529+ if(zErr) {
7530+ sqlite3_result_error(context, zErr, -1);
7531+ sqlite3_free((void *) zErr);
7532+ }else {
7533+ sqlite3_result_error_nomem(context);
7534+ }
7535+ }
7536+ }
7537+
7538+ return SQLITE_OK;
7539+}
7540+
7541+static int vec0Column(sqlite3_vtab_cursor *cur, sqlite3_context *context,
7542+ int i) {
7543+ vec0_cursor *pCur = (vec0_cursor *)cur;
7544+ vec0_vtab *pVtab = (vec0_vtab *)cur->pVtab;
7545+ switch (pCur->query_plan) {
7546+ case VEC0_QUERY_PLAN_FULLSCAN: {
7547+ return vec0Column_fullscan(pVtab, pCur, context, i);
7548+ }
7549+ case VEC0_QUERY_PLAN_KNN: {
7550+ return vec0Column_knn(pVtab, pCur, context, i);
7551+ }
7552+ case VEC0_QUERY_PLAN_POINT: {
7553+ return vec0Column_point(pVtab, pCur, context, i);
7554+ }
7555+ }
7556+ return SQLITE_OK;
7557+}
7558+
7559+/**
7560+ * @brief Handles the "insert rowid" step of a row insert operation of a vec0
7561+ * table.
7562+ *
7563+ * This function will insert a new row into the _rowids vec0 shadow table.
7564+ *
7565+ * @param p: virtual table
7566+ * @param idValue: Value containing the inserted rowid/id value.
7567+ * @param rowid: Output rowid, will point to the "real" i64 rowid
7568+ * value that was inserted
7569+ * @return int SQLITE_OK on success, error code on failure
7570+ */
7571+int vec0Update_InsertRowidStep(vec0_vtab *p, sqlite3_value *idValue,
7572+ i64 *rowid) {
7573+
7574+ /**
7575+ * An insert into a vec0 table can happen a few different ways:
7576+ * 1) With default INTEGER primary key: With a supplied i64 rowid
7577+ * 2) With default INTEGER primary key: WITHOUT a supplied rowid
7578+ * 3) With TEXT primary key: supplied text rowid
7579+ */
7580+
7581+ int rc;
7582+
7583+ // Option 3: vtab has a user-defined TEXT primary key, so ensure a text value
7584+ // is provided.
7585+ if (p->pkIsText) {
7586+ if (sqlite3_value_type(idValue) != SQLITE_TEXT) {
7587+ // IMP: V04200_21039
7588+ vtab_set_error(&p->base,
7589+ "The %s virtual table was declared with a TEXT primary "
7590+ "key, but a non-TEXT value was provided in an INSERT.",
7591+ p->tableName);
7592+ return SQLITE_ERROR;
7593+ }
7594+
7595+ return vec0_rowids_insert_id(p, idValue, rowid);
7596+ }
7597+
7598+ // Option 1: User supplied a i64 rowid
7599+ if (sqlite3_value_type(idValue) == SQLITE_INTEGER) {
7600+ i64 suppliedRowid = sqlite3_value_int64(idValue);
7601+ rc = vec0_rowids_insert_rowid(p, suppliedRowid);
7602+ if (rc == SQLITE_OK) {
7603+ *rowid = suppliedRowid;
7604+ }
7605+ return rc;
7606+ }
7607+
7608+ // Option 2: User did not suppled a rowid
7609+
7610+ if (sqlite3_value_type(idValue) != SQLITE_NULL) {
7611+ // IMP: V30855_14925
7612+ vtab_set_error(&p->base,
7613+ "Only integers are allows for primary key values on %s",
7614+ p->tableName);
7615+ return SQLITE_ERROR;
7616+ }
7617+ // NULL to get next auto-incremented value
7618+ return vec0_rowids_insert_id(p, NULL, rowid);
7619+}
7620+
7621+/**
7622+ * @brief Determines the "next available" chunk position for a newly inserted
7623+ * vec0 row.
7624+ *
7625+ * This operation may insert a new "blank" chunk the _chunks table, if there is
7626+ * no more space in previous chunks.
7627+ *
7628+ * @param p: virtual table
7629+ * @param partitionKeyValues: array of partition key column values, to constrain
7630+ * against any partition key columns.
7631+ * @param chunk_rowid: Output rowid of the chunk in the _chunks virtual table
7632+ * that has the avialabiity.
7633+ * @param chunk_offset: Output the index of the available space insert the
7634+ * chunk, based on the index of the first available validity bit.
7635+ * @param pBlobValidity: Output blob of the validity column of the available
7636+ * chunk. Will be opened with read/write permissions.
7637+ * @param pValidity: Output buffer of the original chunk's validity column.
7638+ * Needs to be cleaned up with sqlite3_free().
7639+ * @return int SQLITE_OK on success, error code on failure
7640+ */
7641+int vec0Update_InsertNextAvailableStep(
7642+ vec0_vtab *p,
7643+ sqlite3_value ** partitionKeyValues,
7644+ i64 *chunk_rowid, i64 *chunk_offset,
7645+ sqlite3_blob **blobChunksValidity,
7646+ const unsigned char **bufferChunksValidity) {
7647+
7648+ int rc;
7649+ i64 validitySize;
7650+ *chunk_offset = -1;
7651+
7652+ rc = vec0_get_latest_chunk_rowid(p, chunk_rowid, partitionKeyValues);
7653+ if(rc == SQLITE_EMPTY) {
7654+ goto done;
7655+ }
7656+ if (rc != SQLITE_OK) {
7657+ goto cleanup;
7658+ }
7659+
7660+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "validity",
7661+ *chunk_rowid, 1, blobChunksValidity);
7662+ if (rc != SQLITE_OK) {
7663+ // IMP: V22053_06123
7664+ vtab_set_error(&p->base,
7665+ VEC_INTERAL_ERROR
7666+ "could not open validity blob on %s.%s.%lld",
7667+ p->schemaName, p->shadowChunksName, *chunk_rowid);
7668+ goto cleanup;
7669+ }
7670+
7671+ validitySize = sqlite3_blob_bytes(*blobChunksValidity);
7672+ if (validitySize != p->chunk_size / CHAR_BIT) {
7673+ // IMP: V29362_13432
7674+ vtab_set_error(&p->base,
7675+ VEC_INTERAL_ERROR
7676+ "validity blob size mismatch on "
7677+ "%s.%s.%lld, expected %lld but received %lld.",
7678+ p->schemaName, p->shadowChunksName, *chunk_rowid,
7679+ (i64)(p->chunk_size / CHAR_BIT), validitySize);
7680+ rc = SQLITE_ERROR;
7681+ goto cleanup;
7682+ }
7683+
7684+ *bufferChunksValidity = sqlite3_malloc(validitySize);
7685+ if (!(*bufferChunksValidity)) {
7686+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
7687+ "Could not allocate memory for validity bitmap");
7688+ rc = SQLITE_NOMEM;
7689+ goto cleanup;
7690+ }
7691+
7692+ rc = sqlite3_blob_read(*blobChunksValidity, (void *)*bufferChunksValidity,
7693+ validitySize, 0);
7694+
7695+ if (rc != SQLITE_OK) {
7696+ vtab_set_error(&p->base,
7697+ VEC_INTERAL_ERROR
7698+ "Could not read validity bitmap for %s.%s.%lld",
7699+ p->schemaName, p->shadowChunksName, *chunk_rowid);
7700+ goto cleanup;
7701+ }
7702+
7703+ // find the next available offset, ie first `0` in the bitmap.
7704+ for (int i = 0; i < validitySize; i++) {
7705+ if ((*bufferChunksValidity)[i] == 0b11111111)
7706+ continue;
7707+ for (int j = 0; j < CHAR_BIT; j++) {
7708+ if (((((*bufferChunksValidity)[i] >> j) & 1) == 0)) {
7709+ *chunk_offset = (i * CHAR_BIT) + j;
7710+ goto done;
7711+ }
7712+ }
7713+ }
7714+
7715+done:
7716+ // latest chunk was full, so need to create a new one
7717+ if (*chunk_offset == -1) {
7718+ rc = vec0_new_chunk(p, partitionKeyValues, chunk_rowid);
7719+ if (rc != SQLITE_OK) {
7720+ // IMP: V08441_25279
7721+ vtab_set_error(&p->base,
7722+ VEC_INTERAL_ERROR "Could not insert a new vector chunk");
7723+ rc = SQLITE_ERROR; // otherwise raises a DatabaseError and not operational
7724+ // error?
7725+ goto cleanup;
7726+ }
7727+ *chunk_offset = 0;
7728+
7729+ // blobChunksValidity and pValidity are stale, pointing to the previous
7730+ // (full) chunk. to re-assign them
7731+ rc = sqlite3_blob_close(*blobChunksValidity);
7732+ sqlite3_free((void *)*bufferChunksValidity);
7733+ *blobChunksValidity = NULL;
7734+ *bufferChunksValidity = NULL;
7735+ if (rc != SQLITE_OK) {
7736+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
7737+ "unknown error, blobChunksValidity could not be closed, "
7738+ "please file an issue.");
7739+ rc = SQLITE_ERROR;
7740+ goto cleanup;
7741+ }
7742+
7743+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName,
7744+ "validity", *chunk_rowid, 1, blobChunksValidity);
7745+ if (rc != SQLITE_OK) {
7746+ vtab_set_error(
7747+ &p->base,
7748+ VEC_INTERAL_ERROR
7749+ "Could not open validity blob for newly created chunk %s.%s.%lld",
7750+ p->schemaName, p->shadowChunksName, *chunk_rowid);
7751+ goto cleanup;
7752+ }
7753+ validitySize = sqlite3_blob_bytes(*blobChunksValidity);
7754+ if (validitySize != p->chunk_size / CHAR_BIT) {
7755+ vtab_set_error(&p->base,
7756+ VEC_INTERAL_ERROR
7757+ "validity blob size mismatch for newly created chunk "
7758+ "%s.%s.%lld. Exepcted %lld, got %lld",
7759+ p->schemaName, p->shadowChunksName, *chunk_rowid,
7760+ p->chunk_size / CHAR_BIT, validitySize);
7761+ goto cleanup;
7762+ }
7763+ *bufferChunksValidity = sqlite3_malloc(validitySize);
7764+ rc = sqlite3_blob_read(*blobChunksValidity, (void *)*bufferChunksValidity,
7765+ validitySize, 0);
7766+ if (rc != SQLITE_OK) {
7767+ vtab_set_error(&p->base,
7768+ VEC_INTERAL_ERROR
7769+ "could not read validity blob newly created chunk "
7770+ "%s.%s.%lld",
7771+ p->schemaName, p->shadowChunksName, *chunk_rowid);
7772+ goto cleanup;
7773+ }
7774+ }
7775+
7776+ rc = SQLITE_OK;
7777+
7778+cleanup:
7779+ return rc;
7780+}
7781+
7782+/**
7783+ * @brief Write the vector data into the provided vector blob at the given
7784+ * offset
7785+ *
7786+ * @param blobVectors SQLite BLOB to write to
7787+ * @param chunk_offset the "offset" (ie validity bitmap position) to write the
7788+ * vector to
7789+ * @param bVector pointer to the vector containing data
7790+ * @param dimensions how many dimensions the vector has
7791+ * @param element_type the vector type
7792+ * @return result of sqlite3_blob_write, SQLITE_OK on success, otherwise failure
7793+ */
7794+static int
7795+vec0_write_vector_to_vector_blob(sqlite3_blob *blobVectors, i64 chunk_offset,
7796+ const void *bVector, size_t dimensions,
7797+ enum VectorElementType element_type) {
7798+ int n;
7799+ int offset;
7800+
7801+ switch (element_type) {
7802+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32:
7803+ n = dimensions * sizeof(f32);
7804+ offset = chunk_offset * dimensions * sizeof(f32);
7805+ break;
7806+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
7807+ n = dimensions * sizeof(i8);
7808+ offset = chunk_offset * dimensions * sizeof(i8);
7809+ break;
7810+ case SQLITE_VEC_ELEMENT_TYPE_BIT:
7811+ n = dimensions / CHAR_BIT;
7812+ offset = chunk_offset * dimensions / CHAR_BIT;
7813+ break;
7814+ }
7815+
7816+ return sqlite3_blob_write(blobVectors, bVector, n, offset);
7817+}
7818+
7819+/**
7820+ * @brief
7821+ *
7822+ * @param p vec0 virtual table
7823+ * @param chunk_rowid: which chunk to write to
7824+ * @param chunk_offset: the offset inside the chunk to write the vector to.
7825+ * @param rowid: the rowid of the inserting row
7826+ * @param vectorDatas: array of the vector data to insert
7827+ * @param blobValidity: writeable validity blob of the row's assigned chunk.
7828+ * @param validity: snapshot buffer of the valdity column from the row's
7829+ * assigned chunk.
7830+ * @return int SQLITE_OK on success, error code on failure
7831+ */
7832+int vec0Update_InsertWriteFinalStep(vec0_vtab *p, i64 chunk_rowid,
7833+ i64 chunk_offset, i64 rowid,
7834+ void *vectorDatas[],
7835+ sqlite3_blob *blobChunksValidity,
7836+ const unsigned char *bufferChunksValidity) {
7837+ int rc, brc;
7838+ sqlite3_blob *blobChunksRowids = NULL;
7839+
7840+ // mark the validity bit for this row in the chunk's validity bitmap
7841+ // Get the byte offset of the bitmap
7842+ char unsigned bx = bufferChunksValidity[chunk_offset / CHAR_BIT];
7843+ // set the bit at the chunk_offset position inside that byte
7844+ bx = bx | (1 << (chunk_offset % CHAR_BIT));
7845+ // write that 1 byte
7846+ rc = sqlite3_blob_write(blobChunksValidity, &bx, 1, chunk_offset / CHAR_BIT);
7847+ if (rc != SQLITE_OK) {
7848+ vtab_set_error(&p->base, VEC_INTERAL_ERROR "could not mark validity bit ");
7849+ return rc;
7850+ }
7851+
7852+ // Go insert the vector data into the vector chunk shadow tables
7853+ for (int i = 0; i < p->numVectorColumns; i++) {
7854+ sqlite3_blob *blobVectors;
7855+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowVectorChunksNames[i],
7856+ "vectors", chunk_rowid, 1, &blobVectors);
7857+ if (rc != SQLITE_OK) {
7858+ vtab_set_error(&p->base, "Error opening vector blob at %s.%s.%lld",
7859+ p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid);
7860+ goto cleanup;
7861+ }
7862+
7863+ i64 expected =
7864+ p->chunk_size * vector_column_byte_size(p->vector_columns[i]);
7865+ i64 actual = sqlite3_blob_bytes(blobVectors);
7866+
7867+ if (actual != expected) {
7868+ // IMP: V16386_00456
7869+ vtab_set_error(
7870+ &p->base,
7871+ VEC_INTERAL_ERROR
7872+ "vector blob size mismatch on %s.%s.%lld. Expected %lld, actual %lld",
7873+ p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid, expected,
7874+ actual);
7875+ rc = SQLITE_ERROR;
7876+ // already error, can ignore result code
7877+ sqlite3_blob_close(blobVectors);
7878+ goto cleanup;
7879+ };
7880+
7881+ rc = vec0_write_vector_to_vector_blob(
7882+ blobVectors, chunk_offset, vectorDatas[i],
7883+ p->vector_columns[i].dimensions, p->vector_columns[i].element_type);
7884+ if (rc != SQLITE_OK) {
7885+ vtab_set_error(&p->base,
7886+ VEC_INTERAL_ERROR
7887+ "could not write vector blob on %s.%s.%lld",
7888+ p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid);
7889+ rc = SQLITE_ERROR;
7890+ // already error, can ignore result code
7891+ sqlite3_blob_close(blobVectors);
7892+ goto cleanup;
7893+ }
7894+ rc = sqlite3_blob_close(blobVectors);
7895+ if (rc != SQLITE_OK) {
7896+ vtab_set_error(&p->base,
7897+ VEC_INTERAL_ERROR
7898+ "could not close vector blob on %s.%s.%lld",
7899+ p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid);
7900+ rc = SQLITE_ERROR;
7901+ goto cleanup;
7902+ }
7903+ }
7904+
7905+ // write the new rowid to the rowids column of the _chunks table
7906+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "rowids",
7907+ chunk_rowid, 1, &blobChunksRowids);
7908+ if (rc != SQLITE_OK) {
7909+ // IMP: V09221_26060
7910+ vtab_set_error(&p->base,
7911+ VEC_INTERAL_ERROR "could not open rowids blob on %s.%s.%lld",
7912+ p->schemaName, p->shadowChunksName, chunk_rowid);
7913+ goto cleanup;
7914+ }
7915+ i64 expected = p->chunk_size * sizeof(i64);
7916+ i64 actual = sqlite3_blob_bytes(blobChunksRowids);
7917+ if (expected != actual) {
7918+ // IMP: V12779_29618
7919+ vtab_set_error(
7920+ &p->base,
7921+ VEC_INTERAL_ERROR
7922+ "rowids blob size mismatch on %s.%s.%lld. Expected %lld, actual %lld",
7923+ p->schemaName, p->shadowChunksName, chunk_rowid, expected, actual);
7924+ rc = SQLITE_ERROR;
7925+ goto cleanup;
7926+ }
7927+ rc = sqlite3_blob_write(blobChunksRowids, &rowid, sizeof(i64),
7928+ chunk_offset * sizeof(i64));
7929+ if (rc != SQLITE_OK) {
7930+ vtab_set_error(
7931+ &p->base, VEC_INTERAL_ERROR "could not write rowids blob on %s.%s.%lld",
7932+ p->schemaName, p->shadowChunksName, chunk_rowid);
7933+ rc = SQLITE_ERROR;
7934+ goto cleanup;
7935+ }
7936+
7937+ // Now with all the vectors inserted, go back and update the _rowids table
7938+ // with the new chunk_rowid/chunk_offset values
7939+ rc = vec0_rowids_update_position(p, rowid, chunk_rowid, chunk_offset);
7940+
7941+cleanup:
7942+ brc = sqlite3_blob_close(blobChunksRowids);
7943+ if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) {
7944+ vtab_set_error(
7945+ &p->base, VEC_INTERAL_ERROR "could not close rowids blob on %s.%s.%lld",
7946+ p->schemaName, p->shadowChunksName, chunk_rowid);
7947+ return brc;
7948+ }
7949+ return rc;
7950+}
7951+
7952+int vec0_write_metadata_value(vec0_vtab *p, int metadata_column_idx, i64 rowid, i64 chunk_id, i64 chunk_offset, sqlite3_value * v, int isupdate) {
7953+ int rc;
7954+ struct Vec0MetadataColumnDefinition * metadata_column = &p->metadata_columns[metadata_column_idx];
7955+ vec0_metadata_column_kind kind = metadata_column->kind;
7956+
7957+ // verify input value matches column type
7958+ switch(kind) {
7959+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
7960+ if(sqlite3_value_type(v) != SQLITE_INTEGER || ((sqlite3_value_int(v) != 0) && (sqlite3_value_int(v) != 1))) {
7961+ rc = SQLITE_ERROR;
7962+ vtab_set_error(&p->base, "Expected 0 or 1 for BOOLEAN metadata column %.*s", metadata_column->name_length, metadata_column->name);
7963+ goto done;
7964+ }
7965+ break;
7966+ }
7967+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
7968+ if(sqlite3_value_type(v) != SQLITE_INTEGER) {
7969+ rc = SQLITE_ERROR;
7970+ vtab_set_error(&p->base, "Expected integer for INTEGER metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v)));
7971+ goto done;
7972+ }
7973+ break;
7974+ }
7975+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
7976+ if(sqlite3_value_type(v) != SQLITE_FLOAT) {
7977+ rc = SQLITE_ERROR;
7978+ vtab_set_error(&p->base, "Expected float for FLOAT metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v)));
7979+ goto done;
7980+ }
7981+ break;
7982+ }
7983+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
7984+ if(sqlite3_value_type(v) != SQLITE_TEXT) {
7985+ rc = SQLITE_ERROR;
7986+ vtab_set_error(&p->base, "Expected text for TEXT metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v)));
7987+ goto done;
7988+ }
7989+ break;
7990+ }
7991+ }
7992+
7993+ sqlite3_blob * blobValue = NULL;
7994+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_column_idx], "data", chunk_id, 1, &blobValue);
7995+ if(rc != SQLITE_OK) {
7996+ goto done;
7997+ }
7998+
7999+ switch(kind) {
8000+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
8001+ u8 block;
8002+ int value = sqlite3_value_int(v);
8003+ rc = sqlite3_blob_read(blobValue, &block, sizeof(u8), (int) (chunk_offset / CHAR_BIT));
8004+ if(rc != SQLITE_OK) {
8005+ goto done;
8006+ }
8007+
8008+ if (value) {
8009+ block |= 1 << (chunk_offset % CHAR_BIT);
8010+ } else {
8011+ block &= ~(1 << (chunk_offset % CHAR_BIT));
8012+ }
8013+
8014+ rc = sqlite3_blob_write(blobValue, &block, sizeof(u8), chunk_offset / CHAR_BIT);
8015+ break;
8016+ }
8017+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
8018+ i64 value = sqlite3_value_int64(v);
8019+ rc = sqlite3_blob_write(blobValue, &value, sizeof(value), chunk_offset * sizeof(i64));
8020+ break;
8021+ }
8022+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
8023+ double value = sqlite3_value_double(v);
8024+ rc = sqlite3_blob_write(blobValue, &value, sizeof(value), chunk_offset * sizeof(double));
8025+ break;
8026+ }
8027+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
8028+ int prev_n;
8029+ rc = sqlite3_blob_read(blobValue, &prev_n, sizeof(int), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8030+ if(rc != SQLITE_OK) {
8031+ goto done;
8032+ }
8033+
8034+ const char * s = (const char *) sqlite3_value_text(v);
8035+ int n = sqlite3_value_bytes(v);
8036+ u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
8037+ memset(view, 0, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8038+ memcpy(view, &n, sizeof(int));
8039+ memcpy(view+4, s, min(n, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH-4));
8040+
8041+ rc = sqlite3_blob_write(blobValue, &view, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH, chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8042+ if(n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
8043+ const char * zSql;
8044+
8045+ if(isupdate && (prev_n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH)) {
8046+ zSql = sqlite3_mprintf("UPDATE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " SET data = ?2 WHERE rowid = ?1", p->schemaName, p->tableName, metadata_column_idx);
8047+ }else {
8048+ zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " (rowid, data) VALUES (?1, ?2)", p->schemaName, p->tableName, metadata_column_idx);
8049+ }
8050+ if(!zSql) {
8051+ rc = SQLITE_NOMEM;
8052+ goto done;
8053+ }
8054+ sqlite3_stmt * stmt;
8055+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8056+ if(rc != SQLITE_OK) {
8057+ goto done;
8058+ }
8059+ sqlite3_bind_int64(stmt, 1, rowid);
8060+ sqlite3_bind_text(stmt, 2, s, n, SQLITE_STATIC);
8061+ rc = sqlite3_step(stmt);
8062+ sqlite3_finalize(stmt);
8063+
8064+ if(rc != SQLITE_DONE) {
8065+ rc = SQLITE_ERROR;
8066+ goto done;
8067+ }
8068+ }
8069+ else if(prev_n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
8070+ const char * zSql = sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_column_idx);
8071+ if(!zSql) {
8072+ rc = SQLITE_NOMEM;
8073+ goto done;
8074+ }
8075+ sqlite3_stmt * stmt;
8076+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8077+ if(rc != SQLITE_OK) {
8078+ goto done;
8079+ }
8080+ sqlite3_bind_int64(stmt, 1, rowid);
8081+ rc = sqlite3_step(stmt);
8082+ sqlite3_finalize(stmt);
8083+
8084+ if(rc != SQLITE_DONE) {
8085+ rc = SQLITE_ERROR;
8086+ goto done;
8087+ }
8088+ }
8089+ break;
8090+ }
8091+ }
8092+
8093+ if(rc != SQLITE_OK) {
8094+
8095+ }
8096+ rc = sqlite3_blob_close(blobValue);
8097+ if(rc != SQLITE_OK) {
8098+ goto done;
8099+ }
8100+
8101+ done:
8102+ return rc;
8103+}
8104+
8105+
8106+/**
8107+ * @brief Handles INSERT INTO operations on a vec0 table.
8108+ *
8109+ * @return int SQLITE_OK on success, otherwise error code on failure
8110+ */
8111+int vec0Update_Insert(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv,
8112+ sqlite_int64 *pRowid) {
8113+ UNUSED_PARAMETER(argc);
8114+ vec0_vtab *p = (vec0_vtab *)pVTab;
8115+ int rc;
8116+ // Rowid for the inserted row, deterimined by the inserted ID + _rowids shadow
8117+ // table
8118+ i64 rowid;
8119+
8120+ // Array to hold the vector data of the inserted row. Individual elements will
8121+ // have a lifetime bound to the argv[..] values.
8122+ void *vectorDatas[VEC0_MAX_VECTOR_COLUMNS];
8123+ // Array to hold cleanup functions for vectorDatas[]
8124+ vector_cleanup cleanups[VEC0_MAX_VECTOR_COLUMNS];
8125+
8126+ sqlite3_value * partitionKeyValues[VEC0_MAX_PARTITION_COLUMNS];
8127+
8128+ // Rowid of the chunk in the _chunks shadow table that the row will be a part
8129+ // of.
8130+ i64 chunk_rowid;
8131+ // offset within the chunk where the rowid belongs
8132+ i64 chunk_offset;
8133+
8134+ // a write-able blob of the validity column for the given chunk. Used to mark
8135+ // validity bit
8136+ sqlite3_blob *blobChunksValidity = NULL;
8137+ // buffer for the valididty column for the given chunk. Maybe not needed here?
8138+ const unsigned char *bufferChunksValidity = NULL;
8139+ int numReadVectors = 0;
8140+
8141+ // Read all provided partition key values into partitionKeyValues
8142+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8143+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_PARTITION) {
8144+ continue;
8145+ }
8146+ int partition_key_idx = p->user_column_idxs[i];
8147+ partitionKeyValues[partition_key_idx] = argv[2+VEC0_COLUMN_USERN_START + i];
8148+
8149+ int new_value_type = sqlite3_value_type(partitionKeyValues[partition_key_idx]);
8150+ if((new_value_type != SQLITE_NULL) && (new_value_type != p->paritition_columns[partition_key_idx].type)) {
8151+ // IMP: V11454_28292
8152+ vtab_set_error(
8153+ pVTab,
8154+ "Parition key type mismatch: The partition key column %.*s has type %s, but %s was provided.",
8155+ p->paritition_columns[partition_key_idx].name_length,
8156+ p->paritition_columns[partition_key_idx].name,
8157+ type_name(p->paritition_columns[partition_key_idx].type),
8158+ type_name(new_value_type)
8159+ );
8160+ rc = SQLITE_ERROR;
8161+ goto cleanup;
8162+ }
8163+ }
8164+
8165+ // read all the inserted vectors into vectorDatas, validate their lengths.
8166+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8167+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) {
8168+ continue;
8169+ }
8170+ int vector_column_idx = p->user_column_idxs[i];
8171+ sqlite3_value *valueVector = argv[2 + VEC0_COLUMN_USERN_START + i];
8172+ size_t dimensions;
8173+
8174+ char *pzError;
8175+ enum VectorElementType elementType;
8176+ rc = vector_from_value(valueVector, &vectorDatas[vector_column_idx], &dimensions,
8177+ &elementType, &cleanups[vector_column_idx], &pzError);
8178+ if (rc != SQLITE_OK) {
8179+ // IMP: V06519_23358
8180+ vtab_set_error(
8181+ pVTab, "Inserted vector for the \"%.*s\" column is invalid: %z",
8182+ p->vector_columns[vector_column_idx].name_length, p->vector_columns[vector_column_idx].name, pzError);
8183+ rc = SQLITE_ERROR;
8184+ goto cleanup;
8185+ }
8186+
8187+ numReadVectors++;
8188+ if (elementType != p->vector_columns[vector_column_idx].element_type) {
8189+ // IMP: V08221_25059
8190+ vtab_set_error(
8191+ pVTab,
8192+ "Inserted vector for the \"%.*s\" column is expected to be of type "
8193+ "%s, but a %s vector was provided.",
8194+ p->vector_columns[i].name_length, p->vector_columns[i].name,
8195+ vector_subtype_name(p->vector_columns[i].element_type),
8196+ vector_subtype_name(elementType));
8197+ rc = SQLITE_ERROR;
8198+ goto cleanup;
8199+ }
8200+
8201+ if (dimensions != p->vector_columns[vector_column_idx].dimensions) {
8202+ // IMP: V01145_17984
8203+ vtab_set_error(
8204+ pVTab,
8205+ "Dimension mismatch for inserted vector for the \"%.*s\" column. "
8206+ "Expected %d dimensions but received %d.",
8207+ p->vector_columns[vector_column_idx].name_length, p->vector_columns[vector_column_idx].name,
8208+ p->vector_columns[vector_column_idx].dimensions, dimensions);
8209+ rc = SQLITE_ERROR;
8210+ goto cleanup;
8211+ }
8212+ }
8213+
8214+ // Cannot insert a value in the hidden "distance" column
8215+ if (sqlite3_value_type(argv[2 + vec0_column_distance_idx(p)]) !=
8216+ SQLITE_NULL) {
8217+ // IMP: V24228_08298
8218+ vtab_set_error(pVTab,
8219+ "A value was provided for the hidden \"distance\" column.");
8220+ rc = SQLITE_ERROR;
8221+ goto cleanup;
8222+ }
8223+ // Cannot insert a value in the hidden "k" column
8224+ if (sqlite3_value_type(argv[2 + vec0_column_k_idx(p)]) != SQLITE_NULL) {
8225+ // IMP: V11875_28713
8226+ vtab_set_error(pVTab, "A value was provided for the hidden \"k\" column.");
8227+ rc = SQLITE_ERROR;
8228+ goto cleanup;
8229+ }
8230+
8231+ // Step #1: Insert/get a rowid for this row, from the _rowids table.
8232+ rc = vec0Update_InsertRowidStep(p, argv[2 + VEC0_COLUMN_ID], &rowid);
8233+ if (rc != SQLITE_OK) {
8234+ goto cleanup;
8235+ }
8236+
8237+ // Step #2: Find the next "available" position in the _chunks table for this
8238+ // row.
8239+ rc = vec0Update_InsertNextAvailableStep(p, partitionKeyValues,
8240+ &chunk_rowid, &chunk_offset,
8241+ &blobChunksValidity,
8242+ &bufferChunksValidity);
8243+ if (rc != SQLITE_OK) {
8244+ goto cleanup;
8245+ }
8246+
8247+ // Step #3: With the next available chunk position, write out all the vectors
8248+ // to their specified location.
8249+ rc = vec0Update_InsertWriteFinalStep(p, chunk_rowid, chunk_offset, rowid,
8250+ vectorDatas, blobChunksValidity,
8251+ bufferChunksValidity);
8252+ if (rc != SQLITE_OK) {
8253+ goto cleanup;
8254+ }
8255+
8256+ if(p->numAuxiliaryColumns > 0) {
8257+ sqlite3_stmt *stmt;
8258+ sqlite3_str * s = sqlite3_str_new(NULL);
8259+ sqlite3_str_appendf(s, "INSERT INTO " VEC0_SHADOW_AUXILIARY_NAME "(rowid ", p->schemaName, p->tableName);
8260+ for(int i = 0; i < p->numAuxiliaryColumns; i++) {
8261+ sqlite3_str_appendf(s, ", value%02d", i);
8262+ }
8263+ sqlite3_str_appendall(s, ") VALUES (? ");
8264+ for(int i = 0; i < p->numAuxiliaryColumns; i++) {
8265+ sqlite3_str_appendall(s, ", ?");
8266+ }
8267+ sqlite3_str_appendall(s, ")");
8268+ char * zSql = sqlite3_str_finish(s);
8269+ // TODO double check error handling ehre
8270+ if(!zSql) {
8271+ rc = SQLITE_NOMEM;
8272+ goto cleanup;
8273+ }
8274+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8275+ if(rc != SQLITE_OK) {
8276+ goto cleanup;
8277+ }
8278+ sqlite3_bind_int64(stmt, 1, rowid);
8279+
8280+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8281+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY) {
8282+ continue;
8283+ }
8284+ int auxiliary_key_idx = p->user_column_idxs[i];
8285+ sqlite3_value * v = argv[2+VEC0_COLUMN_USERN_START + i];
8286+ int v_type = sqlite3_value_type(v);
8287+ if(v_type != SQLITE_NULL && (v_type != p->auxiliary_columns[auxiliary_key_idx].type)) {
8288+ sqlite3_finalize(stmt);
8289+ rc = SQLITE_CONSTRAINT;
8290+ vtab_set_error(
8291+ pVTab,
8292+ "Auxiliary column type mismatch: The auxiliary column %.*s has type %s, but %s was provided.",
8293+ p->auxiliary_columns[auxiliary_key_idx].name_length,
8294+ p->auxiliary_columns[auxiliary_key_idx].name,
8295+ type_name(p->auxiliary_columns[auxiliary_key_idx].type),
8296+ type_name(v_type)
8297+ );
8298+ goto cleanup;
8299+ }
8300+ // first 1 is for 1-based indexing on sqlite3_bind_*, second 1 is to account for initial rowid parameter
8301+ sqlite3_bind_value(stmt, 1 + 1 + auxiliary_key_idx, v);
8302+ }
8303+
8304+ rc = sqlite3_step(stmt);
8305+ if(rc != SQLITE_DONE) {
8306+ sqlite3_finalize(stmt);
8307+ rc = SQLITE_ERROR;
8308+ goto cleanup;
8309+ }
8310+ sqlite3_finalize(stmt);
8311+ }
8312+
8313+
8314+ for(int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8315+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) {
8316+ continue;
8317+ }
8318+ int metadata_idx = p->user_column_idxs[i];
8319+ sqlite3_value *v = argv[2 + VEC0_COLUMN_USERN_START + i];
8320+ rc = vec0_write_metadata_value(p, metadata_idx, rowid, chunk_rowid, chunk_offset, v, 0);
8321+ if(rc != SQLITE_OK) {
8322+ goto cleanup;
8323+ }
8324+ }
8325+
8326+ *pRowid = rowid;
8327+ rc = SQLITE_OK;
8328+
8329+cleanup:
8330+ for (int i = 0; i < numReadVectors; i++) {
8331+ cleanups[i](vectorDatas[i]);
8332+ }
8333+ sqlite3_free((void *)bufferChunksValidity);
8334+ int brc = sqlite3_blob_close(blobChunksValidity);
8335+ if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) {
8336+ vtab_set_error(&p->base,
8337+ VEC_INTERAL_ERROR "unknown error, blobChunksValidity could "
8338+ "not be closed, please file an issue");
8339+ return brc;
8340+ }
8341+ return rc;
8342+}
8343+
8344+int vec0Update_Delete_ClearValidity(vec0_vtab *p, i64 chunk_id,
8345+ u64 chunk_offset) {
8346+ int rc, brc;
8347+ sqlite3_blob *blobChunksValidity = NULL;
8348+ char unsigned bx;
8349+ int validityOffset = chunk_offset / CHAR_BIT;
8350+
8351+ // 2. ensure chunks.validity bit is 1, then set to 0
8352+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "validity",
8353+ chunk_id, 1, &blobChunksValidity);
8354+ if (rc != SQLITE_OK) {
8355+ // IMP: V26002_10073
8356+ vtab_set_error(&p->base, "could not open validity blob for %s.%s.%lld",
8357+ p->schemaName, p->shadowChunksName, chunk_id);
8358+ return SQLITE_ERROR;
8359+ }
8360+ // will skip the sqlite3_blob_bytes(blobChunksValidity) check for now,
8361+ // the read below would catch it
8362+
8363+ rc = sqlite3_blob_read(blobChunksValidity, &bx, sizeof(bx), validityOffset);
8364+ if (rc != SQLITE_OK) {
8365+ // IMP: V21193_05263
8366+ vtab_set_error(
8367+ &p->base, "could not read validity blob for %s.%s.%lld at %d",
8368+ p->schemaName, p->shadowChunksName, chunk_id, validityOffset);
8369+ goto cleanup;
8370+ }
8371+ if (!(bx >> (chunk_offset % CHAR_BIT))) {
8372+ // IMP: V21193_05263
8373+ rc = SQLITE_ERROR;
8374+ vtab_set_error(
8375+ &p->base,
8376+ "vec0 deletion error: validity bit is not set for %s.%s.%lld at %d",
8377+ p->schemaName, p->shadowChunksName, chunk_id, validityOffset);
8378+ goto cleanup;
8379+ }
8380+ char unsigned mask = ~(1 << (chunk_offset % CHAR_BIT));
8381+ char result = bx & mask;
8382+ rc = sqlite3_blob_write(blobChunksValidity, &result, sizeof(bx),
8383+ validityOffset);
8384+ if (rc != SQLITE_OK) {
8385+ vtab_set_error(
8386+ &p->base, "could not write to validity blob for %s.%s.%lld at %d",
8387+ p->schemaName, p->shadowChunksName, chunk_id, validityOffset);
8388+ goto cleanup;
8389+ }
8390+
8391+cleanup:
8392+
8393+ brc = sqlite3_blob_close(blobChunksValidity);
8394+ if (rc != SQLITE_OK)
8395+ return rc;
8396+ if (brc != SQLITE_OK) {
8397+ vtab_set_error(&p->base,
8398+ "vec0 deletion error: Error commiting validity blob "
8399+ "transaction on %s.%s.%lld at %d",
8400+ p->schemaName, p->shadowChunksName, chunk_id,
8401+ validityOffset);
8402+ return brc;
8403+ }
8404+ return SQLITE_OK;
8405+}
8406+
8407+int vec0Update_Delete_DeleteRowids(vec0_vtab *p, i64 rowid) {
8408+ int rc;
8409+ sqlite3_stmt *stmt = NULL;
8410+
8411+ char *zSql =
8412+ sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE rowid = ?",
8413+ p->schemaName, p->tableName);
8414+ if (!zSql) {
8415+ return SQLITE_NOMEM;
8416+ }
8417+
8418+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8419+ sqlite3_free(zSql);
8420+ if (rc != SQLITE_OK) {
8421+ goto cleanup;
8422+ }
8423+ sqlite3_bind_int64(stmt, 1, rowid);
8424+ rc = sqlite3_step(stmt);
8425+ if (rc != SQLITE_DONE) {
8426+ goto cleanup;
8427+ }
8428+ rc = SQLITE_OK;
8429+
8430+cleanup:
8431+ sqlite3_finalize(stmt);
8432+ return rc;
8433+}
8434+
8435+int vec0Update_Delete_DeleteAux(vec0_vtab *p, i64 rowid) {
8436+ int rc;
8437+ sqlite3_stmt *stmt = NULL;
8438+
8439+ char *zSql =
8440+ sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_AUXILIARY_NAME " WHERE rowid = ?",
8441+ p->schemaName, p->tableName);
8442+ if (!zSql) {
8443+ return SQLITE_NOMEM;
8444+ }
8445+
8446+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8447+ sqlite3_free(zSql);
8448+ if (rc != SQLITE_OK) {
8449+ goto cleanup;
8450+ }
8451+ sqlite3_bind_int64(stmt, 1, rowid);
8452+ rc = sqlite3_step(stmt);
8453+ if (rc != SQLITE_DONE) {
8454+ goto cleanup;
8455+ }
8456+ rc = SQLITE_OK;
8457+
8458+cleanup:
8459+ sqlite3_finalize(stmt);
8460+ return rc;
8461+}
8462+
8463+int vec0Update_Delete_ClearMetadata(vec0_vtab *p, int metadata_idx, i64 rowid, i64 chunk_id,
8464+ u64 chunk_offset) {
8465+ int rc;
8466+ sqlite3_blob * blobValue;
8467+ vec0_metadata_column_kind kind = p->metadata_columns[metadata_idx].kind;
8468+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 1, &blobValue);
8469+ if(rc != SQLITE_OK) {
8470+ return rc;
8471+ }
8472+
8473+ switch(kind) {
8474+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
8475+ u8 block;
8476+ rc = sqlite3_blob_read(blobValue, &block, sizeof(u8), (int) (chunk_offset / CHAR_BIT));
8477+ if(rc != SQLITE_OK) {
8478+ goto done;
8479+ }
8480+
8481+ block &= ~(1 << (chunk_offset % CHAR_BIT));
8482+ rc = sqlite3_blob_write(blobValue, &block, sizeof(u8), chunk_offset / CHAR_BIT);
8483+ break;
8484+ }
8485+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
8486+ i64 v = 0;
8487+ rc = sqlite3_blob_write(blobValue, &v, sizeof(v), chunk_offset * sizeof(i64));
8488+ break;
8489+ }
8490+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
8491+ double v = 0;
8492+ rc = sqlite3_blob_write(blobValue, &v, sizeof(v), chunk_offset * sizeof(double));
8493+ break;
8494+ }
8495+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
8496+ int n;
8497+ rc = sqlite3_blob_read(blobValue, &n, sizeof(int), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8498+ if(rc != SQLITE_OK) {
8499+ goto done;
8500+ }
8501+
8502+ u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
8503+ memset(view, 0, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8504+ rc = sqlite3_blob_write(blobValue, &view, sizeof(view), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8505+ if(rc != SQLITE_OK) {
8506+ goto done;
8507+ }
8508+
8509+ if(n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
8510+ const char * zSql = sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_idx);
8511+ if(!zSql) {
8512+ rc = SQLITE_NOMEM;
8513+ goto done;
8514+ }
8515+ sqlite3_stmt * stmt;
8516+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8517+ if(rc != SQLITE_OK) {
8518+ goto done;
8519+ }
8520+ sqlite3_bind_int64(stmt, 1, rowid);
8521+ rc = sqlite3_step(stmt);
8522+ if(rc != SQLITE_DONE) {
8523+ rc = SQLITE_ERROR;
8524+ goto done;
8525+ }
8526+ sqlite3_finalize(stmt);
8527+ }
8528+ break;
8529+ }
8530+ }
8531+ int rc2;
8532+ done:
8533+ rc2 = sqlite3_blob_close(blobValue);
8534+ if(rc == SQLITE_OK) {
8535+ return rc2;
8536+ }
8537+ return rc;
8538+}
8539+
8540+int vec0Update_Delete(sqlite3_vtab *pVTab, sqlite3_value *idValue) {
8541+ vec0_vtab *p = (vec0_vtab *)pVTab;
8542+ int rc;
8543+ i64 rowid;
8544+ i64 chunk_id;
8545+ i64 chunk_offset;
8546+
8547+ if (p->pkIsText) {
8548+ rc = vec0_rowid_from_id(p, idValue, &rowid);
8549+ if (rc != SQLITE_OK) {
8550+ return rc;
8551+ }
8552+ } else {
8553+ rowid = sqlite3_value_int64(idValue);
8554+ }
8555+
8556+ // 1. Find chunk position for given rowid
8557+ // 2. Ensure that validity bit for position is 1, then set to 0
8558+ // 3. Zero out rowid in chunks.rowid
8559+ // 4. Zero out vector data in all vector column chunks
8560+ // 5. Delete value in _rowids table
8561+
8562+ // 1. get chunk_id and chunk_offset from _rowids
8563+ rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset);
8564+ if (rc != SQLITE_OK) {
8565+ return rc;
8566+ }
8567+
8568+ rc = vec0Update_Delete_ClearValidity(p, chunk_id, chunk_offset);
8569+ if (rc != SQLITE_OK) {
8570+ return rc;
8571+ }
8572+
8573+ // 3. zero out rowid in chunks.rowids
8574+ // https://github.com/asg017/sqlite-vec/issues/54
8575+
8576+ // 4. zero out any data in vector chunks tables
8577+ // https://github.com/asg017/sqlite-vec/issues/54
8578+
8579+ // 5. delete from _rowids table
8580+ rc = vec0Update_Delete_DeleteRowids(p, rowid);
8581+ if (rc != SQLITE_OK) {
8582+ return rc;
8583+ }
8584+
8585+ // 6. delete any auxiliary rows
8586+ if(p->numAuxiliaryColumns > 0) {
8587+ rc = vec0Update_Delete_DeleteAux(p, rowid);
8588+ if (rc != SQLITE_OK) {
8589+ return rc;
8590+ }
8591+ }
8592+
8593+ // 6. delete metadata
8594+ for(int i = 0; i < p->numMetadataColumns; i++) {
8595+ rc = vec0Update_Delete_ClearMetadata(p, i, rowid, chunk_id, chunk_offset);
8596+ }
8597+
8598+ return SQLITE_OK;
8599+}
8600+
8601+int vec0Update_UpdateAuxColumn(vec0_vtab *p, int auxiliary_column_idx, sqlite3_value * value, i64 rowid) {
8602+ int rc;
8603+ sqlite3_stmt *stmt;
8604+ const char * zSql = sqlite3_mprintf("UPDATE " VEC0_SHADOW_AUXILIARY_NAME " SET value%02d = ? WHERE rowid = ?", p->schemaName, p->tableName, auxiliary_column_idx);
8605+ if(!zSql) {
8606+ return SQLITE_NOMEM;
8607+ }
8608+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8609+ if(rc != SQLITE_OK) {
8610+ return rc;
8611+ }
8612+ sqlite3_bind_value(stmt, 1, value);
8613+ sqlite3_bind_int64(stmt, 2, rowid);
8614+ rc = sqlite3_step(stmt);
8615+ if(rc != SQLITE_DONE) {
8616+ sqlite3_finalize(stmt);
8617+ return SQLITE_ERROR;
8618+ }
8619+ sqlite3_finalize(stmt);
8620+ return SQLITE_OK;
8621+}
8622+
8623+int vec0Update_UpdateVectorColumn(vec0_vtab *p, i64 chunk_id, i64 chunk_offset,
8624+ int i, sqlite3_value *valueVector) {
8625+ int rc;
8626+
8627+ sqlite3_blob *blobVectors = NULL;
8628+
8629+ char *pzError;
8630+ size_t dimensions;
8631+ enum VectorElementType elementType;
8632+ void *vector;
8633+ vector_cleanup cleanup = vector_cleanup_noop;
8634+ // https://github.com/asg017/sqlite-vec/issues/53
8635+ rc = vector_from_value(valueVector, &vector, &dimensions, &elementType,
8636+ &cleanup, &pzError);
8637+ if (rc != SQLITE_OK) {
8638+ // IMP: V15203_32042
8639+ vtab_set_error(
8640+ &p->base, "Updated vector for the \"%.*s\" column is invalid: %z",
8641+ p->vector_columns[i].name_length, p->vector_columns[i].name, pzError);
8642+ rc = SQLITE_ERROR;
8643+ goto cleanup;
8644+ }
8645+ if (elementType != p->vector_columns[i].element_type) {
8646+ // IMP: V03643_20481
8647+ vtab_set_error(
8648+ &p->base,
8649+ "Updated vector for the \"%.*s\" column is expected to be of type "
8650+ "%s, but a %s vector was provided.",
8651+ p->vector_columns[i].name_length, p->vector_columns[i].name,
8652+ vector_subtype_name(p->vector_columns[i].element_type),
8653+ vector_subtype_name(elementType));
8654+ rc = SQLITE_ERROR;
8655+ goto cleanup;
8656+ }
8657+ if (dimensions != p->vector_columns[i].dimensions) {
8658+ // IMP: V25739_09810
8659+ vtab_set_error(
8660+ &p->base,
8661+ "Dimension mismatch for new updated vector for the \"%.*s\" column. "
8662+ "Expected %d dimensions but received %d.",
8663+ p->vector_columns[i].name_length, p->vector_columns[i].name,
8664+ p->vector_columns[i].dimensions, dimensions);
8665+ rc = SQLITE_ERROR;
8666+ goto cleanup;
8667+ }
8668+
8669+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowVectorChunksNames[i],
8670+ "vectors", chunk_id, 1, &blobVectors);
8671+ if (rc != SQLITE_OK) {
8672+ vtab_set_error(&p->base, "Could not open vectors blob for %s.%s.%lld",
8673+ p->schemaName, p->shadowVectorChunksNames[i], chunk_id);
8674+ goto cleanup;
8675+ }
8676+ rc = vec0_write_vector_to_vector_blob(blobVectors, chunk_offset, vector,
8677+ p->vector_columns[i].dimensions,
8678+ p->vector_columns[i].element_type);
8679+ if (rc != SQLITE_OK) {
8680+ vtab_set_error(&p->base, "Could not write to vectors blob for %s.%s.%lld",
8681+ p->schemaName, p->shadowVectorChunksNames[i], chunk_id);
8682+ goto cleanup;
8683+ }
8684+
8685+cleanup:
8686+ cleanup(vector);
8687+ int brc = sqlite3_blob_close(blobVectors);
8688+ if (rc != SQLITE_OK) {
8689+ return rc;
8690+ }
8691+ if (brc != SQLITE_OK) {
8692+ vtab_set_error(
8693+ &p->base,
8694+ "Could not commit blob transaction for vectors blob for %s.%s.%lld",
8695+ p->schemaName, p->shadowVectorChunksNames[i], chunk_id);
8696+ return brc;
8697+ }
8698+ return SQLITE_OK;
8699+}
8700+
8701+int vec0Update_Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv) {
8702+ UNUSED_PARAMETER(argc);
8703+ vec0_vtab *p = (vec0_vtab *)pVTab;
8704+ int rc;
8705+ i64 chunk_id;
8706+ i64 chunk_offset;
8707+
8708+ i64 rowid;
8709+ if (p->pkIsText) {
8710+ const char *a = (const char *)sqlite3_value_text(argv[0]);
8711+ const char *b = (const char *)sqlite3_value_text(argv[1]);
8712+ // IMP: V08886_25725
8713+ if ((sqlite3_value_bytes(argv[0]) != sqlite3_value_bytes(argv[1])) ||
8714+ strncmp(a, b, sqlite3_value_bytes(argv[0])) != 0) {
8715+ vtab_set_error(pVTab,
8716+ "UPDATEs on vec0 primary key values are not allowed.");
8717+ return SQLITE_ERROR;
8718+ }
8719+ rc = vec0_rowid_from_id(p, argv[0], &rowid);
8720+ if (rc != SQLITE_OK) {
8721+ return rc;
8722+ }
8723+ } else {
8724+ rowid = sqlite3_value_int64(argv[0]);
8725+ }
8726+
8727+ // 1) get chunk_id and chunk_offset from _rowids
8728+ rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset);
8729+ if (rc != SQLITE_OK) {
8730+ return rc;
8731+ }
8732+
8733+ // 2) update any partition key values
8734+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8735+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_PARTITION) {
8736+ continue;
8737+ }
8738+ sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i];
8739+ if(sqlite3_value_nochange(value)) {
8740+ continue;
8741+ }
8742+ vtab_set_error(pVTab, "UPDATE on partition key columns are not supported yet. ");
8743+ return SQLITE_ERROR;
8744+ }
8745+
8746+ // 3) handle auxiliary column updates
8747+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8748+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY) {
8749+ continue;
8750+ }
8751+ int auxiliary_column_idx = p->user_column_idxs[i];
8752+ sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i];
8753+ if(sqlite3_value_nochange(value)) {
8754+ continue;
8755+ }
8756+ rc = vec0Update_UpdateAuxColumn(p, auxiliary_column_idx, value, rowid);
8757+ if(rc != SQLITE_OK) {
8758+ return SQLITE_ERROR;
8759+ }
8760+ }
8761+
8762+ // 4) handle metadata column updates
8763+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8764+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) {
8765+ continue;
8766+ }
8767+ int metadata_column_idx = p->user_column_idxs[i];
8768+ sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i];
8769+ if(sqlite3_value_nochange(value)) {
8770+ continue;
8771+ }
8772+ rc = vec0_write_metadata_value(p, metadata_column_idx, rowid, chunk_id, chunk_offset, value, 1);
8773+ if(rc != SQLITE_OK) {
8774+ return rc;
8775+ }
8776+ }
8777+
8778+ // 5) iterate over all new vectors, update the vectors
8779+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8780+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) {
8781+ continue;
8782+ }
8783+ int vector_idx = p->user_column_idxs[i];
8784+ sqlite3_value *valueVector = argv[2 + VEC0_COLUMN_USERN_START + i];
8785+ // in vec0Column, we check sqlite3_vtab_nochange() on vector columns.
8786+ // If the vector column isn't being changed, we return NULL;
8787+ // That's not great, that means vector columns can never be NULLABLE
8788+ // (bc we cant distinguish if an updated vector is truly NULL or nochange).
8789+ // Also it means that if someone tries to run `UPDATE v SET X = NULL`,
8790+ // we can't effectively detect and raise an error.
8791+ // A better solution would be to use a custom result_type for "empty",
8792+ // but subtypes don't appear to survive xColumn -> xUpdate, it's always 0.
8793+ // So for now, we'll just use NULL and warn people to not SET X = NULL
8794+ // in the docs.
8795+ if (sqlite3_value_type(valueVector) == SQLITE_NULL) {
8796+ continue;
8797+ }
8798+
8799+ rc = vec0Update_UpdateVectorColumn(p, chunk_id, chunk_offset, vector_idx,
8800+ valueVector);
8801+ if (rc != SQLITE_OK) {
8802+ return SQLITE_ERROR;
8803+ }
8804+ }
8805+
8806+ return SQLITE_OK;
8807+}
8808+
8809+static int vec0Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv,
8810+ sqlite_int64 *pRowid) {
8811+ // DELETE operation
8812+ if (argc == 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
8813+ return vec0Update_Delete(pVTab, argv[0]);
8814+ }
8815+ // INSERT operation
8816+ else if (argc > 1 && sqlite3_value_type(argv[0]) == SQLITE_NULL) {
8817+ return vec0Update_Insert(pVTab, argc, argv, pRowid);
8818+ }
8819+ // UPDATE operation
8820+ else if (argc > 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
8821+ return vec0Update_Update(pVTab, argc, argv);
8822+ } else {
8823+ vtab_set_error(pVTab, "Unrecognized xUpdate operation provided for vec0.");
8824+ return SQLITE_ERROR;
8825+ }
8826+}
8827+
8828+static int vec0ShadowName(const char *zName) {
8829+ static const char *azName[] = {
8830+ "rowids", "chunks", "auxiliary", "info",
8831+
8832+ // Up to VEC0_MAX_METADATA_COLUMNS
8833+ // TODO be smarter about this man
8834+ "metadatachunks00",
8835+ "metadatachunks01",
8836+ "metadatachunks02",
8837+ "metadatachunks03",
8838+ "metadatachunks04",
8839+ "metadatachunks05",
8840+ "metadatachunks06",
8841+ "metadatachunks07",
8842+ "metadatachunks08",
8843+ "metadatachunks09",
8844+ "metadatachunks10",
8845+ "metadatachunks11",
8846+ "metadatachunks12",
8847+ "metadatachunks13",
8848+ "metadatachunks14",
8849+ "metadatachunks15",
8850+
8851+ // Up to
8852+ "metadatatext00",
8853+ "metadatatext01",
8854+ "metadatatext02",
8855+ "metadatatext03",
8856+ "metadatatext04",
8857+ "metadatatext05",
8858+ "metadatatext06",
8859+ "metadatatext07",
8860+ "metadatatext08",
8861+ "metadatatext09",
8862+ "metadatatext10",
8863+ "metadatatext11",
8864+ "metadatatext12",
8865+ "metadatatext13",
8866+ "metadatatext14",
8867+ "metadatatext15",
8868+ };
8869+
8870+ for (size_t i = 0; i < sizeof(azName) / sizeof(azName[0]); i++) {
8871+ if (sqlite3_stricmp(zName, azName[i]) == 0)
8872+ return 1;
8873+ }
8874+ //for(size_t i = 0; i < )"vector_chunks", "metadatachunks"
8875+ return 0;
8876+}
8877+
8878+static int vec0Begin(sqlite3_vtab *pVTab) {
8879+ UNUSED_PARAMETER(pVTab);
8880+ return SQLITE_OK;
8881+}
8882+static int vec0Sync(sqlite3_vtab *pVTab) {
8883+ UNUSED_PARAMETER(pVTab);
8884+ vec0_vtab *p = (vec0_vtab *)pVTab;
8885+ if (p->stmtLatestChunk) {
8886+ sqlite3_finalize(p->stmtLatestChunk);
8887+ p->stmtLatestChunk = NULL;
8888+ }
8889+ if (p->stmtRowidsInsertRowid) {
8890+ sqlite3_finalize(p->stmtRowidsInsertRowid);
8891+ p->stmtRowidsInsertRowid = NULL;
8892+ }
8893+ if (p->stmtRowidsInsertId) {
8894+ sqlite3_finalize(p->stmtRowidsInsertId);
8895+ p->stmtRowidsInsertId = NULL;
8896+ }
8897+ if (p->stmtRowidsUpdatePosition) {
8898+ sqlite3_finalize(p->stmtRowidsUpdatePosition);
8899+ p->stmtRowidsUpdatePosition = NULL;
8900+ }
8901+ if (p->stmtRowidsGetChunkPosition) {
8902+ sqlite3_finalize(p->stmtRowidsGetChunkPosition);
8903+ p->stmtRowidsGetChunkPosition = NULL;
8904+ }
8905+ return SQLITE_OK;
8906+}
8907+static int vec0Commit(sqlite3_vtab *pVTab) {
8908+ UNUSED_PARAMETER(pVTab);
8909+ return SQLITE_OK;
8910+}
8911+static int vec0Rollback(sqlite3_vtab *pVTab) {
8912+ UNUSED_PARAMETER(pVTab);
8913+ return SQLITE_OK;
8914+}
8915+
8916+static sqlite3_module vec0Module = {
8917+ /* iVersion */ 3,
8918+ /* xCreate */ vec0Create,
8919+ /* xConnect */ vec0Connect,
8920+ /* xBestIndex */ vec0BestIndex,
8921+ /* xDisconnect */ vec0Disconnect,
8922+ /* xDestroy */ vec0Destroy,
8923+ /* xOpen */ vec0Open,
8924+ /* xClose */ vec0Close,
8925+ /* xFilter */ vec0Filter,
8926+ /* xNext */ vec0Next,
8927+ /* xEof */ vec0Eof,
8928+ /* xColumn */ vec0Column,
8929+ /* xRowid */ vec0Rowid,
8930+ /* xUpdate */ vec0Update,
8931+ /* xBegin */ vec0Begin,
8932+ /* xSync */ vec0Sync,
8933+ /* xCommit */ vec0Commit,
8934+ /* xRollback */ vec0Rollback,
8935+ /* xFindFunction */ 0,
8936+ /* xRename */ 0, // https://github.com/asg017/sqlite-vec/issues/43
8937+ /* xSavepoint */ 0,
8938+ /* xRelease */ 0,
8939+ /* xRollbackTo */ 0,
8940+ /* xShadowName */ vec0ShadowName,
8941+#if SQLITE_VERSION_NUMBER >= 3044000
8942+ /* xIntegrity */ 0, // https://github.com/asg017/sqlite-vec/issues/44
8943+#endif
8944+};
8945+#pragma endregion
8946+
8947+static char *POINTER_NAME_STATIC_BLOB_DEF = "vec0-static_blob_def";
8948+struct static_blob_definition {
8949+ void *p;
8950+ size_t dimensions;
8951+ size_t nvectors;
8952+ enum VectorElementType element_type;
8953+};
8954+static void vec_static_blob_from_raw(sqlite3_context *context, int argc,
8955+ sqlite3_value **argv) {
8956+
8957+ assert(argc == 4);
8958+ struct static_blob_definition *p;
8959+ p = sqlite3_malloc(sizeof(*p));
8960+ if (!p) {
8961+ sqlite3_result_error_nomem(context);
8962+ return;
8963+ }
8964+ memset(p, 0, sizeof(*p));
8965+ p->p = (void *)sqlite3_value_int64(argv[0]);
8966+ p->element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
8967+ p->dimensions = sqlite3_value_int64(argv[2]);
8968+ p->nvectors = sqlite3_value_int64(argv[3]);
8969+ sqlite3_result_pointer(context, p, POINTER_NAME_STATIC_BLOB_DEF,
8970+ sqlite3_free);
8971+}
8972+#pragma region vec_static_blobs() table function
8973+
8974+#define MAX_STATIC_BLOBS 16
8975+
8976+typedef struct static_blob static_blob;
8977+struct static_blob {
8978+ char *name;
8979+ void *p;
8980+ size_t dimensions;
8981+ size_t nvectors;
8982+ enum VectorElementType element_type;
8983+};
8984+
8985+typedef struct vec_static_blob_data vec_static_blob_data;
8986+struct vec_static_blob_data {
8987+ static_blob static_blobs[MAX_STATIC_BLOBS];
8988+};
8989+
8990+typedef struct vec_static_blobs_vtab vec_static_blobs_vtab;
8991+struct vec_static_blobs_vtab {
8992+ sqlite3_vtab base;
8993+ vec_static_blob_data *data;
8994+};
8995+
8996+typedef struct vec_static_blobs_cursor vec_static_blobs_cursor;
8997+struct vec_static_blobs_cursor {
8998+ sqlite3_vtab_cursor base;
8999+ sqlite3_int64 iRowid;
9000+};
9001+
9002+static int vec_static_blobsConnect(sqlite3 *db, void *pAux, int argc,
9003+ const char *const *argv,
9004+ sqlite3_vtab **ppVtab, char **pzErr) {
9005+ UNUSED_PARAMETER(argc);
9006+ UNUSED_PARAMETER(argv);
9007+ UNUSED_PARAMETER(pzErr);
9008+
9009+ vec_static_blobs_vtab *pNew;
9010+#define VEC_STATIC_BLOBS_NAME 0
9011+#define VEC_STATIC_BLOBS_DATA 1
9012+#define VEC_STATIC_BLOBS_DIMENSIONS 2
9013+#define VEC_STATIC_BLOBS_COUNT 3
9014+ int rc = sqlite3_declare_vtab(
9015+ db, "CREATE TABLE x(name, data, dimensions hidden, count hidden)");
9016+ if (rc == SQLITE_OK) {
9017+ pNew = sqlite3_malloc(sizeof(*pNew));
9018+ *ppVtab = (sqlite3_vtab *)pNew;
9019+ if (pNew == 0)
9020+ return SQLITE_NOMEM;
9021+ memset(pNew, 0, sizeof(*pNew));
9022+ pNew->data = pAux;
9023+ }
9024+ return rc;
9025+}
9026+
9027+static int vec_static_blobsDisconnect(sqlite3_vtab *pVtab) {
9028+ vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pVtab;
9029+ sqlite3_free(p);
9030+ return SQLITE_OK;
9031+}
9032+
9033+static int vec_static_blobsUpdate(sqlite3_vtab *pVTab, int argc,
9034+ sqlite3_value **argv, sqlite_int64 *pRowid) {
9035+ UNUSED_PARAMETER(pRowid);
9036+ vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pVTab;
9037+ // DELETE operation
9038+ if (argc == 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
9039+ return SQLITE_ERROR;
9040+ }
9041+ // INSERT operation
9042+ else if (argc > 1 && sqlite3_value_type(argv[0]) == SQLITE_NULL) {
9043+ const char *key =
9044+ (const char *)sqlite3_value_text(argv[2 + VEC_STATIC_BLOBS_NAME]);
9045+ int idx = -1;
9046+ for (int i = 0; i < MAX_STATIC_BLOBS; i++) {
9047+ if (!p->data->static_blobs[i].name) {
9048+ p->data->static_blobs[i].name = sqlite3_mprintf("%s", key);
9049+ idx = i;
9050+ break;
9051+ }
9052+ }
9053+ if (idx < 0)
9054+ abort();
9055+ struct static_blob_definition *def = sqlite3_value_pointer(
9056+ argv[2 + VEC_STATIC_BLOBS_DATA], POINTER_NAME_STATIC_BLOB_DEF);
9057+ p->data->static_blobs[idx].p = def->p;
9058+ p->data->static_blobs[idx].dimensions = def->dimensions;
9059+ p->data->static_blobs[idx].nvectors = def->nvectors;
9060+ p->data->static_blobs[idx].element_type = def->element_type;
9061+
9062+ return SQLITE_OK;
9063+ }
9064+ // UPDATE operation
9065+ else if (argc > 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
9066+ return SQLITE_ERROR;
9067+ }
9068+ return SQLITE_ERROR;
9069+}
9070+
9071+static int vec_static_blobsOpen(sqlite3_vtab *p,
9072+ sqlite3_vtab_cursor **ppCursor) {
9073+ UNUSED_PARAMETER(p);
9074+ vec_static_blobs_cursor *pCur;
9075+ pCur = sqlite3_malloc(sizeof(*pCur));
9076+ if (pCur == 0)
9077+ return SQLITE_NOMEM;
9078+ memset(pCur, 0, sizeof(*pCur));
9079+ *ppCursor = &pCur->base;
9080+ return SQLITE_OK;
9081+}
9082+
9083+static int vec_static_blobsClose(sqlite3_vtab_cursor *cur) {
9084+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9085+ sqlite3_free(pCur);
9086+ return SQLITE_OK;
9087+}
9088+
9089+static int vec_static_blobsBestIndex(sqlite3_vtab *pVTab,
9090+ sqlite3_index_info *pIdxInfo) {
9091+ UNUSED_PARAMETER(pVTab);
9092+ pIdxInfo->idxNum = 1;
9093+ pIdxInfo->estimatedCost = (double)10;
9094+ pIdxInfo->estimatedRows = 10;
9095+ return SQLITE_OK;
9096+}
9097+
9098+static int vec_static_blobsNext(sqlite3_vtab_cursor *cur);
9099+static int vec_static_blobsFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
9100+ const char *idxStr, int argc,
9101+ sqlite3_value **argv) {
9102+ UNUSED_PARAMETER(idxNum);
9103+ UNUSED_PARAMETER(idxStr);
9104+ UNUSED_PARAMETER(argc);
9105+ UNUSED_PARAMETER(argv);
9106+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)pVtabCursor;
9107+ pCur->iRowid = -1;
9108+ vec_static_blobsNext(pVtabCursor);
9109+ return SQLITE_OK;
9110+}
9111+
9112+static int vec_static_blobsRowid(sqlite3_vtab_cursor *cur,
9113+ sqlite_int64 *pRowid) {
9114+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9115+ *pRowid = pCur->iRowid;
9116+ return SQLITE_OK;
9117+}
9118+
9119+static int vec_static_blobsNext(sqlite3_vtab_cursor *cur) {
9120+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9121+ vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pCur->base.pVtab;
9122+ pCur->iRowid++;
9123+ while (pCur->iRowid < MAX_STATIC_BLOBS) {
9124+ if (p->data->static_blobs[pCur->iRowid].name) {
9125+ return SQLITE_OK;
9126+ }
9127+ pCur->iRowid++;
9128+ }
9129+ return SQLITE_OK;
9130+}
9131+
9132+static int vec_static_blobsEof(sqlite3_vtab_cursor *cur) {
9133+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9134+ return pCur->iRowid >= MAX_STATIC_BLOBS;
9135+}
9136+
9137+static int vec_static_blobsColumn(sqlite3_vtab_cursor *cur,
9138+ sqlite3_context *context, int i) {
9139+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9140+ vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)cur->pVtab;
9141+ switch (i) {
9142+ case VEC_STATIC_BLOBS_NAME:
9143+ sqlite3_result_text(context, p->data->static_blobs[pCur->iRowid].name, -1,
9144+ SQLITE_TRANSIENT);
9145+ break;
9146+ case VEC_STATIC_BLOBS_DATA:
9147+ sqlite3_result_null(context);
9148+ break;
9149+ case VEC_STATIC_BLOBS_DIMENSIONS:
9150+ sqlite3_result_int64(context,
9151+ p->data->static_blobs[pCur->iRowid].dimensions);
9152+ break;
9153+ case VEC_STATIC_BLOBS_COUNT:
9154+ sqlite3_result_int64(context, p->data->static_blobs[pCur->iRowid].nvectors);
9155+ break;
9156+ }
9157+ return SQLITE_OK;
9158+}
9159+
9160+static sqlite3_module vec_static_blobsModule = {
9161+ /* iVersion */ 3,
9162+ /* xCreate */ 0,
9163+ /* xConnect */ vec_static_blobsConnect,
9164+ /* xBestIndex */ vec_static_blobsBestIndex,
9165+ /* xDisconnect */ vec_static_blobsDisconnect,
9166+ /* xDestroy */ 0,
9167+ /* xOpen */ vec_static_blobsOpen,
9168+ /* xClose */ vec_static_blobsClose,
9169+ /* xFilter */ vec_static_blobsFilter,
9170+ /* xNext */ vec_static_blobsNext,
9171+ /* xEof */ vec_static_blobsEof,
9172+ /* xColumn */ vec_static_blobsColumn,
9173+ /* xRowid */ vec_static_blobsRowid,
9174+ /* xUpdate */ vec_static_blobsUpdate,
9175+ /* xBegin */ 0,
9176+ /* xSync */ 0,
9177+ /* xCommit */ 0,
9178+ /* xRollback */ 0,
9179+ /* xFindMethod */ 0,
9180+ /* xRename */ 0,
9181+ /* xSavepoint */ 0,
9182+ /* xRelease */ 0,
9183+ /* xRollbackTo */ 0,
9184+ /* xShadowName */ 0,
9185+#if SQLITE_VERSION_NUMBER >= 3044000
9186+ /* xIntegrity */ 0
9187+#endif
9188+};
9189+#pragma endregion
9190+
9191+#pragma region vec_static_blob_entries() table function
9192+
9193+typedef struct vec_static_blob_entries_vtab vec_static_blob_entries_vtab;
9194+struct vec_static_blob_entries_vtab {
9195+ sqlite3_vtab base;
9196+ static_blob *blob;
9197+};
9198+typedef enum {
9199+ VEC_SBE__QUERYPLAN_FULLSCAN = 1,
9200+ VEC_SBE__QUERYPLAN_KNN = 2
9201+} vec_sbe_query_plan;
9202+
9203+struct sbe_query_knn_data {
9204+ i64 k;
9205+ i64 k_used;
9206+ // Array of rowids of size k. Must be freed with sqlite3_free().
9207+ i32 *rowids;
9208+ // Array of distances of size k. Must be freed with sqlite3_free().
9209+ f32 *distances;
9210+ i64 current_idx;
9211+};
9212+void sbe_query_knn_data_clear(struct sbe_query_knn_data *knn_data) {
9213+ if (!knn_data)
9214+ return;
9215+
9216+ if (knn_data->rowids) {
9217+ sqlite3_free(knn_data->rowids);
9218+ knn_data->rowids = NULL;
9219+ }
9220+ if (knn_data->distances) {
9221+ sqlite3_free(knn_data->distances);
9222+ knn_data->distances = NULL;
9223+ }
9224+}
9225+
9226+typedef struct vec_static_blob_entries_cursor vec_static_blob_entries_cursor;
9227+struct vec_static_blob_entries_cursor {
9228+ sqlite3_vtab_cursor base;
9229+ sqlite3_int64 iRowid;
9230+ vec_sbe_query_plan query_plan;
9231+ struct sbe_query_knn_data *knn_data;
9232+};
9233+
9234+static int vec_static_blob_entriesConnect(sqlite3 *db, void *pAux, int argc,
9235+ const char *const *argv,
9236+ sqlite3_vtab **ppVtab, char **pzErr) {
9237+ UNUSED_PARAMETER(argc);
9238+ UNUSED_PARAMETER(argv);
9239+ UNUSED_PARAMETER(pzErr);
9240+ vec_static_blob_data *blob_data = pAux;
9241+ int idx = -1;
9242+ for (int i = 0; i < MAX_STATIC_BLOBS; i++) {
9243+ if (!blob_data->static_blobs[i].name)
9244+ continue;
9245+ if (strncmp(blob_data->static_blobs[i].name, argv[3],
9246+ strlen(blob_data->static_blobs[i].name)) == 0) {
9247+ idx = i;
9248+ break;
9249+ }
9250+ }
9251+ if (idx < 0)
9252+ abort();
9253+ vec_static_blob_entries_vtab *pNew;
9254+#define VEC_STATIC_BLOB_ENTRIES_VECTOR 0
9255+#define VEC_STATIC_BLOB_ENTRIES_DISTANCE 1
9256+#define VEC_STATIC_BLOB_ENTRIES_K 2
9257+ int rc = sqlite3_declare_vtab(
9258+ db, "CREATE TABLE x(vector, distance hidden, k hidden)");
9259+ if (rc == SQLITE_OK) {
9260+ pNew = sqlite3_malloc(sizeof(*pNew));
9261+ *ppVtab = (sqlite3_vtab *)pNew;
9262+ if (pNew == 0)
9263+ return SQLITE_NOMEM;
9264+ memset(pNew, 0, sizeof(*pNew));
9265+ pNew->blob = &blob_data->static_blobs[idx];
9266+ }
9267+ return rc;
9268+}
9269+
9270+static int vec_static_blob_entriesCreate(sqlite3 *db, void *pAux, int argc,
9271+ const char *const *argv,
9272+ sqlite3_vtab **ppVtab, char **pzErr) {
9273+ return vec_static_blob_entriesConnect(db, pAux, argc, argv, ppVtab, pzErr);
9274+}
9275+
9276+static int vec_static_blob_entriesDisconnect(sqlite3_vtab *pVtab) {
9277+ vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVtab;
9278+ sqlite3_free(p);
9279+ return SQLITE_OK;
9280+}
9281+
9282+static int vec_static_blob_entriesOpen(sqlite3_vtab *p,
9283+ sqlite3_vtab_cursor **ppCursor) {
9284+ UNUSED_PARAMETER(p);
9285+ vec_static_blob_entries_cursor *pCur;
9286+ pCur = sqlite3_malloc(sizeof(*pCur));
9287+ if (pCur == 0)
9288+ return SQLITE_NOMEM;
9289+ memset(pCur, 0, sizeof(*pCur));
9290+ *ppCursor = &pCur->base;
9291+ return SQLITE_OK;
9292+}
9293+
9294+static int vec_static_blob_entriesClose(sqlite3_vtab_cursor *cur) {
9295+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9296+ sqlite3_free(pCur->knn_data);
9297+ sqlite3_free(pCur);
9298+ return SQLITE_OK;
9299+}
9300+
9301+static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab,
9302+ sqlite3_index_info *pIdxInfo) {
9303+ vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVTab;
9304+ int iMatchTerm = -1;
9305+ int iLimitTerm = -1;
9306+ // int iRowidTerm = -1; // https://github.com/asg017/sqlite-vec/issues/47
9307+ int iKTerm = -1;
9308+
9309+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
9310+ if (!pIdxInfo->aConstraint[i].usable)
9311+ continue;
9312+
9313+ int iColumn = pIdxInfo->aConstraint[i].iColumn;
9314+ int op = pIdxInfo->aConstraint[i].op;
9315+ if (op == SQLITE_INDEX_CONSTRAINT_MATCH &&
9316+ iColumn == VEC_STATIC_BLOB_ENTRIES_VECTOR) {
9317+ if (iMatchTerm > -1) {
9318+ // https://github.com/asg017/sqlite-vec/issues/51
9319+ return SQLITE_ERROR;
9320+ }
9321+ iMatchTerm = i;
9322+ }
9323+ if (op == SQLITE_INDEX_CONSTRAINT_LIMIT) {
9324+ iLimitTerm = i;
9325+ }
9326+ if (op == SQLITE_INDEX_CONSTRAINT_EQ &&
9327+ iColumn == VEC_STATIC_BLOB_ENTRIES_K) {
9328+ iKTerm = i;
9329+ }
9330+ }
9331+ if (iMatchTerm >= 0) {
9332+ if (iLimitTerm < 0 && iKTerm < 0) {
9333+ // https://github.com/asg017/sqlite-vec/issues/51
9334+ return SQLITE_ERROR;
9335+ }
9336+ if (iLimitTerm >= 0 && iKTerm >= 0) {
9337+ return SQLITE_ERROR; // limit or k, not both
9338+ }
9339+ if (pIdxInfo->nOrderBy < 1) {
9340+ vtab_set_error(pVTab, "ORDER BY distance required");
9341+ return SQLITE_CONSTRAINT;
9342+ }
9343+ if (pIdxInfo->nOrderBy > 1) {
9344+ // https://github.com/asg017/sqlite-vec/issues/51
9345+ vtab_set_error(pVTab, "more than 1 ORDER BY clause provided");
9346+ return SQLITE_CONSTRAINT;
9347+ }
9348+ if (pIdxInfo->aOrderBy[0].iColumn != VEC_STATIC_BLOB_ENTRIES_DISTANCE) {
9349+ vtab_set_error(pVTab, "ORDER BY must be on the distance column");
9350+ return SQLITE_CONSTRAINT;
9351+ }
9352+ if (pIdxInfo->aOrderBy[0].desc) {
9353+ vtab_set_error(pVTab,
9354+ "Only ascending in ORDER BY distance clause is supported, "
9355+ "DESC is not supported yet.");
9356+ return SQLITE_CONSTRAINT;
9357+ }
9358+
9359+ pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_KNN;
9360+ pIdxInfo->estimatedCost = (double)10;
9361+ pIdxInfo->estimatedRows = 10;
9362+
9363+ pIdxInfo->orderByConsumed = 1;
9364+ pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = 1;
9365+ pIdxInfo->aConstraintUsage[iMatchTerm].omit = 1;
9366+ if (iLimitTerm >= 0) {
9367+ pIdxInfo->aConstraintUsage[iLimitTerm].argvIndex = 2;
9368+ pIdxInfo->aConstraintUsage[iLimitTerm].omit = 1;
9369+ } else {
9370+ pIdxInfo->aConstraintUsage[iKTerm].argvIndex = 2;
9371+ pIdxInfo->aConstraintUsage[iKTerm].omit = 1;
9372+ }
9373+
9374+ } else {
9375+ pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_FULLSCAN;
9376+ pIdxInfo->estimatedCost = (double)p->blob->nvectors;
9377+ pIdxInfo->estimatedRows = p->blob->nvectors;
9378+ }
9379+ return SQLITE_OK;
9380+}
9381+
9382+static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor,
9383+ int idxNum, const char *idxStr,
9384+ int argc, sqlite3_value **argv) {
9385+ UNUSED_PARAMETER(idxStr);
9386+ assert(argc >= 0 && argc <= 3);
9387+ vec_static_blob_entries_cursor *pCur =
9388+ (vec_static_blob_entries_cursor *)pVtabCursor;
9389+ vec_static_blob_entries_vtab *p =
9390+ (vec_static_blob_entries_vtab *)pCur->base.pVtab;
9391+
9392+ if (idxNum == VEC_SBE__QUERYPLAN_KNN) {
9393+ assert(argc == 2);
9394+ pCur->query_plan = VEC_SBE__QUERYPLAN_KNN;
9395+ struct sbe_query_knn_data *knn_data;
9396+ knn_data = sqlite3_malloc(sizeof(*knn_data));
9397+ if (!knn_data) {
9398+ return SQLITE_NOMEM;
9399+ }
9400+ memset(knn_data, 0, sizeof(*knn_data));
9401+
9402+ void *queryVector;
9403+ size_t dimensions;
9404+ enum VectorElementType elementType;
9405+ vector_cleanup cleanup;
9406+ char *err;
9407+ int rc = vector_from_value(argv[0], &queryVector, &dimensions, &elementType,
9408+ &cleanup, &err);
9409+ if (rc != SQLITE_OK) {
9410+ return SQLITE_ERROR;
9411+ }
9412+ if (elementType != p->blob->element_type) {
9413+ return SQLITE_ERROR;
9414+ }
9415+ if (dimensions != p->blob->dimensions) {
9416+ return SQLITE_ERROR;
9417+ }
9418+
9419+ i64 k = min(sqlite3_value_int64(argv[1]), (i64)p->blob->nvectors);
9420+ if (k < 0) {
9421+ // HANDLE https://github.com/asg017/sqlite-vec/issues/55
9422+ return SQLITE_ERROR;
9423+ }
9424+ if (k == 0) {
9425+ knn_data->k = 0;
9426+ pCur->knn_data = knn_data;
9427+ return SQLITE_OK;
9428+ }
9429+
9430+ size_t bsize = (p->blob->nvectors + 7) & ~7;
9431+
9432+ i32 *topk_rowids = sqlite3_malloc(k * sizeof(i32));
9433+ if (!topk_rowids) {
9434+ // HANDLE https://github.com/asg017/sqlite-vec/issues/55
9435+ return SQLITE_ERROR;
9436+ }
9437+ f32 *distances = sqlite3_malloc(bsize * sizeof(f32));
9438+ if (!distances) {
9439+ // HANDLE https://github.com/asg017/sqlite-vec/issues/55
9440+ return SQLITE_ERROR;
9441+ }
9442+
9443+ for (size_t i = 0; i < p->blob->nvectors; i++) {
9444+ // https://github.com/asg017/sqlite-vec/issues/52
9445+ float *v = ((float *)p->blob->p) + (i * p->blob->dimensions);
9446+ distances[i] =
9447+ distance_l2_sqr_float(v, (float *)queryVector, &p->blob->dimensions);
9448+ }
9449+ u8 *candidates = bitmap_new(bsize);
9450+ assert(candidates);
9451+
9452+ u8 *taken = bitmap_new(bsize);
9453+ assert(taken);
9454+
9455+ bitmap_fill(candidates, bsize);
9456+ for (size_t i = bsize; i >= p->blob->nvectors; i--) {
9457+ bitmap_set(candidates, i, 0);
9458+ }
9459+ i32 k_used = 0;
9460+ min_idx(distances, bsize, candidates, topk_rowids, k, taken, &k_used);
9461+ knn_data->current_idx = 0;
9462+ knn_data->distances = distances;
9463+ knn_data->k = k;
9464+ knn_data->rowids = topk_rowids;
9465+
9466+ pCur->knn_data = knn_data;
9467+ } else {
9468+ pCur->query_plan = VEC_SBE__QUERYPLAN_FULLSCAN;
9469+ pCur->iRowid = 0;
9470+ }
9471+
9472+ return SQLITE_OK;
9473+}
9474+
9475+static int vec_static_blob_entriesRowid(sqlite3_vtab_cursor *cur,
9476+ sqlite_int64 *pRowid) {
9477+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9478+ switch (pCur->query_plan) {
9479+ case VEC_SBE__QUERYPLAN_FULLSCAN: {
9480+ *pRowid = pCur->iRowid;
9481+ return SQLITE_OK;
9482+ }
9483+ case VEC_SBE__QUERYPLAN_KNN: {
9484+ i32 rowid = ((i32 *)pCur->knn_data->rowids)[pCur->knn_data->current_idx];
9485+ *pRowid = (sqlite3_int64)rowid;
9486+ return SQLITE_OK;
9487+ }
9488+ }
9489+ return SQLITE_ERROR;
9490+}
9491+
9492+static int vec_static_blob_entriesNext(sqlite3_vtab_cursor *cur) {
9493+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9494+ switch (pCur->query_plan) {
9495+ case VEC_SBE__QUERYPLAN_FULLSCAN: {
9496+ pCur->iRowid++;
9497+ return SQLITE_OK;
9498+ }
9499+ case VEC_SBE__QUERYPLAN_KNN: {
9500+ pCur->knn_data->current_idx++;
9501+ return SQLITE_OK;
9502+ }
9503+ }
9504+ return SQLITE_ERROR;
9505+}
9506+
9507+static int vec_static_blob_entriesEof(sqlite3_vtab_cursor *cur) {
9508+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9509+ vec_static_blob_entries_vtab *p =
9510+ (vec_static_blob_entries_vtab *)pCur->base.pVtab;
9511+ switch (pCur->query_plan) {
9512+ case VEC_SBE__QUERYPLAN_FULLSCAN: {
9513+ return (size_t)pCur->iRowid >= p->blob->nvectors;
9514+ }
9515+ case VEC_SBE__QUERYPLAN_KNN: {
9516+ return pCur->knn_data->current_idx >= pCur->knn_data->k;
9517+ }
9518+ }
9519+ return SQLITE_ERROR;
9520+}
9521+
9522+static int vec_static_blob_entriesColumn(sqlite3_vtab_cursor *cur,
9523+ sqlite3_context *context, int i) {
9524+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9525+ vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)cur->pVtab;
9526+
9527+ switch (pCur->query_plan) {
9528+ case VEC_SBE__QUERYPLAN_FULLSCAN: {
9529+ switch (i) {
9530+ case VEC_STATIC_BLOB_ENTRIES_VECTOR:
9531+
9532+ sqlite3_result_blob(
9533+ context,
9534+ ((unsigned char *)p->blob->p) +
9535+ (pCur->iRowid * p->blob->dimensions * sizeof(float)),
9536+ p->blob->dimensions * sizeof(float), SQLITE_TRANSIENT);
9537+ sqlite3_result_subtype(context, p->blob->element_type);
9538+ break;
9539+ }
9540+ return SQLITE_OK;
9541+ }
9542+ case VEC_SBE__QUERYPLAN_KNN: {
9543+ switch (i) {
9544+ case VEC_STATIC_BLOB_ENTRIES_VECTOR: {
9545+ i32 rowid = ((i32 *)pCur->knn_data->rowids)[pCur->knn_data->current_idx];
9546+ sqlite3_result_blob(context,
9547+ ((unsigned char *)p->blob->p) +
9548+ (rowid * p->blob->dimensions * sizeof(float)),
9549+ p->blob->dimensions * sizeof(float),
9550+ SQLITE_TRANSIENT);
9551+ sqlite3_result_subtype(context, p->blob->element_type);
9552+ break;
9553+ }
9554+ }
9555+ return SQLITE_OK;
9556+ }
9557+ }
9558+ return SQLITE_ERROR;
9559+}
9560+
9561+static sqlite3_module vec_static_blob_entriesModule = {
9562+ /* iVersion */ 3,
9563+ /* xCreate */
9564+ vec_static_blob_entriesCreate, // handle rm?
9565+ // https://github.com/asg017/sqlite-vec/issues/55
9566+ /* xConnect */ vec_static_blob_entriesConnect,
9567+ /* xBestIndex */ vec_static_blob_entriesBestIndex,
9568+ /* xDisconnect */ vec_static_blob_entriesDisconnect,
9569+ /* xDestroy */ vec_static_blob_entriesDisconnect,
9570+ /* xOpen */ vec_static_blob_entriesOpen,
9571+ /* xClose */ vec_static_blob_entriesClose,
9572+ /* xFilter */ vec_static_blob_entriesFilter,
9573+ /* xNext */ vec_static_blob_entriesNext,
9574+ /* xEof */ vec_static_blob_entriesEof,
9575+ /* xColumn */ vec_static_blob_entriesColumn,
9576+ /* xRowid */ vec_static_blob_entriesRowid,
9577+ /* xUpdate */ 0,
9578+ /* xBegin */ 0,
9579+ /* xSync */ 0,
9580+ /* xCommit */ 0,
9581+ /* xRollback */ 0,
9582+ /* xFindMethod */ 0,
9583+ /* xRename */ 0,
9584+ /* xSavepoint */ 0,
9585+ /* xRelease */ 0,
9586+ /* xRollbackTo */ 0,
9587+ /* xShadowName */ 0,
9588+#if SQLITE_VERSION_NUMBER >= 3044000
9589+ /* xIntegrity */ 0
9590+#endif
9591+};
9592+#pragma endregion
9593+
9594+#ifdef SQLITE_VEC_ENABLE_AVX
9595+#define SQLITE_VEC_DEBUG_BUILD_AVX "avx"
9596+#else
9597+#define SQLITE_VEC_DEBUG_BUILD_AVX ""
9598+#endif
9599+#ifdef SQLITE_VEC_ENABLE_NEON
9600+#define SQLITE_VEC_DEBUG_BUILD_NEON "neon"
9601+#else
9602+#define SQLITE_VEC_DEBUG_BUILD_NEON ""
9603+#endif
9604+
9605+#define SQLITE_VEC_DEBUG_BUILD \
9606+ SQLITE_VEC_DEBUG_BUILD_AVX " " SQLITE_VEC_DEBUG_BUILD_NEON
9607+
9608+#define SQLITE_VEC_DEBUG_STRING \
9609+ "Version: " SQLITE_VEC_VERSION "\n" \
9610+ "Date: " SQLITE_VEC_DATE "\n" \
9611+ "Commit: " SQLITE_VEC_SOURCE "\n" \
9612+ "Build flags: " SQLITE_VEC_DEBUG_BUILD
9613+
9614+SQLITE_VEC_API int sqlite3_vec_init(sqlite3 *db, char **pzErrMsg,
9615+ const sqlite3_api_routines *pApi) {
9616+#ifndef SQLITE_CORE
9617+ SQLITE_EXTENSION_INIT2(pApi);
9618+#endif
9619+ int rc = SQLITE_OK;
9620+
9621+#define DEFAULT_FLAGS (SQLITE_UTF8 | SQLITE_INNOCUOUS | SQLITE_DETERMINISTIC)
9622+
9623+ rc = sqlite3_create_function_v2(db, "vec_version", 0, DEFAULT_FLAGS,
9624+ SQLITE_VEC_VERSION, _static_text_func, NULL,
9625+ NULL, NULL);
9626+ if (rc != SQLITE_OK) {
9627+ return rc;
9628+ }
9629+ rc = sqlite3_create_function_v2(db, "vec_debug", 0, DEFAULT_FLAGS,
9630+ SQLITE_VEC_DEBUG_STRING, _static_text_func,
9631+ NULL, NULL, NULL);
9632+ if (rc != SQLITE_OK) {
9633+ return rc;
9634+ }
9635+ static struct {
9636+ const char *zFName;
9637+ void (*xFunc)(sqlite3_context *, int, sqlite3_value **);
9638+ int nArg;
9639+ int flags;
9640+ } aFunc[] = {
9641+ // clang-format off
9642+ //{"vec_version", _static_text_func, 0, DEFAULT_FLAGS, (void *) SQLITE_VEC_VERSION },
9643+ //{"vec_debug", _static_text_func, 0, DEFAULT_FLAGS, (void *) SQLITE_VEC_DEBUG_STRING },
9644+ {"vec_distance_l2", vec_distance_l2, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9645+ {"vec_distance_l1", vec_distance_l1, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9646+ {"vec_distance_hamming",vec_distance_hamming, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9647+ {"vec_distance_cosine", vec_distance_cosine, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9648+ {"vec_length", vec_length, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9649+ {"vec_type", vec_type, 1, DEFAULT_FLAGS, },
9650+ {"vec_to_json", vec_to_json, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9651+ {"vec_add", vec_add, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9652+ {"vec_sub", vec_sub, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9653+ {"vec_slice", vec_slice, 3, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9654+ {"vec_normalize", vec_normalize, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9655+ {"vec_f32", vec_f32, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9656+ {"vec_bit", vec_bit, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9657+ {"vec_int8", vec_int8, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9658+ {"vec_quantize_int8", vec_quantize_int8, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9659+ {"vec_quantize_binary", vec_quantize_binary, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9660+ // clang-format on
9661+ };
9662+
9663+ static struct {
9664+ char *name;
9665+ const sqlite3_module *module;
9666+ void *p;
9667+ void (*xDestroy)(void *);
9668+ } aMod[] = {
9669+ // clang-format off
9670+ {"vec0", &vec0Module, NULL, NULL},
9671+ {"vec_each", &vec_eachModule, NULL, NULL},
9672+ // clang-format on
9673+ };
9674+
9675+ for (unsigned long i = 0; i < countof(aFunc) && rc == SQLITE_OK; i++) {
9676+ rc = sqlite3_create_function_v2(db, aFunc[i].zFName, aFunc[i].nArg,
9677+ aFunc[i].flags, NULL, aFunc[i].xFunc, NULL,
9678+ NULL, NULL);
9679+ if (rc != SQLITE_OK) {
9680+ *pzErrMsg = sqlite3_mprintf("Error creating function %s: %s",
9681+ aFunc[i].zFName, sqlite3_errmsg(db));
9682+ return rc;
9683+ }
9684+ }
9685+
9686+ for (unsigned long i = 0; i < countof(aMod) && rc == SQLITE_OK; i++) {
9687+ rc = sqlite3_create_module_v2(db, aMod[i].name, aMod[i].module, NULL, NULL);
9688+ if (rc != SQLITE_OK) {
9689+ *pzErrMsg = sqlite3_mprintf("Error creating module %s: %s", aMod[i].name,
9690+ sqlite3_errmsg(db));
9691+ return rc;
9692+ }
9693+ }
9694+
9695+ return SQLITE_OK;
9696+}
9697+
9698+#ifndef SQLITE_VEC_OMIT_FS
9699+SQLITE_VEC_API int sqlite3_vec_numpy_init(sqlite3 *db, char **pzErrMsg,
9700+ const sqlite3_api_routines *pApi) {
9701+ UNUSED_PARAMETER(pzErrMsg);
9702+#ifndef SQLITE_CORE
9703+ SQLITE_EXTENSION_INIT2(pApi);
9704+#endif
9705+ int rc = SQLITE_OK;
9706+ rc = sqlite3_create_function_v2(db, "vec_npy_file", 1, SQLITE_RESULT_SUBTYPE,
9707+ NULL, vec_npy_file, NULL, NULL, NULL);
9708+ if(rc != SQLITE_OK) {
9709+ return rc;
9710+ }
9711+ rc = sqlite3_create_module_v2(db, "vec_npy_each", &vec_npy_eachModule, NULL, NULL);
9712+ return rc;
9713+}
9714+#endif
9715+
9716+SQLITE_VEC_API int
9717+sqlite3_vec_static_blobs_init(sqlite3 *db, char **pzErrMsg,
9718+ const sqlite3_api_routines *pApi) {
9719+ UNUSED_PARAMETER(pzErrMsg);
9720+#ifndef SQLITE_CORE
9721+ SQLITE_EXTENSION_INIT2(pApi);
9722+#endif
9723+
9724+ int rc = SQLITE_OK;
9725+ vec_static_blob_data *static_blob_data;
9726+ static_blob_data = sqlite3_malloc(sizeof(*static_blob_data));
9727+ if (!static_blob_data) {
9728+ return SQLITE_NOMEM;
9729+ }
9730+ memset(static_blob_data, 0, sizeof(*static_blob_data));
9731+
9732+ rc = sqlite3_create_function_v2(
9733+ db, "vec_static_blob_from_raw", 4,
9734+ DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, NULL,
9735+ vec_static_blob_from_raw, NULL, NULL, NULL);
9736+ if (rc != SQLITE_OK)
9737+ return rc;
9738+
9739+ rc = sqlite3_create_module_v2(db, "vec_static_blobs", &vec_static_blobsModule,
9740+ static_blob_data, sqlite3_free);
9741+ if (rc != SQLITE_OK)
9742+ return rc;
9743+ rc = sqlite3_create_module_v2(db, "vec_static_blob_entries",
9744+ &vec_static_blob_entriesModule,
9745+ static_blob_data, NULL);
9746+ if (rc != SQLITE_OK)
9747+ return rc;
9748+ return rc;
9749+}
new file mode 100644
@@ -0,0 +1,9749 @@
1+#include "sqlite-vec.h"
2+
3+#include <assert.h>
4+#include <errno.h>
5+#include <float.h>
6+#include <inttypes.h>
7+#include <limits.h>
8+#include <math.h>
9+#include <stdbool.h>
10+#include <stdint.h>
11+#include <stdlib.h>
12+#include <string.h>
13+
14+#ifndef SQLITE_VEC_OMIT_FS
15+#include <stdio.h>
16+#endif
17+
18+#ifndef SQLITE_CORE
19+#include "sqlite3ext.h"
20+SQLITE_EXTENSION_INIT1
21+#else
22+#include "sqlite3.h"
23+#endif
24+
25+#ifndef UINT32_TYPE
26+#ifdef HAVE_UINT32_T
27+#define UINT32_TYPE uint32_t
28+#else
29+#define UINT32_TYPE unsigned int
30+#endif
31+#endif
32+#ifndef UINT16_TYPE
33+#ifdef HAVE_UINT16_T
34+#define UINT16_TYPE uint16_t
35+#else
36+#define UINT16_TYPE unsigned short int
37+#endif
38+#endif
39+#ifndef INT16_TYPE
40+#ifdef HAVE_INT16_T
41+#define INT16_TYPE int16_t
42+#else
43+#define INT16_TYPE short int
44+#endif
45+#endif
46+#ifndef UINT8_TYPE
47+#ifdef HAVE_UINT8_T
48+#define UINT8_TYPE uint8_t
49+#else
50+#define UINT8_TYPE unsigned char
51+#endif
52+#endif
53+#ifndef INT8_TYPE
54+#ifdef HAVE_INT8_T
55+#define INT8_TYPE int8_t
56+#else
57+#define INT8_TYPE signed char
58+#endif
59+#endif
60+#ifndef LONGDOUBLE_TYPE
61+#define LONGDOUBLE_TYPE long double
62+#endif
63+
64+#ifndef _WIN32
65+#ifndef __EMSCRIPTEN__
66+#ifndef __COSMOPOLITAN__
67+#ifndef __wasi__
68+typedef u_int8_t uint8_t;
69+typedef u_int16_t uint16_t;
70+typedef u_int64_t uint64_t;
71+#endif
72+#endif
73+#endif
74+#endif
75+
76+typedef int8_t i8;
77+typedef uint8_t u8;
78+typedef int16_t i16;
79+typedef int32_t i32;
80+typedef sqlite3_int64 i64;
81+typedef uint32_t u32;
82+typedef uint64_t u64;
83+typedef float f32;
84+typedef size_t usize;
85+
86+#ifndef UNUSED_PARAMETER
87+#define UNUSED_PARAMETER(X) (void)(X)
88+#endif
89+
90+// sqlite3_vtab_in() was added in SQLite version 3.38 (2022-02-22)
91+// https://www.sqlite.org/changes.html#version_3_38_0
92+#if SQLITE_VERSION_NUMBER >= 3038000
93+#define COMPILER_SUPPORTS_VTAB_IN 1
94+#endif
95+
96+#ifndef SQLITE_SUBTYPE
97+#define SQLITE_SUBTYPE 0x000100000
98+#endif
99+
100+#ifndef SQLITE_RESULT_SUBTYPE
101+#define SQLITE_RESULT_SUBTYPE 0x001000000
102+#endif
103+
104+#ifndef SQLITE_INDEX_CONSTRAINT_LIMIT
105+#define SQLITE_INDEX_CONSTRAINT_LIMIT 73
106+#endif
107+
108+#ifndef SQLITE_INDEX_CONSTRAINT_OFFSET
109+#define SQLITE_INDEX_CONSTRAINT_OFFSET 74
110+#endif
111+
112+#define countof(x) (sizeof(x) / sizeof((x)[0]))
113+#define min(a, b) (((a) <= (b)) ? (a) : (b))
114+
115+enum VectorElementType {
116+ // clang-format off
117+ SQLITE_VEC_ELEMENT_TYPE_FLOAT32 = 223 + 0,
118+ SQLITE_VEC_ELEMENT_TYPE_BIT = 223 + 1,
119+ SQLITE_VEC_ELEMENT_TYPE_INT8 = 223 + 2,
120+ // clang-format on
121+};
122+
123+#ifdef SQLITE_VEC_ENABLE_AVX
124+#include <immintrin.h>
125+#define PORTABLE_ALIGN32 __attribute__((aligned(32)))
126+#define PORTABLE_ALIGN64 __attribute__((aligned(64)))
127+
128+static f32 l2_sqr_float_avx(const void *pVect1v, const void *pVect2v,
129+ const void *qty_ptr) {
130+ f32 *pVect1 = (f32 *)pVect1v;
131+ f32 *pVect2 = (f32 *)pVect2v;
132+ size_t qty = *((size_t *)qty_ptr);
133+ f32 PORTABLE_ALIGN32 TmpRes[8];
134+ size_t qty16 = qty >> 4;
135+
136+ const f32 *pEnd1 = pVect1 + (qty16 << 4);
137+
138+ __m256 diff, v1, v2;
139+ __m256 sum = _mm256_set1_ps(0);
140+
141+ while (pVect1 < pEnd1) {
142+ v1 = _mm256_loadu_ps(pVect1);
143+ pVect1 += 8;
144+ v2 = _mm256_loadu_ps(pVect2);
145+ pVect2 += 8;
146+ diff = _mm256_sub_ps(v1, v2);
147+ sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff));
148+
149+ v1 = _mm256_loadu_ps(pVect1);
150+ pVect1 += 8;
151+ v2 = _mm256_loadu_ps(pVect2);
152+ pVect2 += 8;
153+ diff = _mm256_sub_ps(v1, v2);
154+ sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff));
155+ }
156+
157+ _mm256_store_ps(TmpRes, sum);
158+ return sqrt(TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] +
159+ TmpRes[5] + TmpRes[6] + TmpRes[7]);
160+}
161+#endif
162+
163+#ifdef SQLITE_VEC_ENABLE_NEON
164+#include <arm_neon.h>
165+
166+#define PORTABLE_ALIGN32 __attribute__((aligned(32)))
167+
168+// thx https://github.com/nmslib/hnswlib/pull/299/files
169+static f32 l2_sqr_float_neon(const void *pVect1v, const void *pVect2v,
170+ const void *qty_ptr) {
171+ f32 *pVect1 = (f32 *)pVect1v;
172+ f32 *pVect2 = (f32 *)pVect2v;
173+ size_t qty = *((size_t *)qty_ptr);
174+ size_t qty16 = qty >> 4;
175+
176+ const f32 *pEnd1 = pVect1 + (qty16 << 4);
177+
178+ float32x4_t diff, v1, v2;
179+ float32x4_t sum0 = vdupq_n_f32(0);
180+ float32x4_t sum1 = vdupq_n_f32(0);
181+ float32x4_t sum2 = vdupq_n_f32(0);
182+ float32x4_t sum3 = vdupq_n_f32(0);
183+
184+ while (pVect1 < pEnd1) {
185+ v1 = vld1q_f32(pVect1);
186+ pVect1 += 4;
187+ v2 = vld1q_f32(pVect2);
188+ pVect2 += 4;
189+ diff = vsubq_f32(v1, v2);
190+ sum0 = vfmaq_f32(sum0, diff, diff);
191+
192+ v1 = vld1q_f32(pVect1);
193+ pVect1 += 4;
194+ v2 = vld1q_f32(pVect2);
195+ pVect2 += 4;
196+ diff = vsubq_f32(v1, v2);
197+ sum1 = vfmaq_f32(sum1, diff, diff);
198+
199+ v1 = vld1q_f32(pVect1);
200+ pVect1 += 4;
201+ v2 = vld1q_f32(pVect2);
202+ pVect2 += 4;
203+ diff = vsubq_f32(v1, v2);
204+ sum2 = vfmaq_f32(sum2, diff, diff);
205+
206+ v1 = vld1q_f32(pVect1);
207+ pVect1 += 4;
208+ v2 = vld1q_f32(pVect2);
209+ pVect2 += 4;
210+ diff = vsubq_f32(v1, v2);
211+ sum3 = vfmaq_f32(sum3, diff, diff);
212+ }
213+
214+ f32 sum_scalar =
215+ vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3)));
216+ const f32 *pEnd2 = pVect1 + (qty - (qty16 << 4));
217+ while (pVect1 < pEnd2) {
218+ f32 diff = *pVect1 - *pVect2;
219+ sum_scalar += diff * diff;
220+ pVect1++;
221+ pVect2++;
222+ }
223+
224+ return sqrt(sum_scalar);
225+}
226+
227+static f32 l2_sqr_int8_neon(const void *pVect1v, const void *pVect2v,
228+ const void *qty_ptr) {
229+ i8 *pVect1 = (i8 *)pVect1v;
230+ i8 *pVect2 = (i8 *)pVect2v;
231+ size_t qty = *((size_t *)qty_ptr);
232+
233+ const i8 *pEnd1 = pVect1 + qty;
234+ i32 sum_scalar = 0;
235+
236+ while (pVect1 < pEnd1 - 7) {
237+ // loading 8 at a time
238+ int8x8_t v1 = vld1_s8(pVect1);
239+ int8x8_t v2 = vld1_s8(pVect2);
240+ pVect1 += 8;
241+ pVect2 += 8;
242+
243+ // widen to protect against overflow
244+ int16x8_t v1_wide = vmovl_s8(v1);
245+ int16x8_t v2_wide = vmovl_s8(v2);
246+
247+ int16x8_t diff = vsubq_s16(v1_wide, v2_wide);
248+ int16x8_t squared_diff = vmulq_s16(diff, diff);
249+ int32x4_t sum = vpaddlq_s16(squared_diff);
250+
251+ sum_scalar += vgetq_lane_s32(sum, 0) + vgetq_lane_s32(sum, 1) +
252+ vgetq_lane_s32(sum, 2) + vgetq_lane_s32(sum, 3);
253+ }
254+
255+ // handle leftovers
256+ while (pVect1 < pEnd1) {
257+ i16 diff = (i16)*pVect1 - (i16)*pVect2;
258+ sum_scalar += diff * diff;
259+ pVect1++;
260+ pVect2++;
261+ }
262+
263+ return sqrtf(sum_scalar);
264+}
265+
266+static i32 l1_int8_neon(const void *pVect1v, const void *pVect2v,
267+ const void *qty_ptr) {
268+ i8 *pVect1 = (i8 *)pVect1v;
269+ i8 *pVect2 = (i8 *)pVect2v;
270+ size_t qty = *((size_t *)qty_ptr);
271+
272+ const int8_t *pEnd1 = pVect1 + qty;
273+
274+ int32x4_t acc1 = vdupq_n_s32(0);
275+ int32x4_t acc2 = vdupq_n_s32(0);
276+ int32x4_t acc3 = vdupq_n_s32(0);
277+ int32x4_t acc4 = vdupq_n_s32(0);
278+
279+ while (pVect1 < pEnd1 - 63) {
280+ int8x16_t v1 = vld1q_s8(pVect1);
281+ int8x16_t v2 = vld1q_s8(pVect2);
282+ int8x16_t diff1 = vabdq_s8(v1, v2);
283+ acc1 = vaddq_s32(acc1, vpaddlq_u16(vpaddlq_u8(diff1)));
284+
285+ v1 = vld1q_s8(pVect1 + 16);
286+ v2 = vld1q_s8(pVect2 + 16);
287+ int8x16_t diff2 = vabdq_s8(v1, v2);
288+ acc2 = vaddq_s32(acc2, vpaddlq_u16(vpaddlq_u8(diff2)));
289+
290+ v1 = vld1q_s8(pVect1 + 32);
291+ v2 = vld1q_s8(pVect2 + 32);
292+ int8x16_t diff3 = vabdq_s8(v1, v2);
293+ acc3 = vaddq_s32(acc3, vpaddlq_u16(vpaddlq_u8(diff3)));
294+
295+ v1 = vld1q_s8(pVect1 + 48);
296+ v2 = vld1q_s8(pVect2 + 48);
297+ int8x16_t diff4 = vabdq_s8(v1, v2);
298+ acc4 = vaddq_s32(acc4, vpaddlq_u16(vpaddlq_u8(diff4)));
299+
300+ pVect1 += 64;
301+ pVect2 += 64;
302+ }
303+
304+ while (pVect1 < pEnd1 - 15) {
305+ int8x16_t v1 = vld1q_s8(pVect1);
306+ int8x16_t v2 = vld1q_s8(pVect2);
307+ int8x16_t diff = vabdq_s8(v1, v2);
308+ acc1 = vaddq_s32(acc1, vpaddlq_u16(vpaddlq_u8(diff)));
309+ pVect1 += 16;
310+ pVect2 += 16;
311+ }
312+
313+ int32x4_t acc = vaddq_s32(vaddq_s32(acc1, acc2), vaddq_s32(acc3, acc4));
314+
315+ int32_t sum = 0;
316+ while (pVect1 < pEnd1) {
317+ int32_t diff = abs((int32_t)*pVect1 - (int32_t)*pVect2);
318+ sum += diff;
319+ pVect1++;
320+ pVect2++;
321+ }
322+
323+ return vaddvq_s32(acc) + sum;
324+}
325+
326+static double l1_f32_neon(const void *pVect1v, const void *pVect2v,
327+ const void *qty_ptr) {
328+ f32 *pVect1 = (f32 *)pVect1v;
329+ f32 *pVect2 = (f32 *)pVect2v;
330+ size_t qty = *((size_t *)qty_ptr);
331+
332+ const f32 *pEnd1 = pVect1 + qty;
333+ float64x2_t acc = vdupq_n_f64(0);
334+
335+ while (pVect1 < pEnd1 - 3) {
336+ float32x4_t v1 = vld1q_f32(pVect1);
337+ float32x4_t v2 = vld1q_f32(pVect2);
338+ pVect1 += 4;
339+ pVect2 += 4;
340+
341+ // f32x4 -> f64x2 pad for overflow
342+ float64x2_t low_diff = vabdq_f64(vcvt_f64_f32(vget_low_f32(v1)),
343+ vcvt_f64_f32(vget_low_f32(v2)));
344+ float64x2_t high_diff =
345+ vabdq_f64(vcvt_high_f64_f32(v1), vcvt_high_f64_f32(v2));
346+
347+ acc = vaddq_f64(acc, vaddq_f64(low_diff, high_diff));
348+ }
349+
350+ double sum = 0;
351+ while (pVect1 < pEnd1) {
352+ sum += fabs((double)*pVect1 - (double)*pVect2);
353+ pVect1++;
354+ pVect2++;
355+ }
356+
357+ return vaddvq_f64(acc) + sum;
358+}
359+#endif
360+
361+static f32 l2_sqr_float(const void *pVect1v, const void *pVect2v,
362+ const void *qty_ptr) {
363+ f32 *pVect1 = (f32 *)pVect1v;
364+ f32 *pVect2 = (f32 *)pVect2v;
365+ size_t qty = *((size_t *)qty_ptr);
366+
367+ f32 res = 0;
368+ for (size_t i = 0; i < qty; i++) {
369+ f32 t = *pVect1 - *pVect2;
370+ pVect1++;
371+ pVect2++;
372+ res += t * t;
373+ }
374+ return sqrt(res);
375+}
376+
377+static f32 l2_sqr_int8(const void *pA, const void *pB, const void *pD) {
378+ i8 *a = (i8 *)pA;
379+ i8 *b = (i8 *)pB;
380+ size_t d = *((size_t *)pD);
381+
382+ f32 res = 0;
383+ for (size_t i = 0; i < d; i++) {
384+ f32 t = *a - *b;
385+ a++;
386+ b++;
387+ res += t * t;
388+ }
389+ return sqrt(res);
390+}
391+
392+static f32 distance_l2_sqr_float(const void *a, const void *b, const void *d) {
393+#ifdef SQLITE_VEC_ENABLE_NEON
394+ if ((*(const size_t *)d) > 16) {
395+ return l2_sqr_float_neon(a, b, d);
396+ }
397+#endif
398+#ifdef SQLITE_VEC_ENABLE_AVX
399+ if (((*(const size_t *)d) % 16 == 0)) {
400+ return l2_sqr_float_avx(a, b, d);
401+ }
402+#endif
403+ return l2_sqr_float(a, b, d);
404+}
405+
406+static f32 distance_l2_sqr_int8(const void *a, const void *b, const void *d) {
407+#ifdef SQLITE_VEC_ENABLE_NEON
408+ if ((*(const size_t *)d) > 7) {
409+ return l2_sqr_int8_neon(a, b, d);
410+ }
411+#endif
412+ return l2_sqr_int8(a, b, d);
413+}
414+
415+static i32 l1_int8(const void *pA, const void *pB, const void *pD) {
416+ i8 *a = (i8 *)pA;
417+ i8 *b = (i8 *)pB;
418+ size_t d = *((size_t *)pD);
419+
420+ i32 res = 0;
421+ for (size_t i = 0; i < d; i++) {
422+ res += abs(*a - *b);
423+ a++;
424+ b++;
425+ }
426+
427+ return res;
428+}
429+
430+static i32 distance_l1_int8(const void *a, const void *b, const void *d) {
431+#ifdef SQLITE_VEC_ENABLE_NEON
432+ if ((*(const size_t *)d) > 15) {
433+ return l1_int8_neon(a, b, d);
434+ }
435+#endif
436+ return l1_int8(a, b, d);
437+}
438+
439+static double l1_f32(const void *pA, const void *pB, const void *pD) {
440+ f32 *a = (f32 *)pA;
441+ f32 *b = (f32 *)pB;
442+ size_t d = *((size_t *)pD);
443+
444+ double res = 0;
445+ for (size_t i = 0; i < d; i++) {
446+ res += fabs((double)*a - (double)*b);
447+ a++;
448+ b++;
449+ }
450+
451+ return res;
452+}
453+
454+static double distance_l1_f32(const void *a, const void *b, const void *d) {
455+#ifdef SQLITE_VEC_ENABLE_NEON
456+ if ((*(const size_t *)d) > 3) {
457+ return l1_f32_neon(a, b, d);
458+ }
459+#endif
460+ return l1_f32(a, b, d);
461+}
462+
463+static f32 distance_cosine_float(const void *pVect1v, const void *pVect2v,
464+ const void *qty_ptr) {
465+ f32 *pVect1 = (f32 *)pVect1v;
466+ f32 *pVect2 = (f32 *)pVect2v;
467+ size_t qty = *((size_t *)qty_ptr);
468+
469+ f32 dot = 0;
470+ f32 aMag = 0;
471+ f32 bMag = 0;
472+ for (size_t i = 0; i < qty; i++) {
473+ dot += *pVect1 * *pVect2;
474+ aMag += *pVect1 * *pVect1;
475+ bMag += *pVect2 * *pVect2;
476+ pVect1++;
477+ pVect2++;
478+ }
479+ return 1 - (dot / (sqrt(aMag) * sqrt(bMag)));
480+}
481+static f32 distance_cosine_int8(const void *pA, const void *pB,
482+ const void *pD) {
483+ i8 *a = (i8 *)pA;
484+ i8 *b = (i8 *)pB;
485+ size_t d = *((size_t *)pD);
486+
487+ f32 dot = 0;
488+ f32 aMag = 0;
489+ f32 bMag = 0;
490+ for (size_t i = 0; i < d; i++) {
491+ dot += *a * *b;
492+ aMag += *a * *a;
493+ bMag += *b * *b;
494+ a++;
495+ b++;
496+ }
497+ return 1 - (dot / (sqrt(aMag) * sqrt(bMag)));
498+}
499+
500+// https://github.com/facebookresearch/faiss/blob/77e2e79cd0a680adc343b9840dd865da724c579e/faiss/utils/hamming_distance/common.h#L34
501+static u8 hamdist_table[256] = {
502+ 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 1, 2, 2, 3, 2, 3, 3, 4,
503+ 2, 3, 3, 4, 3, 4, 4, 5, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
504+ 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 1, 2, 2, 3, 2, 3, 3, 4,
505+ 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
506+ 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6,
507+ 4, 5, 5, 6, 5, 6, 6, 7, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
508+ 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5,
509+ 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
510+ 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6,
511+ 4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
512+ 4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8};
513+
514+static f32 distance_hamming_u8(u8 *a, u8 *b, size_t n) {
515+ int same = 0;
516+ for (unsigned long i = 0; i < n; i++) {
517+ same += hamdist_table[a[i] ^ b[i]];
518+ }
519+ return (f32)same;
520+}
521+
522+#ifdef _MSC_VER
523+#if !defined(__clang__) && (defined(_M_ARM) || defined(_M_ARM64))
524+// From
525+// https://github.com/ngtcp2/ngtcp2/blob/b64f1e77b5e0d880b93d31f474147fae4a1d17cc/lib/ngtcp2_ringbuf.c,
526+// line 34-43
527+static unsigned int __builtin_popcountl(unsigned int x) {
528+ unsigned int c = 0;
529+ for (; x; ++c) {
530+ x &= x - 1;
531+ }
532+ return c;
533+}
534+#else
535+#include <intrin.h>
536+#define __builtin_popcountl __popcnt64
537+#endif
538+#endif
539+
540+static f32 distance_hamming_u64(u64 *a, u64 *b, size_t n) {
541+ int same = 0;
542+ for (unsigned long i = 0; i < n; i++) {
543+ same += __builtin_popcountl(a[i] ^ b[i]);
544+ }
545+ return (f32)same;
546+}
547+
548+/**
549+ * @brief Calculate the hamming distance between two bitvectors.
550+ *
551+ * @param a - first bitvector, MUST have d dimensions
552+ * @param b - second bitvector, MUST have d dimensions
553+ * @param d - pointer to size_t, MUST be divisible by CHAR_BIT
554+ * @return f32
555+ */
556+static f32 distance_hamming(const void *a, const void *b, const void *d) {
557+ size_t dimensions = *((size_t *)d);
558+
559+ if ((dimensions % 64) == 0) {
560+ return distance_hamming_u64((u64 *)a, (u64 *)b, dimensions / 8 / CHAR_BIT);
561+ }
562+ return distance_hamming_u8((u8 *)a, (u8 *)b, dimensions / CHAR_BIT);
563+}
564+
565+// from SQLite source:
566+// https://github.com/sqlite/sqlite/blob/a509a90958ddb234d1785ed7801880ccb18b497e/src/json.c#L153
567+static const char vecJsonIsSpaceX[] = {
568+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0,
569+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
570+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
571+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
572+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
573+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
574+
575+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
576+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
577+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
578+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
579+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
580+ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
581+};
582+
583+#define vecJsonIsspace(x) (vecJsonIsSpaceX[(unsigned char)x])
584+
585+typedef void (*vector_cleanup)(void *p);
586+
587+void vector_cleanup_noop(void *_) { UNUSED_PARAMETER(_); }
588+
589+#define JSON_SUBTYPE 74
590+
591+void vtab_set_error(sqlite3_vtab *pVTab, const char *zFormat, ...) {
592+ va_list args;
593+ sqlite3_free(pVTab->zErrMsg);
594+ va_start(args, zFormat);
595+ pVTab->zErrMsg = sqlite3_vmprintf(zFormat, args);
596+ va_end(args);
597+}
598+struct Array {
599+ size_t element_size;
600+ size_t length;
601+ size_t capacity;
602+ void *z;
603+};
604+
605+/**
606+ * @brief Initial an array with the given element size and capacity.
607+ *
608+ * @param array
609+ * @param element_size
610+ * @param init_capacity
611+ * @return SQLITE_OK on success, error code on failure. Only error is
612+ * SQLITE_NOMEM
613+ */
614+int array_init(struct Array *array, size_t element_size, size_t init_capacity) {
615+ int sz = element_size * init_capacity;
616+ void *z = sqlite3_malloc(sz);
617+ if (!z) {
618+ return SQLITE_NOMEM;
619+ }
620+ memset(z, 0, sz);
621+
622+ array->element_size = element_size;
623+ array->length = 0;
624+ array->capacity = init_capacity;
625+ array->z = z;
626+ return SQLITE_OK;
627+}
628+
629+int array_append(struct Array *array, const void *element) {
630+ if (array->length == array->capacity) {
631+ size_t new_capacity = array->capacity * 2 + 100;
632+ void *z = sqlite3_realloc64(array->z, array->element_size * new_capacity);
633+ if (z) {
634+ array->capacity = new_capacity;
635+ array->z = z;
636+ } else {
637+ return SQLITE_NOMEM;
638+ }
639+ }
640+ memcpy(&((unsigned char *)array->z)[array->length * array->element_size],
641+ element, array->element_size);
642+ array->length++;
643+ return SQLITE_OK;
644+}
645+
646+void array_cleanup(struct Array *array) {
647+ if (!array)
648+ return;
649+ array->element_size = 0;
650+ array->length = 0;
651+ array->capacity = 0;
652+ sqlite3_free(array->z);
653+ array->z = NULL;
654+}
655+
656+char *vector_subtype_name(int subtype) {
657+ switch (subtype) {
658+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32:
659+ return "float32";
660+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
661+ return "int8";
662+ case SQLITE_VEC_ELEMENT_TYPE_BIT:
663+ return "bit";
664+ }
665+ return "";
666+}
667+char *type_name(int type) {
668+ switch (type) {
669+ case SQLITE_INTEGER:
670+ return "INTEGER";
671+ case SQLITE_BLOB:
672+ return "BLOB";
673+ case SQLITE_TEXT:
674+ return "TEXT";
675+ case SQLITE_FLOAT:
676+ return "FLOAT";
677+ case SQLITE_NULL:
678+ return "NULL";
679+ }
680+ return "";
681+}
682+
683+typedef void (*fvec_cleanup)(f32 *vector);
684+
685+void fvec_cleanup_noop(f32 *_) { UNUSED_PARAMETER(_); }
686+
687+static int fvec_from_value(sqlite3_value *value, f32 **vector,
688+ size_t *dimensions, fvec_cleanup *cleanup,
689+ char **pzErr) {
690+ int value_type = sqlite3_value_type(value);
691+
692+ if (value_type == SQLITE_BLOB) {
693+ const void *blob = sqlite3_value_blob(value);
694+ int bytes = sqlite3_value_bytes(value);
695+ if (bytes == 0) {
696+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
697+ return SQLITE_ERROR;
698+ }
699+ if ((bytes % sizeof(f32)) != 0) {
700+ *pzErr = sqlite3_mprintf("invalid float32 vector BLOB length. Must be "
701+ "divisible by %d, found %d",
702+ sizeof(f32), bytes);
703+ return SQLITE_ERROR;
704+ }
705+ *vector = (f32 *)blob;
706+ *dimensions = bytes / sizeof(f32);
707+ *cleanup = fvec_cleanup_noop;
708+ return SQLITE_OK;
709+ }
710+
711+ if (value_type == SQLITE_TEXT) {
712+ const char *source = (const char *)sqlite3_value_text(value);
713+ int source_len = sqlite3_value_bytes(value);
714+ if (source_len == 0) {
715+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
716+ return SQLITE_ERROR;
717+ }
718+ int i = 0;
719+
720+ struct Array x;
721+ int rc = array_init(&x, sizeof(f32), ceil(source_len / 2.0));
722+ if (rc != SQLITE_OK) {
723+ return rc;
724+ }
725+
726+ // advance leading whitespace to first '['
727+ while (i < source_len) {
728+ if (vecJsonIsspace(source[i])) {
729+ i++;
730+ continue;
731+ }
732+ if (source[i] == '[') {
733+ break;
734+ }
735+
736+ *pzErr = sqlite3_mprintf(
737+ "JSON array parsing error: Input does not start with '['");
738+ array_cleanup(&x);
739+ return SQLITE_ERROR;
740+ }
741+ if (source[i] != '[') {
742+ *pzErr = sqlite3_mprintf(
743+ "JSON array parsing error: Input does not start with '['");
744+ array_cleanup(&x);
745+ return SQLITE_ERROR;
746+ }
747+ int offset = i + 1;
748+
749+ while (offset < source_len) {
750+ char *ptr = (char *)&source[offset];
751+ char *endptr;
752+
753+ errno = 0;
754+ double result = strtod(ptr, &endptr);
755+ if ((errno != 0 && result == 0) // some interval error?
756+ || (errno == ERANGE &&
757+ (result == HUGE_VAL || result == -HUGE_VAL)) // too big / smalls
758+ ) {
759+ sqlite3_free(x.z);
760+ *pzErr = sqlite3_mprintf("JSON parsing error");
761+ return SQLITE_ERROR;
762+ }
763+
764+ if (endptr == ptr) {
765+ if (*ptr != ']') {
766+ sqlite3_free(x.z);
767+ *pzErr = sqlite3_mprintf("JSON parsing error");
768+ return SQLITE_ERROR;
769+ }
770+ goto done;
771+ }
772+
773+ f32 res = (f32)result;
774+ array_append(&x, (const void *)&res);
775+
776+ offset += (endptr - ptr);
777+ while (offset < source_len) {
778+ if (vecJsonIsspace(source[offset])) {
779+ offset++;
780+ continue;
781+ }
782+ if (source[offset] == ',') {
783+ offset++;
784+ continue;
785+ }
786+ if (source[offset] == ']')
787+ goto done;
788+ break;
789+ }
790+ }
791+
792+ done:
793+
794+ if (x.length > 0) {
795+ *vector = (f32 *)x.z;
796+ *dimensions = x.length;
797+ *cleanup = (fvec_cleanup)sqlite3_free;
798+ return SQLITE_OK;
799+ }
800+ sqlite3_free(x.z);
801+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
802+ return SQLITE_ERROR;
803+ }
804+
805+ *pzErr = sqlite3_mprintf(
806+ "Input must have type BLOB (compact format) or TEXT (JSON), found %s",
807+ type_name(value_type));
808+ return SQLITE_ERROR;
809+}
810+
811+static int bitvec_from_value(sqlite3_value *value, u8 **vector,
812+ size_t *dimensions, vector_cleanup *cleanup,
813+ char **pzErr) {
814+ int value_type = sqlite3_value_type(value);
815+ if (value_type == SQLITE_BLOB) {
816+ const void *blob = sqlite3_value_blob(value);
817+ int bytes = sqlite3_value_bytes(value);
818+ if (bytes == 0) {
819+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
820+ return SQLITE_ERROR;
821+ }
822+ *vector = (u8 *)blob;
823+ *dimensions = bytes * CHAR_BIT;
824+ *cleanup = vector_cleanup_noop;
825+ return SQLITE_OK;
826+ }
827+ *pzErr = sqlite3_mprintf("Unknown type for bitvector.");
828+ return SQLITE_ERROR;
829+}
830+
831+static int int8_vec_from_value(sqlite3_value *value, i8 **vector,
832+ size_t *dimensions, vector_cleanup *cleanup,
833+ char **pzErr) {
834+ int value_type = sqlite3_value_type(value);
835+ if (value_type == SQLITE_BLOB) {
836+ const void *blob = sqlite3_value_blob(value);
837+ int bytes = sqlite3_value_bytes(value);
838+ if (bytes == 0) {
839+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
840+ return SQLITE_ERROR;
841+ }
842+ *vector = (i8 *)blob;
843+ *dimensions = bytes;
844+ *cleanup = vector_cleanup_noop;
845+ return SQLITE_OK;
846+ }
847+
848+ if (value_type == SQLITE_TEXT) {
849+ const char *source = (const char *)sqlite3_value_text(value);
850+ int source_len = sqlite3_value_bytes(value);
851+ int i = 0;
852+
853+ if (source_len == 0) {
854+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
855+ return SQLITE_ERROR;
856+ }
857+
858+ struct Array x;
859+ int rc = array_init(&x, sizeof(i8), ceil(source_len / 2.0));
860+ if (rc != SQLITE_OK) {
861+ return rc;
862+ }
863+
864+ // advance leading whitespace to first '['
865+ while (i < source_len) {
866+ if (vecJsonIsspace(source[i])) {
867+ i++;
868+ continue;
869+ }
870+ if (source[i] == '[') {
871+ break;
872+ }
873+
874+ *pzErr = sqlite3_mprintf(
875+ "JSON array parsing error: Input does not start with '['");
876+ array_cleanup(&x);
877+ return SQLITE_ERROR;
878+ }
879+ if (source[i] != '[') {
880+ *pzErr = sqlite3_mprintf(
881+ "JSON array parsing error: Input does not start with '['");
882+ array_cleanup(&x);
883+ return SQLITE_ERROR;
884+ }
885+ int offset = i + 1;
886+
887+ while (offset < source_len) {
888+ char *ptr = (char *)&source[offset];
889+ char *endptr;
890+
891+ errno = 0;
892+ long result = strtol(ptr, &endptr, 10);
893+ if ((errno != 0 && result == 0) ||
894+ (errno == ERANGE && (result == LONG_MAX || result == LONG_MIN))) {
895+ sqlite3_free(x.z);
896+ *pzErr = sqlite3_mprintf("JSON parsing error");
897+ return SQLITE_ERROR;
898+ }
899+
900+ if (endptr == ptr) {
901+ if (*ptr != ']') {
902+ sqlite3_free(x.z);
903+ *pzErr = sqlite3_mprintf("JSON parsing error");
904+ return SQLITE_ERROR;
905+ }
906+ goto done;
907+ }
908+
909+ if (result < INT8_MIN || result > INT8_MAX) {
910+ sqlite3_free(x.z);
911+ *pzErr =
912+ sqlite3_mprintf("JSON parsing error: value out of range for int8");
913+ return SQLITE_ERROR;
914+ }
915+
916+ i8 res = (i8)result;
917+ array_append(&x, (const void *)&res);
918+
919+ offset += (endptr - ptr);
920+ while (offset < source_len) {
921+ if (vecJsonIsspace(source[offset])) {
922+ offset++;
923+ continue;
924+ }
925+ if (source[offset] == ',') {
926+ offset++;
927+ continue;
928+ }
929+ if (source[offset] == ']')
930+ goto done;
931+ break;
932+ }
933+ }
934+
935+ done:
936+
937+ if (x.length > 0) {
938+ *vector = (i8 *)x.z;
939+ *dimensions = x.length;
940+ *cleanup = (vector_cleanup)sqlite3_free;
941+ return SQLITE_OK;
942+ }
943+ sqlite3_free(x.z);
944+ *pzErr = sqlite3_mprintf("zero-length vectors are not supported.");
945+ return SQLITE_ERROR;
946+ }
947+
948+ *pzErr = sqlite3_mprintf("Unknown type for int8 vector.");
949+ return SQLITE_ERROR;
950+}
951+
952+/**
953+ * @brief Extract a vector from a sqlite3_value. Can be a float32, int8, or bit
954+ * vector.
955+ *
956+ * @param value: the sqlite3_value to read from.
957+ * @param vector: Output pointer to vector data.
958+ * @param dimensions: Output number of dimensions
959+ * @param dimensions: Output vector element type
960+ * @param cleanup
961+ * @param pzErrorMessage
962+ * @return int SQLITE_OK on success, error code otherwise
963+ */
964+int vector_from_value(sqlite3_value *value, void **vector, size_t *dimensions,
965+ enum VectorElementType *element_type,
966+ vector_cleanup *cleanup, char **pzErrorMessage) {
967+ int subtype = sqlite3_value_subtype(value);
968+ if (!subtype || (subtype == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) ||
969+ (subtype == JSON_SUBTYPE)) {
970+ int rc = fvec_from_value(value, (f32 **)vector, dimensions,
971+ (fvec_cleanup *)cleanup, pzErrorMessage);
972+ if (rc == SQLITE_OK) {
973+ *element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
974+ }
975+ return rc;
976+ }
977+
978+ if (subtype == SQLITE_VEC_ELEMENT_TYPE_BIT) {
979+ int rc = bitvec_from_value(value, (u8 **)vector, dimensions, cleanup,
980+ pzErrorMessage);
981+ if (rc == SQLITE_OK) {
982+ *element_type = SQLITE_VEC_ELEMENT_TYPE_BIT;
983+ }
984+ return rc;
985+ }
986+ if (subtype == SQLITE_VEC_ELEMENT_TYPE_INT8) {
987+ int rc = int8_vec_from_value(value, (i8 **)vector, dimensions, cleanup,
988+ pzErrorMessage);
989+ if (rc == SQLITE_OK) {
990+ *element_type = SQLITE_VEC_ELEMENT_TYPE_INT8;
991+ }
992+ return rc;
993+ }
994+ *pzErrorMessage = sqlite3_mprintf("Unknown subtype: %d", subtype);
995+ return SQLITE_ERROR;
996+}
997+
998+int ensure_vector_match(sqlite3_value *aValue, sqlite3_value *bValue, void **a,
999+ void **b, enum VectorElementType *element_type,
1000+ size_t *dimensions, vector_cleanup *outACleanup,
1001+ vector_cleanup *outBCleanup, char **outError) {
1002+ int rc;
1003+ enum VectorElementType aType, bType;
1004+ size_t aDims, bDims;
1005+ char *error = NULL;
1006+ vector_cleanup aCleanup, bCleanup;
1007+
1008+ rc = vector_from_value(aValue, a, &aDims, &aType, &aCleanup, &error);
1009+ if (rc != SQLITE_OK) {
1010+ *outError = sqlite3_mprintf("Error reading 1st vector: %s", error);
1011+ sqlite3_free(error);
1012+ return SQLITE_ERROR;
1013+ }
1014+
1015+ rc = vector_from_value(bValue, b, &bDims, &bType, &bCleanup, &error);
1016+ if (rc != SQLITE_OK) {
1017+ *outError = sqlite3_mprintf("Error reading 2nd vector: %s", error);
1018+ sqlite3_free(error);
1019+ aCleanup(a);
1020+ return SQLITE_ERROR;
1021+ }
1022+
1023+ if (aType != bType) {
1024+ *outError =
1025+ sqlite3_mprintf("Vector type mistmatch. First vector has type %s, "
1026+ "while the second has type %s.",
1027+ vector_subtype_name(aType), vector_subtype_name(bType));
1028+ aCleanup(*a);
1029+ bCleanup(*b);
1030+ return SQLITE_ERROR;
1031+ }
1032+ if (aDims != bDims) {
1033+ *outError = sqlite3_mprintf(
1034+ "Vector dimension mistmatch. First vector has %ld dimensions, "
1035+ "while the second has %ld dimensions.",
1036+ aDims, bDims);
1037+ aCleanup(*a);
1038+ bCleanup(*b);
1039+ return SQLITE_ERROR;
1040+ }
1041+ *element_type = aType;
1042+ *dimensions = aDims;
1043+ *outACleanup = aCleanup;
1044+ *outBCleanup = bCleanup;
1045+ return SQLITE_OK;
1046+}
1047+
1048+int _cmp(const void *a, const void *b) { return (*(i64 *)a - *(i64 *)b); }
1049+
1050+struct VecNpyFile {
1051+ char *path;
1052+ size_t pathLength;
1053+};
1054+#define SQLITE_VEC_NPY_FILE_NAME "vec0-npy-file"
1055+
1056+#ifndef SQLITE_VEC_OMIT_FS
1057+static void vec_npy_file(sqlite3_context *context, int argc,
1058+ sqlite3_value **argv) {
1059+ assert(argc == 1);
1060+ char *path = (char *)sqlite3_value_text(argv[0]);
1061+ size_t pathLength = sqlite3_value_bytes(argv[0]);
1062+ struct VecNpyFile *f;
1063+
1064+ f = sqlite3_malloc(sizeof(*f));
1065+ if (!f) {
1066+ sqlite3_result_error_nomem(context);
1067+ return;
1068+ }
1069+ memset(f, 0, sizeof(*f));
1070+
1071+ f->path = path;
1072+ f->pathLength = pathLength;
1073+ sqlite3_result_pointer(context, f, SQLITE_VEC_NPY_FILE_NAME, sqlite3_free);
1074+}
1075+#endif
1076+
1077+#pragma region scalar functions
1078+static void vec_f32(sqlite3_context *context, int argc, sqlite3_value **argv) {
1079+ assert(argc == 1);
1080+ int rc;
1081+ f32 *vector = NULL;
1082+ size_t dimensions;
1083+ fvec_cleanup cleanup;
1084+ char *errmsg;
1085+ rc = fvec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg);
1086+ if (rc != SQLITE_OK) {
1087+ sqlite3_result_error(context, errmsg, -1);
1088+ sqlite3_free(errmsg);
1089+ return;
1090+ }
1091+ sqlite3_result_blob(context, vector, dimensions * sizeof(f32),
1092+ (void (*)(void *))cleanup);
1093+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1094+}
1095+
1096+static void vec_bit(sqlite3_context *context, int argc, sqlite3_value **argv) {
1097+ assert(argc == 1);
1098+ int rc;
1099+ u8 *vector;
1100+ size_t dimensions;
1101+ vector_cleanup cleanup;
1102+ char *errmsg;
1103+ rc = bitvec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg);
1104+ if (rc != SQLITE_OK) {
1105+ sqlite3_result_error(context, errmsg, -1);
1106+ sqlite3_free(errmsg);
1107+ return;
1108+ }
1109+ sqlite3_result_blob(context, vector, dimensions / CHAR_BIT, SQLITE_TRANSIENT);
1110+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT);
1111+ cleanup(vector);
1112+}
1113+static void vec_int8(sqlite3_context *context, int argc, sqlite3_value **argv) {
1114+ assert(argc == 1);
1115+ int rc;
1116+ i8 *vector;
1117+ size_t dimensions;
1118+ vector_cleanup cleanup;
1119+ char *errmsg;
1120+ rc = int8_vec_from_value(argv[0], &vector, &dimensions, &cleanup, &errmsg);
1121+ if (rc != SQLITE_OK) {
1122+ sqlite3_result_error(context, errmsg, -1);
1123+ sqlite3_free(errmsg);
1124+ return;
1125+ }
1126+ sqlite3_result_blob(context, vector, dimensions, SQLITE_TRANSIENT);
1127+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1128+ cleanup(vector);
1129+}
1130+
1131+static void vec_length(sqlite3_context *context, int argc,
1132+ sqlite3_value **argv) {
1133+ assert(argc == 1);
1134+ int rc;
1135+ void *vector;
1136+ size_t dimensions;
1137+ vector_cleanup cleanup;
1138+ char *errmsg;
1139+ enum VectorElementType elementType;
1140+ rc = vector_from_value(argv[0], &vector, &dimensions, &elementType, &cleanup,
1141+ &errmsg);
1142+ if (rc != SQLITE_OK) {
1143+ sqlite3_result_error(context, errmsg, -1);
1144+ sqlite3_free(errmsg);
1145+ return;
1146+ }
1147+ sqlite3_result_int64(context, dimensions);
1148+ cleanup(vector);
1149+}
1150+
1151+static void vec_distance_cosine(sqlite3_context *context, int argc,
1152+ sqlite3_value **argv) {
1153+ assert(argc == 2);
1154+ int rc;
1155+ void *a = NULL, *b = NULL;
1156+ size_t dimensions;
1157+ vector_cleanup aCleanup, bCleanup;
1158+ char *error;
1159+ enum VectorElementType elementType;
1160+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1161+ &aCleanup, &bCleanup, &error);
1162+ if (rc != SQLITE_OK) {
1163+ sqlite3_result_error(context, error, -1);
1164+ sqlite3_free(error);
1165+ return;
1166+ }
1167+
1168+ switch (elementType) {
1169+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1170+ sqlite3_result_error(
1171+ context, "Cannot calculate cosine distance between two bitvectors.",
1172+ -1);
1173+ goto finish;
1174+ }
1175+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1176+ f32 result = distance_cosine_float(a, b, &dimensions);
1177+ sqlite3_result_double(context, result);
1178+ goto finish;
1179+ }
1180+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1181+ f32 result = distance_cosine_int8(a, b, &dimensions);
1182+ sqlite3_result_double(context, result);
1183+ goto finish;
1184+ }
1185+ }
1186+
1187+finish:
1188+ aCleanup(a);
1189+ bCleanup(b);
1190+ return;
1191+}
1192+
1193+static void vec_distance_l2(sqlite3_context *context, int argc,
1194+ sqlite3_value **argv) {
1195+ assert(argc == 2);
1196+ int rc;
1197+ void *a = NULL, *b = NULL;
1198+ size_t dimensions;
1199+ vector_cleanup aCleanup, bCleanup;
1200+ char *error;
1201+ enum VectorElementType elementType;
1202+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1203+ &aCleanup, &bCleanup, &error);
1204+ if (rc != SQLITE_OK) {
1205+ sqlite3_result_error(context, error, -1);
1206+ sqlite3_free(error);
1207+ return;
1208+ }
1209+
1210+ switch (elementType) {
1211+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1212+ sqlite3_result_error(
1213+ context, "Cannot calculate L2 distance between two bitvectors.", -1);
1214+ goto finish;
1215+ }
1216+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1217+ f32 result = distance_l2_sqr_float(a, b, &dimensions);
1218+ sqlite3_result_double(context, result);
1219+ goto finish;
1220+ }
1221+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1222+ f32 result = distance_l2_sqr_int8(a, b, &dimensions);
1223+ sqlite3_result_double(context, result);
1224+ goto finish;
1225+ }
1226+ }
1227+
1228+finish:
1229+ aCleanup(a);
1230+ bCleanup(b);
1231+ return;
1232+}
1233+
1234+static void vec_distance_l1(sqlite3_context *context, int argc,
1235+ sqlite3_value **argv) {
1236+ assert(argc == 2);
1237+ int rc;
1238+ void *a, *b;
1239+ size_t dimensions;
1240+ vector_cleanup aCleanup, bCleanup;
1241+ char *error;
1242+ enum VectorElementType elementType;
1243+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1244+ &aCleanup, &bCleanup, &error);
1245+ if (rc != SQLITE_OK) {
1246+ sqlite3_result_error(context, error, -1);
1247+ sqlite3_free(error);
1248+ return;
1249+ }
1250+
1251+ switch (elementType) {
1252+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1253+ sqlite3_result_error(
1254+ context, "Cannot calculate L1 distance between two bitvectors.", -1);
1255+ goto finish;
1256+ }
1257+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1258+ double result = distance_l1_f32(a, b, &dimensions);
1259+ sqlite3_result_double(context, result);
1260+ goto finish;
1261+ }
1262+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1263+ i64 result = distance_l1_int8(a, b, &dimensions);
1264+ sqlite3_result_int(context, result);
1265+ goto finish;
1266+ }
1267+ }
1268+
1269+finish:
1270+ aCleanup(a);
1271+ bCleanup(b);
1272+ return;
1273+}
1274+
1275+static void vec_distance_hamming(sqlite3_context *context, int argc,
1276+ sqlite3_value **argv) {
1277+ assert(argc == 2);
1278+ int rc;
1279+ void *a = NULL, *b = NULL;
1280+ size_t dimensions;
1281+ vector_cleanup aCleanup, bCleanup;
1282+ char *error;
1283+ enum VectorElementType elementType;
1284+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1285+ &aCleanup, &bCleanup, &error);
1286+ if (rc != SQLITE_OK) {
1287+ sqlite3_result_error(context, error, -1);
1288+ sqlite3_free(error);
1289+ return;
1290+ }
1291+
1292+ switch (elementType) {
1293+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1294+ sqlite3_result_double(context, distance_hamming(a, b, &dimensions));
1295+ goto finish;
1296+ }
1297+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1298+ sqlite3_result_error(
1299+ context,
1300+ "Cannot calculate hamming distance between two float32 vectors.", -1);
1301+ goto finish;
1302+ }
1303+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1304+ sqlite3_result_error(
1305+ context, "Cannot calculate hamming distance between two int8 vectors.",
1306+ -1);
1307+ goto finish;
1308+ }
1309+ }
1310+
1311+finish:
1312+ aCleanup(a);
1313+ bCleanup(b);
1314+ return;
1315+}
1316+
1317+char *vec_type_name(enum VectorElementType elementType) {
1318+ switch (elementType) {
1319+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32:
1320+ return "float32";
1321+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
1322+ return "int8";
1323+ case SQLITE_VEC_ELEMENT_TYPE_BIT:
1324+ return "bit";
1325+ }
1326+ return "";
1327+}
1328+
1329+static void vec_type(sqlite3_context *context, int argc, sqlite3_value **argv) {
1330+ assert(argc == 1);
1331+ void *vector;
1332+ size_t dimensions;
1333+ vector_cleanup cleanup;
1334+ char *pzError;
1335+ enum VectorElementType elementType;
1336+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1337+ &cleanup, &pzError);
1338+ if (rc != SQLITE_OK) {
1339+ sqlite3_result_error(context, pzError, -1);
1340+ sqlite3_free(pzError);
1341+ return;
1342+ }
1343+ sqlite3_result_text(context, vec_type_name(elementType), -1, SQLITE_STATIC);
1344+ cleanup(vector);
1345+}
1346+static void vec_quantize_binary(sqlite3_context *context, int argc,
1347+ sqlite3_value **argv) {
1348+ assert(argc == 1);
1349+ void *vector;
1350+ size_t dimensions;
1351+ vector_cleanup vectorCleanup;
1352+ char *pzError;
1353+ enum VectorElementType elementType;
1354+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1355+ &vectorCleanup, &pzError);
1356+ if (rc != SQLITE_OK) {
1357+ sqlite3_result_error(context, pzError, -1);
1358+ sqlite3_free(pzError);
1359+ return;
1360+ }
1361+
1362+ if (dimensions <= 0) {
1363+ sqlite3_result_error(context, "Zero length vectors are not supported.", -1);
1364+ goto cleanup;
1365+ return;
1366+ }
1367+ if ((dimensions % CHAR_BIT) != 0) {
1368+ sqlite3_result_error(
1369+ context,
1370+ "Binary quantization requires vectors with a length divisible by 8",
1371+ -1);
1372+ goto cleanup;
1373+ return;
1374+ }
1375+
1376+ int sz = dimensions / CHAR_BIT;
1377+ u8 *out = sqlite3_malloc(sz);
1378+ if (!out) {
1379+ sqlite3_result_error_code(context, SQLITE_NOMEM);
1380+ goto cleanup;
1381+ return;
1382+ }
1383+ memset(out, 0, sz);
1384+
1385+ switch (elementType) {
1386+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1387+
1388+ for (size_t i = 0; i < dimensions; i++) {
1389+ int res = ((f32 *)vector)[i] > 0.0;
1390+ out[i / 8] |= (res << (i % 8));
1391+ }
1392+ break;
1393+ }
1394+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1395+ for (size_t i = 0; i < dimensions; i++) {
1396+ int res = ((i8 *)vector)[i] > 0;
1397+ out[i / 8] |= (res << (i % 8));
1398+ }
1399+ break;
1400+ }
1401+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1402+ sqlite3_result_error(context,
1403+ "Can only binary quantize float or int8 vectors", -1);
1404+ sqlite3_free(out);
1405+ return;
1406+ }
1407+ }
1408+ sqlite3_result_blob(context, out, sz, sqlite3_free);
1409+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT);
1410+
1411+cleanup:
1412+ vectorCleanup(vector);
1413+}
1414+
1415+static void vec_quantize_int8(sqlite3_context *context, int argc,
1416+ sqlite3_value **argv) {
1417+ assert(argc == 2);
1418+ f32 *srcVector;
1419+ size_t dimensions;
1420+ fvec_cleanup srcCleanup;
1421+ char *err;
1422+ i8 *out = NULL;
1423+ int rc = fvec_from_value(argv[0], &srcVector, &dimensions, &srcCleanup, &err);
1424+ if (rc != SQLITE_OK) {
1425+ sqlite3_result_error(context, err, -1);
1426+ sqlite3_free(err);
1427+ return;
1428+ }
1429+
1430+ int sz = dimensions * sizeof(i8);
1431+ out = sqlite3_malloc(sz);
1432+ if (!out) {
1433+ sqlite3_result_error_nomem(context);
1434+ goto cleanup;
1435+ }
1436+ memset(out, 0, sz);
1437+
1438+ if ((sqlite3_value_type(argv[1]) != SQLITE_TEXT) ||
1439+ (sqlite3_value_bytes(argv[1]) != strlen("unit")) ||
1440+ (sqlite3_stricmp((const char *)sqlite3_value_text(argv[1]), "unit") !=
1441+ 0)) {
1442+ sqlite3_result_error(
1443+ context, "2nd argument to vec_quantize_int8() must be 'unit'.", -1);
1444+ sqlite3_free(out);
1445+ goto cleanup;
1446+ }
1447+ f32 step = (1.0 - (-1.0)) / 255;
1448+ for (size_t i = 0; i < dimensions; i++) {
1449+ out[i] = ((srcVector[i] - (-1.0)) / step) - 128;
1450+ }
1451+
1452+ sqlite3_result_blob(context, out, dimensions * sizeof(i8), sqlite3_free);
1453+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1454+
1455+cleanup:
1456+ srcCleanup(srcVector);
1457+}
1458+
1459+static void vec_add(sqlite3_context *context, int argc, sqlite3_value **argv) {
1460+ assert(argc == 2);
1461+ int rc;
1462+ void *a = NULL, *b = NULL;
1463+ size_t dimensions;
1464+ vector_cleanup aCleanup, bCleanup;
1465+ char *error;
1466+ enum VectorElementType elementType;
1467+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1468+ &aCleanup, &bCleanup, &error);
1469+ if (rc != SQLITE_OK) {
1470+ sqlite3_result_error(context, error, -1);
1471+ sqlite3_free(error);
1472+ return;
1473+ }
1474+
1475+ switch (elementType) {
1476+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1477+ sqlite3_result_error(context, "Cannot add two bitvectors together.", -1);
1478+ goto finish;
1479+ }
1480+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1481+ size_t outSize = dimensions * sizeof(f32);
1482+ f32 *out = sqlite3_malloc(outSize);
1483+ if (!out) {
1484+ sqlite3_result_error_nomem(context);
1485+ goto finish;
1486+ }
1487+ memset(out, 0, outSize);
1488+ for (size_t i = 0; i < dimensions; i++) {
1489+ out[i] = ((f32 *)a)[i] + ((f32 *)b)[i];
1490+ }
1491+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1492+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1493+ goto finish;
1494+ }
1495+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1496+ size_t outSize = dimensions * sizeof(i8);
1497+ i8 *out = sqlite3_malloc(outSize);
1498+ if (!out) {
1499+ sqlite3_result_error_nomem(context);
1500+ goto finish;
1501+ }
1502+ memset(out, 0, outSize);
1503+ for (size_t i = 0; i < dimensions; i++) {
1504+ out[i] = ((i8 *)a)[i] + ((i8 *)b)[i];
1505+ }
1506+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1507+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1508+ goto finish;
1509+ }
1510+ }
1511+finish:
1512+ aCleanup(a);
1513+ bCleanup(b);
1514+ return;
1515+}
1516+static void vec_sub(sqlite3_context *context, int argc, sqlite3_value **argv) {
1517+ assert(argc == 2);
1518+ int rc;
1519+ void *a = NULL, *b = NULL;
1520+ size_t dimensions;
1521+ vector_cleanup aCleanup, bCleanup;
1522+ char *error;
1523+ enum VectorElementType elementType;
1524+ rc = ensure_vector_match(argv[0], argv[1], &a, &b, &elementType, &dimensions,
1525+ &aCleanup, &bCleanup, &error);
1526+ if (rc != SQLITE_OK) {
1527+ sqlite3_result_error(context, error, -1);
1528+ sqlite3_free(error);
1529+ return;
1530+ }
1531+
1532+ switch (elementType) {
1533+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1534+ sqlite3_result_error(context, "Cannot subtract two bitvectors together.",
1535+ -1);
1536+ goto finish;
1537+ }
1538+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1539+ size_t outSize = dimensions * sizeof(f32);
1540+ f32 *out = sqlite3_malloc(outSize);
1541+ if (!out) {
1542+ sqlite3_result_error_nomem(context);
1543+ goto finish;
1544+ }
1545+ memset(out, 0, outSize);
1546+ for (size_t i = 0; i < dimensions; i++) {
1547+ out[i] = ((f32 *)a)[i] - ((f32 *)b)[i];
1548+ }
1549+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1550+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1551+ goto finish;
1552+ }
1553+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1554+ size_t outSize = dimensions * sizeof(i8);
1555+ i8 *out = sqlite3_malloc(outSize);
1556+ if (!out) {
1557+ sqlite3_result_error_nomem(context);
1558+ goto finish;
1559+ }
1560+ memset(out, 0, outSize);
1561+ for (size_t i = 0; i < dimensions; i++) {
1562+ out[i] = ((i8 *)a)[i] - ((i8 *)b)[i];
1563+ }
1564+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1565+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1566+ goto finish;
1567+ }
1568+ }
1569+finish:
1570+ aCleanup(a);
1571+ bCleanup(b);
1572+ return;
1573+}
1574+static void vec_slice(sqlite3_context *context, int argc,
1575+ sqlite3_value **argv) {
1576+ assert(argc == 3);
1577+
1578+ void *vector;
1579+ size_t dimensions;
1580+ vector_cleanup cleanup;
1581+ char *err;
1582+ enum VectorElementType elementType;
1583+
1584+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1585+ &cleanup, &err);
1586+ if (rc != SQLITE_OK) {
1587+ sqlite3_result_error(context, err, -1);
1588+ sqlite3_free(err);
1589+ return;
1590+ }
1591+
1592+ int start = sqlite3_value_int(argv[1]);
1593+ int end = sqlite3_value_int(argv[2]);
1594+
1595+ if (start < 0) {
1596+ sqlite3_result_error(context,
1597+ "slice 'start' index must be a postive number.", -1);
1598+ goto done;
1599+ }
1600+ if (end < 0) {
1601+ sqlite3_result_error(context, "slice 'end' index must be a postive number.",
1602+ -1);
1603+ goto done;
1604+ }
1605+ if (((size_t)start) > dimensions) {
1606+ sqlite3_result_error(
1607+ context, "slice 'start' index is greater than the number of dimensions",
1608+ -1);
1609+ goto done;
1610+ }
1611+ if (((size_t)end) > dimensions) {
1612+ sqlite3_result_error(
1613+ context, "slice 'end' index is greater than the number of dimensions",
1614+ -1);
1615+ goto done;
1616+ }
1617+ if (start > end) {
1618+ sqlite3_result_error(context,
1619+ "slice 'start' index is greater than 'end' index", -1);
1620+ goto done;
1621+ }
1622+ if (start == end) {
1623+ sqlite3_result_error(context,
1624+ "slice 'start' index is equal to the 'end' index, "
1625+ "vectors must have non-zero length",
1626+ -1);
1627+ goto done;
1628+ }
1629+ size_t n = end - start;
1630+
1631+ switch (elementType) {
1632+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
1633+ int outSize = n * sizeof(f32);
1634+ f32 *out = sqlite3_malloc(outSize);
1635+ if (!out) {
1636+ sqlite3_result_error_nomem(context);
1637+ goto done;
1638+ }
1639+ memset(out, 0, outSize);
1640+ for (size_t i = 0; i < n; i++) {
1641+ out[i] = ((f32 *)vector)[start + i];
1642+ }
1643+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1644+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1645+ goto done;
1646+ }
1647+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
1648+ int outSize = n * sizeof(i8);
1649+ i8 *out = sqlite3_malloc(outSize);
1650+ if (!out) {
1651+ sqlite3_result_error_nomem(context);
1652+ return;
1653+ }
1654+ memset(out, 0, outSize);
1655+ for (size_t i = 0; i < n; i++) {
1656+ out[i] = ((i8 *)vector)[start + i];
1657+ }
1658+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1659+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_INT8);
1660+ goto done;
1661+ }
1662+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
1663+ if ((start % CHAR_BIT) != 0) {
1664+ sqlite3_result_error(context, "start index must be divisible by 8.", -1);
1665+ goto done;
1666+ }
1667+ if ((end % CHAR_BIT) != 0) {
1668+ sqlite3_result_error(context, "end index must be divisible by 8.", -1);
1669+ goto done;
1670+ }
1671+ int outSize = n / CHAR_BIT;
1672+ u8 *out = sqlite3_malloc(outSize);
1673+ if (!out) {
1674+ sqlite3_result_error_nomem(context);
1675+ return;
1676+ }
1677+ memset(out, 0, outSize);
1678+ for (size_t i = 0; i < n / CHAR_BIT; i++) {
1679+ out[i] = ((u8 *)vector)[(start / CHAR_BIT) + i];
1680+ }
1681+ sqlite3_result_blob(context, out, outSize, sqlite3_free);
1682+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT);
1683+ goto done;
1684+ }
1685+ }
1686+done:
1687+ cleanup(vector);
1688+}
1689+
1690+static void vec_to_json(sqlite3_context *context, int argc,
1691+ sqlite3_value **argv) {
1692+ assert(argc == 1);
1693+ void *vector;
1694+ size_t dimensions;
1695+ vector_cleanup cleanup;
1696+ char *err;
1697+ enum VectorElementType elementType;
1698+
1699+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1700+ &cleanup, &err);
1701+ if (rc != SQLITE_OK) {
1702+ sqlite3_result_error(context, err, -1);
1703+ sqlite3_free(err);
1704+ return;
1705+ }
1706+
1707+ sqlite3_str *str = sqlite3_str_new(sqlite3_context_db_handle(context));
1708+ sqlite3_str_appendall(str, "[");
1709+ for (size_t i = 0; i < dimensions; i++) {
1710+ if (i != 0) {
1711+ sqlite3_str_appendall(str, ",");
1712+ }
1713+ if (elementType == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
1714+ f32 value = ((f32 *)vector)[i];
1715+ if (isnan(value)) {
1716+ sqlite3_str_appendall(str, "null");
1717+ } else {
1718+ sqlite3_str_appendf(str, "%f", value);
1719+ }
1720+
1721+ } else if (elementType == SQLITE_VEC_ELEMENT_TYPE_INT8) {
1722+ sqlite3_str_appendf(str, "%d", ((i8 *)vector)[i]);
1723+ } else if (elementType == SQLITE_VEC_ELEMENT_TYPE_BIT) {
1724+ u8 b = (((u8 *)vector)[i / 8] >> (i % CHAR_BIT)) & 1;
1725+ sqlite3_str_appendf(str, "%d", b);
1726+ }
1727+ }
1728+ sqlite3_str_appendall(str, "]");
1729+ int len = sqlite3_str_length(str);
1730+ char *s = sqlite3_str_finish(str);
1731+ if (s) {
1732+ sqlite3_result_text(context, s, len, sqlite3_free);
1733+ sqlite3_result_subtype(context, JSON_SUBTYPE);
1734+ } else {
1735+ sqlite3_result_error_nomem(context);
1736+ }
1737+ cleanup(vector);
1738+}
1739+
1740+static void vec_normalize(sqlite3_context *context, int argc,
1741+ sqlite3_value **argv) {
1742+ assert(argc == 1);
1743+ void *vector;
1744+ size_t dimensions;
1745+ vector_cleanup cleanup;
1746+ char *err;
1747+ enum VectorElementType elementType;
1748+
1749+ int rc = vector_from_value(argv[0], &vector, &dimensions, &elementType,
1750+ &cleanup, &err);
1751+ if (rc != SQLITE_OK) {
1752+ sqlite3_result_error(context, err, -1);
1753+ sqlite3_free(err);
1754+ return;
1755+ }
1756+
1757+ if (elementType != SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
1758+ sqlite3_result_error(
1759+ context, "only float32 vectors are supported when normalizing", -1);
1760+ cleanup(vector);
1761+ return;
1762+ }
1763+
1764+ int outSize = dimensions * sizeof(f32);
1765+ f32 *out = sqlite3_malloc(outSize);
1766+ if (!out) {
1767+ cleanup(vector);
1768+ sqlite3_result_error_code(context, SQLITE_NOMEM);
1769+ return;
1770+ }
1771+ memset(out, 0, outSize);
1772+
1773+ f32 *v = (f32 *)vector;
1774+
1775+ f32 norm = 0;
1776+ for (size_t i = 0; i < dimensions; i++) {
1777+ norm += v[i] * v[i];
1778+ }
1779+ norm = sqrt(norm);
1780+ for (size_t i = 0; i < dimensions; i++) {
1781+ out[i] = v[i] / norm;
1782+ }
1783+
1784+ sqlite3_result_blob(context, out, dimensions * sizeof(f32), sqlite3_free);
1785+ sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
1786+ cleanup(vector);
1787+}
1788+
1789+static void _static_text_func(sqlite3_context *context, int argc,
1790+ sqlite3_value **argv) {
1791+ UNUSED_PARAMETER(argc);
1792+ UNUSED_PARAMETER(argv);
1793+ sqlite3_result_text(context, sqlite3_user_data(context), -1, SQLITE_STATIC);
1794+}
1795+
1796+#pragma endregion
1797+
1798+enum Vec0TokenType {
1799+ TOKEN_TYPE_IDENTIFIER,
1800+ TOKEN_TYPE_DIGIT,
1801+ TOKEN_TYPE_LBRACKET,
1802+ TOKEN_TYPE_RBRACKET,
1803+ TOKEN_TYPE_PLUS,
1804+ TOKEN_TYPE_EQ,
1805+};
1806+struct Vec0Token {
1807+ enum Vec0TokenType token_type;
1808+ char *start;
1809+ char *end;
1810+};
1811+
1812+int is_alpha(char x) {
1813+ return (x >= 'a' && x <= 'z') || (x >= 'A' && x <= 'Z');
1814+}
1815+int is_digit(char x) { return (x >= '0' && x <= '9'); }
1816+int is_whitespace(char x) {
1817+ return x == ' ' || x == '\t' || x == '\n' || x == '\r';
1818+}
1819+
1820+#define VEC0_TOKEN_RESULT_EOF 1
1821+#define VEC0_TOKEN_RESULT_SOME 2
1822+#define VEC0_TOKEN_RESULT_ERROR 3
1823+
1824+int vec0_token_next(char *start, char *end, struct Vec0Token *out) {
1825+ char *ptr = start;
1826+ while (ptr < end) {
1827+ char curr = *ptr;
1828+ if (is_whitespace(curr)) {
1829+ ptr++;
1830+ continue;
1831+ } else if (curr == '+') {
1832+ ptr++;
1833+ out->start = ptr;
1834+ out->end = ptr;
1835+ out->token_type = TOKEN_TYPE_PLUS;
1836+ return VEC0_TOKEN_RESULT_SOME;
1837+ } else if (curr == '[') {
1838+ ptr++;
1839+ out->start = ptr;
1840+ out->end = ptr;
1841+ out->token_type = TOKEN_TYPE_LBRACKET;
1842+ return VEC0_TOKEN_RESULT_SOME;
1843+ } else if (curr == ']') {
1844+ ptr++;
1845+ out->start = ptr;
1846+ out->end = ptr;
1847+ out->token_type = TOKEN_TYPE_RBRACKET;
1848+ return VEC0_TOKEN_RESULT_SOME;
1849+ } else if (curr == '=') {
1850+ ptr++;
1851+ out->start = ptr;
1852+ out->end = ptr;
1853+ out->token_type = TOKEN_TYPE_EQ;
1854+ return VEC0_TOKEN_RESULT_SOME;
1855+ } else if (is_alpha(curr)) {
1856+ char *start = ptr;
1857+ while (ptr < end && (is_alpha(*ptr) || is_digit(*ptr) || *ptr == '_')) {
1858+ ptr++;
1859+ }
1860+ out->start = start;
1861+ out->end = ptr;
1862+ out->token_type = TOKEN_TYPE_IDENTIFIER;
1863+ return VEC0_TOKEN_RESULT_SOME;
1864+ } else if (is_digit(curr)) {
1865+ char *start = ptr;
1866+ while (ptr < end && (is_digit(*ptr))) {
1867+ ptr++;
1868+ }
1869+ out->start = start;
1870+ out->end = ptr;
1871+ out->token_type = TOKEN_TYPE_DIGIT;
1872+ return VEC0_TOKEN_RESULT_SOME;
1873+ } else {
1874+ return VEC0_TOKEN_RESULT_ERROR;
1875+ }
1876+ }
1877+ return VEC0_TOKEN_RESULT_EOF;
1878+}
1879+
1880+struct Vec0Scanner {
1881+ char *start;
1882+ char *end;
1883+ char *ptr;
1884+};
1885+
1886+void vec0_scanner_init(struct Vec0Scanner *scanner, const char *source,
1887+ int source_length) {
1888+ scanner->start = (char *)source;
1889+ scanner->end = (char *)source + source_length;
1890+ scanner->ptr = (char *)source;
1891+}
1892+int vec0_scanner_next(struct Vec0Scanner *scanner, struct Vec0Token *out) {
1893+ int rc = vec0_token_next(scanner->start, scanner->end, out);
1894+ if (rc == VEC0_TOKEN_RESULT_SOME) {
1895+ scanner->start = out->end;
1896+ }
1897+ return rc;
1898+}
1899+
1900+int vec0_parse_table_option(const char *source, int source_length,
1901+ char **out_key, int *out_key_length,
1902+ char **out_value, int *out_value_length) {
1903+ int rc;
1904+ struct Vec0Scanner scanner;
1905+ struct Vec0Token token;
1906+ char *key;
1907+ char *value;
1908+ int keyLength, valueLength;
1909+
1910+ vec0_scanner_init(&scanner, source, source_length);
1911+
1912+ rc = vec0_scanner_next(&scanner, &token);
1913+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1914+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
1915+ return SQLITE_EMPTY;
1916+ }
1917+ key = token.start;
1918+ keyLength = token.end - token.start;
1919+
1920+ rc = vec0_scanner_next(&scanner, &token);
1921+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_EQ) {
1922+ return SQLITE_EMPTY;
1923+ }
1924+
1925+ rc = vec0_scanner_next(&scanner, &token);
1926+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1927+ !((token.token_type == TOKEN_TYPE_IDENTIFIER) ||
1928+ (token.token_type == TOKEN_TYPE_DIGIT))) {
1929+ return SQLITE_ERROR;
1930+ }
1931+ value = token.start;
1932+ valueLength = token.end - token.start;
1933+
1934+ rc = vec0_scanner_next(&scanner, &token);
1935+ if (rc == VEC0_TOKEN_RESULT_EOF) {
1936+ *out_key = key;
1937+ *out_key_length = keyLength;
1938+ *out_value = value;
1939+ *out_value_length = valueLength;
1940+ return SQLITE_OK;
1941+ }
1942+ return SQLITE_ERROR;
1943+}
1944+/**
1945+ * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if
1946+ * it's a PARTITION KEY definition.
1947+ *
1948+ * @param source: argv[i] source string
1949+ * @param source_length: length of the source string
1950+ * @param out_column_name: If it is a partition key, the output column name. Same lifetime
1951+ * as source, points to specific char *
1952+ * @param out_column_name_length: Length of out_column_name in bytes
1953+ * @param out_column_type: SQLITE_TEXT or SQLITE_INTEGER.
1954+ * @return int: SQLITE_EMPTY if not a PK, SQLITE_OK if it is.
1955+ */
1956+int vec0_parse_partition_key_definition(const char *source, int source_length,
1957+ char **out_column_name,
1958+ int *out_column_name_length,
1959+ int *out_column_type) {
1960+ struct Vec0Scanner scanner;
1961+ struct Vec0Token token;
1962+ char *column_name;
1963+ int column_name_length;
1964+ int column_type;
1965+ vec0_scanner_init(&scanner, source, source_length);
1966+
1967+ // Check first token is identifier, will be the column name
1968+ int rc = vec0_scanner_next(&scanner, &token);
1969+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1970+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
1971+ return SQLITE_EMPTY;
1972+ }
1973+
1974+ column_name = token.start;
1975+ column_name_length = token.end - token.start;
1976+
1977+ // Check the next token matches "text" or "integer", as column type
1978+ rc = vec0_scanner_next(&scanner, &token);
1979+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1980+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
1981+ return SQLITE_EMPTY;
1982+ }
1983+ if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) {
1984+ column_type = SQLITE_TEXT;
1985+ } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) ==
1986+ 0 ||
1987+ sqlite3_strnicmp(token.start, "integer",
1988+ token.end - token.start) == 0) {
1989+ column_type = SQLITE_INTEGER;
1990+ } else {
1991+ return SQLITE_EMPTY;
1992+ }
1993+
1994+ // Check the next token is identifier and matches "partition"
1995+ rc = vec0_scanner_next(&scanner, &token);
1996+ if (rc != VEC0_TOKEN_RESULT_SOME &&
1997+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
1998+ return SQLITE_EMPTY;
1999+ }
2000+ if (sqlite3_strnicmp(token.start, "partition", token.end - token.start) != 0) {
2001+ return SQLITE_EMPTY;
2002+ }
2003+
2004+ // Check the next token is identifier and matches "key"
2005+ rc = vec0_scanner_next(&scanner, &token);
2006+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2007+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2008+ return SQLITE_EMPTY;
2009+ }
2010+ if (sqlite3_strnicmp(token.start, "key", token.end - token.start) != 0) {
2011+ return SQLITE_EMPTY;
2012+ }
2013+
2014+ *out_column_name = column_name;
2015+ *out_column_name_length = column_name_length;
2016+ *out_column_type = column_type;
2017+
2018+ return SQLITE_OK;
2019+}
2020+
2021+/**
2022+ * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if
2023+ * it's an auxiliar column definition, ie `+[name] [type]` like `+contents text`
2024+ *
2025+ * @param source: argv[i] source string
2026+ * @param source_length: length of the source string
2027+ * @param out_column_name: If it is a partition key, the output column name. Same lifetime
2028+ * as source, points to specific char *
2029+ * @param out_column_name_length: Length of out_column_name in bytes
2030+ * @param out_column_type: SQLITE_TEXT, SQLITE_INTEGER, SQLITE_FLOAT, or SQLITE_BLOB.
2031+ * @return int: SQLITE_EMPTY if not an aux column, SQLITE_OK if it is.
2032+ */
2033+int vec0_parse_auxiliary_column_definition(const char *source, int source_length,
2034+ char **out_column_name,
2035+ int *out_column_name_length,
2036+ int *out_column_type) {
2037+ struct Vec0Scanner scanner;
2038+ struct Vec0Token token;
2039+ char *column_name;
2040+ int column_name_length;
2041+ int column_type;
2042+ vec0_scanner_init(&scanner, source, source_length);
2043+
2044+ // Check first token is '+', which denotes aux columns
2045+ int rc = vec0_scanner_next(&scanner, &token);
2046+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2047+ token.token_type != TOKEN_TYPE_PLUS) {
2048+ return SQLITE_EMPTY;
2049+ }
2050+
2051+ rc = vec0_scanner_next(&scanner, &token);
2052+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2053+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2054+ return SQLITE_EMPTY;
2055+ }
2056+
2057+ column_name = token.start;
2058+ column_name_length = token.end - token.start;
2059+
2060+ // Check the next token matches "text" or "integer", as column type
2061+ rc = vec0_scanner_next(&scanner, &token);
2062+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2063+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2064+ return SQLITE_EMPTY;
2065+ }
2066+ if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) {
2067+ column_type = SQLITE_TEXT;
2068+ } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) ==
2069+ 0 ||
2070+ sqlite3_strnicmp(token.start, "integer",
2071+ token.end - token.start) == 0) {
2072+ column_type = SQLITE_INTEGER;
2073+ } else if (sqlite3_strnicmp(token.start, "float", token.end - token.start) ==
2074+ 0 ||
2075+ sqlite3_strnicmp(token.start, "double",
2076+ token.end - token.start) == 0) {
2077+ column_type = SQLITE_FLOAT;
2078+ } else if (sqlite3_strnicmp(token.start, "blob", token.end - token.start) ==0) {
2079+ column_type = SQLITE_BLOB;
2080+ } else {
2081+ return SQLITE_EMPTY;
2082+ }
2083+
2084+ *out_column_name = column_name;
2085+ *out_column_name_length = column_name_length;
2086+ *out_column_type = column_type;
2087+
2088+ return SQLITE_OK;
2089+}
2090+
2091+typedef enum {
2092+ VEC0_METADATA_COLUMN_KIND_BOOLEAN,
2093+ VEC0_METADATA_COLUMN_KIND_INTEGER,
2094+ VEC0_METADATA_COLUMN_KIND_FLOAT,
2095+ VEC0_METADATA_COLUMN_KIND_TEXT,
2096+ // future: blob, date, datetime
2097+} vec0_metadata_column_kind;
2098+
2099+/**
2100+ * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if
2101+ * it's an metadata column definition, ie `[name] [type]` like `is_released boolean`
2102+ *
2103+ * @param source: argv[i] source string
2104+ * @param source_length: length of the source string
2105+ * @param out_column_name: If it is a metadata column, the output column name. Same lifetime
2106+ * as source, points to specific char *
2107+ * @param out_column_name_length: Length of out_column_name in bytes
2108+ * @param out_column_type: one of vec0_metadata_column_kind
2109+ * @return int: SQLITE_EMPTY if not an metadata column, SQLITE_OK if it is.
2110+ */
2111+int vec0_parse_metadata_column_definition(const char *source, int source_length,
2112+ char **out_column_name,
2113+ int *out_column_name_length,
2114+ vec0_metadata_column_kind *out_column_type) {
2115+ struct Vec0Scanner scanner;
2116+ struct Vec0Token token;
2117+ char *column_name;
2118+ int column_name_length;
2119+ vec0_metadata_column_kind column_type;
2120+ int rc;
2121+ vec0_scanner_init(&scanner, source, source_length);
2122+
2123+ rc = vec0_scanner_next(&scanner, &token);
2124+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2125+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2126+ return SQLITE_EMPTY;
2127+ }
2128+
2129+ column_name = token.start;
2130+ column_name_length = token.end - token.start;
2131+
2132+ // Check the next token matches a valid metadata type
2133+ rc = vec0_scanner_next(&scanner, &token);
2134+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2135+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2136+ return SQLITE_EMPTY;
2137+ }
2138+ char * t = token.start;
2139+ int n = token.end - token.start;
2140+ if (sqlite3_strnicmp(t, "boolean", n) == 0 || sqlite3_strnicmp(t, "bool", n) == 0) {
2141+ column_type = VEC0_METADATA_COLUMN_KIND_BOOLEAN;
2142+ }else if (sqlite3_strnicmp(t, "int64", n) == 0 || sqlite3_strnicmp(t, "integer64", n) == 0 || sqlite3_strnicmp(t, "integer", n) == 0 || sqlite3_strnicmp(t, "int", n) == 0) {
2143+ column_type = VEC0_METADATA_COLUMN_KIND_INTEGER;
2144+ }else if (sqlite3_strnicmp(t, "float", n) == 0 || sqlite3_strnicmp(t, "double", n) == 0 || sqlite3_strnicmp(t, "float64", n) == 0 || sqlite3_strnicmp(t, "f64", n) == 0) {
2145+ column_type = VEC0_METADATA_COLUMN_KIND_FLOAT;
2146+ } else if (sqlite3_strnicmp(t, "text", n) == 0) {
2147+ column_type = VEC0_METADATA_COLUMN_KIND_TEXT;
2148+ } else {
2149+ return SQLITE_EMPTY;
2150+ }
2151+
2152+ *out_column_name = column_name;
2153+ *out_column_name_length = column_name_length;
2154+ *out_column_type = column_type;
2155+
2156+ return SQLITE_OK;
2157+}
2158+
2159+/**
2160+ * @brief Parse an argv[i] entry of a vec0 virtual table definition, and see if
2161+ * it's a PRIMARY KEY definition.
2162+ *
2163+ * @param source: argv[i] source string
2164+ * @param source_length: length of the source string
2165+ * @param out_column_name: If it is a PK, the output column name. Same lifetime
2166+ * as source, points to specific char *
2167+ * @param out_column_name_length: Length of out_column_name in bytes
2168+ * @param out_column_type: SQLITE_TEXT or SQLITE_INTEGER.
2169+ * @return int: SQLITE_EMPTY if not a PK, SQLITE_OK if it is.
2170+ */
2171+int vec0_parse_primary_key_definition(const char *source, int source_length,
2172+ char **out_column_name,
2173+ int *out_column_name_length,
2174+ int *out_column_type) {
2175+ struct Vec0Scanner scanner;
2176+ struct Vec0Token token;
2177+ char *column_name;
2178+ int column_name_length;
2179+ int column_type;
2180+ vec0_scanner_init(&scanner, source, source_length);
2181+
2182+ // Check first token is identifier, will be the column name
2183+ int rc = vec0_scanner_next(&scanner, &token);
2184+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2185+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2186+ return SQLITE_EMPTY;
2187+ }
2188+
2189+ column_name = token.start;
2190+ column_name_length = token.end - token.start;
2191+
2192+ // Check the next token matches "text" or "integer", as column type
2193+ rc = vec0_scanner_next(&scanner, &token);
2194+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2195+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2196+ return SQLITE_EMPTY;
2197+ }
2198+ if (sqlite3_strnicmp(token.start, "text", token.end - token.start) == 0) {
2199+ column_type = SQLITE_TEXT;
2200+ } else if (sqlite3_strnicmp(token.start, "int", token.end - token.start) ==
2201+ 0 ||
2202+ sqlite3_strnicmp(token.start, "integer",
2203+ token.end - token.start) == 0) {
2204+ column_type = SQLITE_INTEGER;
2205+ } else {
2206+ return SQLITE_EMPTY;
2207+ }
2208+
2209+ // Check the next token is identifier and matches "primary"
2210+ rc = vec0_scanner_next(&scanner, &token);
2211+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2212+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2213+ return SQLITE_EMPTY;
2214+ }
2215+ if (sqlite3_strnicmp(token.start, "primary", token.end - token.start) != 0) {
2216+ return SQLITE_EMPTY;
2217+ }
2218+
2219+ // Check the next token is identifier and matches "key"
2220+ rc = vec0_scanner_next(&scanner, &token);
2221+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2222+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2223+ return SQLITE_EMPTY;
2224+ }
2225+ if (sqlite3_strnicmp(token.start, "key", token.end - token.start) != 0) {
2226+ return SQLITE_EMPTY;
2227+ }
2228+
2229+ *out_column_name = column_name;
2230+ *out_column_name_length = column_name_length;
2231+ *out_column_type = column_type;
2232+
2233+ return SQLITE_OK;
2234+}
2235+
2236+enum Vec0DistanceMetrics {
2237+ VEC0_DISTANCE_METRIC_L2 = 1,
2238+ VEC0_DISTANCE_METRIC_COSINE = 2,
2239+ VEC0_DISTANCE_METRIC_L1 = 3,
2240+};
2241+
2242+struct VectorColumnDefinition {
2243+ char *name;
2244+ int name_length;
2245+ size_t dimensions;
2246+ enum VectorElementType element_type;
2247+ enum Vec0DistanceMetrics distance_metric;
2248+};
2249+
2250+struct Vec0PartitionColumnDefinition {
2251+ int type;
2252+ char * name;
2253+ int name_length;
2254+};
2255+
2256+struct Vec0AuxiliaryColumnDefinition {
2257+ int type;
2258+ char * name;
2259+ int name_length;
2260+};
2261+struct Vec0MetadataColumnDefinition {
2262+ vec0_metadata_column_kind kind;
2263+ char * name;
2264+ int name_length;
2265+};
2266+
2267+size_t vector_byte_size(enum VectorElementType element_type,
2268+ size_t dimensions) {
2269+ switch (element_type) {
2270+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32:
2271+ return dimensions * sizeof(f32);
2272+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
2273+ return dimensions * sizeof(i8);
2274+ case SQLITE_VEC_ELEMENT_TYPE_BIT:
2275+ return dimensions / CHAR_BIT;
2276+ }
2277+ return 0;
2278+}
2279+
2280+size_t vector_column_byte_size(struct VectorColumnDefinition column) {
2281+ return vector_byte_size(column.element_type, column.dimensions);
2282+}
2283+
2284+/**
2285+ * @brief Parse an vec0 vtab argv[i] column definition and see if
2286+ * it's a vector column defintion, ex `contents_embedding float[768]`.
2287+ *
2288+ * @param source vec0 argv[i] item
2289+ * @param source_length length of source in bytes
2290+ * @param outColumn Output the parse vector column to this struct, if success
2291+ * @return int SQLITE_OK on success, SQLITE_EMPTY is it's not a vector column
2292+ * definition, SQLITE_ERROR on error.
2293+ */
2294+int vec0_parse_vector_column(const char *source, int source_length,
2295+ struct VectorColumnDefinition *outColumn) {
2296+ // parses a vector column definition like so:
2297+ // "abc float[123]", "abc_123 bit[1234]", eetc.
2298+ // https://github.com/asg017/sqlite-vec/issues/46
2299+ int rc;
2300+ struct Vec0Scanner scanner;
2301+ struct Vec0Token token;
2302+
2303+ char *name;
2304+ int nameLength;
2305+ enum VectorElementType elementType;
2306+ enum Vec0DistanceMetrics distanceMetric = VEC0_DISTANCE_METRIC_L2;
2307+ int dimensions;
2308+
2309+ vec0_scanner_init(&scanner, source, source_length);
2310+
2311+ // starts with an identifier
2312+ rc = vec0_scanner_next(&scanner, &token);
2313+
2314+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2315+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2316+ return SQLITE_EMPTY;
2317+ }
2318+
2319+ name = token.start;
2320+ nameLength = token.end - token.start;
2321+
2322+ // vector column type comes next: float, int, or bit
2323+ rc = vec0_scanner_next(&scanner, &token);
2324+
2325+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2326+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2327+ return SQLITE_EMPTY;
2328+ }
2329+ if (sqlite3_strnicmp(token.start, "float", 5) == 0 ||
2330+ sqlite3_strnicmp(token.start, "f32", 3) == 0) {
2331+ elementType = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
2332+ } else if (sqlite3_strnicmp(token.start, "int8", 4) == 0 ||
2333+ sqlite3_strnicmp(token.start, "i8", 2) == 0) {
2334+ elementType = SQLITE_VEC_ELEMENT_TYPE_INT8;
2335+ } else if (sqlite3_strnicmp(token.start, "bit", 3) == 0) {
2336+ elementType = SQLITE_VEC_ELEMENT_TYPE_BIT;
2337+ } else {
2338+ return SQLITE_EMPTY;
2339+ }
2340+
2341+ // left '[' bracket
2342+ rc = vec0_scanner_next(&scanner, &token);
2343+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_LBRACKET) {
2344+ return SQLITE_EMPTY;
2345+ }
2346+
2347+ // digit, for vector dimension length
2348+ rc = vec0_scanner_next(&scanner, &token);
2349+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_DIGIT) {
2350+ return SQLITE_ERROR;
2351+ }
2352+ dimensions = atoi(token.start);
2353+ if (dimensions <= 0) {
2354+ return SQLITE_ERROR;
2355+ }
2356+
2357+ // // right ']' bracket
2358+ rc = vec0_scanner_next(&scanner, &token);
2359+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_RBRACKET) {
2360+ return SQLITE_ERROR;
2361+ }
2362+
2363+ // any other tokens left should be column-level options , ex `key=value`
2364+ // ex `distance_metric=L2 distance_metric=cosine` should error
2365+ while (1) {
2366+ // should be EOF or identifier (option key)
2367+ rc = vec0_scanner_next(&scanner, &token);
2368+ if (rc == VEC0_TOKEN_RESULT_EOF) {
2369+ break;
2370+ }
2371+
2372+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2373+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2374+ return SQLITE_ERROR;
2375+ }
2376+
2377+ char *key = token.start;
2378+ int keyLength = token.end - token.start;
2379+
2380+ if (sqlite3_strnicmp(key, "distance_metric", keyLength) == 0) {
2381+
2382+ if (elementType == SQLITE_VEC_ELEMENT_TYPE_BIT) {
2383+ return SQLITE_ERROR;
2384+ }
2385+ // ensure equal sign after distance_metric
2386+ rc = vec0_scanner_next(&scanner, &token);
2387+ if (rc != VEC0_TOKEN_RESULT_SOME && token.token_type != TOKEN_TYPE_EQ) {
2388+ return SQLITE_ERROR;
2389+ }
2390+
2391+ // distance_metric value, an identifier (L2, cosine, etc)
2392+ rc = vec0_scanner_next(&scanner, &token);
2393+ if (rc != VEC0_TOKEN_RESULT_SOME &&
2394+ token.token_type != TOKEN_TYPE_IDENTIFIER) {
2395+ return SQLITE_ERROR;
2396+ }
2397+
2398+ char *value = token.start;
2399+ int valueLength = token.end - token.start;
2400+ if (sqlite3_strnicmp(value, "l2", valueLength) == 0) {
2401+ distanceMetric = VEC0_DISTANCE_METRIC_L2;
2402+ } else if (sqlite3_strnicmp(value, "l1", valueLength) == 0) {
2403+ distanceMetric = VEC0_DISTANCE_METRIC_L1;
2404+ } else if (sqlite3_strnicmp(value, "cosine", valueLength) == 0) {
2405+ distanceMetric = VEC0_DISTANCE_METRIC_COSINE;
2406+ } else {
2407+ return SQLITE_ERROR;
2408+ }
2409+ }
2410+ // unknown key
2411+ else {
2412+ return SQLITE_ERROR;
2413+ }
2414+ }
2415+
2416+ outColumn->name = sqlite3_mprintf("%.*s", nameLength, name);
2417+ if (!outColumn->name) {
2418+ return SQLITE_ERROR;
2419+ }
2420+ outColumn->name_length = nameLength;
2421+ outColumn->distance_metric = distanceMetric;
2422+ outColumn->element_type = elementType;
2423+ outColumn->dimensions = dimensions;
2424+ return SQLITE_OK;
2425+}
2426+
2427+#pragma region vec_each table function
2428+
2429+typedef struct vec_each_vtab vec_each_vtab;
2430+struct vec_each_vtab {
2431+ sqlite3_vtab base;
2432+};
2433+
2434+typedef struct vec_each_cursor vec_each_cursor;
2435+struct vec_each_cursor {
2436+ sqlite3_vtab_cursor base;
2437+ i64 iRowid;
2438+ enum VectorElementType vector_type;
2439+ void *vector;
2440+ size_t dimensions;
2441+ vector_cleanup cleanup;
2442+};
2443+
2444+static int vec_eachConnect(sqlite3 *db, void *pAux, int argc,
2445+ const char *const *argv, sqlite3_vtab **ppVtab,
2446+ char **pzErr) {
2447+ UNUSED_PARAMETER(pAux);
2448+ UNUSED_PARAMETER(argc);
2449+ UNUSED_PARAMETER(argv);
2450+ UNUSED_PARAMETER(pzErr);
2451+ vec_each_vtab *pNew;
2452+ int rc;
2453+
2454+ rc = sqlite3_declare_vtab(db, "CREATE TABLE x(value, vector hidden)");
2455+#define VEC_EACH_COLUMN_VALUE 0
2456+#define VEC_EACH_COLUMN_VECTOR 1
2457+ if (rc == SQLITE_OK) {
2458+ pNew = sqlite3_malloc(sizeof(*pNew));
2459+ *ppVtab = (sqlite3_vtab *)pNew;
2460+ if (pNew == 0)
2461+ return SQLITE_NOMEM;
2462+ memset(pNew, 0, sizeof(*pNew));
2463+ }
2464+ return rc;
2465+}
2466+
2467+static int vec_eachDisconnect(sqlite3_vtab *pVtab) {
2468+ vec_each_vtab *p = (vec_each_vtab *)pVtab;
2469+ sqlite3_free(p);
2470+ return SQLITE_OK;
2471+}
2472+
2473+static int vec_eachOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) {
2474+ UNUSED_PARAMETER(p);
2475+ vec_each_cursor *pCur;
2476+ pCur = sqlite3_malloc(sizeof(*pCur));
2477+ if (pCur == 0)
2478+ return SQLITE_NOMEM;
2479+ memset(pCur, 0, sizeof(*pCur));
2480+ *ppCursor = &pCur->base;
2481+ return SQLITE_OK;
2482+}
2483+
2484+static int vec_eachClose(sqlite3_vtab_cursor *cur) {
2485+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2486+ pCur->cleanup(pCur->vector);
2487+ sqlite3_free(pCur);
2488+ return SQLITE_OK;
2489+}
2490+
2491+static int vec_eachBestIndex(sqlite3_vtab *pVTab,
2492+ sqlite3_index_info *pIdxInfo) {
2493+ UNUSED_PARAMETER(pVTab);
2494+ int hasVector = 0;
2495+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
2496+ const struct sqlite3_index_constraint *pCons = &pIdxInfo->aConstraint[i];
2497+ // printf("i=%d iColumn=%d, op=%d, usable=%d\n", i, pCons->iColumn,
2498+ // pCons->op, pCons->usable);
2499+ switch (pCons->iColumn) {
2500+ case VEC_EACH_COLUMN_VECTOR: {
2501+ if (pCons->op == SQLITE_INDEX_CONSTRAINT_EQ && pCons->usable) {
2502+ hasVector = 1;
2503+ pIdxInfo->aConstraintUsage[i].argvIndex = 1;
2504+ pIdxInfo->aConstraintUsage[i].omit = 1;
2505+ }
2506+ break;
2507+ }
2508+ }
2509+ }
2510+ if (!hasVector) {
2511+ return SQLITE_CONSTRAINT;
2512+ }
2513+
2514+ pIdxInfo->estimatedCost = (double)100000;
2515+ pIdxInfo->estimatedRows = 100000;
2516+
2517+ return SQLITE_OK;
2518+}
2519+
2520+static int vec_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
2521+ const char *idxStr, int argc, sqlite3_value **argv) {
2522+ UNUSED_PARAMETER(idxNum);
2523+ UNUSED_PARAMETER(idxStr);
2524+ assert(argc == 1);
2525+ vec_each_cursor *pCur = (vec_each_cursor *)pVtabCursor;
2526+
2527+ if (pCur->vector) {
2528+ pCur->cleanup(pCur->vector);
2529+ pCur->vector = NULL;
2530+ }
2531+
2532+ char *pzErrMsg;
2533+ int rc = vector_from_value(argv[0], &pCur->vector, &pCur->dimensions,
2534+ &pCur->vector_type, &pCur->cleanup, &pzErrMsg);
2535+ if (rc != SQLITE_OK) {
2536+ return SQLITE_ERROR;
2537+ }
2538+ pCur->iRowid = 0;
2539+ return SQLITE_OK;
2540+}
2541+
2542+static int vec_eachRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) {
2543+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2544+ *pRowid = pCur->iRowid;
2545+ return SQLITE_OK;
2546+}
2547+
2548+static int vec_eachEof(sqlite3_vtab_cursor *cur) {
2549+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2550+ return pCur->iRowid >= (i64)pCur->dimensions;
2551+}
2552+
2553+static int vec_eachNext(sqlite3_vtab_cursor *cur) {
2554+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2555+ pCur->iRowid++;
2556+ return SQLITE_OK;
2557+}
2558+
2559+static int vec_eachColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context,
2560+ int i) {
2561+ vec_each_cursor *pCur = (vec_each_cursor *)cur;
2562+ switch (i) {
2563+ case VEC_EACH_COLUMN_VALUE:
2564+ switch (pCur->vector_type) {
2565+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
2566+ sqlite3_result_double(context, ((f32 *)pCur->vector)[pCur->iRowid]);
2567+ break;
2568+ }
2569+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
2570+ u8 x = ((u8 *)pCur->vector)[pCur->iRowid / CHAR_BIT];
2571+ sqlite3_result_int(context,
2572+ (x & (0b10000000 >> ((pCur->iRowid % CHAR_BIT)))) > 0);
2573+ break;
2574+ }
2575+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
2576+ sqlite3_result_int(context, ((i8 *)pCur->vector)[pCur->iRowid]);
2577+ break;
2578+ }
2579+ }
2580+
2581+ break;
2582+ }
2583+ return SQLITE_OK;
2584+}
2585+
2586+static sqlite3_module vec_eachModule = {
2587+ /* iVersion */ 0,
2588+ /* xCreate */ 0,
2589+ /* xConnect */ vec_eachConnect,
2590+ /* xBestIndex */ vec_eachBestIndex,
2591+ /* xDisconnect */ vec_eachDisconnect,
2592+ /* xDestroy */ 0,
2593+ /* xOpen */ vec_eachOpen,
2594+ /* xClose */ vec_eachClose,
2595+ /* xFilter */ vec_eachFilter,
2596+ /* xNext */ vec_eachNext,
2597+ /* xEof */ vec_eachEof,
2598+ /* xColumn */ vec_eachColumn,
2599+ /* xRowid */ vec_eachRowid,
2600+ /* xUpdate */ 0,
2601+ /* xBegin */ 0,
2602+ /* xSync */ 0,
2603+ /* xCommit */ 0,
2604+ /* xRollback */ 0,
2605+ /* xFindMethod */ 0,
2606+ /* xRename */ 0,
2607+ /* xSavepoint */ 0,
2608+ /* xRelease */ 0,
2609+ /* xRollbackTo */ 0,
2610+ /* xShadowName */ 0,
2611+#if SQLITE_VERSION_NUMBER >= 3044000
2612+ /* xIntegrity */ 0
2613+#endif
2614+};
2615+
2616+#pragma endregion
2617+
2618+#pragma region vec_npy_each table function
2619+
2620+enum NpyTokenType {
2621+ NPY_TOKEN_TYPE_IDENTIFIER,
2622+ NPY_TOKEN_TYPE_NUMBER,
2623+ NPY_TOKEN_TYPE_LPAREN,
2624+ NPY_TOKEN_TYPE_RPAREN,
2625+ NPY_TOKEN_TYPE_LBRACE,
2626+ NPY_TOKEN_TYPE_RBRACE,
2627+ NPY_TOKEN_TYPE_COLON,
2628+ NPY_TOKEN_TYPE_COMMA,
2629+ NPY_TOKEN_TYPE_STRING,
2630+ NPY_TOKEN_TYPE_FALSE,
2631+};
2632+
2633+struct NpyToken {
2634+ enum NpyTokenType token_type;
2635+ unsigned char *start;
2636+ unsigned char *end;
2637+};
2638+
2639+int npy_token_next(unsigned char *start, unsigned char *end,
2640+ struct NpyToken *out) {
2641+ unsigned char *ptr = start;
2642+ while (ptr < end) {
2643+ unsigned char curr = *ptr;
2644+ if (is_whitespace(curr)) {
2645+ ptr++;
2646+ continue;
2647+ } else if (curr == '(') {
2648+ out->start = ptr++;
2649+ out->end = ptr;
2650+ out->token_type = NPY_TOKEN_TYPE_LPAREN;
2651+ return VEC0_TOKEN_RESULT_SOME;
2652+ } else if (curr == ')') {
2653+ out->start = ptr++;
2654+ out->end = ptr;
2655+ out->token_type = NPY_TOKEN_TYPE_RPAREN;
2656+ return VEC0_TOKEN_RESULT_SOME;
2657+ } else if (curr == '{') {
2658+ out->start = ptr++;
2659+ out->end = ptr;
2660+ out->token_type = NPY_TOKEN_TYPE_LBRACE;
2661+ return VEC0_TOKEN_RESULT_SOME;
2662+ } else if (curr == '}') {
2663+ out->start = ptr++;
2664+ out->end = ptr;
2665+ out->token_type = NPY_TOKEN_TYPE_RBRACE;
2666+ return VEC0_TOKEN_RESULT_SOME;
2667+ } else if (curr == ':') {
2668+ out->start = ptr++;
2669+ out->end = ptr;
2670+ out->token_type = NPY_TOKEN_TYPE_COLON;
2671+ return VEC0_TOKEN_RESULT_SOME;
2672+ } else if (curr == ',') {
2673+ out->start = ptr++;
2674+ out->end = ptr;
2675+ out->token_type = NPY_TOKEN_TYPE_COMMA;
2676+ return VEC0_TOKEN_RESULT_SOME;
2677+ } else if (curr == '\'') {
2678+ unsigned char *start = ptr;
2679+ ptr++;
2680+ while (ptr < end) {
2681+ if ((*ptr) == '\'') {
2682+ break;
2683+ }
2684+ ptr++;
2685+ }
2686+ if ((*ptr) != '\'') {
2687+ return VEC0_TOKEN_RESULT_ERROR;
2688+ }
2689+ out->start = start;
2690+ out->end = ++ptr;
2691+ out->token_type = NPY_TOKEN_TYPE_STRING;
2692+ return VEC0_TOKEN_RESULT_SOME;
2693+ } else if (curr == 'F' &&
2694+ strncmp((char *)ptr, "False", strlen("False")) == 0) {
2695+ out->start = ptr;
2696+ out->end = (ptr + (int)strlen("False"));
2697+ ptr = out->end;
2698+ out->token_type = NPY_TOKEN_TYPE_FALSE;
2699+ return VEC0_TOKEN_RESULT_SOME;
2700+ } else if (is_digit(curr)) {
2701+ unsigned char *start = ptr;
2702+ while (ptr < end && (is_digit(*ptr))) {
2703+ ptr++;
2704+ }
2705+ out->start = start;
2706+ out->end = ptr;
2707+ out->token_type = NPY_TOKEN_TYPE_NUMBER;
2708+ return VEC0_TOKEN_RESULT_SOME;
2709+ } else {
2710+ return VEC0_TOKEN_RESULT_ERROR;
2711+ }
2712+ }
2713+ return VEC0_TOKEN_RESULT_ERROR;
2714+}
2715+
2716+struct NpyScanner {
2717+ unsigned char *start;
2718+ unsigned char *end;
2719+ unsigned char *ptr;
2720+};
2721+
2722+void npy_scanner_init(struct NpyScanner *scanner, const unsigned char *source,
2723+ int source_length) {
2724+ scanner->start = (unsigned char *)source;
2725+ scanner->end = (unsigned char *)source + source_length;
2726+ scanner->ptr = (unsigned char *)source;
2727+}
2728+
2729+int npy_scanner_next(struct NpyScanner *scanner, struct NpyToken *out) {
2730+ int rc = npy_token_next(scanner->start, scanner->end, out);
2731+ if (rc == VEC0_TOKEN_RESULT_SOME) {
2732+ scanner->start = out->end;
2733+ }
2734+ return rc;
2735+}
2736+
2737+#define NPY_PARSE_ERROR "Error parsing numpy array: "
2738+int parse_npy_header(sqlite3_vtab *pVTab, const unsigned char *header,
2739+ size_t headerLength,
2740+ enum VectorElementType *out_element_type,
2741+ int *fortran_order, size_t *numElements,
2742+ size_t *numDimensions) {
2743+
2744+ struct NpyScanner scanner;
2745+ struct NpyToken token;
2746+ int rc;
2747+ npy_scanner_init(&scanner, header, headerLength);
2748+
2749+ if (npy_scanner_next(&scanner, &token) != VEC0_TOKEN_RESULT_SOME &&
2750+ token.token_type != NPY_TOKEN_TYPE_LBRACE) {
2751+ vtab_set_error(pVTab,
2752+ NPY_PARSE_ERROR "numpy header did not start with '{'");
2753+ return SQLITE_ERROR;
2754+ }
2755+ while (1) {
2756+ rc = npy_scanner_next(&scanner, &token);
2757+ if (rc != VEC0_TOKEN_RESULT_SOME) {
2758+ vtab_set_error(pVTab, NPY_PARSE_ERROR "expected key in numpy header");
2759+ return SQLITE_ERROR;
2760+ }
2761+
2762+ if (token.token_type == NPY_TOKEN_TYPE_RBRACE) {
2763+ break;
2764+ }
2765+ if (token.token_type != NPY_TOKEN_TYPE_STRING) {
2766+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2767+ "expected a string as key in numpy header");
2768+ return SQLITE_ERROR;
2769+ }
2770+ unsigned char *key = token.start;
2771+
2772+ rc = npy_scanner_next(&scanner, &token);
2773+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2774+ (token.token_type != NPY_TOKEN_TYPE_COLON)) {
2775+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2776+ "expected a ':' after key in numpy header");
2777+ return SQLITE_ERROR;
2778+ }
2779+
2780+ if (strncmp((char *)key, "'descr'", strlen("'descr'")) == 0) {
2781+ rc = npy_scanner_next(&scanner, &token);
2782+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2783+ (token.token_type != NPY_TOKEN_TYPE_STRING)) {
2784+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2785+ "expected a string value after 'descr' key");
2786+ return SQLITE_ERROR;
2787+ }
2788+ if (strncmp((char *)token.start, "'<f4'", strlen("'<f4'")) != 0) {
2789+ vtab_set_error(
2790+ pVTab, NPY_PARSE_ERROR
2791+ "Only '<f4' values are supported in sqlite-vec numpy functions");
2792+ return SQLITE_ERROR;
2793+ }
2794+ *out_element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
2795+ } else if (strncmp((char *)key, "'fortran_order'",
2796+ strlen("'fortran_order'")) == 0) {
2797+ rc = npy_scanner_next(&scanner, &token);
2798+ if (rc != VEC0_TOKEN_RESULT_SOME ||
2799+ token.token_type != NPY_TOKEN_TYPE_FALSE) {
2800+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2801+ "Only fortran_order = False is supported in sqlite-vec "
2802+ "numpy functions");
2803+ return SQLITE_ERROR;
2804+ }
2805+ *fortran_order = 0;
2806+ } else if (strncmp((char *)key, "'shape'", strlen("'shape'")) == 0) {
2807+ // "(xxx, xxx)" OR (xxx,)
2808+ size_t first;
2809+ rc = npy_scanner_next(&scanner, &token);
2810+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2811+ (token.token_type != NPY_TOKEN_TYPE_LPAREN)) {
2812+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2813+ "Expected left parenthesis '(' after shape key");
2814+ return SQLITE_ERROR;
2815+ }
2816+
2817+ rc = npy_scanner_next(&scanner, &token);
2818+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2819+ (token.token_type != NPY_TOKEN_TYPE_NUMBER)) {
2820+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2821+ "Expected an initial number in shape value");
2822+ return SQLITE_ERROR;
2823+ }
2824+ first = strtol((char *)token.start, NULL, 10);
2825+
2826+ rc = npy_scanner_next(&scanner, &token);
2827+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2828+ (token.token_type != NPY_TOKEN_TYPE_COMMA)) {
2829+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2830+ "Expected comma after first shape value");
2831+ return SQLITE_ERROR;
2832+ }
2833+
2834+ rc = npy_scanner_next(&scanner, &token);
2835+ if (rc != VEC0_TOKEN_RESULT_SOME) {
2836+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2837+ "unexpected header EOF while parsing shape");
2838+ return SQLITE_ERROR;
2839+ }
2840+ if (token.token_type == NPY_TOKEN_TYPE_NUMBER) {
2841+ *numElements = first;
2842+ *numDimensions = strtol((char *)token.start, NULL, 10);
2843+ rc = npy_scanner_next(&scanner, &token);
2844+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2845+ (token.token_type != NPY_TOKEN_TYPE_RPAREN)) {
2846+ vtab_set_error(pVTab, NPY_PARSE_ERROR
2847+ "expected right parenthesis after shape value");
2848+ return SQLITE_ERROR;
2849+ }
2850+ } else if (token.token_type == NPY_TOKEN_TYPE_RPAREN) {
2851+ // '(0,)' means an empty array!
2852+ *numElements = first ? 1 : 0;
2853+ *numDimensions = first;
2854+ } else {
2855+ vtab_set_error(pVTab, NPY_PARSE_ERROR "unknown type in shape value");
2856+ return SQLITE_ERROR;
2857+ }
2858+ } else {
2859+ vtab_set_error(pVTab, NPY_PARSE_ERROR "unknown key in numpy header");
2860+ return SQLITE_ERROR;
2861+ }
2862+
2863+ rc = npy_scanner_next(&scanner, &token);
2864+ if ((rc != VEC0_TOKEN_RESULT_SOME) ||
2865+ (token.token_type != NPY_TOKEN_TYPE_COMMA)) {
2866+ vtab_set_error(pVTab, NPY_PARSE_ERROR "unknown extra token after value");
2867+ return SQLITE_ERROR;
2868+ }
2869+ }
2870+
2871+ return SQLITE_OK;
2872+}
2873+
2874+typedef struct vec_npy_each_vtab vec_npy_each_vtab;
2875+struct vec_npy_each_vtab {
2876+ sqlite3_vtab base;
2877+};
2878+
2879+typedef enum {
2880+ VEC_NPY_EACH_INPUT_BUFFER,
2881+ VEC_NPY_EACH_INPUT_FILE,
2882+} vec_npy_each_input_type;
2883+
2884+typedef struct vec_npy_each_cursor vec_npy_each_cursor;
2885+struct vec_npy_each_cursor {
2886+ sqlite3_vtab_cursor base;
2887+ i64 iRowid;
2888+ // sqlite-vec compatible type of vector
2889+ enum VectorElementType elementType;
2890+ // number of vectors in the npy array
2891+ size_t nElements;
2892+ // number of dimensions each vector has
2893+ size_t nDimensions;
2894+
2895+ vec_npy_each_input_type input_type;
2896+
2897+ // when input_type == VEC_NPY_EACH_INPUT_BUFFER
2898+
2899+ // Buffer containing the vector data, when reading from an in-memory buffer.
2900+ // Size: nElements * nDimensions * element_size
2901+ // Clean up with sqlite3_free() once complete
2902+ void *vector;
2903+
2904+ // when input_type == VEC_NPY_EACH_INPUT_FILE
2905+
2906+ // Opened npy file, when reading from a file.
2907+ // fclose() when complete.
2908+#ifndef SQLITE_VEC_OMIT_FS
2909+ FILE *file;
2910+#endif
2911+
2912+ // an in-memory buffer containing a portion of the npy array.
2913+ // Used for faster reading, instead of calling fread a lot.
2914+ // Will have a byte-size of fileBufferSize
2915+ void *chunksBuffer;
2916+ // size of allocated fileBuffer in bytes
2917+ size_t chunksBufferSize;
2918+ //// Maximum length of the buffer, in terms of number of vectors.
2919+ size_t maxChunks;
2920+
2921+ // Counter index of the current vector into of fileBuffer to yield.
2922+ // Starts at 0 once fileBuffer is read, and iterates to bufferLength.
2923+ // Resets to 0 once that "buffer" is yielded and a new one is read.
2924+ size_t currentChunkIndex;
2925+ size_t currentChunkSize;
2926+
2927+ // 0 when there are still more elements to read/yield, 1 when complete.
2928+ int eof;
2929+};
2930+
2931+static unsigned char NPY_MAGIC[6] = "\x93NUMPY";
2932+
2933+#ifndef SQLITE_VEC_OMIT_FS
2934+int parse_npy_file(sqlite3_vtab *pVTab, FILE *file, vec_npy_each_cursor *pCur) {
2935+ int n;
2936+ fseek(file, 0, SEEK_END);
2937+ long fileSize = ftell(file);
2938+
2939+ fseek(file, 0L, SEEK_SET);
2940+
2941+ unsigned char header[10];
2942+ n = fread(&header, sizeof(unsigned char), 10, file);
2943+ if (n != 10) {
2944+ vtab_set_error(pVTab, "numpy array file too short");
2945+ return SQLITE_ERROR;
2946+ }
2947+
2948+ if (memcmp(NPY_MAGIC, header, sizeof(NPY_MAGIC)) != 0) {
2949+ vtab_set_error(pVTab,
2950+ "numpy array file does not contain the 'magic' header");
2951+ return SQLITE_ERROR;
2952+ }
2953+
2954+ u8 major = header[6];
2955+ u8 minor = header[7];
2956+ uint16_t headerLength = 0;
2957+ memcpy(&headerLength, &header[8], sizeof(uint16_t));
2958+
2959+ size_t totalHeaderLength = sizeof(NPY_MAGIC) + sizeof(major) + sizeof(minor) +
2960+ sizeof(headerLength) + headerLength;
2961+ i32 dataSize = fileSize - totalHeaderLength;
2962+ if (dataSize < 0) {
2963+ vtab_set_error(pVTab, "numpy array file header length is invalid");
2964+ return SQLITE_ERROR;
2965+ }
2966+
2967+ unsigned char *headerX = sqlite3_malloc(headerLength);
2968+ if (headerLength && !headerX) {
2969+ return SQLITE_NOMEM;
2970+ }
2971+
2972+ n = fread(headerX, sizeof(char), headerLength, file);
2973+ if (n != headerLength) {
2974+ sqlite3_free(headerX);
2975+ vtab_set_error(pVTab, "numpy array file header length is invalid");
2976+ return SQLITE_ERROR;
2977+ }
2978+
2979+ int fortran_order;
2980+ enum VectorElementType element_type;
2981+ size_t numElements;
2982+ size_t numDimensions;
2983+ int rc = parse_npy_header(pVTab, headerX, headerLength, &element_type,
2984+ &fortran_order, &numElements, &numDimensions);
2985+ sqlite3_free(headerX);
2986+ if (rc != SQLITE_OK) {
2987+ // parse_npy_header already attackes an error emssage
2988+ return rc;
2989+ }
2990+
2991+ i32 expectedDataSize =
2992+ numElements * vector_byte_size(element_type, numDimensions);
2993+ if (expectedDataSize != dataSize) {
2994+ vtab_set_error(
2995+ pVTab, "numpy array file error: Expected a data size of %d, found %d",
2996+ expectedDataSize, dataSize);
2997+ return SQLITE_ERROR;
2998+ }
2999+
3000+ pCur->maxChunks = 1024;
3001+ pCur->chunksBufferSize =
3002+ (vector_byte_size(element_type, numDimensions)) * pCur->maxChunks;
3003+ pCur->chunksBuffer = sqlite3_malloc(pCur->chunksBufferSize);
3004+ if (pCur->chunksBufferSize && !pCur->chunksBuffer) {
3005+ return SQLITE_NOMEM;
3006+ }
3007+
3008+ pCur->currentChunkSize =
3009+ fread(pCur->chunksBuffer, vector_byte_size(element_type, numDimensions),
3010+ pCur->maxChunks, file);
3011+
3012+ pCur->currentChunkIndex = 0;
3013+ pCur->elementType = element_type;
3014+ pCur->nElements = numElements;
3015+ pCur->nDimensions = numDimensions;
3016+ pCur->input_type = VEC_NPY_EACH_INPUT_FILE;
3017+
3018+ pCur->eof = pCur->currentChunkSize == 0;
3019+ pCur->file = file;
3020+ return SQLITE_OK;
3021+}
3022+#endif
3023+
3024+int parse_npy_buffer(sqlite3_vtab *pVTab, const unsigned char *buffer,
3025+ int bufferLength, void **data, size_t *numElements,
3026+ size_t *numDimensions,
3027+ enum VectorElementType *element_type) {
3028+
3029+ if (bufferLength < 10) {
3030+ // IMP: V03312_20150
3031+ vtab_set_error(pVTab, "numpy array too short");
3032+ return SQLITE_ERROR;
3033+ }
3034+ if (memcmp(NPY_MAGIC, buffer, sizeof(NPY_MAGIC)) != 0) {
3035+ // V11954_28792
3036+ vtab_set_error(pVTab, "numpy array does not contain the 'magic' header");
3037+ return SQLITE_ERROR;
3038+ }
3039+
3040+ u8 major = buffer[6];
3041+ u8 minor = buffer[7];
3042+ uint16_t headerLength = 0;
3043+ memcpy(&headerLength, &buffer[8], sizeof(uint16_t));
3044+
3045+ i32 totalHeaderLength = sizeof(NPY_MAGIC) + sizeof(major) + sizeof(minor) +
3046+ sizeof(headerLength) + headerLength;
3047+ i32 dataSize = bufferLength - totalHeaderLength;
3048+
3049+ if (dataSize < 0) {
3050+ vtab_set_error(pVTab, "numpy array header length is invalid");
3051+ return SQLITE_ERROR;
3052+ }
3053+
3054+ const unsigned char *header = &buffer[10];
3055+ int fortran_order;
3056+
3057+ int rc = parse_npy_header(pVTab, header, headerLength, element_type,
3058+ &fortran_order, numElements, numDimensions);
3059+ if (rc != SQLITE_OK) {
3060+ return rc;
3061+ }
3062+
3063+ i32 expectedDataSize =
3064+ (*numElements * vector_byte_size(*element_type, *numDimensions));
3065+ if (expectedDataSize != dataSize) {
3066+ vtab_set_error(pVTab,
3067+ "numpy array error: Expected a data size of %d, found %d",
3068+ expectedDataSize, dataSize);
3069+ return SQLITE_ERROR;
3070+ }
3071+
3072+ *data = (void *)&buffer[totalHeaderLength];
3073+ return SQLITE_OK;
3074+}
3075+
3076+static int vec_npy_eachConnect(sqlite3 *db, void *pAux, int argc,
3077+ const char *const *argv, sqlite3_vtab **ppVtab,
3078+ char **pzErr) {
3079+ UNUSED_PARAMETER(pAux);
3080+ UNUSED_PARAMETER(argc);
3081+ UNUSED_PARAMETER(argv);
3082+ UNUSED_PARAMETER(pzErr);
3083+ vec_npy_each_vtab *pNew;
3084+ int rc;
3085+
3086+ rc = sqlite3_declare_vtab(db, "CREATE TABLE x(vector, input hidden)");
3087+#define VEC_NPY_EACH_COLUMN_VECTOR 0
3088+#define VEC_NPY_EACH_COLUMN_INPUT 1
3089+ if (rc == SQLITE_OK) {
3090+ pNew = sqlite3_malloc(sizeof(*pNew));
3091+ *ppVtab = (sqlite3_vtab *)pNew;
3092+ if (pNew == 0)
3093+ return SQLITE_NOMEM;
3094+ memset(pNew, 0, sizeof(*pNew));
3095+ }
3096+ return rc;
3097+}
3098+
3099+static int vec_npy_eachDisconnect(sqlite3_vtab *pVtab) {
3100+ vec_npy_each_vtab *p = (vec_npy_each_vtab *)pVtab;
3101+ sqlite3_free(p);
3102+ return SQLITE_OK;
3103+}
3104+
3105+static int vec_npy_eachOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) {
3106+ UNUSED_PARAMETER(p);
3107+ vec_npy_each_cursor *pCur;
3108+ pCur = sqlite3_malloc(sizeof(*pCur));
3109+ if (pCur == 0)
3110+ return SQLITE_NOMEM;
3111+ memset(pCur, 0, sizeof(*pCur));
3112+ *ppCursor = &pCur->base;
3113+ return SQLITE_OK;
3114+}
3115+
3116+static int vec_npy_eachClose(sqlite3_vtab_cursor *cur) {
3117+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3118+#ifndef SQLITE_VEC_OMIT_FS
3119+ if (pCur->file) {
3120+ fclose(pCur->file);
3121+ pCur->file = NULL;
3122+ }
3123+#endif
3124+ if (pCur->chunksBuffer) {
3125+ sqlite3_free(pCur->chunksBuffer);
3126+ pCur->chunksBuffer = NULL;
3127+ }
3128+ if (pCur->vector) {
3129+ pCur->vector = NULL;
3130+ }
3131+ sqlite3_free(pCur);
3132+ return SQLITE_OK;
3133+}
3134+
3135+static int vec_npy_eachBestIndex(sqlite3_vtab *pVTab,
3136+ sqlite3_index_info *pIdxInfo) {
3137+ int hasInput;
3138+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
3139+ const struct sqlite3_index_constraint *pCons = &pIdxInfo->aConstraint[i];
3140+ // printf("i=%d iColumn=%d, op=%d, usable=%d\n", i, pCons->iColumn,
3141+ // pCons->op, pCons->usable);
3142+ switch (pCons->iColumn) {
3143+ case VEC_NPY_EACH_COLUMN_INPUT: {
3144+ if (pCons->op == SQLITE_INDEX_CONSTRAINT_EQ && pCons->usable) {
3145+ hasInput = 1;
3146+ pIdxInfo->aConstraintUsage[i].argvIndex = 1;
3147+ pIdxInfo->aConstraintUsage[i].omit = 1;
3148+ }
3149+ break;
3150+ }
3151+ }
3152+ }
3153+ if (!hasInput) {
3154+ pVTab->zErrMsg = sqlite3_mprintf("input argument is required");
3155+ return SQLITE_ERROR;
3156+ }
3157+
3158+ pIdxInfo->estimatedCost = (double)100000;
3159+ pIdxInfo->estimatedRows = 100000;
3160+
3161+ return SQLITE_OK;
3162+}
3163+
3164+static int vec_npy_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
3165+ const char *idxStr, int argc,
3166+ sqlite3_value **argv) {
3167+ UNUSED_PARAMETER(idxNum);
3168+ UNUSED_PARAMETER(idxStr);
3169+ assert(argc == 1);
3170+ int rc;
3171+
3172+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)pVtabCursor;
3173+
3174+#ifndef SQLITE_VEC_OMIT_FS
3175+ if (pCur->file) {
3176+ fclose(pCur->file);
3177+ pCur->file = NULL;
3178+ }
3179+#endif
3180+ if (pCur->chunksBuffer) {
3181+ sqlite3_free(pCur->chunksBuffer);
3182+ pCur->chunksBuffer = NULL;
3183+ }
3184+ if (pCur->vector) {
3185+ pCur->vector = NULL;
3186+ }
3187+
3188+#ifndef SQLITE_VEC_OMIT_FS
3189+ struct VecNpyFile *f = NULL;
3190+ if ((f = sqlite3_value_pointer(argv[0], SQLITE_VEC_NPY_FILE_NAME))) {
3191+ FILE *file = fopen(f->path, "r");
3192+ if (!file) {
3193+ vtab_set_error(pVtabCursor->pVtab, "Could not open numpy file");
3194+ return SQLITE_ERROR;
3195+ }
3196+
3197+ rc = parse_npy_file(pVtabCursor->pVtab, file, pCur);
3198+ if (rc != SQLITE_OK) {
3199+#ifndef SQLITE_VEC_OMIT_FS
3200+ fclose(file);
3201+#endif
3202+ return rc;
3203+ }
3204+
3205+ } else
3206+#endif
3207+ {
3208+
3209+ const unsigned char *input = sqlite3_value_blob(argv[0]);
3210+ int inputLength = sqlite3_value_bytes(argv[0]);
3211+ void *data;
3212+ size_t numElements;
3213+ size_t numDimensions;
3214+ enum VectorElementType element_type;
3215+
3216+ rc = parse_npy_buffer(pVtabCursor->pVtab, input, inputLength, &data,
3217+ &numElements, &numDimensions, &element_type);
3218+ if (rc != SQLITE_OK) {
3219+ return rc;
3220+ }
3221+
3222+ pCur->vector = data;
3223+ pCur->elementType = element_type;
3224+ pCur->nElements = numElements;
3225+ pCur->nDimensions = numDimensions;
3226+ pCur->input_type = VEC_NPY_EACH_INPUT_BUFFER;
3227+ }
3228+
3229+ pCur->iRowid = 0;
3230+ return SQLITE_OK;
3231+}
3232+
3233+static int vec_npy_eachRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) {
3234+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3235+ *pRowid = pCur->iRowid;
3236+ return SQLITE_OK;
3237+}
3238+
3239+static int vec_npy_eachEof(sqlite3_vtab_cursor *cur) {
3240+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3241+ if (pCur->input_type == VEC_NPY_EACH_INPUT_BUFFER) {
3242+ return (!pCur->nElements) || (size_t)pCur->iRowid >= pCur->nElements;
3243+ }
3244+ return pCur->eof;
3245+}
3246+
3247+static int vec_npy_eachNext(sqlite3_vtab_cursor *cur) {
3248+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3249+ pCur->iRowid++;
3250+ if (pCur->input_type == VEC_NPY_EACH_INPUT_BUFFER) {
3251+ return SQLITE_OK;
3252+ }
3253+
3254+#ifndef SQLITE_VEC_OMIT_FS
3255+ // else: input is a file
3256+ pCur->currentChunkIndex++;
3257+ if (pCur->currentChunkIndex >= pCur->currentChunkSize) {
3258+ pCur->currentChunkSize =
3259+ fread(pCur->chunksBuffer,
3260+ vector_byte_size(pCur->elementType, pCur->nDimensions),
3261+ pCur->maxChunks, pCur->file);
3262+ if (!pCur->currentChunkSize) {
3263+ pCur->eof = 1;
3264+ }
3265+ pCur->currentChunkIndex = 0;
3266+ }
3267+#endif
3268+ return SQLITE_OK;
3269+}
3270+
3271+static int vec_npy_eachColumnBuffer(vec_npy_each_cursor *pCur,
3272+ sqlite3_context *context, int i) {
3273+ switch (i) {
3274+ case VEC_NPY_EACH_COLUMN_VECTOR: {
3275+ sqlite3_result_subtype(context, pCur->elementType);
3276+ switch (pCur->elementType) {
3277+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
3278+ sqlite3_result_blob(
3279+ context,
3280+ &((unsigned char *)
3281+ pCur->vector)[pCur->iRowid * pCur->nDimensions * sizeof(f32)],
3282+ pCur->nDimensions * sizeof(f32), SQLITE_TRANSIENT);
3283+
3284+ break;
3285+ }
3286+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
3287+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
3288+ // https://github.com/asg017/sqlite-vec/issues/42
3289+ sqlite3_result_error(context,
3290+ "vec_npy_each only supports float32 vectors", -1);
3291+ break;
3292+ }
3293+ }
3294+
3295+ break;
3296+ }
3297+ }
3298+ return SQLITE_OK;
3299+}
3300+static int vec_npy_eachColumnFile(vec_npy_each_cursor *pCur,
3301+ sqlite3_context *context, int i) {
3302+ switch (i) {
3303+ case VEC_NPY_EACH_COLUMN_VECTOR: {
3304+ switch (pCur->elementType) {
3305+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
3306+ sqlite3_result_blob(
3307+ context,
3308+ &((unsigned char *)
3309+ pCur->chunksBuffer)[pCur->currentChunkIndex *
3310+ pCur->nDimensions * sizeof(f32)],
3311+ pCur->nDimensions * sizeof(f32), SQLITE_TRANSIENT);
3312+ break;
3313+ }
3314+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
3315+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
3316+ // https://github.com/asg017/sqlite-vec/issues/42
3317+ sqlite3_result_error(context,
3318+ "vec_npy_each only supports float32 vectors", -1);
3319+ break;
3320+ }
3321+ }
3322+ break;
3323+ }
3324+ }
3325+ return SQLITE_OK;
3326+}
3327+static int vec_npy_eachColumn(sqlite3_vtab_cursor *cur,
3328+ sqlite3_context *context, int i) {
3329+ vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
3330+ switch (pCur->input_type) {
3331+ case VEC_NPY_EACH_INPUT_BUFFER:
3332+ return vec_npy_eachColumnBuffer(pCur, context, i);
3333+ case VEC_NPY_EACH_INPUT_FILE:
3334+ return vec_npy_eachColumnFile(pCur, context, i);
3335+ }
3336+ return SQLITE_ERROR;
3337+}
3338+
3339+static sqlite3_module vec_npy_eachModule = {
3340+ /* iVersion */ 0,
3341+ /* xCreate */ 0,
3342+ /* xConnect */ vec_npy_eachConnect,
3343+ /* xBestIndex */ vec_npy_eachBestIndex,
3344+ /* xDisconnect */ vec_npy_eachDisconnect,
3345+ /* xDestroy */ 0,
3346+ /* xOpen */ vec_npy_eachOpen,
3347+ /* xClose */ vec_npy_eachClose,
3348+ /* xFilter */ vec_npy_eachFilter,
3349+ /* xNext */ vec_npy_eachNext,
3350+ /* xEof */ vec_npy_eachEof,
3351+ /* xColumn */ vec_npy_eachColumn,
3352+ /* xRowid */ vec_npy_eachRowid,
3353+ /* xUpdate */ 0,
3354+ /* xBegin */ 0,
3355+ /* xSync */ 0,
3356+ /* xCommit */ 0,
3357+ /* xRollback */ 0,
3358+ /* xFindMethod */ 0,
3359+ /* xRename */ 0,
3360+ /* xSavepoint */ 0,
3361+ /* xRelease */ 0,
3362+ /* xRollbackTo */ 0,
3363+ /* xShadowName */ 0,
3364+#if SQLITE_VERSION_NUMBER >= 3044000
3365+ /* xIntegrity */ 0,
3366+#endif
3367+};
3368+
3369+#pragma endregion
3370+
3371+#pragma region vec0 virtual table
3372+
3373+#define VEC0_COLUMN_ID 0
3374+#define VEC0_COLUMN_USERN_START 1
3375+#define VEC0_COLUMN_OFFSET_DISTANCE 1
3376+#define VEC0_COLUMN_OFFSET_K 2
3377+
3378+#define VEC0_SHADOW_INFO_NAME "\"%w\".\"%w_info\""
3379+
3380+#define VEC0_SHADOW_CHUNKS_NAME "\"%w\".\"%w_chunks\""
3381+/// 1) schema, 2) original vtab table name
3382+#define VEC0_SHADOW_CHUNKS_CREATE \
3383+ "CREATE TABLE " VEC0_SHADOW_CHUNKS_NAME "(" \
3384+ "chunk_id INTEGER PRIMARY KEY AUTOINCREMENT," \
3385+ "size INTEGER NOT NULL," \
3386+ "validity BLOB NOT NULL," \
3387+ "rowids BLOB NOT NULL" \
3388+ ");"
3389+
3390+#define VEC0_SHADOW_ROWIDS_NAME "\"%w\".\"%w_rowids\""
3391+/// 1) schema, 2) original vtab table name
3392+#define VEC0_SHADOW_ROWIDS_CREATE_BASIC \
3393+ "CREATE TABLE " VEC0_SHADOW_ROWIDS_NAME "(" \
3394+ "rowid INTEGER PRIMARY KEY AUTOINCREMENT," \
3395+ "id," \
3396+ "chunk_id INTEGER," \
3397+ "chunk_offset INTEGER" \
3398+ ");"
3399+
3400+// vec0 tables with a text primary keys are still backed by int64 primary keys,
3401+// since a fixed-length rowid is required for vec0 chunks. But we add a new 'id
3402+// text unique' column to emulate a text primary key interface.
3403+#define VEC0_SHADOW_ROWIDS_CREATE_PK_TEXT \
3404+ "CREATE TABLE " VEC0_SHADOW_ROWIDS_NAME "(" \
3405+ "rowid INTEGER PRIMARY KEY AUTOINCREMENT," \
3406+ "id TEXT UNIQUE NOT NULL," \
3407+ "chunk_id INTEGER," \
3408+ "chunk_offset INTEGER" \
3409+ ");"
3410+
3411+/// 1) schema, 2) original vtab table name
3412+#define VEC0_SHADOW_VECTOR_N_NAME "\"%w\".\"%w_vector_chunks%02d\""
3413+
3414+/// 1) schema, 2) original vtab table name
3415+#define VEC0_SHADOW_VECTOR_N_CREATE \
3416+ "CREATE TABLE " VEC0_SHADOW_VECTOR_N_NAME "(" \
3417+ "rowid PRIMARY KEY," \
3418+ "vectors BLOB NOT NULL" \
3419+ ");"
3420+
3421+#define VEC0_SHADOW_AUXILIARY_NAME "\"%w\".\"%w_auxiliary\""
3422+
3423+#define VEC0_SHADOW_METADATA_N_NAME "\"%w\".\"%w_metadatachunks%02d\""
3424+#define VEC0_SHADOW_METADATA_TEXT_DATA_NAME "\"%w\".\"%w_metadatatext%02d\""
3425+
3426+#define VEC_INTERAL_ERROR "Internal sqlite-vec error: "
3427+#define REPORT_URL "https://github.com/asg017/sqlite-vec/issues/new"
3428+
3429+typedef struct vec0_vtab vec0_vtab;
3430+
3431+#define VEC0_MAX_VECTOR_COLUMNS 16
3432+#define VEC0_MAX_PARTITION_COLUMNS 4
3433+#define VEC0_MAX_AUXILIARY_COLUMNS 16
3434+#define VEC0_MAX_METADATA_COLUMNS 16
3435+
3436+#define SQLITE_VEC_VEC0_MAX_DIMENSIONS 8192
3437+#define VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH 16
3438+#define VEC0_METADATA_TEXT_VIEW_DATA_LENGTH 12
3439+
3440+typedef enum {
3441+ // vector column, ie "contents_embedding float[1024]"
3442+ SQLITE_VEC0_USER_COLUMN_KIND_VECTOR = 1,
3443+
3444+ // partition key column, ie "user_id integer partition key"
3445+ SQLITE_VEC0_USER_COLUMN_KIND_PARTITION = 2,
3446+
3447+ //
3448+ SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY = 3,
3449+
3450+ // metadata column that can be filtered, ie "genre text"
3451+ SQLITE_VEC0_USER_COLUMN_KIND_METADATA = 4,
3452+} vec0_user_column_kind;
3453+
3454+struct vec0_vtab {
3455+ sqlite3_vtab base;
3456+
3457+ // the SQLite connection of the host database
3458+ sqlite3 *db;
3459+
3460+ // True if the primary key of the vec0 table has a column type TEXT.
3461+ // Will change the schema of the _rowids table, and insert/query logic.
3462+ int pkIsText;
3463+
3464+ // number of defined vector columns.
3465+ int numVectorColumns;
3466+
3467+ // number of defined PARTITION KEY columns.
3468+ int numPartitionColumns;
3469+
3470+ // number of defined auxiliary columns
3471+ int numAuxiliaryColumns;
3472+
3473+ // number of defined metadata columns
3474+ int numMetadataColumns;
3475+
3476+
3477+ // Name of the schema the table exists on.
3478+ // Must be freed with sqlite3_free()
3479+ char *schemaName;
3480+
3481+ // Name of the table the table exists on.
3482+ // Must be freed with sqlite3_free()
3483+ char *tableName;
3484+
3485+ // Name of the _rowids shadow table.
3486+ // Must be freed with sqlite3_free()
3487+ char *shadowRowidsName;
3488+
3489+ // Name of the _chunks shadow table.
3490+ // Must be freed with sqlite3_free()
3491+ char *shadowChunksName;
3492+
3493+ // contains enum vec0_user_column_kind values for up to
3494+ // numVectorColumns + numPartitionColumns entries
3495+ vec0_user_column_kind user_column_kinds[VEC0_MAX_VECTOR_COLUMNS + VEC0_MAX_PARTITION_COLUMNS + VEC0_MAX_AUXILIARY_COLUMNS + VEC0_MAX_METADATA_COLUMNS];
3496+
3497+ uint8_t user_column_idxs[VEC0_MAX_VECTOR_COLUMNS + VEC0_MAX_PARTITION_COLUMNS + VEC0_MAX_AUXILIARY_COLUMNS + VEC0_MAX_METADATA_COLUMNS];
3498+
3499+
3500+ // Name of all the vector chunk shadow tables.
3501+ // Ex '_vector_chunks00'
3502+ // Only the first numVectorColumns entries will be available.
3503+ // The first numVectorColumns entries must be freed with sqlite3_free()
3504+ char *shadowVectorChunksNames[VEC0_MAX_VECTOR_COLUMNS];
3505+
3506+ // Name of all metadata chunk shadow tables, ie `_metadatachunks00`
3507+ // Only the first numMetadataColumns entries will be available.
3508+ // The first numMetadataColumns entries must be freed with sqlite3_free()
3509+ char *shadowMetadataChunksNames[VEC0_MAX_METADATA_COLUMNS];
3510+
3511+ struct VectorColumnDefinition vector_columns[VEC0_MAX_VECTOR_COLUMNS];
3512+ struct Vec0PartitionColumnDefinition paritition_columns[VEC0_MAX_PARTITION_COLUMNS];
3513+ struct Vec0AuxiliaryColumnDefinition auxiliary_columns[VEC0_MAX_AUXILIARY_COLUMNS];
3514+ struct Vec0MetadataColumnDefinition metadata_columns[VEC0_MAX_METADATA_COLUMNS];
3515+
3516+ int chunk_size;
3517+
3518+ // select latest chunk from _chunks, getting chunk_id
3519+ sqlite3_stmt *stmtLatestChunk;
3520+
3521+ /**
3522+ * Statement to insert a row into the _rowids table, with a rowid.
3523+ * Parameters:
3524+ * 1: int64, rowid to insert
3525+ * Result columns: none
3526+ * SQL: "INSERT INTO _rowids(rowid) VALUES (?)"
3527+ *
3528+ * Must be cleaned up with sqlite3_finalize().
3529+ */
3530+ sqlite3_stmt *stmtRowidsInsertRowid;
3531+
3532+ /**
3533+ * Statement to insert a row into the _rowids table, with an id.
3534+ * The id column isn't a tradition primary key, but instead a unique
3535+ * column to handle "text primary key" vec0 tables. The true int64 rowid
3536+ * can be retrieved after inserting with sqlite3_last_rowid().
3537+ *
3538+ * Parameters:
3539+ * 1: text or null, id to insert
3540+ * Result columns: none
3541+ *
3542+ * Must be cleaned up with sqlite3_finalize().
3543+ */
3544+ sqlite3_stmt *stmtRowidsInsertId;
3545+
3546+ /**
3547+ * Statement to update the "position" columns chunk_id and chunk_offset for
3548+ * a given _rowids row. Used when the "next available" chunk position is found
3549+ * for a vector.
3550+ *
3551+ * Parameters:
3552+ * 1: int64, chunk_id value
3553+ * 2: int64, chunk_offset value
3554+ * 3: int64, rowid value
3555+ * Result columns: none
3556+ *
3557+ * Must be cleaned up with sqlite3_finalize().
3558+ */
3559+ sqlite3_stmt *stmtRowidsUpdatePosition;
3560+
3561+ /**
3562+ * Statement to quickly find the chunk_id + chunk_offset of a given row.
3563+ * Parameters:
3564+ * 1: rowid of the row/vector to lookup
3565+ * Result columns:
3566+ * 0: chunk_id (i64)
3567+ * 1: chunk_offset (i64)
3568+ * SQL: "SELECT id, chunk_id, chunk_offset FROM _rowids WHERE rowid = ?""
3569+ *
3570+ * Must be cleaned up with sqlite3_finalize().
3571+ */
3572+ sqlite3_stmt *stmtRowidsGetChunkPosition;
3573+};
3574+
3575+/**
3576+ * @brief Finalize all the sqlite3_stmt members in a vec0_vtab.
3577+ *
3578+ * @param p vec0_vtab pointer
3579+ */
3580+void vec0_free_resources(vec0_vtab *p) {
3581+ sqlite3_finalize(p->stmtLatestChunk);
3582+ p->stmtLatestChunk = NULL;
3583+ sqlite3_finalize(p->stmtRowidsInsertRowid);
3584+ p->stmtRowidsInsertRowid = NULL;
3585+ sqlite3_finalize(p->stmtRowidsInsertId);
3586+ p->stmtRowidsInsertId = NULL;
3587+ sqlite3_finalize(p->stmtRowidsUpdatePosition);
3588+ p->stmtRowidsUpdatePosition = NULL;
3589+ sqlite3_finalize(p->stmtRowidsGetChunkPosition);
3590+ p->stmtRowidsGetChunkPosition = NULL;
3591+}
3592+
3593+/**
3594+ * @brief Free all memory and sqlite3_stmt members of a vec0_vtab
3595+ *
3596+ * @param p vec0_vtab pointer
3597+ */
3598+void vec0_free(vec0_vtab *p) {
3599+ vec0_free_resources(p);
3600+
3601+ sqlite3_free(p->schemaName);
3602+ p->schemaName = NULL;
3603+ sqlite3_free(p->tableName);
3604+ p->tableName = NULL;
3605+ sqlite3_free(p->shadowChunksName);
3606+ p->shadowChunksName = NULL;
3607+ sqlite3_free(p->shadowRowidsName);
3608+ p->shadowRowidsName = NULL;
3609+
3610+ for (int i = 0; i < p->numVectorColumns; i++) {
3611+ sqlite3_free(p->shadowVectorChunksNames[i]);
3612+ p->shadowVectorChunksNames[i] = NULL;
3613+
3614+ sqlite3_free(p->vector_columns[i].name);
3615+ p->vector_columns[i].name = NULL;
3616+ }
3617+}
3618+
3619+int vec0_num_defined_user_columns(vec0_vtab *p) {
3620+ return p->numVectorColumns + p->numPartitionColumns + p->numAuxiliaryColumns + p->numMetadataColumns;
3621+}
3622+
3623+/**
3624+ * @brief Returns the index of the distance hidden column for the given vec0
3625+ * table.
3626+ *
3627+ * @param p vec0 table
3628+ * @return int
3629+ */
3630+int vec0_column_distance_idx(vec0_vtab *p) {
3631+ return VEC0_COLUMN_USERN_START + (vec0_num_defined_user_columns(p) - 1) +
3632+ VEC0_COLUMN_OFFSET_DISTANCE;
3633+}
3634+
3635+/**
3636+ * @brief Returns the index of the k hidden column for the given vec0 table.
3637+ *
3638+ * @param p vec0 table
3639+ * @return int k column index
3640+ */
3641+int vec0_column_k_idx(vec0_vtab *p) {
3642+ return VEC0_COLUMN_USERN_START + (vec0_num_defined_user_columns(p) - 1) +
3643+ VEC0_COLUMN_OFFSET_K;
3644+}
3645+
3646+/**
3647+ * Returns 1 if the given column-based index is a valid vector column,
3648+ * 0 otherwise.
3649+ */
3650+int vec0_column_idx_is_vector(vec0_vtab *pVtab, int column_idx) {
3651+ return column_idx >= VEC0_COLUMN_USERN_START &&
3652+ column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) &&
3653+ pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_VECTOR;
3654+}
3655+
3656+/**
3657+ * Returns the vector index of the given user column index.
3658+ * ONLY call if validated with vec0_column_idx_is_vector before
3659+ */
3660+int vec0_column_idx_to_vector_idx(vec0_vtab *pVtab, int column_idx) {
3661+ UNUSED_PARAMETER(pVtab);
3662+ return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START];
3663+}
3664+/**
3665+ * Returns 1 if the given column-based index is a "partition key" column,
3666+ * 0 otherwise.
3667+ */
3668+int vec0_column_idx_is_partition(vec0_vtab *pVtab, int column_idx) {
3669+ return column_idx >= VEC0_COLUMN_USERN_START &&
3670+ column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) &&
3671+ pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_PARTITION;
3672+}
3673+
3674+/**
3675+ * Returns the partition column index of the given user column index.
3676+ * ONLY call if validated with vec0_column_idx_is_vector before
3677+ */
3678+int vec0_column_idx_to_partition_idx(vec0_vtab *pVtab, int column_idx) {
3679+ UNUSED_PARAMETER(pVtab);
3680+ return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START];
3681+}
3682+
3683+/**
3684+ * Returns 1 if the given column-based index is a auxiliary column,
3685+ * 0 otherwise.
3686+ */
3687+int vec0_column_idx_is_auxiliary(vec0_vtab *pVtab, int column_idx) {
3688+ return column_idx >= VEC0_COLUMN_USERN_START &&
3689+ column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) &&
3690+ pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY;
3691+}
3692+
3693+/**
3694+ * Returns the auxiliary column index of the given user column index.
3695+ * ONLY call if validated with vec0_column_idx_to_partition_idx before
3696+ */
3697+int vec0_column_idx_to_auxiliary_idx(vec0_vtab *pVtab, int column_idx) {
3698+ UNUSED_PARAMETER(pVtab);
3699+ return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START];
3700+}
3701+
3702+/**
3703+ * Returns 1 if the given column-based index is a metadata column,
3704+ * 0 otherwise.
3705+ */
3706+int vec0_column_idx_is_metadata(vec0_vtab *pVtab, int column_idx) {
3707+ return column_idx >= VEC0_COLUMN_USERN_START &&
3708+ column_idx <= (VEC0_COLUMN_USERN_START + vec0_num_defined_user_columns(pVtab) - 1) &&
3709+ pVtab->user_column_kinds[column_idx - VEC0_COLUMN_USERN_START] == SQLITE_VEC0_USER_COLUMN_KIND_METADATA;
3710+}
3711+
3712+/**
3713+ * Returns the metadata column index of the given user column index.
3714+ * ONLY call if validated with vec0_column_idx_is_metadata before
3715+ */
3716+int vec0_column_idx_to_metadata_idx(vec0_vtab *pVtab, int column_idx) {
3717+ UNUSED_PARAMETER(pVtab);
3718+ return pVtab->user_column_idxs[column_idx - VEC0_COLUMN_USERN_START];
3719+}
3720+
3721+/**
3722+ * @brief Retrieve the chunk_id, chunk_offset, and possible "id" value
3723+ * of a vec0_vtab row with the provided rowid
3724+ *
3725+ * @param p vec0_vtab
3726+ * @param rowid the rowid of the row to query
3727+ * @param id output, optional sqlite3_value to provide the id.
3728+ * Useful for text PK rows. Must be freed with sqlite3_value_free()
3729+ * @param chunk_id output, the chunk_id the row belongs to
3730+ * @param chunk_offset output, the offset within the chunk the row belongs to
3731+ * @return SQLITE_ROW on success, error code otherwise. SQLITE_EMPTY if row DNE
3732+ */
3733+int vec0_get_chunk_position(vec0_vtab *p, i64 rowid, sqlite3_value **id,
3734+ i64 *chunk_id, i64 *chunk_offset) {
3735+ int rc;
3736+
3737+ if (!p->stmtRowidsGetChunkPosition) {
3738+ const char *zSql =
3739+ sqlite3_mprintf("SELECT id, chunk_id, chunk_offset "
3740+ "FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE rowid = ?",
3741+ p->schemaName, p->tableName);
3742+ if (!zSql) {
3743+ rc = SQLITE_NOMEM;
3744+ goto cleanup;
3745+ }
3746+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsGetChunkPosition, 0);
3747+ sqlite3_free((void *)zSql);
3748+ if (rc != SQLITE_OK) {
3749+ vtab_set_error(
3750+ &p->base, VEC_INTERAL_ERROR
3751+ "could not initialize 'rowids get chunk position' statement");
3752+ goto cleanup;
3753+ }
3754+ }
3755+
3756+ sqlite3_bind_int64(p->stmtRowidsGetChunkPosition, 1, rowid);
3757+ rc = sqlite3_step(p->stmtRowidsGetChunkPosition);
3758+ // special case: when no results, return SQLITE_EMPTY to convey "that chunk
3759+ // position doesnt exist"
3760+ if (rc == SQLITE_DONE) {
3761+ rc = SQLITE_EMPTY;
3762+ goto cleanup;
3763+ }
3764+ if (rc != SQLITE_ROW) {
3765+ goto cleanup;
3766+ }
3767+
3768+ if (id) {
3769+ sqlite3_value *value =
3770+ sqlite3_column_value(p->stmtRowidsGetChunkPosition, 0);
3771+ *id = sqlite3_value_dup(value);
3772+ if (!*id) {
3773+ rc = SQLITE_NOMEM;
3774+ goto cleanup;
3775+ }
3776+ }
3777+
3778+ if (chunk_id) {
3779+ *chunk_id = sqlite3_column_int64(p->stmtRowidsGetChunkPosition, 1);
3780+ }
3781+ if (chunk_offset) {
3782+ *chunk_offset = sqlite3_column_int64(p->stmtRowidsGetChunkPosition, 2);
3783+ }
3784+
3785+ rc = SQLITE_OK;
3786+
3787+cleanup:
3788+ sqlite3_reset(p->stmtRowidsGetChunkPosition);
3789+ sqlite3_clear_bindings(p->stmtRowidsGetChunkPosition);
3790+ return rc;
3791+}
3792+
3793+/**
3794+ * @brief Return the id value from the _rowids table where _rowids.rowid =
3795+ * rowid.
3796+ *
3797+ * @param pVtab: vec0 table to query
3798+ * @param rowid: rowid of the row to query.
3799+ * @param out: A dup'ed sqlite3_value of the id column. Might be null.
3800+ * Must be cleaned up with sqlite3_value_free().
3801+ * @returns SQLITE_OK on success, error code on failure
3802+ */
3803+int vec0_get_id_value_from_rowid(vec0_vtab *pVtab, i64 rowid,
3804+ sqlite3_value **out) {
3805+ // PERF: different strategy than get_chunk_position?
3806+ return vec0_get_chunk_position((vec0_vtab *)pVtab, rowid, out, NULL, NULL);
3807+}
3808+
3809+int vec0_rowid_from_id(vec0_vtab *p, sqlite3_value *valueId, i64 *rowid) {
3810+ sqlite3_stmt *stmt = NULL;
3811+ int rc;
3812+ char *zSql;
3813+ zSql = sqlite3_mprintf("SELECT rowid"
3814+ " FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE id = ?",
3815+ p->schemaName, p->tableName);
3816+ if (!zSql) {
3817+ rc = SQLITE_NOMEM;
3818+ goto cleanup;
3819+ }
3820+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
3821+ sqlite3_free(zSql);
3822+ if (rc != SQLITE_OK) {
3823+ goto cleanup;
3824+ }
3825+ sqlite3_bind_value(stmt, 1, valueId);
3826+ rc = sqlite3_step(stmt);
3827+ if (rc == SQLITE_DONE) {
3828+ rc = SQLITE_EMPTY;
3829+ goto cleanup;
3830+ }
3831+ if (rc != SQLITE_ROW) {
3832+ goto cleanup;
3833+ }
3834+ *rowid = sqlite3_column_int64(stmt, 0);
3835+ rc = sqlite3_step(stmt);
3836+ if (rc != SQLITE_DONE) {
3837+ goto cleanup;
3838+ }
3839+
3840+ rc = SQLITE_OK;
3841+
3842+cleanup:
3843+ sqlite3_finalize(stmt);
3844+ return rc;
3845+}
3846+
3847+int vec0_result_id(vec0_vtab *p, sqlite3_context *context, i64 rowid) {
3848+ if (!p->pkIsText) {
3849+ sqlite3_result_int64(context, rowid);
3850+ return SQLITE_OK;
3851+ }
3852+ sqlite3_value *valueId;
3853+ int rc = vec0_get_id_value_from_rowid(p, rowid, &valueId);
3854+ if (rc != SQLITE_OK) {
3855+ return rc;
3856+ }
3857+ if (!valueId) {
3858+ sqlite3_result_error_nomem(context);
3859+ } else {
3860+ sqlite3_result_value(context, valueId);
3861+ sqlite3_value_free(valueId);
3862+ }
3863+ return SQLITE_OK;
3864+}
3865+
3866+/**
3867+ * @brief
3868+ *
3869+ * @param pVtab: virtual table to query
3870+ * @param rowid: row to lookup
3871+ * @param vector_column_idx: which vector column to query
3872+ * @param outVector: Output pointer to the vector buffer.
3873+ * Must be sqlite3_free()'ed.
3874+ * @param outVectorSize: Pointer to a int where the size of outVector
3875+ * will be stored.
3876+ * @return int SQLITE_OK on success.
3877+ */
3878+int vec0_get_vector_data(vec0_vtab *pVtab, i64 rowid, int vector_column_idx,
3879+ void **outVector, int *outVectorSize) {
3880+ vec0_vtab *p = pVtab;
3881+ int rc, brc;
3882+ i64 chunk_id;
3883+ i64 chunk_offset;
3884+ size_t size;
3885+ void *buf = NULL;
3886+ int blobOffset;
3887+ sqlite3_blob *vectorBlob = NULL;
3888+ assert((vector_column_idx >= 0) &&
3889+ (vector_column_idx < pVtab->numVectorColumns));
3890+
3891+ rc = vec0_get_chunk_position(pVtab, rowid, NULL, &chunk_id, &chunk_offset);
3892+ if (rc == SQLITE_EMPTY) {
3893+ vtab_set_error(&pVtab->base, "Could not find a row with rowid %lld", rowid);
3894+ goto cleanup;
3895+ }
3896+ if (rc != SQLITE_OK) {
3897+ goto cleanup;
3898+ }
3899+
3900+ rc = sqlite3_blob_open(p->db, p->schemaName,
3901+ p->shadowVectorChunksNames[vector_column_idx],
3902+ "vectors", chunk_id, 0, &vectorBlob);
3903+
3904+ if (rc != SQLITE_OK) {
3905+ vtab_set_error(&pVtab->base,
3906+ "Could not fetch vector data for %lld, opening blob failed",
3907+ rowid);
3908+ rc = SQLITE_ERROR;
3909+ goto cleanup;
3910+ }
3911+
3912+ size = vector_column_byte_size(pVtab->vector_columns[vector_column_idx]);
3913+ blobOffset = chunk_offset * size;
3914+
3915+ buf = sqlite3_malloc(size);
3916+ if (!buf) {
3917+ rc = SQLITE_NOMEM;
3918+ goto cleanup;
3919+ }
3920+
3921+ rc = sqlite3_blob_read(vectorBlob, buf, size, blobOffset);
3922+ if (rc != SQLITE_OK) {
3923+ sqlite3_free(buf);
3924+ buf = NULL;
3925+ vtab_set_error(
3926+ &pVtab->base,
3927+ "Could not fetch vector data for %lld, reading from blob failed",
3928+ rowid);
3929+ rc = SQLITE_ERROR;
3930+ goto cleanup;
3931+ }
3932+
3933+ *outVector = buf;
3934+ if (outVectorSize) {
3935+ *outVectorSize = size;
3936+ }
3937+ rc = SQLITE_OK;
3938+
3939+cleanup:
3940+ brc = sqlite3_blob_close(vectorBlob);
3941+ if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) {
3942+ vtab_set_error(
3943+ &p->base, VEC_INTERAL_ERROR
3944+ "unknown error, could not close vector blob, please file an issue");
3945+ return brc;
3946+ }
3947+
3948+ return rc;
3949+}
3950+
3951+/**
3952+ * @brief Retrieve the sqlite3_value of the i'th partition value for the given row.
3953+ *
3954+ * @param pVtab - the vec0_vtab in questions
3955+ * @param rowid - rowid of target row
3956+ * @param partition_idx - which partition column to retrieve
3957+ * @param outValue - output sqlite3_value
3958+ * @return int - SQLITE_OK on success, otherwise error code
3959+ */
3960+int vec0_get_partition_value_for_rowid(vec0_vtab *pVtab, i64 rowid, int partition_idx, sqlite3_value ** outValue) {
3961+ int rc;
3962+ i64 chunk_id;
3963+ i64 chunk_offset;
3964+ rc = vec0_get_chunk_position(pVtab, rowid, NULL, &chunk_id, &chunk_offset);
3965+ if(rc != SQLITE_OK) {
3966+ return rc;
3967+ }
3968+ sqlite3_stmt * stmt = NULL;
3969+ char * zSql = sqlite3_mprintf("SELECT partition%02d FROM " VEC0_SHADOW_CHUNKS_NAME " WHERE chunk_id = ?", partition_idx, pVtab->schemaName, pVtab->tableName);
3970+ if(!zSql) {
3971+ return SQLITE_NOMEM;
3972+ }
3973+ rc = sqlite3_prepare_v2(pVtab->db, zSql, -1, &stmt, NULL);
3974+ sqlite3_free(zSql);
3975+ if(rc != SQLITE_OK) {
3976+ return rc;
3977+ }
3978+ sqlite3_bind_int64(stmt, 1, chunk_id);
3979+ rc = sqlite3_step(stmt);
3980+ if(rc != SQLITE_ROW) {
3981+ rc = SQLITE_ERROR;
3982+ goto done;
3983+ }
3984+ *outValue = sqlite3_value_dup(sqlite3_column_value(stmt, 0));
3985+ if(!*outValue) {
3986+ rc = SQLITE_NOMEM;
3987+ goto done;
3988+ }
3989+ rc = SQLITE_OK;
3990+
3991+ done:
3992+ sqlite3_finalize(stmt);
3993+ return rc;
3994+
3995+}
3996+
3997+/**
3998+ * @brief Get the value of an auxiliary column for the given rowid
3999+ *
4000+ * @param pVtab vec0_vtab
4001+ * @param rowid the rowid of the row to lookup
4002+ * @param auxiliary_idx aux index of the column we care about
4003+ * @param outValue Output sqlite3_value to store
4004+ * @return int SQLITE_OK on success, error code otherwise
4005+ */
4006+int vec0_get_auxiliary_value_for_rowid(vec0_vtab *pVtab, i64 rowid, int auxiliary_idx, sqlite3_value ** outValue) {
4007+ int rc;
4008+ sqlite3_stmt * stmt = NULL;
4009+ char * zSql = sqlite3_mprintf("SELECT value%02d FROM " VEC0_SHADOW_AUXILIARY_NAME " WHERE rowid = ?", auxiliary_idx, pVtab->schemaName, pVtab->tableName);
4010+ if(!zSql) {
4011+ return SQLITE_NOMEM;
4012+ }
4013+ rc = sqlite3_prepare_v2(pVtab->db, zSql, -1, &stmt, NULL);
4014+ sqlite3_free(zSql);
4015+ if(rc != SQLITE_OK) {
4016+ return rc;
4017+ }
4018+ sqlite3_bind_int64(stmt, 1, rowid);
4019+ rc = sqlite3_step(stmt);
4020+ if(rc != SQLITE_ROW) {
4021+ rc = SQLITE_ERROR;
4022+ goto done;
4023+ }
4024+ *outValue = sqlite3_value_dup(sqlite3_column_value(stmt, 0));
4025+ if(!*outValue) {
4026+ rc = SQLITE_NOMEM;
4027+ goto done;
4028+ }
4029+ rc = SQLITE_OK;
4030+
4031+ done:
4032+ sqlite3_finalize(stmt);
4033+ return rc;
4034+}
4035+
4036+/**
4037+ * @brief Result the given metadata value for the given row and metadata column index.
4038+ * Will traverse the metadatachunksNN table with BLOB I/0 for the given rowid.
4039+ *
4040+ * @param p
4041+ * @param rowid
4042+ * @param metadata_idx
4043+ * @param context
4044+ * @return int
4045+ */
4046+int vec0_result_metadata_value_for_rowid(vec0_vtab *p, i64 rowid, int metadata_idx, sqlite3_context * context) {
4047+ int rc;
4048+ i64 chunk_id;
4049+ i64 chunk_offset;
4050+ rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset);
4051+ if(rc != SQLITE_OK) {
4052+ return rc;
4053+ }
4054+ sqlite3_blob * blobValue;
4055+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 0, &blobValue);
4056+ if(rc != SQLITE_OK) {
4057+ return rc;
4058+ }
4059+
4060+ switch(p->metadata_columns[metadata_idx].kind) {
4061+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
4062+ u8 block;
4063+ rc = sqlite3_blob_read(blobValue, &block, sizeof(block), chunk_offset / CHAR_BIT);
4064+ if(rc != SQLITE_OK) {
4065+ goto done;
4066+ }
4067+ int value = block >> ((chunk_offset % CHAR_BIT)) & 1;
4068+ sqlite3_result_int(context, value);
4069+ break;
4070+ }
4071+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
4072+ i64 value;
4073+ rc = sqlite3_blob_read(blobValue, &value, sizeof(value), chunk_offset * sizeof(i64));
4074+ if(rc != SQLITE_OK) {
4075+ goto done;
4076+ }
4077+ sqlite3_result_int64(context, value);
4078+ break;
4079+ }
4080+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
4081+ double value;
4082+ rc = sqlite3_blob_read(blobValue, &value, sizeof(value), chunk_offset * sizeof(double));
4083+ if(rc != SQLITE_OK) {
4084+ goto done;
4085+ }
4086+ sqlite3_result_double(context, value);
4087+ break;
4088+ }
4089+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
4090+ u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
4091+ rc = sqlite3_blob_read(blobValue, &view, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH, chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
4092+ if(rc != SQLITE_OK) {
4093+ goto done;
4094+ }
4095+ int length = ((int *)view)[0];
4096+ if(length <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
4097+ sqlite3_result_text(context, (const char*) (view + 4), length, SQLITE_TRANSIENT);
4098+ }
4099+ else {
4100+ sqlite3_stmt * stmt;
4101+ const char * zSql = sqlite3_mprintf("SELECT data FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_idx);
4102+ if(!zSql) {
4103+ rc = SQLITE_ERROR;
4104+ goto done;
4105+ }
4106+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
4107+ sqlite3_free((void *) zSql);
4108+ if(rc != SQLITE_OK) {
4109+ goto done;
4110+ }
4111+ sqlite3_bind_int64(stmt, 1, rowid);
4112+ rc = sqlite3_step(stmt);
4113+ if(rc != SQLITE_ROW) {
4114+ sqlite3_finalize(stmt);
4115+ rc = SQLITE_ERROR;
4116+ goto done;
4117+ }
4118+ sqlite3_result_value(context, sqlite3_column_value(stmt, 0));
4119+ sqlite3_finalize(stmt);
4120+ rc = SQLITE_OK;
4121+ }
4122+ break;
4123+ }
4124+ }
4125+ done:
4126+ // blobValue is read-only, will not fail on close
4127+ sqlite3_blob_close(blobValue);
4128+ return rc;
4129+
4130+}
4131+
4132+int vec0_get_latest_chunk_rowid(vec0_vtab *p, i64 *chunk_rowid, sqlite3_value ** partitionKeyValues) {
4133+ int rc;
4134+ const char *zSql;
4135+ // lazy initialize stmtLatestChunk when needed. May be cleared during xSync()
4136+ if (!p->stmtLatestChunk) {
4137+ if(p->numPartitionColumns > 0) {
4138+ sqlite3_str * s = sqlite3_str_new(NULL);
4139+ sqlite3_str_appendf(s, "SELECT max(rowid) FROM " VEC0_SHADOW_CHUNKS_NAME " WHERE ",
4140+ p->schemaName, p->tableName);
4141+
4142+ for(int i = 0; i < p->numPartitionColumns; i++) {
4143+ if(i != 0) {
4144+ sqlite3_str_appendall(s, " AND ");
4145+ }
4146+ sqlite3_str_appendf(s, " partition%02d = ? ", i);
4147+ }
4148+ zSql = sqlite3_str_finish(s);
4149+ }else {
4150+ zSql = sqlite3_mprintf("SELECT max(rowid) FROM " VEC0_SHADOW_CHUNKS_NAME,
4151+ p->schemaName, p->tableName);
4152+ }
4153+
4154+ if (!zSql) {
4155+ rc = SQLITE_NOMEM;
4156+ goto cleanup;
4157+ }
4158+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtLatestChunk, 0);
4159+ sqlite3_free((void *)zSql);
4160+ if (rc != SQLITE_OK) {
4161+ // IMP: V21406_05476
4162+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
4163+ "could not initialize 'latest chunk' statement");
4164+ goto cleanup;
4165+ }
4166+ }
4167+
4168+ for(int i = 0; i < p->numPartitionColumns; i++) {
4169+ sqlite3_bind_value(p->stmtLatestChunk, i+1, (partitionKeyValues[i]));
4170+ }
4171+
4172+ rc = sqlite3_step(p->stmtLatestChunk);
4173+ if (rc != SQLITE_ROW) {
4174+ // IMP: V31559_15629
4175+ vtab_set_error(&p->base, VEC_INTERAL_ERROR "Could not find latest chunk");
4176+ rc = SQLITE_ERROR;
4177+ goto cleanup;
4178+ }
4179+ if(sqlite3_column_type(p->stmtLatestChunk, 0) == SQLITE_NULL){
4180+ rc = SQLITE_EMPTY;
4181+ goto cleanup;
4182+ }
4183+ *chunk_rowid = sqlite3_column_int64(p->stmtLatestChunk, 0);
4184+ rc = sqlite3_step(p->stmtLatestChunk);
4185+ if (rc != SQLITE_DONE) {
4186+ vtab_set_error(&p->base,
4187+ VEC_INTERAL_ERROR
4188+ "unknown result code when closing out stmtLatestChunk. "
4189+ "Please file an issue: " REPORT_URL,
4190+ p->schemaName, p->shadowChunksName);
4191+ goto cleanup;
4192+ }
4193+ rc = SQLITE_OK;
4194+
4195+cleanup:
4196+ if (p->stmtLatestChunk) {
4197+ sqlite3_reset(p->stmtLatestChunk);
4198+ sqlite3_clear_bindings(p->stmtLatestChunk);
4199+ }
4200+ return rc;
4201+}
4202+
4203+int vec0_rowids_insert_rowid(vec0_vtab *p, i64 rowid) {
4204+ int rc = SQLITE_OK;
4205+ int entered = 0;
4206+ UNUSED_PARAMETER(entered); // temporary
4207+ if (!p->stmtRowidsInsertRowid) {
4208+ const char *zSql =
4209+ sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_ROWIDS_NAME "(rowid)"
4210+ "VALUES (?);",
4211+ p->schemaName, p->tableName);
4212+ if (!zSql) {
4213+ rc = SQLITE_NOMEM;
4214+ goto cleanup;
4215+ }
4216+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsInsertRowid, 0);
4217+ sqlite3_free((void *)zSql);
4218+ if (rc != SQLITE_OK) {
4219+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
4220+ "could not initialize 'insert rowids' statement");
4221+ goto cleanup;
4222+ }
4223+ }
4224+
4225+#if SQLITE_THREADSAFE
4226+ if (sqlite3_mutex_enter) {
4227+ sqlite3_mutex_enter(sqlite3_db_mutex(p->db));
4228+ entered = 1;
4229+ }
4230+#endif
4231+ sqlite3_bind_int64(p->stmtRowidsInsertRowid, 1, rowid);
4232+ rc = sqlite3_step(p->stmtRowidsInsertRowid);
4233+
4234+ if (rc != SQLITE_DONE) {
4235+ if (sqlite3_extended_errcode(p->db) == SQLITE_CONSTRAINT_PRIMARYKEY) {
4236+ // IMP: V17090_01160
4237+ vtab_set_error(&p->base, "UNIQUE constraint failed on %s primary key",
4238+ p->tableName);
4239+ } else {
4240+ // IMP: V04679_21517
4241+ vtab_set_error(&p->base,
4242+ "Error inserting rowid into rowids shadow table: %s",
4243+ sqlite3_errmsg(sqlite3_db_handle(p->stmtRowidsInsertId)));
4244+ }
4245+ rc = SQLITE_ERROR;
4246+ goto cleanup;
4247+ }
4248+
4249+ rc = SQLITE_OK;
4250+
4251+cleanup:
4252+ if (p->stmtRowidsInsertRowid) {
4253+ sqlite3_reset(p->stmtRowidsInsertRowid);
4254+ sqlite3_clear_bindings(p->stmtRowidsInsertRowid);
4255+ }
4256+
4257+#if SQLITE_THREADSAFE
4258+ if (sqlite3_mutex_leave && entered) {
4259+ sqlite3_mutex_leave(sqlite3_db_mutex(p->db));
4260+ }
4261+#endif
4262+ return rc;
4263+}
4264+
4265+int vec0_rowids_insert_id(vec0_vtab *p, sqlite3_value *idValue, i64 *rowid) {
4266+ int rc = SQLITE_OK;
4267+ int entered = 0;
4268+ UNUSED_PARAMETER(entered); // temporary
4269+ if (!p->stmtRowidsInsertId) {
4270+ const char *zSql =
4271+ sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_ROWIDS_NAME "(id)"
4272+ "VALUES (?);",
4273+ p->schemaName, p->tableName);
4274+ if (!zSql) {
4275+ rc = SQLITE_NOMEM;
4276+ goto complete;
4277+ }
4278+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsInsertId, 0);
4279+ sqlite3_free((void *)zSql);
4280+ if (rc != SQLITE_OK) {
4281+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
4282+ "could not initialize 'insert rowids id' statement");
4283+ goto complete;
4284+ }
4285+ }
4286+
4287+#if SQLITE_THREADSAFE
4288+ if (sqlite3_mutex_enter) {
4289+ sqlite3_mutex_enter(sqlite3_db_mutex(p->db));
4290+ entered = 1;
4291+ }
4292+#endif
4293+
4294+ if (idValue) {
4295+ sqlite3_bind_value(p->stmtRowidsInsertId, 1, idValue);
4296+ }
4297+ rc = sqlite3_step(p->stmtRowidsInsertId);
4298+
4299+ if (rc != SQLITE_DONE) {
4300+ if (sqlite3_extended_errcode(p->db) == SQLITE_CONSTRAINT_UNIQUE) {
4301+ // IMP: V20497_04568
4302+ vtab_set_error(&p->base, "UNIQUE constraint failed on %s primary key",
4303+ p->tableName);
4304+ } else {
4305+ // IMP: V24016_08086
4306+ // IMP: V15177_32015
4307+ vtab_set_error(&p->base,
4308+ "Error inserting id into rowids shadow table: %s",
4309+ sqlite3_errmsg(sqlite3_db_handle(p->stmtRowidsInsertId)));
4310+ }
4311+ rc = SQLITE_ERROR;
4312+ goto complete;
4313+ }
4314+
4315+ *rowid = sqlite3_last_insert_rowid(p->db);
4316+ rc = SQLITE_OK;
4317+
4318+complete:
4319+ if (p->stmtRowidsInsertId) {
4320+ sqlite3_reset(p->stmtRowidsInsertId);
4321+ sqlite3_clear_bindings(p->stmtRowidsInsertId);
4322+ }
4323+
4324+#if SQLITE_THREADSAFE
4325+ if (sqlite3_mutex_leave && entered) {
4326+ sqlite3_mutex_leave(sqlite3_db_mutex(p->db));
4327+ }
4328+#endif
4329+ return rc;
4330+}
4331+
4332+int vec0_metadata_chunk_size(vec0_metadata_column_kind kind, int chunk_size) {
4333+ switch(kind) {
4334+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN:
4335+ return chunk_size / 8;
4336+ case VEC0_METADATA_COLUMN_KIND_INTEGER:
4337+ return chunk_size * sizeof(i64);
4338+ case VEC0_METADATA_COLUMN_KIND_FLOAT:
4339+ return chunk_size * sizeof(double);
4340+ case VEC0_METADATA_COLUMN_KIND_TEXT:
4341+ return chunk_size * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH;
4342+ }
4343+ return 0;
4344+}
4345+
4346+int vec0_rowids_update_position(vec0_vtab *p, i64 rowid, i64 chunk_rowid,
4347+ i64 chunk_offset) {
4348+ int rc = SQLITE_OK;
4349+
4350+ if (!p->stmtRowidsUpdatePosition) {
4351+ const char *zSql = sqlite3_mprintf(" UPDATE " VEC0_SHADOW_ROWIDS_NAME
4352+ " SET chunk_id = ?, chunk_offset = ?"
4353+ " WHERE rowid = ?",
4354+ p->schemaName, p->tableName);
4355+ if (!zSql) {
4356+ rc = SQLITE_NOMEM;
4357+ goto cleanup;
4358+ }
4359+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &p->stmtRowidsUpdatePosition, 0);
4360+ sqlite3_free((void *)zSql);
4361+ if (rc != SQLITE_OK) {
4362+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
4363+ "could not initialize 'update rowids position' statement");
4364+ goto cleanup;
4365+ }
4366+ }
4367+
4368+ sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 1, chunk_rowid);
4369+ sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 2, chunk_offset);
4370+ sqlite3_bind_int64(p->stmtRowidsUpdatePosition, 3, rowid);
4371+
4372+ rc = sqlite3_step(p->stmtRowidsUpdatePosition);
4373+ if (rc != SQLITE_DONE) {
4374+ // IMP: V21925_05995
4375+ vtab_set_error(&p->base,
4376+ VEC_INTERAL_ERROR
4377+ "could not update rowids position for rowid=%lld, "
4378+ "chunk_rowid=%lld, chunk_offset=%lld",
4379+ rowid, chunk_rowid, chunk_offset);
4380+ rc = SQLITE_ERROR;
4381+ goto cleanup;
4382+ }
4383+ rc = SQLITE_OK;
4384+
4385+cleanup:
4386+ if (p->stmtRowidsUpdatePosition) {
4387+ sqlite3_reset(p->stmtRowidsUpdatePosition);
4388+ sqlite3_clear_bindings(p->stmtRowidsUpdatePosition);
4389+ }
4390+
4391+ return rc;
4392+}
4393+
4394+/**
4395+ * @brief Adds a new chunk for the vec0 table, and the corresponding vector
4396+ * chunks.
4397+ *
4398+ * Inserts a new row into the _chunks table, with blank data, and uses that new
4399+ * rowid to insert new blank rows into _vector_chunksXX tables.
4400+ *
4401+ * @param p: vec0 table to add new chunk
4402+ * @param paritionKeyValues: Array of partition key valeus for the new chunk, if available
4403+ * @param chunk_rowid: Output pointer, if not NULL, then will be filled with the
4404+ * new chunk rowid.
4405+ * @return int SQLITE_OK on success, error code otherwise.
4406+ */
4407+int vec0_new_chunk(vec0_vtab *p, sqlite3_value ** partitionKeyValues, i64 *chunk_rowid) {
4408+ int rc;
4409+ char *zSql;
4410+ sqlite3_stmt *stmt;
4411+ i64 rowid;
4412+
4413+ // Step 1: Insert a new row in _chunks, capture that new rowid
4414+ if(p->numPartitionColumns > 0) {
4415+ sqlite3_str * s = sqlite3_str_new(NULL);
4416+ sqlite3_str_appendf(s, "INSERT INTO " VEC0_SHADOW_CHUNKS_NAME, p->schemaName, p->tableName);
4417+ sqlite3_str_appendall(s, "(size, validity, rowids");
4418+ for(int i = 0; i < p->numPartitionColumns; i++) {
4419+ sqlite3_str_appendf(s, ", partition%02d", i);
4420+ }
4421+ sqlite3_str_appendall(s, ") VALUES (?, ?, ?");
4422+ for(int i = 0; i < p->numPartitionColumns; i++) {
4423+ sqlite3_str_appendall(s, ", ?");
4424+ }
4425+ sqlite3_str_appendall(s, ")");
4426+
4427+ zSql = sqlite3_str_finish(s);
4428+ }else {
4429+ zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_CHUNKS_NAME
4430+ "(size, validity, rowids) "
4431+ "VALUES (?, ?, ?);",
4432+ p->schemaName, p->tableName);
4433+ }
4434+
4435+ if (!zSql) {
4436+ return SQLITE_NOMEM;
4437+ }
4438+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
4439+ sqlite3_free(zSql);
4440+ if (rc != SQLITE_OK) {
4441+ sqlite3_finalize(stmt);
4442+ return rc;
4443+ }
4444+
4445+#if SQLITE_THREADSAFE
4446+ if (sqlite3_mutex_enter) {
4447+ sqlite3_mutex_enter(sqlite3_db_mutex(p->db));
4448+ }
4449+#endif
4450+
4451+ sqlite3_bind_int64(stmt, 1, p->chunk_size); // size
4452+ sqlite3_bind_zeroblob(stmt, 2, p->chunk_size / CHAR_BIT); // validity bitmap
4453+ sqlite3_bind_zeroblob(stmt, 3, p->chunk_size * sizeof(i64)); // rowids
4454+
4455+ for(int i = 0; i < p->numPartitionColumns; i++) {
4456+ sqlite3_bind_value(stmt, 4 + i, partitionKeyValues[i]);
4457+ }
4458+
4459+ rc = sqlite3_step(stmt);
4460+ int failed = rc != SQLITE_DONE;
4461+ rowid = sqlite3_last_insert_rowid(p->db);
4462+#if SQLITE_THREADSAFE
4463+ if (sqlite3_mutex_leave) {
4464+ sqlite3_mutex_leave(sqlite3_db_mutex(p->db));
4465+ }
4466+#endif
4467+ sqlite3_finalize(stmt);
4468+ if (failed) {
4469+ return SQLITE_ERROR;
4470+ }
4471+
4472+ // Step 2: Create new vector chunks for each vector column, with
4473+ // that new chunk_rowid.
4474+
4475+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
4476+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) {
4477+ continue;
4478+ }
4479+ int vector_column_idx = p->user_column_idxs[i];
4480+ i64 vectorsSize =
4481+ p->chunk_size * vector_column_byte_size(p->vector_columns[vector_column_idx]);
4482+
4483+ zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_VECTOR_N_NAME
4484+ "(rowid, vectors)"
4485+ "VALUES (?, ?)",
4486+ p->schemaName, p->tableName, vector_column_idx);
4487+ if (!zSql) {
4488+ return SQLITE_NOMEM;
4489+ }
4490+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
4491+ sqlite3_free(zSql);
4492+
4493+ if (rc != SQLITE_OK) {
4494+ sqlite3_finalize(stmt);
4495+ return rc;
4496+ }
4497+
4498+ sqlite3_bind_int64(stmt, 1, rowid);
4499+ sqlite3_bind_zeroblob64(stmt, 2, vectorsSize);
4500+
4501+ rc = sqlite3_step(stmt);
4502+ sqlite3_finalize(stmt);
4503+ if (rc != SQLITE_DONE) {
4504+ return rc;
4505+ }
4506+ }
4507+
4508+ // Step 3: Create new metadata chunks for each metadata column
4509+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
4510+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) {
4511+ continue;
4512+ }
4513+ int metadata_column_idx = p->user_column_idxs[i];
4514+ zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_METADATA_N_NAME
4515+ "(rowid, data)"
4516+ "VALUES (?, ?)",
4517+ p->schemaName, p->tableName, metadata_column_idx);
4518+ if (!zSql) {
4519+ return SQLITE_NOMEM;
4520+ }
4521+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
4522+ sqlite3_free(zSql);
4523+
4524+ if (rc != SQLITE_OK) {
4525+ sqlite3_finalize(stmt);
4526+ return rc;
4527+ }
4528+
4529+ sqlite3_bind_int64(stmt, 1, rowid);
4530+ sqlite3_bind_zeroblob64(stmt, 2, vec0_metadata_chunk_size(p->metadata_columns[metadata_column_idx].kind, p->chunk_size));
4531+
4532+ rc = sqlite3_step(stmt);
4533+ sqlite3_finalize(stmt);
4534+ if (rc != SQLITE_DONE) {
4535+ return rc;
4536+ }
4537+ }
4538+
4539+
4540+ if (chunk_rowid) {
4541+ *chunk_rowid = rowid;
4542+ }
4543+
4544+ return SQLITE_OK;
4545+}
4546+
4547+struct vec0_query_fullscan_data {
4548+ sqlite3_stmt *rowids_stmt;
4549+ i8 done;
4550+};
4551+void vec0_query_fullscan_data_clear(
4552+ struct vec0_query_fullscan_data *fullscan_data) {
4553+ if (!fullscan_data)
4554+ return;
4555+
4556+ if (fullscan_data->rowids_stmt) {
4557+ sqlite3_finalize(fullscan_data->rowids_stmt);
4558+ fullscan_data->rowids_stmt = NULL;
4559+ }
4560+}
4561+
4562+struct vec0_query_knn_data {
4563+ i64 k;
4564+ i64 k_used;
4565+ // Array of rowids of size k. Must be freed with sqlite3_free().
4566+ i64 *rowids;
4567+ // Array of distances of size k. Must be freed with sqlite3_free().
4568+ f32 *distances;
4569+ i64 current_idx;
4570+};
4571+void vec0_query_knn_data_clear(struct vec0_query_knn_data *knn_data) {
4572+ if (!knn_data)
4573+ return;
4574+
4575+ if (knn_data->rowids) {
4576+ sqlite3_free(knn_data->rowids);
4577+ knn_data->rowids = NULL;
4578+ }
4579+ if (knn_data->distances) {
4580+ sqlite3_free(knn_data->distances);
4581+ knn_data->distances = NULL;
4582+ }
4583+}
4584+
4585+struct vec0_query_point_data {
4586+ i64 rowid;
4587+ void *vectors[VEC0_MAX_VECTOR_COLUMNS];
4588+ int done;
4589+};
4590+void vec0_query_point_data_clear(struct vec0_query_point_data *point_data) {
4591+ if (!point_data)
4592+ return;
4593+ for (int i = 0; i < VEC0_MAX_VECTOR_COLUMNS; i++) {
4594+ sqlite3_free(point_data->vectors[i]);
4595+ point_data->vectors[i] = NULL;
4596+ }
4597+}
4598+
4599+typedef enum {
4600+ // If any values are updated, please update the ARCHITECTURE.md docs accordingly!
4601+
4602+ VEC0_QUERY_PLAN_FULLSCAN = '1',
4603+ VEC0_QUERY_PLAN_POINT = '2',
4604+ VEC0_QUERY_PLAN_KNN = '3',
4605+} vec0_query_plan;
4606+
4607+typedef struct vec0_cursor vec0_cursor;
4608+struct vec0_cursor {
4609+ sqlite3_vtab_cursor base;
4610+
4611+ vec0_query_plan query_plan;
4612+ struct vec0_query_fullscan_data *fullscan_data;
4613+ struct vec0_query_knn_data *knn_data;
4614+ struct vec0_query_point_data *point_data;
4615+};
4616+
4617+void vec0_cursor_clear(vec0_cursor *pCur) {
4618+ if (pCur->fullscan_data) {
4619+ vec0_query_fullscan_data_clear(pCur->fullscan_data);
4620+ sqlite3_free(pCur->fullscan_data);
4621+ pCur->fullscan_data = NULL;
4622+ }
4623+ if (pCur->knn_data) {
4624+ vec0_query_knn_data_clear(pCur->knn_data);
4625+ sqlite3_free(pCur->knn_data);
4626+ pCur->knn_data = NULL;
4627+ }
4628+ if (pCur->point_data) {
4629+ vec0_query_point_data_clear(pCur->point_data);
4630+ sqlite3_free(pCur->point_data);
4631+ pCur->point_data = NULL;
4632+ }
4633+}
4634+
4635+#define VEC_CONSTRUCTOR_ERROR "vec0 constructor error: "
4636+static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv,
4637+ sqlite3_vtab **ppVtab, char **pzErr, bool isCreate) {
4638+ UNUSED_PARAMETER(pAux);
4639+ vec0_vtab *pNew;
4640+ int rc;
4641+ const char *zSql;
4642+
4643+ pNew = sqlite3_malloc(sizeof(*pNew));
4644+ if (pNew == 0)
4645+ return SQLITE_NOMEM;
4646+ memset(pNew, 0, sizeof(*pNew));
4647+
4648+ // Declared chunk_size=N for entire table.
4649+ // -1 to use the defualt, otherwise will get re-assigned on `chunk_size=N`
4650+ // option
4651+ int chunk_size = -1;
4652+ int numVectorColumns = 0;
4653+ int numPartitionColumns = 0;
4654+ int numAuxiliaryColumns = 0;
4655+ int numMetadataColumns = 0;
4656+ int user_column_idx = 0;
4657+
4658+ // track if a "primary key" column is defined
4659+ char *pkColumnName = NULL;
4660+ int pkColumnNameLength;
4661+ int pkColumnType = SQLITE_INTEGER;
4662+
4663+ for (int i = 3; i < argc; i++) {
4664+ struct VectorColumnDefinition vecColumn;
4665+ struct Vec0PartitionColumnDefinition partitionColumn;
4666+ struct Vec0AuxiliaryColumnDefinition auxColumn;
4667+ struct Vec0MetadataColumnDefinition metadataColumn;
4668+ char *cName = NULL;
4669+ int cNameLength;
4670+ int cType;
4671+
4672+ // Scenario #1: Constructor argument is a vector column definition, ie `foo float[1024]`
4673+ rc = vec0_parse_vector_column(argv[i], strlen(argv[i]), &vecColumn);
4674+ if (rc == SQLITE_ERROR) {
4675+ *pzErr = sqlite3_mprintf(
4676+ VEC_CONSTRUCTOR_ERROR "could not parse vector column '%s'", argv[i]);
4677+ goto error;
4678+ }
4679+ if (rc == SQLITE_OK) {
4680+ if (numVectorColumns >= VEC0_MAX_VECTOR_COLUMNS) {
4681+ sqlite3_free(vecColumn.name);
4682+ *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4683+ "Too many provided vector columns, maximum %d",
4684+ VEC0_MAX_VECTOR_COLUMNS);
4685+ goto error;
4686+ }
4687+
4688+ if (vecColumn.dimensions > SQLITE_VEC_VEC0_MAX_DIMENSIONS) {
4689+ sqlite3_free(vecColumn.name);
4690+ *pzErr = sqlite3_mprintf(
4691+ VEC_CONSTRUCTOR_ERROR
4692+ "Dimension on vector column too large, provided %lld, maximum %lld",
4693+ (i64)vecColumn.dimensions, SQLITE_VEC_VEC0_MAX_DIMENSIONS);
4694+ goto error;
4695+ }
4696+ pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_VECTOR;
4697+ pNew->user_column_idxs[user_column_idx] = numVectorColumns;
4698+ memcpy(&pNew->vector_columns[numVectorColumns], &vecColumn, sizeof(vecColumn));
4699+ numVectorColumns++;
4700+ user_column_idx++;
4701+
4702+ continue;
4703+ }
4704+
4705+ // Scenario #2: Constructor argument is a partition key column definition, ie `user_id text partition key`
4706+ rc = vec0_parse_partition_key_definition(argv[i], strlen(argv[i]), &cName,
4707+ &cNameLength, &cType);
4708+ if (rc == SQLITE_OK) {
4709+ if (numPartitionColumns >= VEC0_MAX_PARTITION_COLUMNS) {
4710+ *pzErr = sqlite3_mprintf(
4711+ VEC_CONSTRUCTOR_ERROR
4712+ "More than %d partition key columns were provided",
4713+ VEC0_MAX_PARTITION_COLUMNS);
4714+ goto error;
4715+ }
4716+ partitionColumn.type = cType;
4717+ partitionColumn.name_length = cNameLength;
4718+ partitionColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName);
4719+ if(!partitionColumn.name) {
4720+ rc = SQLITE_NOMEM;
4721+ goto error;
4722+ }
4723+
4724+ pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_PARTITION;
4725+ pNew->user_column_idxs[user_column_idx] = numPartitionColumns;
4726+ memcpy(&pNew->paritition_columns[numPartitionColumns], &partitionColumn, sizeof(partitionColumn));
4727+ numPartitionColumns++;
4728+ user_column_idx++;
4729+ continue;
4730+ }
4731+
4732+ // Scenario #3: Constructor argument is a primary key column definition, ie `article_id text primary key`
4733+ rc = vec0_parse_primary_key_definition(argv[i], strlen(argv[i]), &cName,
4734+ &cNameLength, &cType);
4735+ if (rc == SQLITE_OK) {
4736+ if (pkColumnName) {
4737+ *pzErr = sqlite3_mprintf(
4738+ VEC_CONSTRUCTOR_ERROR
4739+ "More than one primary key definition was provided, vec0 only "
4740+ "suports a single primary key column",
4741+ argv[i]);
4742+ goto error;
4743+ }
4744+ pkColumnName = cName;
4745+ pkColumnNameLength = cNameLength;
4746+ pkColumnType = cType;
4747+ continue;
4748+ }
4749+
4750+ // Scenario #4: Constructor argument is a auxiliary column definition, ie `+contents text`
4751+ rc = vec0_parse_auxiliary_column_definition(argv[i], strlen(argv[i]), &cName,
4752+ &cNameLength, &cType);
4753+ if(rc == SQLITE_OK) {
4754+ if (numAuxiliaryColumns >= VEC0_MAX_AUXILIARY_COLUMNS) {
4755+ *pzErr = sqlite3_mprintf(
4756+ VEC_CONSTRUCTOR_ERROR
4757+ "More than %d auxiliary columns were provided",
4758+ VEC0_MAX_AUXILIARY_COLUMNS);
4759+ goto error;
4760+ }
4761+ auxColumn.type = cType;
4762+ auxColumn.name_length = cNameLength;
4763+ auxColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName);
4764+ if(!auxColumn.name) {
4765+ rc = SQLITE_NOMEM;
4766+ goto error;
4767+ }
4768+
4769+ pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY;
4770+ pNew->user_column_idxs[user_column_idx] = numAuxiliaryColumns;
4771+ memcpy(&pNew->auxiliary_columns[numAuxiliaryColumns], &auxColumn, sizeof(auxColumn));
4772+ numAuxiliaryColumns++;
4773+ user_column_idx++;
4774+ continue;
4775+ }
4776+
4777+ vec0_metadata_column_kind kind;
4778+ rc = vec0_parse_metadata_column_definition(argv[i], strlen(argv[i]), &cName,
4779+ &cNameLength, &kind);
4780+ if(rc == SQLITE_OK) {
4781+ if (numMetadataColumns >= VEC0_MAX_METADATA_COLUMNS) {
4782+ *pzErr = sqlite3_mprintf(
4783+ VEC_CONSTRUCTOR_ERROR
4784+ "More than %d metadata columns were provided",
4785+ VEC0_MAX_METADATA_COLUMNS);
4786+ goto error;
4787+ }
4788+ metadataColumn.kind = kind;
4789+ metadataColumn.name_length = cNameLength;
4790+ metadataColumn.name = sqlite3_mprintf("%.*s", cNameLength, cName);
4791+ if(!metadataColumn.name) {
4792+ rc = SQLITE_NOMEM;
4793+ goto error;
4794+ }
4795+
4796+ pNew->user_column_kinds[user_column_idx] = SQLITE_VEC0_USER_COLUMN_KIND_METADATA;
4797+ pNew->user_column_idxs[user_column_idx] = numMetadataColumns;
4798+ memcpy(&pNew->metadata_columns[numMetadataColumns], &metadataColumn, sizeof(metadataColumn));
4799+ numMetadataColumns++;
4800+ user_column_idx++;
4801+ continue;
4802+ }
4803+
4804+ // Scenario #4: Constructor argument is a table-level option, ie `chunk_size`
4805+
4806+ char *key;
4807+ char *value;
4808+ int keyLength, valueLength;
4809+ rc = vec0_parse_table_option(argv[i], strlen(argv[i]), &key, &keyLength,
4810+ &value, &valueLength);
4811+ if (rc == SQLITE_ERROR) {
4812+ *pzErr = sqlite3_mprintf(
4813+ VEC_CONSTRUCTOR_ERROR "could not parse table option '%s'", argv[i]);
4814+ goto error;
4815+ }
4816+ if (rc == SQLITE_OK) {
4817+ if (sqlite3_strnicmp(key, "chunk_size", keyLength) == 0) {
4818+ chunk_size = atoi(value);
4819+ if (chunk_size <= 0) {
4820+ // IMP: V01931_18769
4821+ *pzErr =
4822+ sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4823+ "chunk_size must be a non-zero positive integer");
4824+ goto error;
4825+ }
4826+ if ((chunk_size % 8) != 0) {
4827+ // IMP: V14110_30948
4828+ *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4829+ "chunk_size must be divisible by 8");
4830+ goto error;
4831+ }
4832+#define SQLITE_VEC_CHUNK_SIZE_MAX 4096
4833+ if (chunk_size > SQLITE_VEC_CHUNK_SIZE_MAX) {
4834+ *pzErr =
4835+ sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR "chunk_size too large");
4836+ goto error;
4837+ }
4838+ } else {
4839+ // IMP: V27642_11712
4840+ *pzErr = sqlite3_mprintf(
4841+ VEC_CONSTRUCTOR_ERROR "Unknown table option: %.*s", keyLength, key);
4842+ goto error;
4843+ }
4844+ continue;
4845+ }
4846+
4847+ // Scenario #5: Unknown constructor argument
4848+ *pzErr =
4849+ sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR "Could not parse '%s'", argv[i]);
4850+ goto error;
4851+ }
4852+
4853+ if (chunk_size < 0) {
4854+ chunk_size = 1024;
4855+ }
4856+
4857+ if (numVectorColumns <= 0) {
4858+ *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4859+ "At least one vector column is required");
4860+ goto error;
4861+ }
4862+
4863+ sqlite3_str *createStr = sqlite3_str_new(NULL);
4864+ sqlite3_str_appendall(createStr, "CREATE TABLE x(");
4865+ if (pkColumnName) {
4866+ sqlite3_str_appendf(createStr, "\"%.*w\" primary key, ", pkColumnNameLength,
4867+ pkColumnName);
4868+ } else {
4869+ sqlite3_str_appendall(createStr, "rowid, ");
4870+ }
4871+ for (int i = 0; i < numVectorColumns + numPartitionColumns + numAuxiliaryColumns + numMetadataColumns; i++) {
4872+ switch(pNew->user_column_kinds[i]) {
4873+ case SQLITE_VEC0_USER_COLUMN_KIND_VECTOR: {
4874+ int vector_idx = pNew->user_column_idxs[i];
4875+ sqlite3_str_appendf(createStr, "\"%.*w\", ",
4876+ pNew->vector_columns[vector_idx].name_length,
4877+ pNew->vector_columns[vector_idx].name);
4878+ break;
4879+ }
4880+ case SQLITE_VEC0_USER_COLUMN_KIND_PARTITION: {
4881+ int partition_idx = pNew->user_column_idxs[i];
4882+ sqlite3_str_appendf(createStr, "\"%.*w\", ",
4883+ pNew->paritition_columns[partition_idx].name_length,
4884+ pNew->paritition_columns[partition_idx].name);
4885+ break;
4886+ }
4887+ case SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY: {
4888+ int auxiliary_idx = pNew->user_column_idxs[i];
4889+ sqlite3_str_appendf(createStr, "\"%.*w\", ",
4890+ pNew->auxiliary_columns[auxiliary_idx].name_length,
4891+ pNew->auxiliary_columns[auxiliary_idx].name);
4892+ break;
4893+ }
4894+ case SQLITE_VEC0_USER_COLUMN_KIND_METADATA: {
4895+ int metadata_idx = pNew->user_column_idxs[i];
4896+ sqlite3_str_appendf(createStr, "\"%.*w\", ",
4897+ pNew->metadata_columns[metadata_idx].name_length,
4898+ pNew->metadata_columns[metadata_idx].name);
4899+ break;
4900+ }
4901+ }
4902+
4903+ }
4904+ sqlite3_str_appendall(createStr, " distance hidden, k hidden) ");
4905+ if (pkColumnName) {
4906+ sqlite3_str_appendall(createStr, "without rowid ");
4907+ }
4908+ zSql = sqlite3_str_finish(createStr);
4909+ if (!zSql) {
4910+ goto error;
4911+ }
4912+ rc = sqlite3_declare_vtab(db, zSql);
4913+ sqlite3_free((void *)zSql);
4914+ if (rc != SQLITE_OK) {
4915+ *pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
4916+ "could not declare virtual table, '%s'",
4917+ sqlite3_errmsg(db));
4918+ goto error;
4919+ }
4920+
4921+ const char *schemaName = argv[1];
4922+ const char *tableName = argv[2];
4923+
4924+ pNew->db = db;
4925+ pNew->pkIsText = pkColumnType == SQLITE_TEXT;
4926+ pNew->schemaName = sqlite3_mprintf("%s", schemaName);
4927+ if (!pNew->schemaName) {
4928+ goto error;
4929+ }
4930+ pNew->tableName = sqlite3_mprintf("%s", tableName);
4931+ if (!pNew->tableName) {
4932+ goto error;
4933+ }
4934+ pNew->shadowRowidsName = sqlite3_mprintf("%s_rowids", tableName);
4935+ if (!pNew->shadowRowidsName) {
4936+ goto error;
4937+ }
4938+ pNew->shadowChunksName = sqlite3_mprintf("%s_chunks", tableName);
4939+ if (!pNew->shadowChunksName) {
4940+ goto error;
4941+ }
4942+ pNew->numVectorColumns = numVectorColumns;
4943+ pNew->numPartitionColumns = numPartitionColumns;
4944+ pNew->numAuxiliaryColumns = numAuxiliaryColumns;
4945+ pNew->numMetadataColumns = numMetadataColumns;
4946+
4947+ for (int i = 0; i < pNew->numVectorColumns; i++) {
4948+ pNew->shadowVectorChunksNames[i] =
4949+ sqlite3_mprintf("%s_vector_chunks%02d", tableName, i);
4950+ if (!pNew->shadowVectorChunksNames[i]) {
4951+ goto error;
4952+ }
4953+ }
4954+ for (int i = 0; i < pNew->numMetadataColumns; i++) {
4955+ pNew->shadowMetadataChunksNames[i] =
4956+ sqlite3_mprintf("%s_metadatachunks%02d", tableName, i);
4957+ if (!pNew->shadowMetadataChunksNames[i]) {
4958+ goto error;
4959+ }
4960+ }
4961+ pNew->chunk_size = chunk_size;
4962+
4963+ // if xCreate, then create the necessary shadow tables
4964+ if (isCreate) {
4965+ sqlite3_stmt *stmt;
4966+ int rc;
4967+
4968+ char * zCreateInfo = sqlite3_mprintf("CREATE TABLE "VEC0_SHADOW_INFO_NAME " (key text primary key, value any)", pNew->schemaName, pNew->tableName);
4969+ if(!zCreateInfo) {
4970+ goto error;
4971+ }
4972+ rc = sqlite3_prepare_v2(db, zCreateInfo, -1, &stmt, NULL);
4973+
4974+ sqlite3_free((void *) zCreateInfo);
4975+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
4976+ // TODO(IMP)
4977+ sqlite3_finalize(stmt);
4978+ *pzErr = sqlite3_mprintf("Could not create '_info' shadow table: %s",
4979+ sqlite3_errmsg(db));
4980+ goto error;
4981+ }
4982+ sqlite3_finalize(stmt);
4983+
4984+ char * zSeedInfo = sqlite3_mprintf(
4985+ "INSERT INTO "VEC0_SHADOW_INFO_NAME "(key, value) VALUES "
4986+ "(?1, ?2), (?3, ?4), (?5, ?6), (?7, ?8) ",
4987+ pNew->schemaName, pNew->tableName
4988+ );
4989+ if(!zSeedInfo) {
4990+ goto error;
4991+ }
4992+ rc = sqlite3_prepare_v2(db, zSeedInfo, -1, &stmt, NULL);
4993+ sqlite3_free((void *) zSeedInfo);
4994+ if (rc != SQLITE_OK) {
4995+ // TODO(IMP)
4996+ sqlite3_finalize(stmt);
4997+ *pzErr = sqlite3_mprintf("Could not seed '_info' shadow table: %s",
4998+ sqlite3_errmsg(db));
4999+ goto error;
5000+ }
5001+ sqlite3_bind_text(stmt, 1, "CREATE_VERSION", -1, SQLITE_STATIC);
5002+ sqlite3_bind_text(stmt, 2, SQLITE_VEC_VERSION, -1, SQLITE_STATIC);
5003+ sqlite3_bind_text(stmt, 3, "CREATE_VERSION_MAJOR", -1, SQLITE_STATIC);
5004+ sqlite3_bind_int(stmt, 4, SQLITE_VEC_VERSION_MAJOR);
5005+ sqlite3_bind_text(stmt, 5, "CREATE_VERSION_MINOR", -1, SQLITE_STATIC);
5006+ sqlite3_bind_int(stmt, 6, SQLITE_VEC_VERSION_MINOR);
5007+ sqlite3_bind_text(stmt, 7, "CREATE_VERSION_PATCH", -1, SQLITE_STATIC);
5008+ sqlite3_bind_int(stmt, 8, SQLITE_VEC_VERSION_PATCH);
5009+
5010+ if(sqlite3_step(stmt) != SQLITE_DONE) {
5011+ // TODO(IMP)
5012+ sqlite3_finalize(stmt);
5013+ *pzErr = sqlite3_mprintf("Could not seed '_info' shadow table: %s",
5014+ sqlite3_errmsg(db));
5015+ goto error;
5016+ }
5017+ sqlite3_finalize(stmt);
5018+
5019+
5020+
5021+ // create the _chunks shadow table
5022+ char *zCreateShadowChunks = NULL;
5023+ if(pNew->numPartitionColumns) {
5024+ sqlite3_str * s = sqlite3_str_new(NULL);
5025+ sqlite3_str_appendf(s, "CREATE TABLE " VEC0_SHADOW_CHUNKS_NAME "(", pNew->schemaName, pNew->tableName);
5026+ sqlite3_str_appendall(s, "chunk_id INTEGER PRIMARY KEY AUTOINCREMENT," "size INTEGER NOT NULL,");
5027+ sqlite3_str_appendall(s, "sequence_id integer,");
5028+ for(int i = 0; i < pNew->numPartitionColumns;i++) {
5029+ sqlite3_str_appendf(s, "partition%02d,", i);
5030+ }
5031+ sqlite3_str_appendall(s, "validity BLOB NOT NULL, rowids BLOB NOT NULL);");
5032+ zCreateShadowChunks = sqlite3_str_finish(s);
5033+ }else {
5034+ zCreateShadowChunks = sqlite3_mprintf(VEC0_SHADOW_CHUNKS_CREATE,
5035+ pNew->schemaName, pNew->tableName);
5036+ }
5037+ if (!zCreateShadowChunks) {
5038+ goto error;
5039+ }
5040+ rc = sqlite3_prepare_v2(db, zCreateShadowChunks, -1, &stmt, 0);
5041+ sqlite3_free((void *)zCreateShadowChunks);
5042+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5043+ // IMP: V17740_01811
5044+ sqlite3_finalize(stmt);
5045+ *pzErr = sqlite3_mprintf("Could not create '_chunks' shadow table: %s",
5046+ sqlite3_errmsg(db));
5047+ goto error;
5048+ }
5049+ sqlite3_finalize(stmt);
5050+
5051+ // create the _rowids shadow table
5052+ char *zCreateShadowRowids;
5053+ if (pNew->pkIsText) {
5054+ // adds a "text unique not null" constraint to the id column
5055+ zCreateShadowRowids = sqlite3_mprintf(VEC0_SHADOW_ROWIDS_CREATE_PK_TEXT,
5056+ pNew->schemaName, pNew->tableName);
5057+ } else {
5058+ zCreateShadowRowids = sqlite3_mprintf(VEC0_SHADOW_ROWIDS_CREATE_BASIC,
5059+ pNew->schemaName, pNew->tableName);
5060+ }
5061+ if (!zCreateShadowRowids) {
5062+ goto error;
5063+ }
5064+ rc = sqlite3_prepare_v2(db, zCreateShadowRowids, -1, &stmt, 0);
5065+ sqlite3_free((void *)zCreateShadowRowids);
5066+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5067+ // IMP: V11631_28470
5068+ sqlite3_finalize(stmt);
5069+ *pzErr = sqlite3_mprintf("Could not create '_rowids' shadow table: %s",
5070+ sqlite3_errmsg(db));
5071+ goto error;
5072+ }
5073+ sqlite3_finalize(stmt);
5074+
5075+ for (int i = 0; i < pNew->numVectorColumns; i++) {
5076+ char *zSql = sqlite3_mprintf(VEC0_SHADOW_VECTOR_N_CREATE,
5077+ pNew->schemaName, pNew->tableName, i);
5078+ if (!zSql) {
5079+ goto error;
5080+ }
5081+ rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0);
5082+ sqlite3_free((void *)zSql);
5083+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5084+ // IMP: V25919_09989
5085+ sqlite3_finalize(stmt);
5086+ *pzErr = sqlite3_mprintf(
5087+ "Could not create '_vector_chunks%02d' shadow table: %s", i,
5088+ sqlite3_errmsg(db));
5089+ goto error;
5090+ }
5091+ sqlite3_finalize(stmt);
5092+ }
5093+
5094+ for (int i = 0; i < pNew->numMetadataColumns; i++) {
5095+ char *zSql = sqlite3_mprintf("CREATE TABLE " VEC0_SHADOW_METADATA_N_NAME "(rowid PRIMARY KEY, data BLOB NOT NULL);",
5096+ pNew->schemaName, pNew->tableName, i);
5097+ if (!zSql) {
5098+ goto error;
5099+ }
5100+ rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0);
5101+ sqlite3_free((void *)zSql);
5102+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5103+ sqlite3_finalize(stmt);
5104+ *pzErr = sqlite3_mprintf(
5105+ "Could not create '_metata_chunks%02d' shadow table: %s", i,
5106+ sqlite3_errmsg(db));
5107+ goto error;
5108+ }
5109+ sqlite3_finalize(stmt);
5110+
5111+ if(pNew->metadata_columns[i].kind == VEC0_METADATA_COLUMN_KIND_TEXT) {
5112+ char *zSql = sqlite3_mprintf("CREATE TABLE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME "(rowid PRIMARY KEY, data TEXT);",
5113+ pNew->schemaName, pNew->tableName, i);
5114+ if (!zSql) {
5115+ goto error;
5116+ }
5117+ rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, 0);
5118+ sqlite3_free((void *)zSql);
5119+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5120+ sqlite3_finalize(stmt);
5121+ *pzErr = sqlite3_mprintf(
5122+ "Could not create '_metadatatext%02d' shadow table: %s", i,
5123+ sqlite3_errmsg(db));
5124+ goto error;
5125+ }
5126+ sqlite3_finalize(stmt);
5127+
5128+ }
5129+ }
5130+
5131+ if(pNew->numAuxiliaryColumns > 0) {
5132+ sqlite3_stmt * stmt;
5133+ sqlite3_str * s = sqlite3_str_new(NULL);
5134+ sqlite3_str_appendf(s, "CREATE TABLE " VEC0_SHADOW_AUXILIARY_NAME "( rowid integer PRIMARY KEY ", pNew->schemaName, pNew->tableName);
5135+ for(int i = 0; i < pNew->numAuxiliaryColumns; i++) {
5136+ sqlite3_str_appendf(s, ", value%02d", i);
5137+ }
5138+ sqlite3_str_appendall(s, ")");
5139+ char *zSql = sqlite3_str_finish(s);
5140+ if(!zSql) {
5141+ goto error;
5142+ }
5143+ rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, NULL);
5144+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5145+ sqlite3_finalize(stmt);
5146+ *pzErr = sqlite3_mprintf(
5147+ "Could not create auxiliary shadow table: %s",
5148+ sqlite3_errmsg(db));
5149+
5150+ goto error;
5151+ }
5152+ sqlite3_finalize(stmt);
5153+ }
5154+ }
5155+
5156+ *ppVtab = (sqlite3_vtab *)pNew;
5157+ return SQLITE_OK;
5158+
5159+error:
5160+ vec0_free(pNew);
5161+ return SQLITE_ERROR;
5162+}
5163+
5164+static int vec0Create(sqlite3 *db, void *pAux, int argc,
5165+ const char *const *argv, sqlite3_vtab **ppVtab,
5166+ char **pzErr) {
5167+ return vec0_init(db, pAux, argc, argv, ppVtab, pzErr, true);
5168+}
5169+static int vec0Connect(sqlite3 *db, void *pAux, int argc,
5170+ const char *const *argv, sqlite3_vtab **ppVtab,
5171+ char **pzErr) {
5172+ return vec0_init(db, pAux, argc, argv, ppVtab, pzErr, false);
5173+}
5174+
5175+static int vec0Disconnect(sqlite3_vtab *pVtab) {
5176+ vec0_vtab *p = (vec0_vtab *)pVtab;
5177+ vec0_free(p);
5178+ sqlite3_free(p);
5179+ return SQLITE_OK;
5180+}
5181+static int vec0Destroy(sqlite3_vtab *pVtab) {
5182+ vec0_vtab *p = (vec0_vtab *)pVtab;
5183+ sqlite3_stmt *stmt;
5184+ int rc;
5185+ const char *zSql;
5186+
5187+ // Free up any sqlite3_stmt, otherwise DROPs on those tables will fail
5188+ vec0_free_resources(p);
5189+
5190+ // TODO(test) later: can't evidence-of here, bc always gives "SQL logic error" instead of
5191+ // provided error
5192+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_CHUNKS_NAME, p->schemaName,
5193+ p->tableName);
5194+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5195+ sqlite3_free((void *)zSql);
5196+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5197+ rc = SQLITE_ERROR;
5198+ vtab_set_error(pVtab, "could not drop chunks shadow table");
5199+ goto done;
5200+ }
5201+ sqlite3_finalize(stmt);
5202+
5203+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_INFO_NAME, p->schemaName,
5204+ p->tableName);
5205+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5206+ sqlite3_free((void *)zSql);
5207+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5208+ rc = SQLITE_ERROR;
5209+ vtab_set_error(pVtab, "could not drop info shadow table");
5210+ goto done;
5211+ }
5212+ sqlite3_finalize(stmt);
5213+
5214+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_ROWIDS_NAME, p->schemaName,
5215+ p->tableName);
5216+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5217+ sqlite3_free((void *)zSql);
5218+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5219+ rc = SQLITE_ERROR;
5220+ goto done;
5221+ }
5222+ sqlite3_finalize(stmt);
5223+
5224+ for (int i = 0; i < p->numVectorColumns; i++) {
5225+ zSql = sqlite3_mprintf("DROP TABLE \"%w\".\"%w\"", p->schemaName,
5226+ p->shadowVectorChunksNames[i]);
5227+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5228+ sqlite3_free((void *)zSql);
5229+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5230+ rc = SQLITE_ERROR;
5231+ goto done;
5232+ }
5233+ sqlite3_finalize(stmt);
5234+ }
5235+
5236+ if(p->numAuxiliaryColumns > 0) {
5237+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_AUXILIARY_NAME, p->schemaName, p->tableName);
5238+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5239+ sqlite3_free((void *)zSql);
5240+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5241+ rc = SQLITE_ERROR;
5242+ goto done;
5243+ }
5244+ sqlite3_finalize(stmt);
5245+ }
5246+
5247+
5248+ for (int i = 0; i < p->numMetadataColumns; i++) {
5249+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_METADATA_N_NAME, p->schemaName,p->tableName, i);
5250+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5251+ sqlite3_free((void *)zSql);
5252+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5253+ rc = SQLITE_ERROR;
5254+ goto done;
5255+ }
5256+ sqlite3_finalize(stmt);
5257+
5258+ if(p->metadata_columns[i].kind == VEC0_METADATA_COLUMN_KIND_TEXT) {
5259+ zSql = sqlite3_mprintf("DROP TABLE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME, p->schemaName,p->tableName, i);
5260+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, 0);
5261+ sqlite3_free((void *)zSql);
5262+ if ((rc != SQLITE_OK) || (sqlite3_step(stmt) != SQLITE_DONE)) {
5263+ rc = SQLITE_ERROR;
5264+ goto done;
5265+ }
5266+ sqlite3_finalize(stmt);
5267+ }
5268+ }
5269+
5270+ stmt = NULL;
5271+ rc = SQLITE_OK;
5272+
5273+done:
5274+ sqlite3_finalize(stmt);
5275+ vec0_free(p);
5276+ // If there was an error
5277+ if (rc == SQLITE_OK) {
5278+ sqlite3_free(p);
5279+ }
5280+ return rc;
5281+}
5282+
5283+static int vec0Open(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) {
5284+ UNUSED_PARAMETER(p);
5285+ vec0_cursor *pCur;
5286+ pCur = sqlite3_malloc(sizeof(*pCur));
5287+ if (pCur == 0)
5288+ return SQLITE_NOMEM;
5289+ memset(pCur, 0, sizeof(*pCur));
5290+ *ppCursor = &pCur->base;
5291+ return SQLITE_OK;
5292+}
5293+
5294+static int vec0Close(sqlite3_vtab_cursor *cur) {
5295+ vec0_cursor *pCur = (vec0_cursor *)cur;
5296+ vec0_cursor_clear(pCur);
5297+ sqlite3_free(pCur);
5298+ return SQLITE_OK;
5299+}
5300+
5301+// All the different type of "values" provided to argv/argc in vec0Filter.
5302+// These enums denote the use and purpose of all of them.
5303+typedef enum {
5304+ // If any values are updated, please update the ARCHITECTURE.md docs accordingly!
5305+
5306+ VEC0_IDXSTR_KIND_KNN_MATCH = '{',
5307+ VEC0_IDXSTR_KIND_KNN_K = '}',
5308+ VEC0_IDXSTR_KIND_KNN_ROWID_IN = '[',
5309+ VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT = ']',
5310+ VEC0_IDXSTR_KIND_POINT_ID = '!',
5311+ VEC0_IDXSTR_KIND_METADATA_CONSTRAINT = '&',
5312+} vec0_idxstr_kind;
5313+
5314+// The different SQLITE_INDEX_CONSTRAINT values that vec0 partition key columns
5315+// support, but as characters that fit nicely in idxstr.
5316+typedef enum {
5317+ // If any values are updated, please update the ARCHITECTURE.md docs accordingly!
5318+
5319+ VEC0_PARTITION_OPERATOR_EQ = 'a',
5320+ VEC0_PARTITION_OPERATOR_GT = 'b',
5321+ VEC0_PARTITION_OPERATOR_LE = 'c',
5322+ VEC0_PARTITION_OPERATOR_LT = 'd',
5323+ VEC0_PARTITION_OPERATOR_GE = 'e',
5324+ VEC0_PARTITION_OPERATOR_NE = 'f',
5325+} vec0_partition_operator;
5326+typedef enum {
5327+ VEC0_METADATA_OPERATOR_EQ = 'a',
5328+ VEC0_METADATA_OPERATOR_GT = 'b',
5329+ VEC0_METADATA_OPERATOR_LE = 'c',
5330+ VEC0_METADATA_OPERATOR_LT = 'd',
5331+ VEC0_METADATA_OPERATOR_GE = 'e',
5332+ VEC0_METADATA_OPERATOR_NE = 'f',
5333+ VEC0_METADATA_OPERATOR_IN = 'g',
5334+} vec0_metadata_operator;
5335+
5336+static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
5337+ vec0_vtab *p = (vec0_vtab *)pVTab;
5338+ /**
5339+ * Possible query plans are:
5340+ * 1. KNN when:
5341+ * a) An `MATCH` op on vector column
5342+ * b) ORDER BY on distance column
5343+ * c) LIMIT
5344+ * d) rowid in (...) OPTIONAL
5345+ * 2. Point when:
5346+ * a) An `EQ` op on rowid column
5347+ * 3. else: fullscan
5348+ *
5349+ */
5350+ int iMatchTerm = -1;
5351+ int iMatchVectorTerm = -1;
5352+ int iLimitTerm = -1;
5353+ int iRowidTerm = -1;
5354+ int iKTerm = -1;
5355+ int iRowidInTerm = -1;
5356+ int hasAuxConstraint = 0;
5357+
5358+#ifdef SQLITE_VEC_DEBUG
5359+ printf("pIdxInfo->nOrderBy=%d, pIdxInfo->nConstraint=%d\n", pIdxInfo->nOrderBy, pIdxInfo->nConstraint);
5360+#endif
5361+
5362+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
5363+ u8 vtabIn = 0;
5364+
5365+#if COMPILER_SUPPORTS_VTAB_IN
5366+ if (sqlite3_libversion_number() >= 3038000) {
5367+ vtabIn = sqlite3_vtab_in(pIdxInfo, i, -1);
5368+ }
5369+#endif
5370+
5371+#ifdef SQLITE_VEC_DEBUG
5372+ printf("xBestIndex [%d] usable=%d iColumn=%d op=%d vtabin=%d\n", i,
5373+ pIdxInfo->aConstraint[i].usable, pIdxInfo->aConstraint[i].iColumn,
5374+ pIdxInfo->aConstraint[i].op, vtabIn);
5375+#endif
5376+ if (!pIdxInfo->aConstraint[i].usable)
5377+ continue;
5378+
5379+ int iColumn = pIdxInfo->aConstraint[i].iColumn;
5380+ int op = pIdxInfo->aConstraint[i].op;
5381+
5382+ if (op == SQLITE_INDEX_CONSTRAINT_LIMIT) {
5383+ iLimitTerm = i;
5384+ }
5385+ if (op == SQLITE_INDEX_CONSTRAINT_MATCH &&
5386+ vec0_column_idx_is_vector(p, iColumn)) {
5387+ if (iMatchTerm > -1) {
5388+ vtab_set_error(
5389+ pVTab, "only 1 MATCH operator is allowed in a single vec0 query");
5390+ return SQLITE_ERROR;
5391+ }
5392+ iMatchTerm = i;
5393+ iMatchVectorTerm = vec0_column_idx_to_vector_idx(p, iColumn);
5394+ }
5395+ if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == VEC0_COLUMN_ID) {
5396+ if (vtabIn) {
5397+ if (iRowidInTerm != -1) {
5398+ vtab_set_error(pVTab, "only 1 'rowid in (..)' operator is allowed in "
5399+ "a single vec0 query");
5400+ return SQLITE_ERROR;
5401+ }
5402+ iRowidInTerm = i;
5403+
5404+ } else {
5405+ iRowidTerm = i;
5406+ }
5407+ }
5408+ if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == vec0_column_k_idx(p)) {
5409+ iKTerm = i;
5410+ }
5411+ if(
5412+ (op != SQLITE_INDEX_CONSTRAINT_LIMIT && op != SQLITE_INDEX_CONSTRAINT_OFFSET)
5413+ && vec0_column_idx_is_auxiliary(p, iColumn)) {
5414+ hasAuxConstraint = 1;
5415+ }
5416+ }
5417+
5418+ sqlite3_str *idxStr = sqlite3_str_new(NULL);
5419+ int rc;
5420+
5421+ if (iMatchTerm >= 0) {
5422+ if (iLimitTerm < 0 && iKTerm < 0) {
5423+ vtab_set_error(
5424+ pVTab,
5425+ "A LIMIT or 'k = ?' constraint is required on vec0 knn queries.");
5426+ rc = SQLITE_ERROR;
5427+ goto done;
5428+ }
5429+ if (iLimitTerm >= 0 && iKTerm >= 0) {
5430+ vtab_set_error(pVTab, "Only LIMIT or 'k =?' can be provided, not both");
5431+ rc = SQLITE_ERROR;
5432+ goto done;
5433+ }
5434+
5435+ if (pIdxInfo->nOrderBy) {
5436+ if (pIdxInfo->nOrderBy > 1) {
5437+ vtab_set_error(pVTab, "Only a single 'ORDER BY distance' clause is "
5438+ "allowed on vec0 KNN queries");
5439+ rc = SQLITE_ERROR;
5440+ goto done;
5441+ }
5442+ if (pIdxInfo->aOrderBy[0].iColumn != vec0_column_distance_idx(p)) {
5443+ vtab_set_error(pVTab,
5444+ "Only a single 'ORDER BY distance' clause is allowed on "
5445+ "vec0 KNN queries, not on other columns");
5446+ rc = SQLITE_ERROR;
5447+ goto done;
5448+ }
5449+ if (pIdxInfo->aOrderBy[0].desc) {
5450+ vtab_set_error(
5451+ pVTab, "Only ascending in ORDER BY distance clause is supported, "
5452+ "DESC is not supported yet.");
5453+ rc = SQLITE_ERROR;
5454+ goto done;
5455+ }
5456+ }
5457+
5458+ if(hasAuxConstraint) {
5459+ // IMP: V25623_09693
5460+ vtab_set_error(pVTab, "An illegal WHERE constraint was provided on a vec0 auxiliary column in a KNN query.");
5461+ rc = SQLITE_ERROR;
5462+ goto done;
5463+ }
5464+
5465+ sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_KNN);
5466+
5467+ int argvIndex = 1;
5468+ pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = argvIndex++;
5469+ pIdxInfo->aConstraintUsage[iMatchTerm].omit = 1;
5470+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_MATCH);
5471+ sqlite3_str_appendchar(idxStr, 3, '_');
5472+
5473+ if (iLimitTerm >= 0) {
5474+ pIdxInfo->aConstraintUsage[iLimitTerm].argvIndex = argvIndex++;
5475+ pIdxInfo->aConstraintUsage[iLimitTerm].omit = 1;
5476+ } else {
5477+ pIdxInfo->aConstraintUsage[iKTerm].argvIndex = argvIndex++;
5478+ pIdxInfo->aConstraintUsage[iKTerm].omit = 1;
5479+ }
5480+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_K);
5481+ sqlite3_str_appendchar(idxStr, 3, '_');
5482+
5483+#if COMPILER_SUPPORTS_VTAB_IN
5484+ if (iRowidInTerm >= 0) {
5485+ // already validated as >= SQLite 3.38 bc iRowidInTerm is only >= 0 when
5486+ // vtabIn == 1
5487+ sqlite3_vtab_in(pIdxInfo, iRowidInTerm, 1);
5488+ pIdxInfo->aConstraintUsage[iRowidInTerm].argvIndex = argvIndex++;
5489+ pIdxInfo->aConstraintUsage[iRowidInTerm].omit = 1;
5490+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_ROWID_IN);
5491+ sqlite3_str_appendchar(idxStr, 3, '_');
5492+ }
5493+#endif
5494+
5495+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
5496+ if (!pIdxInfo->aConstraint[i].usable)
5497+ continue;
5498+
5499+ int iColumn = pIdxInfo->aConstraint[i].iColumn;
5500+ int op = pIdxInfo->aConstraint[i].op;
5501+ if(op == SQLITE_INDEX_CONSTRAINT_LIMIT || op == SQLITE_INDEX_CONSTRAINT_OFFSET) {
5502+ continue;
5503+ }
5504+ if(!vec0_column_idx_is_partition(p, iColumn)) {
5505+ continue;
5506+ }
5507+
5508+ int partition_idx = vec0_column_idx_to_partition_idx(p, iColumn);
5509+ char value = 0;
5510+
5511+ switch(op) {
5512+ case SQLITE_INDEX_CONSTRAINT_EQ: {
5513+ value = VEC0_PARTITION_OPERATOR_EQ;
5514+ break;
5515+ }
5516+ case SQLITE_INDEX_CONSTRAINT_GT: {
5517+ value = VEC0_PARTITION_OPERATOR_GT;
5518+ break;
5519+ }
5520+ case SQLITE_INDEX_CONSTRAINT_LE: {
5521+ value = VEC0_PARTITION_OPERATOR_LE;
5522+ break;
5523+ }
5524+ case SQLITE_INDEX_CONSTRAINT_LT: {
5525+ value = VEC0_PARTITION_OPERATOR_LT;
5526+ break;
5527+ }
5528+ case SQLITE_INDEX_CONSTRAINT_GE: {
5529+ value = VEC0_PARTITION_OPERATOR_GE;
5530+ break;
5531+ }
5532+ case SQLITE_INDEX_CONSTRAINT_NE: {
5533+ value = VEC0_PARTITION_OPERATOR_NE;
5534+ break;
5535+ }
5536+ }
5537+
5538+ if(value) {
5539+ pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++;
5540+ pIdxInfo->aConstraintUsage[i].omit = 1;
5541+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT);
5542+ sqlite3_str_appendchar(idxStr, 1, 'A' + partition_idx);
5543+ sqlite3_str_appendchar(idxStr, 1, value);
5544+ sqlite3_str_appendchar(idxStr, 1, '_');
5545+ }
5546+
5547+ }
5548+
5549+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
5550+ if (!pIdxInfo->aConstraint[i].usable)
5551+ continue;
5552+
5553+ int iColumn = pIdxInfo->aConstraint[i].iColumn;
5554+ int op = pIdxInfo->aConstraint[i].op;
5555+ if(op == SQLITE_INDEX_CONSTRAINT_LIMIT || op == SQLITE_INDEX_CONSTRAINT_OFFSET) {
5556+ continue;
5557+ }
5558+ if(!vec0_column_idx_is_metadata(p, iColumn)) {
5559+ continue;
5560+ }
5561+
5562+ int metadata_idx = vec0_column_idx_to_metadata_idx(p, iColumn);
5563+ char value = 0;
5564+
5565+ switch(op) {
5566+ case SQLITE_INDEX_CONSTRAINT_EQ: {
5567+ int vtabIn = 0;
5568+ #if COMPILER_SUPPORTS_VTAB_IN
5569+ if (sqlite3_libversion_number() >= 3038000) {
5570+ vtabIn = sqlite3_vtab_in(pIdxInfo, i, -1);
5571+ }
5572+ if(vtabIn) {
5573+ switch(p->metadata_columns[metadata_idx].kind) {
5574+ case VEC0_METADATA_COLUMN_KIND_FLOAT:
5575+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
5576+ // IMP: V15248_32086
5577+ rc = SQLITE_ERROR;
5578+ vtab_set_error(pVTab, "'xxx in (...)' is only available on INTEGER or TEXT metadata columns.");
5579+ goto done;
5580+ break;
5581+ }
5582+ case VEC0_METADATA_COLUMN_KIND_INTEGER:
5583+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
5584+ break;
5585+ }
5586+ }
5587+ value = VEC0_METADATA_OPERATOR_IN;
5588+ sqlite3_vtab_in(pIdxInfo, i, 1);
5589+ }else
5590+ #endif
5591+ {
5592+ value = VEC0_PARTITION_OPERATOR_EQ;
5593+ }
5594+ break;
5595+ }
5596+ case SQLITE_INDEX_CONSTRAINT_GT: {
5597+ value = VEC0_METADATA_OPERATOR_GT;
5598+ break;
5599+ }
5600+ case SQLITE_INDEX_CONSTRAINT_LE: {
5601+ value = VEC0_METADATA_OPERATOR_LE;
5602+ break;
5603+ }
5604+ case SQLITE_INDEX_CONSTRAINT_LT: {
5605+ value = VEC0_METADATA_OPERATOR_LT;
5606+ break;
5607+ }
5608+ case SQLITE_INDEX_CONSTRAINT_GE: {
5609+ value = VEC0_METADATA_OPERATOR_GE;
5610+ break;
5611+ }
5612+ case SQLITE_INDEX_CONSTRAINT_NE: {
5613+ value = VEC0_METADATA_OPERATOR_NE;
5614+ break;
5615+ }
5616+ default: {
5617+ // IMP: V16511_00582
5618+ rc = SQLITE_ERROR;
5619+ vtab_set_error(pVTab,
5620+ "An illegal WHERE constraint was provided on a vec0 metadata column in a KNN query. "
5621+ "Only one of EQUALS, GREATER_THAN, LESS_THAN_OR_EQUAL, LESS_THAN, GREATER_THAN_OR_EQUAL, NOT_EQUALS is allowed."
5622+ );
5623+ goto done;
5624+ }
5625+ }
5626+
5627+ if(p->metadata_columns[metadata_idx].kind == VEC0_METADATA_COLUMN_KIND_BOOLEAN) {
5628+ if(!(value == VEC0_METADATA_OPERATOR_EQ || value == VEC0_METADATA_OPERATOR_NE)) {
5629+ // IMP: V10145_26984
5630+ rc = SQLITE_ERROR;
5631+ vtab_set_error(pVTab, "ONLY EQUALS (=) or NOT_EQUALS (!=) operators are allowed on boolean metadata columns.");
5632+ goto done;
5633+ }
5634+ }
5635+
5636+ pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++;
5637+ pIdxInfo->aConstraintUsage[i].omit = 1;
5638+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_METADATA_CONSTRAINT);
5639+ sqlite3_str_appendchar(idxStr, 1, 'A' + metadata_idx);
5640+ sqlite3_str_appendchar(idxStr, 1, value);
5641+ sqlite3_str_appendchar(idxStr, 1, '_');
5642+
5643+ }
5644+
5645+
5646+
5647+ pIdxInfo->idxNum = iMatchVectorTerm;
5648+ pIdxInfo->estimatedCost = 30.0;
5649+ pIdxInfo->estimatedRows = 10;
5650+
5651+ } else if (iRowidTerm >= 0) {
5652+ sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_POINT);
5653+ pIdxInfo->aConstraintUsage[iRowidTerm].argvIndex = 1;
5654+ pIdxInfo->aConstraintUsage[iRowidTerm].omit = 1;
5655+ sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_POINT_ID);
5656+ sqlite3_str_appendchar(idxStr, 3, '_');
5657+ pIdxInfo->idxNum = pIdxInfo->colUsed;
5658+ pIdxInfo->estimatedCost = 10.0;
5659+ pIdxInfo->estimatedRows = 1;
5660+ } else {
5661+ sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_FULLSCAN);
5662+ pIdxInfo->estimatedCost = 3000000.0;
5663+ pIdxInfo->estimatedRows = 100000;
5664+ }
5665+ pIdxInfo->idxStr = sqlite3_str_finish(idxStr);
5666+ idxStr = NULL;
5667+ if (!pIdxInfo->idxStr) {
5668+ rc = SQLITE_OK;
5669+ goto done;
5670+ }
5671+ pIdxInfo->needToFreeIdxStr = 1;
5672+
5673+
5674+ rc = SQLITE_OK;
5675+
5676+ done:
5677+ if(idxStr) {
5678+ sqlite3_str_finish(idxStr);
5679+ }
5680+ return rc;
5681+}
5682+
5683+// forward delcaration bc vec0Filter uses it
5684+static int vec0Next(sqlite3_vtab_cursor *cur);
5685+
5686+void merge_sorted_lists(f32 *a, i64 *a_rowids, i64 a_length, f32 *b,
5687+ i64 *b_rowids, i32 *b_top_idxs, i64 b_length, f32 *out,
5688+ i64 *out_rowids, i64 out_length, i64 *out_used) {
5689+ // assert((a_length >= out_length) || (b_length >= out_length));
5690+ i64 ptrA = 0;
5691+ i64 ptrB = 0;
5692+ for (int i = 0; i < out_length; i++) {
5693+ if ((ptrA >= a_length) && (ptrB >= b_length)) {
5694+ *out_used = i;
5695+ return;
5696+ }
5697+ if (ptrA >= a_length) {
5698+ out[i] = b[b_top_idxs[ptrB]];
5699+ out_rowids[i] = b_rowids[b_top_idxs[ptrB]];
5700+ ptrB++;
5701+ } else if (ptrB >= b_length) {
5702+ out[i] = a[ptrA];
5703+ out_rowids[i] = a_rowids[ptrA];
5704+ ptrA++;
5705+ } else {
5706+ if (a[ptrA] <= b[b_top_idxs[ptrB]]) {
5707+ out[i] = a[ptrA];
5708+ out_rowids[i] = a_rowids[ptrA];
5709+ ptrA++;
5710+ } else {
5711+ out[i] = b[b_top_idxs[ptrB]];
5712+ out_rowids[i] = b_rowids[b_top_idxs[ptrB]];
5713+ ptrB++;
5714+ }
5715+ }
5716+ }
5717+
5718+ *out_used = out_length;
5719+}
5720+
5721+u8 *bitmap_new(i32 n) {
5722+ assert(n % 8 == 0);
5723+ u8 *p = sqlite3_malloc(n * sizeof(u8) / CHAR_BIT);
5724+ if (p) {
5725+ memset(p, 0, n * sizeof(u8) / CHAR_BIT);
5726+ }
5727+ return p;
5728+}
5729+u8 *bitmap_new_from(i32 n, u8 *from) {
5730+ assert(n % 8 == 0);
5731+ u8 *p = sqlite3_malloc(n * sizeof(u8) / CHAR_BIT);
5732+ if (p) {
5733+ memcpy(p, from, n / CHAR_BIT);
5734+ }
5735+ return p;
5736+}
5737+
5738+void bitmap_copy(u8 *base, u8 *from, i32 n) {
5739+ assert(n % 8 == 0);
5740+ memcpy(base, from, n / CHAR_BIT);
5741+}
5742+
5743+void bitmap_and_inplace(u8 *base, u8 *other, i32 n) {
5744+ assert((n % 8) == 0);
5745+ for (int i = 0; i < n / CHAR_BIT; i++) {
5746+ base[i] = base[i] & other[i];
5747+ }
5748+}
5749+
5750+void bitmap_set(u8 *bitmap, i32 position, int value) {
5751+ if (value) {
5752+ bitmap[position / CHAR_BIT] |= 1 << (position % CHAR_BIT);
5753+ } else {
5754+ bitmap[position / CHAR_BIT] &= ~(1 << (position % CHAR_BIT));
5755+ }
5756+}
5757+
5758+int bitmap_get(u8 *bitmap, i32 position) {
5759+ return (((bitmap[position / CHAR_BIT]) >> (position % CHAR_BIT)) & 1);
5760+}
5761+
5762+void bitmap_clear(u8 *bitmap, i32 n) {
5763+ assert((n % 8) == 0);
5764+ memset(bitmap, 0, n / CHAR_BIT);
5765+}
5766+
5767+void bitmap_fill(u8 *bitmap, i32 n) {
5768+ assert((n % 8) == 0);
5769+ memset(bitmap, 0xFF, n / CHAR_BIT);
5770+}
5771+
5772+/**
5773+ * @brief Finds the minimum k items in distances, and writes the indicies to
5774+ * out.
5775+ *
5776+ * @param distances input f32 array of size n, the items to consider.
5777+ * @param n: size of distances array.
5778+ * @param out: Output array of size k, will contain at most k element indicies
5779+ * @param k: Size of output array
5780+ * @return int
5781+ */
5782+int min_idx(const f32 *distances, i32 n, u8 *candidates, i32 *out, i32 k,
5783+ u8 *bTaken, i32 *k_used) {
5784+ assert(k > 0);
5785+ assert(k <= n);
5786+
5787+ bitmap_clear(bTaken, n);
5788+
5789+ for (int ik = 0; ik < k; ik++) {
5790+ int min_idx = 0;
5791+ while (min_idx < n &&
5792+ (bitmap_get(bTaken, min_idx) || !bitmap_get(candidates, min_idx))) {
5793+ min_idx++;
5794+ }
5795+ if (min_idx >= n) {
5796+ *k_used = ik;
5797+ return SQLITE_OK;
5798+ }
5799+
5800+ for (int i = 0; i < n; i++) {
5801+ if (distances[i] <= distances[min_idx] && !bitmap_get(bTaken, i) &&
5802+ (bitmap_get(candidates, i))) {
5803+ min_idx = i;
5804+ }
5805+ }
5806+
5807+ out[ik] = min_idx;
5808+ bitmap_set(bTaken, min_idx, 1);
5809+ }
5810+ *k_used = k;
5811+ return SQLITE_OK;
5812+}
5813+
5814+int vec0_get_metadata_text_long_value(
5815+ vec0_vtab * p,
5816+ sqlite3_stmt ** stmt,
5817+ int metadata_idx,
5818+ i64 rowid,
5819+ int *n,
5820+ char ** s) {
5821+ int rc;
5822+ if(!(*stmt)) {
5823+ const char * zSql = sqlite3_mprintf("select data from " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " where rowid = ?", p->schemaName, p->tableName, metadata_idx);
5824+ if(!zSql) {
5825+ rc = SQLITE_NOMEM;
5826+ goto done;
5827+ }
5828+ rc = sqlite3_prepare_v2(p->db, zSql, -1, stmt, NULL);
5829+ sqlite3_free( (void *) zSql);
5830+ if(rc != SQLITE_OK) {
5831+ goto done;
5832+ }
5833+ }
5834+
5835+ sqlite3_reset(*stmt);
5836+ sqlite3_bind_int64(*stmt, 1, rowid);
5837+ rc = sqlite3_step(*stmt);
5838+ if(rc != SQLITE_ROW) {
5839+ rc = SQLITE_ERROR;
5840+ goto done;
5841+ }
5842+ *s = (char *) sqlite3_column_text(*stmt, 0);
5843+ *n = sqlite3_column_bytes(*stmt, 0);
5844+ rc = SQLITE_OK;
5845+ done:
5846+ return rc;
5847+}
5848+
5849+/**
5850+ * @brief Crete at "iterator" (sqlite3_stmt) of chunks with the given constraints
5851+ *
5852+ * Any VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT values in idxStr/argv will be applied
5853+ * as WHERE constraints in the underlying stmt SQL, and any consumer of the stmt
5854+ * can freely step through the stmt with all constraints satisfied.
5855+ *
5856+ * @param p - vec0_vtab
5857+ * @param idxStr - the xBestIndex/xFilter idxstr containing VEC0_IDXSTR values
5858+ * @param argc - number of argv values from xFilter
5859+ * @param argv - array of sqlite3_value from xFilter
5860+ * @param outStmt - output sqlite3_stmt of chunks with all filters applied
5861+ * @return int SQLITE_OK on success, error code otherwise
5862+ */
5863+int vec0_chunks_iter(vec0_vtab * p, const char * idxStr, int argc, sqlite3_value ** argv, sqlite3_stmt** outStmt) {
5864+ // always null terminated, enforced by SQLite
5865+ int idxStrLength = strlen(idxStr);
5866+ // "1" refers to the initial vec0_query_plan char, 4 is the number of chars per "element"
5867+ int numValueEntries = (idxStrLength-1) / 4;
5868+ assert(argc == numValueEntries);
5869+
5870+ int rc;
5871+ sqlite3_str * s = sqlite3_str_new(NULL);
5872+ sqlite3_str_appendf(s, "select chunk_id, validity, rowids "
5873+ " from " VEC0_SHADOW_CHUNKS_NAME,
5874+ p->schemaName, p->tableName);
5875+
5876+ int appendedWhere = 0;
5877+ for(int i = 0; i < numValueEntries; i++) {
5878+ int idx = 1 + (i * 4);
5879+ char kind = idxStr[idx + 0];
5880+ if(kind != VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT) {
5881+ continue;
5882+ }
5883+
5884+ int partition_idx = idxStr[idx + 1] - 'A';
5885+ int operator = idxStr[idx + 2];
5886+ // idxStr[idx + 3] is just null, a '_' placeholder
5887+
5888+ if(!appendedWhere) {
5889+ sqlite3_str_appendall(s, " WHERE ");
5890+ appendedWhere = 1;
5891+ }else {
5892+ sqlite3_str_appendall(s, " AND ");
5893+ }
5894+ switch(operator) {
5895+ case VEC0_PARTITION_OPERATOR_EQ:
5896+ sqlite3_str_appendf(s, " partition%02d = ? ", partition_idx);
5897+ break;
5898+ case VEC0_PARTITION_OPERATOR_GT:
5899+ sqlite3_str_appendf(s, " partition%02d > ? ", partition_idx);
5900+ break;
5901+ case VEC0_PARTITION_OPERATOR_LE:
5902+ sqlite3_str_appendf(s, " partition%02d <= ? ", partition_idx);
5903+ break;
5904+ case VEC0_PARTITION_OPERATOR_LT:
5905+ sqlite3_str_appendf(s, " partition%02d < ? ", partition_idx);
5906+ break;
5907+ case VEC0_PARTITION_OPERATOR_GE:
5908+ sqlite3_str_appendf(s, " partition%02d >= ? ", partition_idx);
5909+ break;
5910+ case VEC0_PARTITION_OPERATOR_NE:
5911+ sqlite3_str_appendf(s, " partition%02d != ? ", partition_idx);
5912+ break;
5913+ default: {
5914+ char * zSql = sqlite3_str_finish(s);
5915+ sqlite3_free(zSql);
5916+ return SQLITE_ERROR;
5917+ }
5918+
5919+ }
5920+
5921+ }
5922+
5923+ char *zSql = sqlite3_str_finish(s);
5924+ if (!zSql) {
5925+ return SQLITE_NOMEM;
5926+ }
5927+
5928+ rc = sqlite3_prepare_v2(p->db, zSql, -1, outStmt, NULL);
5929+ sqlite3_free(zSql);
5930+ if(rc != SQLITE_OK) {
5931+ return rc;
5932+ }
5933+
5934+ int n = 1;
5935+ for(int i = 0; i < numValueEntries; i++) {
5936+ int idx = 1 + (i * 4);
5937+ char kind = idxStr[idx + 0];
5938+ if(kind != VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT) {
5939+ continue;
5940+ }
5941+ sqlite3_bind_value(*outStmt, n++, argv[i]);
5942+ }
5943+
5944+ return rc;
5945+}
5946+
5947+// a single `xxx in (...)` constraint on a metadata column. TEXT or INTEGER only for now.
5948+struct Vec0MetadataIn{
5949+ // index of argv[i]` the constraint is on
5950+ int argv_idx;
5951+ // metadata column index of the constraint, derived from idxStr + argv_idx
5952+ int metadata_idx;
5953+ // array of the copied `(...)` values from sqlite3_vtab_in_first()/sqlite3_vtab_in_next()
5954+ struct Array array;
5955+};
5956+
5957+// Array elements for `xxx in (...)` values for a text column. basically just a string
5958+struct Vec0MetadataInTextEntry {
5959+ int n;
5960+ char * zString;
5961+};
5962+
5963+
5964+int vec0_metadata_filter_text(vec0_vtab * p, sqlite3_value * value, const void * buffer, int size, vec0_metadata_operator op, u8* b, int metadata_idx, int chunk_rowid, struct Array * aMetadataIn, int argv_idx) {
5965+ int rc;
5966+ sqlite3_stmt * stmt = NULL;
5967+ i64 * rowids = NULL;
5968+ sqlite3_blob * rowidsBlob;
5969+ const char * sTarget = (const char *) sqlite3_value_text(value);
5970+ int nTarget = sqlite3_value_bytes(value);
5971+
5972+
5973+ // TODO(perf): only text metadata news the rowids BLOB. Make it so that
5974+ // rowids BLOB is re-used when multiple fitlers on text columns,
5975+ // ex "name BETWEEN 'a' and 'b'""
5976+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "rowids", chunk_rowid, 0, &rowidsBlob);
5977+ if(rc != SQLITE_OK) {
5978+ return rc;
5979+ }
5980+ assert(sqlite3_blob_bytes(rowidsBlob) % sizeof(i64) == 0);
5981+ assert((sqlite3_blob_bytes(rowidsBlob) / sizeof(i64)) == size);
5982+
5983+ rowids = sqlite3_malloc(sqlite3_blob_bytes(rowidsBlob));
5984+ if(!rowids) {
5985+ sqlite3_blob_close(rowidsBlob);
5986+ return SQLITE_NOMEM;
5987+ }
5988+
5989+ rc = sqlite3_blob_read(rowidsBlob, rowids, sqlite3_blob_bytes(rowidsBlob), 0);
5990+ if(rc != SQLITE_OK) {
5991+ sqlite3_blob_close(rowidsBlob);
5992+ return rc;
5993+ }
5994+ sqlite3_blob_close(rowidsBlob);
5995+
5996+ switch(op) {
5997+ int nPrefix;
5998+ char * sPrefix;
5999+ char *sFull;
6000+ int nFull;
6001+ u8 * view;
6002+ case VEC0_METADATA_OPERATOR_EQ: {
6003+ for(int i = 0; i < size; i++) {
6004+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6005+ nPrefix = ((int*) view)[0];
6006+ sPrefix = (char *) &view[4];
6007+
6008+ // for EQ the text lengths must match
6009+ if(nPrefix != nTarget) {
6010+ bitmap_set(b, i, 0);
6011+ continue;
6012+ }
6013+ int cmpPrefix = strncmp(sPrefix, sTarget, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH));
6014+
6015+ // for short strings, use the prefix comparison direclty
6016+ if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6017+ bitmap_set(b, i, cmpPrefix == 0);
6018+ continue;
6019+ }
6020+ // for EQ on longs strings, the prefix must match
6021+ if(cmpPrefix) {
6022+ bitmap_set(b, i, 0);
6023+ continue;
6024+ }
6025+ // consult the full string
6026+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6027+ if(rc != SQLITE_OK) {
6028+ goto done;
6029+ }
6030+ if(nPrefix != nFull) {
6031+ rc = SQLITE_ERROR;
6032+ goto done;
6033+ }
6034+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) == 0);
6035+ }
6036+ break;
6037+ }
6038+ case VEC0_METADATA_OPERATOR_NE: {
6039+ for(int i = 0; i < size; i++) {
6040+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6041+ nPrefix = ((int*) view)[0];
6042+ sPrefix = (char *) &view[4];
6043+
6044+ // for NE if text lengths dont match, it never will
6045+ if(nPrefix != nTarget) {
6046+ bitmap_set(b, i, 1);
6047+ continue;
6048+ }
6049+
6050+ int cmpPrefix = strncmp(sPrefix, sTarget, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH));
6051+
6052+ // for short strings, use the prefix comparison direclty
6053+ if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6054+ bitmap_set(b, i, cmpPrefix != 0);
6055+ continue;
6056+ }
6057+ // for NE on longs strings, if prefixes dont match, then long string wont
6058+ if(cmpPrefix) {
6059+ bitmap_set(b, i, 1);
6060+ continue;
6061+ }
6062+ // consult the full string
6063+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6064+ if(rc != SQLITE_OK) {
6065+ goto done;
6066+ }
6067+ if(nPrefix != nFull) {
6068+ rc = SQLITE_ERROR;
6069+ goto done;
6070+ }
6071+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) != 0);
6072+ }
6073+ break;
6074+ }
6075+ case VEC0_METADATA_OPERATOR_GT: {
6076+ for(int i = 0; i < size; i++) {
6077+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6078+ nPrefix = ((int*) view)[0];
6079+ sPrefix = (char *) &view[4];
6080+ int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget));
6081+
6082+ if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6083+ // if prefix match, check which is longer
6084+ if(cmpPrefix == 0) {
6085+ bitmap_set(b, i, nPrefix > nTarget);
6086+ }
6087+ else {
6088+ bitmap_set(b, i, cmpPrefix > 0);
6089+ }
6090+ continue;
6091+ }
6092+ // TODO(perf): may not need to compare full text in some cases
6093+
6094+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6095+ if(rc != SQLITE_OK) {
6096+ goto done;
6097+ }
6098+ if(nPrefix != nFull) {
6099+ rc = SQLITE_ERROR;
6100+ goto done;
6101+ }
6102+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) > 0);
6103+ }
6104+ break;
6105+ }
6106+ case VEC0_METADATA_OPERATOR_GE: {
6107+ for(int i = 0; i < size; i++) {
6108+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6109+ nPrefix = ((int*) view)[0];
6110+ sPrefix = (char *) &view[4];
6111+ int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget));
6112+
6113+ if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6114+ // if prefix match, check which is longer
6115+ if(cmpPrefix == 0) {
6116+ bitmap_set(b, i, nPrefix >= nTarget);
6117+ }
6118+ else {
6119+ bitmap_set(b, i, cmpPrefix >= 0);
6120+ }
6121+ continue;
6122+ }
6123+ // TODO(perf): may not need to compare full text in some cases
6124+
6125+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6126+ if(rc != SQLITE_OK) {
6127+ goto done;
6128+ }
6129+ if(nPrefix != nFull) {
6130+ rc = SQLITE_ERROR;
6131+ goto done;
6132+ }
6133+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) >= 0);
6134+ }
6135+ break;
6136+ }
6137+ case VEC0_METADATA_OPERATOR_LE: {
6138+ for(int i = 0; i < size; i++) {
6139+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6140+ nPrefix = ((int*) view)[0];
6141+ sPrefix = (char *) &view[4];
6142+ int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget));
6143+
6144+ if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6145+ // if prefix match, check which is longer
6146+ if(cmpPrefix == 0) {
6147+ bitmap_set(b, i, nPrefix <= nTarget);
6148+ }
6149+ else {
6150+ bitmap_set(b, i, cmpPrefix <= 0);
6151+ }
6152+ continue;
6153+ }
6154+ // TODO(perf): may not need to compare full text in some cases
6155+
6156+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6157+ if(rc != SQLITE_OK) {
6158+ goto done;
6159+ }
6160+ if(nPrefix != nFull) {
6161+ rc = SQLITE_ERROR;
6162+ goto done;
6163+ }
6164+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) <= 0);
6165+ }
6166+ break;
6167+ }
6168+ case VEC0_METADATA_OPERATOR_LT: {
6169+ for(int i = 0; i < size; i++) {
6170+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6171+ nPrefix = ((int*) view)[0];
6172+ sPrefix = (char *) &view[4];
6173+ int cmpPrefix = strncmp(sPrefix, sTarget, min(min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH), nTarget));
6174+
6175+ if(nPrefix < VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6176+ // if prefix match, check which is longer
6177+ if(cmpPrefix == 0) {
6178+ bitmap_set(b, i, nPrefix < nTarget);
6179+ }
6180+ else {
6181+ bitmap_set(b, i, cmpPrefix < 0);
6182+ }
6183+ continue;
6184+ }
6185+ // TODO(perf): may not need to compare full text in some cases
6186+
6187+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6188+ if(rc != SQLITE_OK) {
6189+ goto done;
6190+ }
6191+ if(nPrefix != nFull) {
6192+ rc = SQLITE_ERROR;
6193+ goto done;
6194+ }
6195+ bitmap_set(b, i, strncmp(sFull, sTarget, nFull) < 0);
6196+ }
6197+ break;
6198+ }
6199+
6200+ case VEC0_METADATA_OPERATOR_IN: {
6201+ size_t metadataInIdx = -1;
6202+ for(size_t i = 0; i < aMetadataIn->length; i++) {
6203+ struct Vec0MetadataIn * metadataIn = &(((struct Vec0MetadataIn *) aMetadataIn->z)[i]);
6204+ if(metadataIn->argv_idx == argv_idx) {
6205+ metadataInIdx = i;
6206+ break;
6207+ }
6208+ }
6209+ if(metadataInIdx < 0) {
6210+ rc = SQLITE_ERROR;
6211+ goto done;
6212+ }
6213+
6214+ struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[metadataInIdx];
6215+ struct Array * aTarget = &(metadataIn->array);
6216+
6217+
6218+ int nPrefix;
6219+ char * sPrefix;
6220+ char *sFull;
6221+ int nFull;
6222+ u8 * view;
6223+ for(int i = 0; i < size; i++) {
6224+ view = &((u8*) buffer)[i * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
6225+ nPrefix = ((int*) view)[0];
6226+ sPrefix = (char *) &view[4];
6227+ for(size_t target_idx = 0; target_idx < aTarget->length; target_idx++) {
6228+ struct Vec0MetadataInTextEntry * entry = &(((struct Vec0MetadataInTextEntry*)aTarget->z)[target_idx]);
6229+ if(entry->n != nPrefix) {
6230+ continue;
6231+ }
6232+ int cmpPrefix = strncmp(sPrefix, entry->zString, min(nPrefix, VEC0_METADATA_TEXT_VIEW_DATA_LENGTH));
6233+ if(nPrefix <= VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
6234+ if(cmpPrefix == 0) {
6235+ bitmap_set(b, i, 1);
6236+ break;
6237+ }
6238+ continue;
6239+ }
6240+ if(cmpPrefix) {
6241+ continue;
6242+ }
6243+
6244+ rc = vec0_get_metadata_text_long_value(p, &stmt, metadata_idx, rowids[i], &nFull, &sFull);
6245+ if(rc != SQLITE_OK) {
6246+ goto done;
6247+ }
6248+ if(nPrefix != nFull) {
6249+ rc = SQLITE_ERROR;
6250+ goto done;
6251+ }
6252+ if(strncmp(sFull, entry->zString, nFull) == 0) {
6253+ bitmap_set(b, i, 1);
6254+ break;
6255+ }
6256+ }
6257+ }
6258+ break;
6259+ }
6260+
6261+ }
6262+ rc = SQLITE_OK;
6263+
6264+ done:
6265+ sqlite3_finalize(stmt);
6266+ sqlite3_free(rowids);
6267+ return rc;
6268+
6269+}
6270+
6271+/**
6272+ * @brief Fill in bitmap of chunk values, whether or not the values match a metadata constraint
6273+ *
6274+ * @param p vec0_vtab
6275+ * @param metadata_idx index of the metatadata column to perfrom constraints on
6276+ * @param value sqlite3_value of the constraints value
6277+ * @param blob sqlite3_blob that is already opened on the metdata column's shadow chunk table
6278+ * @param chunk_rowid rowid of the chunk to calculate on
6279+ * @param b pre-allocated and zero'd out bitmap to write results to
6280+ * @param size size of the chunk
6281+ * @return int SQLITE_OK on success, error code otherwise
6282+ */
6283+int vec0_set_metadata_filter_bitmap(
6284+ vec0_vtab *p,
6285+ int metadata_idx,
6286+ vec0_metadata_operator op,
6287+ sqlite3_value * value,
6288+ sqlite3_blob * blob,
6289+ i64 chunk_rowid,
6290+ u8* b,
6291+ int size,
6292+ struct Array * aMetadataIn, int argv_idx) {
6293+ // TODO: shouldn't this skip in-valid entries from the chunk's validity bitmap?
6294+
6295+ int rc;
6296+ rc = sqlite3_blob_reopen(blob, chunk_rowid);
6297+ if(rc != SQLITE_OK) {
6298+ return rc;
6299+ }
6300+
6301+ vec0_metadata_column_kind kind = p->metadata_columns[metadata_idx].kind;
6302+ int szMatch = 0;
6303+ int blobSize = sqlite3_blob_bytes(blob);
6304+ switch(kind) {
6305+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
6306+ szMatch = blobSize == size / CHAR_BIT;
6307+ break;
6308+ }
6309+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
6310+ szMatch = blobSize == size * sizeof(i64);
6311+ break;
6312+ }
6313+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
6314+ szMatch = blobSize == size * sizeof(double);
6315+ break;
6316+ }
6317+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
6318+ szMatch = blobSize == size * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH;
6319+ break;
6320+ }
6321+ }
6322+ if(!szMatch) {
6323+ return SQLITE_ERROR;
6324+ }
6325+ void * buffer = sqlite3_malloc(blobSize);
6326+ if(!buffer) {
6327+ return SQLITE_NOMEM;
6328+ }
6329+ rc = sqlite3_blob_read(blob, buffer, blobSize, 0);
6330+ if(rc != SQLITE_OK) {
6331+ goto done;
6332+ }
6333+ switch(kind) {
6334+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
6335+ int target = sqlite3_value_int(value);
6336+ if( (target && op == VEC0_METADATA_OPERATOR_EQ) || (!target && op == VEC0_METADATA_OPERATOR_NE)) {
6337+ for(int i = 0; i < size; i++) { bitmap_set(b, i, bitmap_get((u8*) buffer, i)); }
6338+ }
6339+ else {
6340+ for(int i = 0; i < size; i++) { bitmap_set(b, i, !bitmap_get((u8*) buffer, i)); }
6341+ }
6342+ break;
6343+ }
6344+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
6345+ i64 * array = (i64*) buffer;
6346+ i64 target = sqlite3_value_int64(value);
6347+ switch(op) {
6348+ case VEC0_METADATA_OPERATOR_EQ: {
6349+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] == target); }
6350+ break;
6351+ }
6352+ case VEC0_METADATA_OPERATOR_GT: {
6353+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] > target); }
6354+ break;
6355+ }
6356+ case VEC0_METADATA_OPERATOR_LE: {
6357+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] <= target); }
6358+ break;
6359+ }
6360+ case VEC0_METADATA_OPERATOR_LT: {
6361+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] < target); }
6362+ break;
6363+ }
6364+ case VEC0_METADATA_OPERATOR_GE: {
6365+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] >= target); }
6366+ break;
6367+ }
6368+ case VEC0_METADATA_OPERATOR_NE: {
6369+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] != target); }
6370+ break;
6371+ }
6372+ case VEC0_METADATA_OPERATOR_IN: {
6373+ int metadataInIdx = -1;
6374+ for(size_t i = 0; i < aMetadataIn->length; i++) {
6375+ struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[i];
6376+ if(metadataIn->argv_idx == argv_idx) {
6377+ metadataInIdx = i;
6378+ break;
6379+ }
6380+ }
6381+ if(metadataInIdx < 0) {
6382+ rc = SQLITE_ERROR;
6383+ goto done;
6384+ }
6385+ struct Vec0MetadataIn * metadataIn = &((struct Vec0MetadataIn *) aMetadataIn->z)[metadataInIdx];
6386+ struct Array * aTarget = &(metadataIn->array);
6387+
6388+ for(int i = 0; i < size; i++) {
6389+ for(size_t target_idx = 0; target_idx < aTarget->length; target_idx++) {
6390+ if( ((i64*)aTarget->z)[target_idx] == array[i]) {
6391+ bitmap_set(b, i, 1);
6392+ break;
6393+ }
6394+ }
6395+ }
6396+ break;
6397+ }
6398+ }
6399+ break;
6400+ }
6401+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
6402+ double * array = (double*) buffer;
6403+ double target = sqlite3_value_double(value);
6404+ switch(op) {
6405+ case VEC0_METADATA_OPERATOR_EQ: {
6406+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] == target); }
6407+ break;
6408+ }
6409+ case VEC0_METADATA_OPERATOR_GT: {
6410+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] > target); }
6411+ break;
6412+ }
6413+ case VEC0_METADATA_OPERATOR_LE: {
6414+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] <= target); }
6415+ break;
6416+ }
6417+ case VEC0_METADATA_OPERATOR_LT: {
6418+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] < target); }
6419+ break;
6420+ }
6421+ case VEC0_METADATA_OPERATOR_GE: {
6422+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] >= target); }
6423+ break;
6424+ }
6425+ case VEC0_METADATA_OPERATOR_NE: {
6426+ for(int i = 0; i < size; i++) { bitmap_set(b, i, array[i] != target); }
6427+ break;
6428+ }
6429+ case VEC0_METADATA_OPERATOR_IN: {
6430+ // should never be reached
6431+ break;
6432+ }
6433+ }
6434+ break;
6435+ }
6436+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
6437+ rc = vec0_metadata_filter_text(p, value, buffer, size, op, b, metadata_idx, chunk_rowid, aMetadataIn, argv_idx);
6438+ if(rc != SQLITE_OK) {
6439+ goto done;
6440+ }
6441+ break;
6442+ }
6443+ }
6444+ done:
6445+ sqlite3_free(buffer);
6446+ return rc;
6447+}
6448+
6449+int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
6450+ struct VectorColumnDefinition *vector_column,
6451+ int vectorColumnIdx, struct Array *arrayRowidsIn,
6452+ struct Array * aMetadataIn,
6453+ const char * idxStr, int argc, sqlite3_value ** argv,
6454+ void *queryVector, i64 k, i64 **out_topk_rowids,
6455+ f32 **out_topk_distances, i64 *out_used) {
6456+ // for each chunk, get top min(k, chunk_size) rowid + distances to query vec.
6457+ // then reconcile all topk_chunks for a true top k.
6458+ // output only rowids + distances for now
6459+
6460+ int rc = SQLITE_OK;
6461+ sqlite3_blob *blobVectors = NULL;
6462+
6463+ void *baseVectors = NULL; // memory: chunk_size * dimensions * element_size
6464+
6465+ // OWNED BY CALLER ON SUCCESS
6466+ i64 *topk_rowids = NULL; // memory: k * 4
6467+ // OWNED BY CALLER ON SUCCESS
6468+ f32 *topk_distances = NULL; // memory: k * 4
6469+
6470+ i64 *tmp_topk_rowids = NULL; // memory: k * 4
6471+ f32 *tmp_topk_distances = NULL; // memory: k * 4
6472+ f32 *chunk_distances = NULL; // memory: chunk_size * 4
6473+ u8 *b = NULL; // memory: chunk_size / 8
6474+ u8 *bTaken = NULL; // memory: chunk_size / 8
6475+ i32 *chunk_topk_idxs = NULL; // memory: k * 4
6476+ u8 *bmRowids = NULL; // memory: chunk_size / 8
6477+ u8 *bmMetadata = NULL; // memory: chunk_size / 8
6478+ // // total: a lot???
6479+
6480+ // 6 * (k * 4) + (k * 2) + (chunk_size / 8) + (chunk_size * dimensions * 4)
6481+
6482+ topk_rowids = sqlite3_malloc(k * sizeof(i64));
6483+ if (!topk_rowids) {
6484+ rc = SQLITE_NOMEM;
6485+ goto cleanup;
6486+ }
6487+ memset(topk_rowids, 0, k * sizeof(i64));
6488+
6489+ topk_distances = sqlite3_malloc(k * sizeof(f32));
6490+ if (!topk_distances) {
6491+ rc = SQLITE_NOMEM;
6492+ goto cleanup;
6493+ }
6494+ memset(topk_distances, 0, k * sizeof(f32));
6495+
6496+ tmp_topk_rowids = sqlite3_malloc(k * sizeof(i64));
6497+ if (!tmp_topk_rowids) {
6498+ rc = SQLITE_NOMEM;
6499+ goto cleanup;
6500+ }
6501+ memset(tmp_topk_rowids, 0, k * sizeof(i64));
6502+
6503+ tmp_topk_distances = sqlite3_malloc(k * sizeof(f32));
6504+ if (!tmp_topk_distances) {
6505+ rc = SQLITE_NOMEM;
6506+ goto cleanup;
6507+ }
6508+ memset(tmp_topk_distances, 0, k * sizeof(f32));
6509+
6510+ i64 k_used = 0;
6511+ i64 baseVectorsSize = p->chunk_size * vector_column_byte_size(*vector_column);
6512+ baseVectors = sqlite3_malloc(baseVectorsSize);
6513+ if (!baseVectors) {
6514+ rc = SQLITE_NOMEM;
6515+ goto cleanup;
6516+ }
6517+
6518+ chunk_distances = sqlite3_malloc(p->chunk_size * sizeof(f32));
6519+ if (!chunk_distances) {
6520+ rc = SQLITE_NOMEM;
6521+ goto cleanup;
6522+ }
6523+
6524+ b = bitmap_new(p->chunk_size);
6525+ if (!b) {
6526+ rc = SQLITE_NOMEM;
6527+ goto cleanup;
6528+ }
6529+
6530+ bTaken = bitmap_new(p->chunk_size);
6531+ if (!bTaken) {
6532+ rc = SQLITE_NOMEM;
6533+ goto cleanup;
6534+ }
6535+
6536+ chunk_topk_idxs = sqlite3_malloc(k * sizeof(i32));
6537+ if (!chunk_topk_idxs) {
6538+ rc = SQLITE_NOMEM;
6539+ goto cleanup;
6540+ }
6541+
6542+ bmRowids = arrayRowidsIn ? bitmap_new(p->chunk_size) : NULL;
6543+ if (arrayRowidsIn && !bmRowids) {
6544+ rc = SQLITE_NOMEM;
6545+ goto cleanup;
6546+ }
6547+
6548+ sqlite3_blob * metadataBlobs[VEC0_MAX_METADATA_COLUMNS];
6549+ memset(metadataBlobs, 0, sizeof(sqlite3_blob*) * VEC0_MAX_METADATA_COLUMNS);
6550+
6551+ bmMetadata = bitmap_new(p->chunk_size);
6552+ if(!bmMetadata) {
6553+ rc = SQLITE_NOMEM;
6554+ goto cleanup;
6555+ }
6556+
6557+ int idxStrLength = strlen(idxStr);
6558+ int numValueEntries = (idxStrLength-1) / 4;
6559+ assert(numValueEntries == argc);
6560+ int hasMetadataFilters = 0;
6561+ for(int i = 0; i < argc; i++) {
6562+ int idx = 1 + (i * 4);
6563+ char kind = idxStr[idx + 0];
6564+ if(kind == VEC0_IDXSTR_KIND_METADATA_CONSTRAINT) {
6565+ hasMetadataFilters = 1;
6566+ break;
6567+ }
6568+ }
6569+
6570+ while (true) {
6571+ rc = sqlite3_step(stmtChunks);
6572+ if (rc == SQLITE_DONE) {
6573+ break;
6574+ }
6575+ if (rc != SQLITE_ROW) {
6576+ vtab_set_error(&p->base, "chunks iter error");
6577+ rc = SQLITE_ERROR;
6578+ goto cleanup;
6579+ }
6580+ memset(chunk_distances, 0, p->chunk_size * sizeof(f32));
6581+ memset(chunk_topk_idxs, 0, k * sizeof(i32));
6582+ bitmap_clear(b, p->chunk_size);
6583+
6584+ i64 chunk_id = sqlite3_column_int64(stmtChunks, 0);
6585+ unsigned char *chunkValidity =
6586+ (unsigned char *)sqlite3_column_blob(stmtChunks, 1);
6587+ i64 validitySize = sqlite3_column_bytes(stmtChunks, 1);
6588+ if (validitySize != p->chunk_size / CHAR_BIT) {
6589+ // IMP: V05271_22109
6590+ vtab_set_error(
6591+ &p->base,
6592+ "chunk validity size doesn't match - expected %lld, found %lld",
6593+ p->chunk_size / CHAR_BIT, validitySize);
6594+ rc = SQLITE_ERROR;
6595+ goto cleanup;
6596+ }
6597+
6598+ i64 *chunkRowids = (i64 *)sqlite3_column_blob(stmtChunks, 2);
6599+ i64 rowidsSize = sqlite3_column_bytes(stmtChunks, 2);
6600+ if (rowidsSize != p->chunk_size * sizeof(i64)) {
6601+ // IMP: V02796_19635
6602+ vtab_set_error(&p->base, "rowids size doesn't match");
6603+ vtab_set_error(
6604+ &p->base,
6605+ "chunk rowids size doesn't match - expected %lld, found %lld",
6606+ p->chunk_size * sizeof(i64), rowidsSize);
6607+ rc = SQLITE_ERROR;
6608+ goto cleanup;
6609+ }
6610+
6611+ // open the vector chunk blob for the current chunk
6612+ rc = sqlite3_blob_open(p->db, p->schemaName,
6613+ p->shadowVectorChunksNames[vectorColumnIdx],
6614+ "vectors", chunk_id, 0, &blobVectors);
6615+ if (rc != SQLITE_OK) {
6616+ vtab_set_error(&p->base, "could not open vectors blob for chunk %lld",
6617+ chunk_id);
6618+ rc = SQLITE_ERROR;
6619+ goto cleanup;
6620+ }
6621+
6622+ i64 currentBaseVectorsSize = sqlite3_blob_bytes(blobVectors);
6623+ i64 expectedBaseVectorsSize =
6624+ p->chunk_size * vector_column_byte_size(*vector_column);
6625+ if (currentBaseVectorsSize != expectedBaseVectorsSize) {
6626+ // IMP: V16465_00535
6627+ vtab_set_error(
6628+ &p->base,
6629+ "vectors blob size doesn't match - expected %lld, found %lld",
6630+ expectedBaseVectorsSize, currentBaseVectorsSize);
6631+ rc = SQLITE_ERROR;
6632+ goto cleanup;
6633+ }
6634+ rc = sqlite3_blob_read(blobVectors, baseVectors, currentBaseVectorsSize, 0);
6635+
6636+ if (rc != SQLITE_OK) {
6637+ vtab_set_error(&p->base, "vectors blob read error for %lld", chunk_id);
6638+ rc = SQLITE_ERROR;
6639+ goto cleanup;
6640+ }
6641+
6642+ bitmap_copy(b, chunkValidity, p->chunk_size);
6643+ if (arrayRowidsIn) {
6644+ bitmap_clear(bmRowids, p->chunk_size);
6645+
6646+ for (int i = 0; i < p->chunk_size; i++) {
6647+ if (!bitmap_get(chunkValidity, i)) {
6648+ continue;
6649+ }
6650+ i64 rowid = chunkRowids[i];
6651+ void *in = bsearch(&rowid, arrayRowidsIn->z, arrayRowidsIn->length,
6652+ sizeof(i64), _cmp);
6653+ bitmap_set(bmRowids, i, in ? 1 : 0);
6654+ }
6655+ bitmap_and_inplace(b, bmRowids, p->chunk_size);
6656+ }
6657+
6658+ if(hasMetadataFilters) {
6659+ for(int i = 0; i < argc; i++) {
6660+ int idx = 1 + (i * 4);
6661+ char kind = idxStr[idx + 0];
6662+ if(kind != VEC0_IDXSTR_KIND_METADATA_CONSTRAINT) {
6663+ continue;
6664+ }
6665+ int metadata_idx = idxStr[idx + 1] - 'A';
6666+ int operator = idxStr[idx + 2];
6667+
6668+ if(!metadataBlobs[metadata_idx]) {
6669+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 0, &metadataBlobs[metadata_idx]);
6670+ vtab_set_error(&p->base, "Could not open metadata blob");
6671+ if(rc != SQLITE_OK) {
6672+ goto cleanup;
6673+ }
6674+ }
6675+
6676+ bitmap_clear(bmMetadata, p->chunk_size);
6677+ rc = vec0_set_metadata_filter_bitmap(p, metadata_idx, operator, argv[i], metadataBlobs[metadata_idx], chunk_id, bmMetadata, p->chunk_size, aMetadataIn, i);
6678+ if(rc != SQLITE_OK) {
6679+ vtab_set_error(&p->base, "Could not filter metadata fields");
6680+ if(rc != SQLITE_OK) {
6681+ goto cleanup;
6682+ }
6683+ }
6684+ bitmap_and_inplace(b, bmMetadata, p->chunk_size);
6685+ }
6686+ }
6687+
6688+
6689+ for (int i = 0; i < p->chunk_size; i++) {
6690+ if (!bitmap_get(b, i)) {
6691+ continue;
6692+ };
6693+
6694+ f32 result;
6695+ switch (vector_column->element_type) {
6696+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: {
6697+ const f32 *base_i =
6698+ ((f32 *)baseVectors) + (i * vector_column->dimensions);
6699+ switch (vector_column->distance_metric) {
6700+ case VEC0_DISTANCE_METRIC_L2: {
6701+ result = distance_l2_sqr_float(base_i, (f32 *)queryVector,
6702+ &vector_column->dimensions);
6703+ break;
6704+ }
6705+ case VEC0_DISTANCE_METRIC_L1: {
6706+ result = distance_l1_f32(base_i, (f32 *)queryVector,
6707+ &vector_column->dimensions);
6708+ break;
6709+ }
6710+ case VEC0_DISTANCE_METRIC_COSINE: {
6711+ result = distance_cosine_float(base_i, (f32 *)queryVector,
6712+ &vector_column->dimensions);
6713+ break;
6714+ }
6715+ }
6716+ break;
6717+ }
6718+ case SQLITE_VEC_ELEMENT_TYPE_INT8: {
6719+ const i8 *base_i =
6720+ ((i8 *)baseVectors) + (i * vector_column->dimensions);
6721+ switch (vector_column->distance_metric) {
6722+ case VEC0_DISTANCE_METRIC_L2: {
6723+ result = distance_l2_sqr_int8(base_i, (i8 *)queryVector,
6724+ &vector_column->dimensions);
6725+ break;
6726+ }
6727+ case VEC0_DISTANCE_METRIC_L1: {
6728+ result = distance_l1_int8(base_i, (i8 *)queryVector,
6729+ &vector_column->dimensions);
6730+ break;
6731+ }
6732+ case VEC0_DISTANCE_METRIC_COSINE: {
6733+ result = distance_cosine_int8(base_i, (i8 *)queryVector,
6734+ &vector_column->dimensions);
6735+ break;
6736+ }
6737+ }
6738+
6739+ break;
6740+ }
6741+ case SQLITE_VEC_ELEMENT_TYPE_BIT: {
6742+ const u8 *base_i =
6743+ ((u8 *)baseVectors) + (i * (vector_column->dimensions / CHAR_BIT));
6744+ result = distance_hamming(base_i, (u8 *)queryVector,
6745+ &vector_column->dimensions);
6746+ break;
6747+ }
6748+ }
6749+
6750+ chunk_distances[i] = result;
6751+ }
6752+
6753+ int used1;
6754+ min_idx(chunk_distances, p->chunk_size, b, chunk_topk_idxs,
6755+ min(k, p->chunk_size), bTaken, &used1);
6756+
6757+ i64 used;
6758+ merge_sorted_lists(topk_distances, topk_rowids, k_used, chunk_distances,
6759+ chunkRowids, chunk_topk_idxs,
6760+ min(min(k, p->chunk_size), used1), tmp_topk_distances,
6761+ tmp_topk_rowids, k, &used);
6762+
6763+ for (int i = 0; i < used; i++) {
6764+ topk_rowids[i] = tmp_topk_rowids[i];
6765+ topk_distances[i] = tmp_topk_distances[i];
6766+ }
6767+ k_used = used;
6768+ // blobVectors is always opened with read-only permissions, so this never
6769+ // fails.
6770+ sqlite3_blob_close(blobVectors);
6771+ blobVectors = NULL;
6772+ }
6773+
6774+ *out_topk_rowids = topk_rowids;
6775+ *out_topk_distances = topk_distances;
6776+ *out_used = k_used;
6777+ rc = SQLITE_OK;
6778+
6779+cleanup:
6780+ if (rc != SQLITE_OK) {
6781+ sqlite3_free(topk_rowids);
6782+ sqlite3_free(topk_distances);
6783+ }
6784+ sqlite3_free(chunk_topk_idxs);
6785+ sqlite3_free(tmp_topk_rowids);
6786+ sqlite3_free(tmp_topk_distances);
6787+ sqlite3_free(b);
6788+ sqlite3_free(bTaken);
6789+ sqlite3_free(bmRowids);
6790+ sqlite3_free(baseVectors);
6791+ sqlite3_free(chunk_distances);
6792+ sqlite3_free(bmMetadata);
6793+ for(int i = 0; i < VEC0_MAX_METADATA_COLUMNS; i++) {
6794+ sqlite3_blob_close(metadataBlobs[i]);
6795+ }
6796+ // blobVectors is always opened with read-only permissions, so this never
6797+ // fails.
6798+ sqlite3_blob_close(blobVectors);
6799+ return rc;
6800+}
6801+
6802+int vec0Filter_knn(vec0_cursor *pCur, vec0_vtab *p, int idxNum,
6803+ const char *idxStr, int argc, sqlite3_value **argv) {
6804+ assert(argc == (strlen(idxStr)-1) / 4);
6805+ int rc;
6806+ struct vec0_query_knn_data *knn_data;
6807+
6808+ int vectorColumnIdx = idxNum;
6809+ struct VectorColumnDefinition *vector_column =
6810+ &p->vector_columns[vectorColumnIdx];
6811+
6812+ struct Array *arrayRowidsIn = NULL;
6813+ sqlite3_stmt *stmtChunks = NULL;
6814+ void *queryVector;
6815+ size_t dimensions;
6816+ enum VectorElementType elementType;
6817+ vector_cleanup queryVectorCleanup = vector_cleanup_noop;
6818+ char *pzError;
6819+ knn_data = sqlite3_malloc(sizeof(*knn_data));
6820+ if (!knn_data) {
6821+ return SQLITE_NOMEM;
6822+ }
6823+ memset(knn_data, 0, sizeof(*knn_data));
6824+ // array of `struct Vec0MetadataIn`, IF there are any `xxx in (...)` metadata constraints
6825+ struct Array * aMetadataIn = NULL;
6826+
6827+ int query_idx =-1;
6828+ int k_idx = -1;
6829+ int rowid_in_idx = -1;
6830+ for(int i = 0; i < argc; i++) {
6831+ if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_MATCH) {
6832+ query_idx = i;
6833+ }
6834+ if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_K) {
6835+ k_idx = i;
6836+ }
6837+ if(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_KNN_ROWID_IN) {
6838+ rowid_in_idx = i;
6839+ }
6840+ }
6841+ assert(query_idx >= 0);
6842+ assert(k_idx >= 0);
6843+
6844+ // make sure the query vector matches the vector column (type dimensions etc.)
6845+ rc = vector_from_value(argv[query_idx], &queryVector, &dimensions, &elementType,
6846+ &queryVectorCleanup, &pzError);
6847+
6848+ if (rc != SQLITE_OK) {
6849+ vtab_set_error(&p->base,
6850+ "Query vector on the \"%.*s\" column is invalid: %z",
6851+ vector_column->name_length, vector_column->name, pzError);
6852+ rc = SQLITE_ERROR;
6853+ goto cleanup;
6854+ }
6855+ if (elementType != vector_column->element_type) {
6856+ vtab_set_error(
6857+ &p->base,
6858+ "Query vector for the \"%.*s\" column is expected to be of type "
6859+ "%s, but a %s vector was provided.",
6860+ vector_column->name_length, vector_column->name,
6861+ vector_subtype_name(vector_column->element_type),
6862+ vector_subtype_name(elementType));
6863+ rc = SQLITE_ERROR;
6864+ goto cleanup;
6865+ }
6866+ if (dimensions != vector_column->dimensions) {
6867+ vtab_set_error(
6868+ &p->base,
6869+ "Dimension mismatch for query vector for the \"%.*s\" column. "
6870+ "Expected %d dimensions but received %d.",
6871+ vector_column->name_length, vector_column->name,
6872+ vector_column->dimensions, dimensions);
6873+ rc = SQLITE_ERROR;
6874+ goto cleanup;
6875+ }
6876+
6877+ i64 k = sqlite3_value_int64(argv[k_idx]);
6878+ if (k < 0) {
6879+ vtab_set_error(
6880+ &p->base, "k value in knn queries must be greater than or equal to 0.");
6881+ rc = SQLITE_ERROR;
6882+ goto cleanup;
6883+ }
6884+#define SQLITE_VEC_VEC0_K_MAX 4096
6885+ if (k > SQLITE_VEC_VEC0_K_MAX) {
6886+ vtab_set_error(
6887+ &p->base,
6888+ "k value in knn query too large, provided %lld and the limit is %lld",
6889+ k, SQLITE_VEC_VEC0_K_MAX);
6890+ rc = SQLITE_ERROR;
6891+ goto cleanup;
6892+ }
6893+
6894+ if (k == 0) {
6895+ knn_data->k = 0;
6896+ pCur->knn_data = knn_data;
6897+ pCur->query_plan = VEC0_QUERY_PLAN_KNN;
6898+ rc = SQLITE_OK;
6899+ goto cleanup;
6900+ }
6901+
6902+// handle when a `rowid in (...)` operation was provided
6903+// Array of all the rowids that appear in any `rowid in (...)` constraint.
6904+// NULL if none were provided, which means a "full" scan.
6905+#if COMPILER_SUPPORTS_VTAB_IN
6906+ if (rowid_in_idx >= 0) {
6907+ sqlite3_value *item;
6908+ int rc;
6909+ arrayRowidsIn = sqlite3_malloc(sizeof(*arrayRowidsIn));
6910+ if (!arrayRowidsIn) {
6911+ rc = SQLITE_NOMEM;
6912+ goto cleanup;
6913+ }
6914+ memset(arrayRowidsIn, 0, sizeof(*arrayRowidsIn));
6915+
6916+ rc = array_init(arrayRowidsIn, sizeof(i64), 32);
6917+ if (rc != SQLITE_OK) {
6918+ goto cleanup;
6919+ }
6920+ for (rc = sqlite3_vtab_in_first(argv[rowid_in_idx], &item); rc == SQLITE_OK && item;
6921+ rc = sqlite3_vtab_in_next(argv[rowid_in_idx], &item)) {
6922+ i64 rowid;
6923+ if (p->pkIsText) {
6924+ rc = vec0_rowid_from_id(p, item, &rowid);
6925+ if (rc != SQLITE_OK) {
6926+ goto cleanup;
6927+ }
6928+ } else {
6929+ rowid = sqlite3_value_int64(item);
6930+ }
6931+ rc = array_append(arrayRowidsIn, &rowid);
6932+ if (rc != SQLITE_OK) {
6933+ goto cleanup;
6934+ }
6935+ }
6936+ if (rc != SQLITE_DONE) {
6937+ vtab_set_error(&p->base, "error processing rowid in (...) array");
6938+ goto cleanup;
6939+ }
6940+ qsort(arrayRowidsIn->z, arrayRowidsIn->length, arrayRowidsIn->element_size,
6941+ _cmp);
6942+ }
6943+#endif
6944+
6945+ #if COMPILER_SUPPORTS_VTAB_IN
6946+ for(int i = 0; i < argc; i++) {
6947+ if(!(idxStr[1 + (i*4)] == VEC0_IDXSTR_KIND_METADATA_CONSTRAINT && idxStr[1 + (i*4) + 2] == VEC0_METADATA_OPERATOR_IN)) {
6948+ continue;
6949+ }
6950+ int metadata_idx = idxStr[1 + (i*4) + 1] - 'A';
6951+ if(!aMetadataIn) {
6952+ aMetadataIn = sqlite3_malloc(sizeof(*aMetadataIn));
6953+ if(!aMetadataIn) {
6954+ rc = SQLITE_NOMEM;
6955+ goto cleanup;
6956+ }
6957+ memset(aMetadataIn, 0, sizeof(*aMetadataIn));
6958+ rc = array_init(aMetadataIn, sizeof(struct Vec0MetadataIn), 8);
6959+ if(rc != SQLITE_OK) {
6960+ goto cleanup;
6961+ }
6962+ }
6963+
6964+ struct Vec0MetadataIn item;
6965+ memset(&item, 0, sizeof(item));
6966+ item.metadata_idx=metadata_idx;
6967+ item.argv_idx = i;
6968+
6969+ switch(p->metadata_columns[metadata_idx].kind) {
6970+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
6971+ rc = array_init(&item.array, sizeof(i64), 16);
6972+ if(rc != SQLITE_OK) {
6973+ goto cleanup;
6974+ }
6975+ sqlite3_value *entry;
6976+ for (rc = sqlite3_vtab_in_first(argv[i], &entry); rc == SQLITE_OK && entry; rc = sqlite3_vtab_in_next(argv[i], &entry)) {
6977+ i64 v = sqlite3_value_int64(entry);
6978+ rc = array_append(&item.array, &v);
6979+ if (rc != SQLITE_OK) {
6980+ goto cleanup;
6981+ }
6982+ }
6983+
6984+ if (rc != SQLITE_DONE) {
6985+ vtab_set_error(&p->base, "Error fetching next value in `x in (...)` integer expression");
6986+ goto cleanup;
6987+ }
6988+
6989+ break;
6990+ }
6991+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
6992+ rc = array_init(&item.array, sizeof(struct Vec0MetadataInTextEntry), 16);
6993+ if(rc != SQLITE_OK) {
6994+ goto cleanup;
6995+ }
6996+ sqlite3_value *entry;
6997+ for (rc = sqlite3_vtab_in_first(argv[i], &entry); rc == SQLITE_OK && entry; rc = sqlite3_vtab_in_next(argv[i], &entry)) {
6998+ const char * s = (const char *) sqlite3_value_text(entry);
6999+ int n = sqlite3_value_bytes(entry);
7000+
7001+ struct Vec0MetadataInTextEntry entry;
7002+ entry.zString = sqlite3_mprintf("%.*s", n, s);
7003+ if(!entry.zString) {
7004+ rc = SQLITE_NOMEM;
7005+ goto cleanup;
7006+ }
7007+ entry.n = n;
7008+ rc = array_append(&item.array, &entry);
7009+ if (rc != SQLITE_OK) {
7010+ goto cleanup;
7011+ }
7012+ }
7013+
7014+ if (rc != SQLITE_DONE) {
7015+ vtab_set_error(&p->base, "Error fetching next value in `x in (...)` text expression");
7016+ goto cleanup;
7017+ }
7018+
7019+ break;
7020+ }
7021+ default: {
7022+ vtab_set_error(&p->base, "Internal sqlite-vec error");
7023+ goto cleanup;
7024+ }
7025+ }
7026+
7027+ rc = array_append(aMetadataIn, &item);
7028+ if(rc != SQLITE_OK) {
7029+ goto cleanup;
7030+ }
7031+ }
7032+ #endif
7033+
7034+ rc = vec0_chunks_iter(p, idxStr, argc, argv, &stmtChunks);
7035+ if (rc != SQLITE_OK) {
7036+ // IMP: V06942_23781
7037+ vtab_set_error(&p->base, "Error preparing stmtChunk: %s",
7038+ sqlite3_errmsg(p->db));
7039+ goto cleanup;
7040+ }
7041+
7042+ i64 *topk_rowids = NULL;
7043+ f32 *topk_distances = NULL;
7044+ i64 k_used = 0;
7045+ rc = vec0Filter_knn_chunks_iter(p, stmtChunks, vector_column, vectorColumnIdx,
7046+ arrayRowidsIn, aMetadataIn, idxStr, argc, argv, queryVector, k, &topk_rowids,
7047+ &topk_distances, &k_used);
7048+ if (rc != SQLITE_OK) {
7049+ goto cleanup;
7050+ }
7051+
7052+ knn_data->current_idx = 0;
7053+ knn_data->k = k;
7054+ knn_data->rowids = topk_rowids;
7055+ knn_data->distances = topk_distances;
7056+ knn_data->k_used = k_used;
7057+
7058+ pCur->knn_data = knn_data;
7059+ pCur->query_plan = VEC0_QUERY_PLAN_KNN;
7060+ rc = SQLITE_OK;
7061+
7062+cleanup:
7063+ sqlite3_finalize(stmtChunks);
7064+ array_cleanup(arrayRowidsIn);
7065+ sqlite3_free(arrayRowidsIn);
7066+ queryVectorCleanup(queryVector);
7067+ if(aMetadataIn) {
7068+ for(size_t i = 0; i < aMetadataIn->length; i++) {
7069+ struct Vec0MetadataIn* item = &((struct Vec0MetadataIn *) aMetadataIn->z)[i];
7070+ for(size_t j = 0; j < item->array.length; j++) {
7071+ if(p->metadata_columns[item->metadata_idx].kind == VEC0_METADATA_COLUMN_KIND_TEXT) {
7072+ struct Vec0MetadataInTextEntry entry = ((struct Vec0MetadataInTextEntry*)item->array.z)[j];
7073+ sqlite3_free(entry.zString);
7074+ }
7075+ }
7076+ array_cleanup(&item->array);
7077+ }
7078+ array_cleanup(aMetadataIn);
7079+ }
7080+
7081+ sqlite3_free(aMetadataIn);
7082+
7083+ return rc;
7084+}
7085+
7086+int vec0Filter_fullscan(vec0_vtab *p, vec0_cursor *pCur) {
7087+ int rc;
7088+ char *zSql;
7089+ struct vec0_query_fullscan_data *fullscan_data;
7090+
7091+ fullscan_data = sqlite3_malloc(sizeof(*fullscan_data));
7092+ if (!fullscan_data) {
7093+ return SQLITE_NOMEM;
7094+ }
7095+ memset(fullscan_data, 0, sizeof(*fullscan_data));
7096+
7097+ zSql = sqlite3_mprintf(" SELECT rowid "
7098+ " FROM " VEC0_SHADOW_ROWIDS_NAME
7099+ " ORDER by chunk_id, chunk_offset ",
7100+ p->schemaName, p->tableName);
7101+ if (!zSql) {
7102+ rc = SQLITE_NOMEM;
7103+ goto error;
7104+ }
7105+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &fullscan_data->rowids_stmt, NULL);
7106+ sqlite3_free(zSql);
7107+ if (rc != SQLITE_OK) {
7108+ // IMP: V09901_26739
7109+ vtab_set_error(&p->base, "Error preparing rowid scan: %s",
7110+ sqlite3_errmsg(p->db));
7111+ goto error;
7112+ }
7113+
7114+ rc = sqlite3_step(fullscan_data->rowids_stmt);
7115+
7116+ // DONE when there's no rowids, ROW when there are, both "success"
7117+ if (!(rc == SQLITE_ROW || rc == SQLITE_DONE)) {
7118+ goto error;
7119+ }
7120+
7121+ fullscan_data->done = rc == SQLITE_DONE;
7122+ pCur->query_plan = VEC0_QUERY_PLAN_FULLSCAN;
7123+ pCur->fullscan_data = fullscan_data;
7124+ return SQLITE_OK;
7125+
7126+error:
7127+ vec0_query_fullscan_data_clear(fullscan_data);
7128+ sqlite3_free(fullscan_data);
7129+ return rc;
7130+}
7131+
7132+int vec0Filter_point(vec0_cursor *pCur, vec0_vtab *p, int argc,
7133+ sqlite3_value **argv) {
7134+ int rc;
7135+ assert(argc == 1);
7136+ i64 rowid;
7137+ struct vec0_query_point_data *point_data = NULL;
7138+
7139+ point_data = sqlite3_malloc(sizeof(*point_data));
7140+ if (!point_data) {
7141+ rc = SQLITE_NOMEM;
7142+ goto error;
7143+ }
7144+ memset(point_data, 0, sizeof(*point_data));
7145+
7146+ if (p->pkIsText) {
7147+ rc = vec0_rowid_from_id(p, argv[0], &rowid);
7148+ if (rc == SQLITE_EMPTY) {
7149+ goto eof;
7150+ }
7151+ if (rc != SQLITE_OK) {
7152+ goto error;
7153+ }
7154+ } else {
7155+ rowid = sqlite3_value_int64(argv[0]);
7156+ }
7157+
7158+ for (int i = 0; i < p->numVectorColumns; i++) {
7159+ rc = vec0_get_vector_data(p, rowid, i, &point_data->vectors[i], NULL);
7160+ if (rc == SQLITE_EMPTY) {
7161+ goto eof;
7162+ }
7163+ if (rc != SQLITE_OK) {
7164+ goto error;
7165+ }
7166+ }
7167+
7168+ point_data->rowid = rowid;
7169+ point_data->done = 0;
7170+ pCur->point_data = point_data;
7171+ pCur->query_plan = VEC0_QUERY_PLAN_POINT;
7172+ return SQLITE_OK;
7173+
7174+eof:
7175+ point_data->rowid = rowid;
7176+ point_data->done = 1;
7177+ pCur->point_data = point_data;
7178+ pCur->query_plan = VEC0_QUERY_PLAN_POINT;
7179+ return SQLITE_OK;
7180+
7181+error:
7182+ vec0_query_point_data_clear(point_data);
7183+ sqlite3_free(point_data);
7184+ return rc;
7185+}
7186+
7187+static int vec0Filter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
7188+ const char *idxStr, int argc, sqlite3_value **argv) {
7189+ vec0_vtab *p = (vec0_vtab *)pVtabCursor->pVtab;
7190+ vec0_cursor *pCur = (vec0_cursor *)pVtabCursor;
7191+ vec0_cursor_clear(pCur);
7192+
7193+ int idxStrLength = strlen(idxStr);
7194+ if(idxStrLength <= 0) {
7195+ return SQLITE_ERROR;
7196+ }
7197+ if((idxStrLength-1) % 4 != 0) {
7198+ return SQLITE_ERROR;
7199+ }
7200+ int numValueEntries = (idxStrLength-1) / 4;
7201+ if(numValueEntries != argc) {
7202+ return SQLITE_ERROR;
7203+ }
7204+
7205+ char query_plan = idxStr[0];
7206+ switch(query_plan) {
7207+ case VEC0_QUERY_PLAN_FULLSCAN:
7208+ return vec0Filter_fullscan(p, pCur);
7209+ case VEC0_QUERY_PLAN_KNN:
7210+ return vec0Filter_knn(pCur, p, idxNum, idxStr, argc, argv);
7211+ case VEC0_QUERY_PLAN_POINT:
7212+ return vec0Filter_point(pCur, p, argc, argv);
7213+ default:
7214+ vtab_set_error(pVtabCursor->pVtab, "unknown idxStr '%s'", idxStr);
7215+ return SQLITE_ERROR;
7216+ }
7217+}
7218+
7219+static int vec0Rowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) {
7220+ vec0_cursor *pCur = (vec0_cursor *)cur;
7221+ switch (pCur->query_plan) {
7222+ case VEC0_QUERY_PLAN_FULLSCAN: {
7223+ *pRowid = sqlite3_column_int64(pCur->fullscan_data->rowids_stmt, 0);
7224+ return SQLITE_OK;
7225+ }
7226+ case VEC0_QUERY_PLAN_POINT: {
7227+ *pRowid = pCur->point_data->rowid;
7228+ return SQLITE_OK;
7229+ }
7230+ case VEC0_QUERY_PLAN_KNN: {
7231+ vtab_set_error(cur->pVtab,
7232+ "Internal sqlite-vec error: expected point query plan in "
7233+ "vec0Rowid, found %d",
7234+ pCur->query_plan);
7235+ return SQLITE_ERROR;
7236+ }
7237+ }
7238+ return SQLITE_ERROR;
7239+}
7240+
7241+static int vec0Next(sqlite3_vtab_cursor *cur) {
7242+ vec0_cursor *pCur = (vec0_cursor *)cur;
7243+ switch (pCur->query_plan) {
7244+ case VEC0_QUERY_PLAN_FULLSCAN: {
7245+ if (!pCur->fullscan_data) {
7246+ return SQLITE_ERROR;
7247+ }
7248+ int rc = sqlite3_step(pCur->fullscan_data->rowids_stmt);
7249+ if (rc == SQLITE_DONE) {
7250+ pCur->fullscan_data->done = 1;
7251+ return SQLITE_OK;
7252+ }
7253+ if (rc == SQLITE_ROW) {
7254+ return SQLITE_OK;
7255+ }
7256+ return SQLITE_ERROR;
7257+ }
7258+ case VEC0_QUERY_PLAN_KNN: {
7259+ if (!pCur->knn_data) {
7260+ return SQLITE_ERROR;
7261+ }
7262+
7263+ pCur->knn_data->current_idx++;
7264+ return SQLITE_OK;
7265+ }
7266+ case VEC0_QUERY_PLAN_POINT: {
7267+ if (!pCur->point_data) {
7268+ return SQLITE_ERROR;
7269+ }
7270+ pCur->point_data->done = 1;
7271+ return SQLITE_OK;
7272+ }
7273+ }
7274+ return SQLITE_ERROR;
7275+}
7276+
7277+static int vec0Eof(sqlite3_vtab_cursor *cur) {
7278+ vec0_cursor *pCur = (vec0_cursor *)cur;
7279+ switch (pCur->query_plan) {
7280+ case VEC0_QUERY_PLAN_FULLSCAN: {
7281+ if (!pCur->fullscan_data) {
7282+ return 1;
7283+ }
7284+ return pCur->fullscan_data->done;
7285+ }
7286+ case VEC0_QUERY_PLAN_KNN: {
7287+ if (!pCur->knn_data) {
7288+ return 1;
7289+ }
7290+ // return (pCur->knn_data->current_idx >= pCur->knn_data->k) ||
7291+ // (pCur->knn_data->distances[pCur->knn_data->current_idx] == FLT_MAX);
7292+ return (pCur->knn_data->current_idx >= pCur->knn_data->k_used);
7293+ }
7294+ case VEC0_QUERY_PLAN_POINT: {
7295+ if (!pCur->point_data) {
7296+ return 1;
7297+ }
7298+ return pCur->point_data->done;
7299+ }
7300+ }
7301+ return 1;
7302+}
7303+
7304+static int vec0Column_fullscan(vec0_vtab *pVtab, vec0_cursor *pCur,
7305+ sqlite3_context *context, int i) {
7306+ if (!pCur->fullscan_data) {
7307+ sqlite3_result_error(
7308+ context, "Internal sqlite-vec error: fullscan_data is NULL.", -1);
7309+ return SQLITE_ERROR;
7310+ }
7311+ i64 rowid = sqlite3_column_int64(pCur->fullscan_data->rowids_stmt, 0);
7312+ if (i == VEC0_COLUMN_ID) {
7313+ return vec0_result_id(pVtab, context, rowid);
7314+ }
7315+ else if (vec0_column_idx_is_vector(pVtab, i)) {
7316+ void *v;
7317+ int sz;
7318+ int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i);
7319+ int rc = vec0_get_vector_data(pVtab, rowid, vector_idx, &v, &sz);
7320+ if (rc != SQLITE_OK) {
7321+ return rc;
7322+ }
7323+ sqlite3_result_blob(context, v, sz, sqlite3_free);
7324+ sqlite3_result_subtype(context,
7325+ pVtab->vector_columns[vector_idx].element_type);
7326+
7327+ }
7328+ else if (i == vec0_column_distance_idx(pVtab)) {
7329+ sqlite3_result_null(context);
7330+ }
7331+ else if(vec0_column_idx_is_partition(pVtab, i)) {
7332+ int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i);
7333+ sqlite3_value * v;
7334+ int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v);
7335+ if(rc == SQLITE_OK) {
7336+ sqlite3_result_value(context, v);
7337+ sqlite3_value_free(v);
7338+ }else {
7339+ sqlite3_result_error_code(context, rc);
7340+ }
7341+ }
7342+ else if(vec0_column_idx_is_auxiliary(pVtab, i)) {
7343+ int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i);
7344+ sqlite3_value * v;
7345+ int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v);
7346+ if(rc == SQLITE_OK) {
7347+ sqlite3_result_value(context, v);
7348+ sqlite3_value_free(v);
7349+ }else {
7350+ sqlite3_result_error_code(context, rc);
7351+ }
7352+ }
7353+
7354+ else if(vec0_column_idx_is_metadata(pVtab, i)) {
7355+ if(sqlite3_vtab_nochange(context)) {
7356+ return SQLITE_OK;
7357+ }
7358+ int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i);
7359+ int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context);
7360+ if(rc != SQLITE_OK) {
7361+ // IMP: V15466_32305
7362+ const char * zErr = sqlite3_mprintf(
7363+ "Could not extract metadata value for column %.*s at rowid %lld",
7364+ pVtab->metadata_columns[metadata_idx].name_length,
7365+ pVtab->metadata_columns[metadata_idx].name, rowid
7366+ );
7367+ if(zErr) {
7368+ sqlite3_result_error(context, zErr, -1);
7369+ sqlite3_free((void *) zErr);
7370+ }else {
7371+ sqlite3_result_error_nomem(context);
7372+ }
7373+ }
7374+ }
7375+
7376+ return SQLITE_OK;
7377+}
7378+
7379+static int vec0Column_point(vec0_vtab *pVtab, vec0_cursor *pCur,
7380+ sqlite3_context *context, int i) {
7381+ if (!pCur->point_data) {
7382+ sqlite3_result_error(context,
7383+ "Internal sqlite-vec error: point_data is NULL.", -1);
7384+ return SQLITE_ERROR;
7385+ }
7386+ if (i == VEC0_COLUMN_ID) {
7387+ return vec0_result_id(pVtab, context, pCur->point_data->rowid);
7388+ }
7389+ else if (i == vec0_column_distance_idx(pVtab)) {
7390+ sqlite3_result_null(context);
7391+ return SQLITE_OK;
7392+ }
7393+ else if (vec0_column_idx_is_vector(pVtab, i)) {
7394+ if (sqlite3_vtab_nochange(context)) {
7395+ sqlite3_result_null(context);
7396+ return SQLITE_OK;
7397+ }
7398+ int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i);
7399+ sqlite3_result_blob(
7400+ context, pCur->point_data->vectors[vector_idx],
7401+ vector_column_byte_size(pVtab->vector_columns[vector_idx]),
7402+ SQLITE_TRANSIENT);
7403+ sqlite3_result_subtype(context,
7404+ pVtab->vector_columns[vector_idx].element_type);
7405+ return SQLITE_OK;
7406+ }
7407+ else if(vec0_column_idx_is_partition(pVtab, i)) {
7408+ if(sqlite3_vtab_nochange(context)) {
7409+ return SQLITE_OK;
7410+ }
7411+ int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i);
7412+ i64 rowid = pCur->point_data->rowid;
7413+ sqlite3_value * v;
7414+ int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v);
7415+ if(rc == SQLITE_OK) {
7416+ sqlite3_result_value(context, v);
7417+ sqlite3_value_free(v);
7418+ }else {
7419+ sqlite3_result_error_code(context, rc);
7420+ }
7421+ }
7422+ else if(vec0_column_idx_is_auxiliary(pVtab, i)) {
7423+ if(sqlite3_vtab_nochange(context)) {
7424+ return SQLITE_OK;
7425+ }
7426+ i64 rowid = pCur->point_data->rowid;
7427+ int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i);
7428+ sqlite3_value * v;
7429+ int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v);
7430+ if(rc == SQLITE_OK) {
7431+ sqlite3_result_value(context, v);
7432+ sqlite3_value_free(v);
7433+ }else {
7434+ sqlite3_result_error_code(context, rc);
7435+ }
7436+ }
7437+
7438+ else if(vec0_column_idx_is_metadata(pVtab, i)) {
7439+ if(sqlite3_vtab_nochange(context)) {
7440+ return SQLITE_OK;
7441+ }
7442+ i64 rowid = pCur->point_data->rowid;
7443+ int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i);
7444+ int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context);
7445+ if(rc != SQLITE_OK) {
7446+ const char * zErr = sqlite3_mprintf(
7447+ "Could not extract metadata value for column %.*s at rowid %lld",
7448+ pVtab->metadata_columns[metadata_idx].name_length,
7449+ pVtab->metadata_columns[metadata_idx].name, rowid
7450+ );
7451+ if(zErr) {
7452+ sqlite3_result_error(context, zErr, -1);
7453+ sqlite3_free((void *) zErr);
7454+ }else {
7455+ sqlite3_result_error_nomem(context);
7456+ }
7457+ }
7458+ }
7459+
7460+ return SQLITE_OK;
7461+}
7462+
7463+static int vec0Column_knn(vec0_vtab *pVtab, vec0_cursor *pCur,
7464+ sqlite3_context *context, int i) {
7465+ if (!pCur->knn_data) {
7466+ sqlite3_result_error(context,
7467+ "Internal sqlite-vec error: knn_data is NULL.", -1);
7468+ return SQLITE_ERROR;
7469+ }
7470+ if (i == VEC0_COLUMN_ID) {
7471+ i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx];
7472+ return vec0_result_id(pVtab, context, rowid);
7473+ }
7474+ else if (i == vec0_column_distance_idx(pVtab)) {
7475+ sqlite3_result_double(
7476+ context, pCur->knn_data->distances[pCur->knn_data->current_idx]);
7477+ return SQLITE_OK;
7478+ }
7479+ else if (vec0_column_idx_is_vector(pVtab, i)) {
7480+ void *out;
7481+ int sz;
7482+ int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i);
7483+ int rc = vec0_get_vector_data(
7484+ pVtab, pCur->knn_data->rowids[pCur->knn_data->current_idx], vector_idx,
7485+ &out, &sz);
7486+ if (rc != SQLITE_OK) {
7487+ return rc;
7488+ }
7489+ sqlite3_result_blob(context, out, sz, sqlite3_free);
7490+ sqlite3_result_subtype(context,
7491+ pVtab->vector_columns[vector_idx].element_type);
7492+ return SQLITE_OK;
7493+ }
7494+ else if(vec0_column_idx_is_partition(pVtab, i)) {
7495+ int partition_idx = vec0_column_idx_to_partition_idx(pVtab, i);
7496+ i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx];
7497+ sqlite3_value * v;
7498+ int rc = vec0_get_partition_value_for_rowid(pVtab, rowid, partition_idx, &v);
7499+ if(rc == SQLITE_OK) {
7500+ sqlite3_result_value(context, v);
7501+ sqlite3_value_free(v);
7502+ }else {
7503+ sqlite3_result_error_code(context, rc);
7504+ }
7505+ }
7506+ else if(vec0_column_idx_is_auxiliary(pVtab, i)) {
7507+ int auxiliary_idx = vec0_column_idx_to_auxiliary_idx(pVtab, i);
7508+ i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx];
7509+ sqlite3_value * v;
7510+ int rc = vec0_get_auxiliary_value_for_rowid(pVtab, rowid, auxiliary_idx, &v);
7511+ if(rc == SQLITE_OK) {
7512+ sqlite3_result_value(context, v);
7513+ sqlite3_value_free(v);
7514+ }else {
7515+ sqlite3_result_error_code(context, rc);
7516+ }
7517+ }
7518+
7519+ else if(vec0_column_idx_is_metadata(pVtab, i)) {
7520+ int metadata_idx = vec0_column_idx_to_metadata_idx(pVtab, i);
7521+ i64 rowid = pCur->knn_data->rowids[pCur->knn_data->current_idx];
7522+ int rc = vec0_result_metadata_value_for_rowid(pVtab, rowid, metadata_idx, context);
7523+ if(rc != SQLITE_OK) {
7524+ const char * zErr = sqlite3_mprintf(
7525+ "Could not extract metadata value for column %.*s at rowid %lld",
7526+ pVtab->metadata_columns[metadata_idx].name_length,
7527+ pVtab->metadata_columns[metadata_idx].name, rowid
7528+ );
7529+ if(zErr) {
7530+ sqlite3_result_error(context, zErr, -1);
7531+ sqlite3_free((void *) zErr);
7532+ }else {
7533+ sqlite3_result_error_nomem(context);
7534+ }
7535+ }
7536+ }
7537+
7538+ return SQLITE_OK;
7539+}
7540+
7541+static int vec0Column(sqlite3_vtab_cursor *cur, sqlite3_context *context,
7542+ int i) {
7543+ vec0_cursor *pCur = (vec0_cursor *)cur;
7544+ vec0_vtab *pVtab = (vec0_vtab *)cur->pVtab;
7545+ switch (pCur->query_plan) {
7546+ case VEC0_QUERY_PLAN_FULLSCAN: {
7547+ return vec0Column_fullscan(pVtab, pCur, context, i);
7548+ }
7549+ case VEC0_QUERY_PLAN_KNN: {
7550+ return vec0Column_knn(pVtab, pCur, context, i);
7551+ }
7552+ case VEC0_QUERY_PLAN_POINT: {
7553+ return vec0Column_point(pVtab, pCur, context, i);
7554+ }
7555+ }
7556+ return SQLITE_OK;
7557+}
7558+
7559+/**
7560+ * @brief Handles the "insert rowid" step of a row insert operation of a vec0
7561+ * table.
7562+ *
7563+ * This function will insert a new row into the _rowids vec0 shadow table.
7564+ *
7565+ * @param p: virtual table
7566+ * @param idValue: Value containing the inserted rowid/id value.
7567+ * @param rowid: Output rowid, will point to the "real" i64 rowid
7568+ * value that was inserted
7569+ * @return int SQLITE_OK on success, error code on failure
7570+ */
7571+int vec0Update_InsertRowidStep(vec0_vtab *p, sqlite3_value *idValue,
7572+ i64 *rowid) {
7573+
7574+ /**
7575+ * An insert into a vec0 table can happen a few different ways:
7576+ * 1) With default INTEGER primary key: With a supplied i64 rowid
7577+ * 2) With default INTEGER primary key: WITHOUT a supplied rowid
7578+ * 3) With TEXT primary key: supplied text rowid
7579+ */
7580+
7581+ int rc;
7582+
7583+ // Option 3: vtab has a user-defined TEXT primary key, so ensure a text value
7584+ // is provided.
7585+ if (p->pkIsText) {
7586+ if (sqlite3_value_type(idValue) != SQLITE_TEXT) {
7587+ // IMP: V04200_21039
7588+ vtab_set_error(&p->base,
7589+ "The %s virtual table was declared with a TEXT primary "
7590+ "key, but a non-TEXT value was provided in an INSERT.",
7591+ p->tableName);
7592+ return SQLITE_ERROR;
7593+ }
7594+
7595+ return vec0_rowids_insert_id(p, idValue, rowid);
7596+ }
7597+
7598+ // Option 1: User supplied a i64 rowid
7599+ if (sqlite3_value_type(idValue) == SQLITE_INTEGER) {
7600+ i64 suppliedRowid = sqlite3_value_int64(idValue);
7601+ rc = vec0_rowids_insert_rowid(p, suppliedRowid);
7602+ if (rc == SQLITE_OK) {
7603+ *rowid = suppliedRowid;
7604+ }
7605+ return rc;
7606+ }
7607+
7608+ // Option 2: User did not suppled a rowid
7609+
7610+ if (sqlite3_value_type(idValue) != SQLITE_NULL) {
7611+ // IMP: V30855_14925
7612+ vtab_set_error(&p->base,
7613+ "Only integers are allows for primary key values on %s",
7614+ p->tableName);
7615+ return SQLITE_ERROR;
7616+ }
7617+ // NULL to get next auto-incremented value
7618+ return vec0_rowids_insert_id(p, NULL, rowid);
7619+}
7620+
7621+/**
7622+ * @brief Determines the "next available" chunk position for a newly inserted
7623+ * vec0 row.
7624+ *
7625+ * This operation may insert a new "blank" chunk the _chunks table, if there is
7626+ * no more space in previous chunks.
7627+ *
7628+ * @param p: virtual table
7629+ * @param partitionKeyValues: array of partition key column values, to constrain
7630+ * against any partition key columns.
7631+ * @param chunk_rowid: Output rowid of the chunk in the _chunks virtual table
7632+ * that has the avialabiity.
7633+ * @param chunk_offset: Output the index of the available space insert the
7634+ * chunk, based on the index of the first available validity bit.
7635+ * @param pBlobValidity: Output blob of the validity column of the available
7636+ * chunk. Will be opened with read/write permissions.
7637+ * @param pValidity: Output buffer of the original chunk's validity column.
7638+ * Needs to be cleaned up with sqlite3_free().
7639+ * @return int SQLITE_OK on success, error code on failure
7640+ */
7641+int vec0Update_InsertNextAvailableStep(
7642+ vec0_vtab *p,
7643+ sqlite3_value ** partitionKeyValues,
7644+ i64 *chunk_rowid, i64 *chunk_offset,
7645+ sqlite3_blob **blobChunksValidity,
7646+ const unsigned char **bufferChunksValidity) {
7647+
7648+ int rc;
7649+ i64 validitySize;
7650+ *chunk_offset = -1;
7651+
7652+ rc = vec0_get_latest_chunk_rowid(p, chunk_rowid, partitionKeyValues);
7653+ if(rc == SQLITE_EMPTY) {
7654+ goto done;
7655+ }
7656+ if (rc != SQLITE_OK) {
7657+ goto cleanup;
7658+ }
7659+
7660+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "validity",
7661+ *chunk_rowid, 1, blobChunksValidity);
7662+ if (rc != SQLITE_OK) {
7663+ // IMP: V22053_06123
7664+ vtab_set_error(&p->base,
7665+ VEC_INTERAL_ERROR
7666+ "could not open validity blob on %s.%s.%lld",
7667+ p->schemaName, p->shadowChunksName, *chunk_rowid);
7668+ goto cleanup;
7669+ }
7670+
7671+ validitySize = sqlite3_blob_bytes(*blobChunksValidity);
7672+ if (validitySize != p->chunk_size / CHAR_BIT) {
7673+ // IMP: V29362_13432
7674+ vtab_set_error(&p->base,
7675+ VEC_INTERAL_ERROR
7676+ "validity blob size mismatch on "
7677+ "%s.%s.%lld, expected %lld but received %lld.",
7678+ p->schemaName, p->shadowChunksName, *chunk_rowid,
7679+ (i64)(p->chunk_size / CHAR_BIT), validitySize);
7680+ rc = SQLITE_ERROR;
7681+ goto cleanup;
7682+ }
7683+
7684+ *bufferChunksValidity = sqlite3_malloc(validitySize);
7685+ if (!(*bufferChunksValidity)) {
7686+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
7687+ "Could not allocate memory for validity bitmap");
7688+ rc = SQLITE_NOMEM;
7689+ goto cleanup;
7690+ }
7691+
7692+ rc = sqlite3_blob_read(*blobChunksValidity, (void *)*bufferChunksValidity,
7693+ validitySize, 0);
7694+
7695+ if (rc != SQLITE_OK) {
7696+ vtab_set_error(&p->base,
7697+ VEC_INTERAL_ERROR
7698+ "Could not read validity bitmap for %s.%s.%lld",
7699+ p->schemaName, p->shadowChunksName, *chunk_rowid);
7700+ goto cleanup;
7701+ }
7702+
7703+ // find the next available offset, ie first `0` in the bitmap.
7704+ for (int i = 0; i < validitySize; i++) {
7705+ if ((*bufferChunksValidity)[i] == 0b11111111)
7706+ continue;
7707+ for (int j = 0; j < CHAR_BIT; j++) {
7708+ if (((((*bufferChunksValidity)[i] >> j) & 1) == 0)) {
7709+ *chunk_offset = (i * CHAR_BIT) + j;
7710+ goto done;
7711+ }
7712+ }
7713+ }
7714+
7715+done:
7716+ // latest chunk was full, so need to create a new one
7717+ if (*chunk_offset == -1) {
7718+ rc = vec0_new_chunk(p, partitionKeyValues, chunk_rowid);
7719+ if (rc != SQLITE_OK) {
7720+ // IMP: V08441_25279
7721+ vtab_set_error(&p->base,
7722+ VEC_INTERAL_ERROR "Could not insert a new vector chunk");
7723+ rc = SQLITE_ERROR; // otherwise raises a DatabaseError and not operational
7724+ // error?
7725+ goto cleanup;
7726+ }
7727+ *chunk_offset = 0;
7728+
7729+ // blobChunksValidity and pValidity are stale, pointing to the previous
7730+ // (full) chunk. to re-assign them
7731+ rc = sqlite3_blob_close(*blobChunksValidity);
7732+ sqlite3_free((void *)*bufferChunksValidity);
7733+ *blobChunksValidity = NULL;
7734+ *bufferChunksValidity = NULL;
7735+ if (rc != SQLITE_OK) {
7736+ vtab_set_error(&p->base, VEC_INTERAL_ERROR
7737+ "unknown error, blobChunksValidity could not be closed, "
7738+ "please file an issue.");
7739+ rc = SQLITE_ERROR;
7740+ goto cleanup;
7741+ }
7742+
7743+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName,
7744+ "validity", *chunk_rowid, 1, blobChunksValidity);
7745+ if (rc != SQLITE_OK) {
7746+ vtab_set_error(
7747+ &p->base,
7748+ VEC_INTERAL_ERROR
7749+ "Could not open validity blob for newly created chunk %s.%s.%lld",
7750+ p->schemaName, p->shadowChunksName, *chunk_rowid);
7751+ goto cleanup;
7752+ }
7753+ validitySize = sqlite3_blob_bytes(*blobChunksValidity);
7754+ if (validitySize != p->chunk_size / CHAR_BIT) {
7755+ vtab_set_error(&p->base,
7756+ VEC_INTERAL_ERROR
7757+ "validity blob size mismatch for newly created chunk "
7758+ "%s.%s.%lld. Exepcted %lld, got %lld",
7759+ p->schemaName, p->shadowChunksName, *chunk_rowid,
7760+ p->chunk_size / CHAR_BIT, validitySize);
7761+ goto cleanup;
7762+ }
7763+ *bufferChunksValidity = sqlite3_malloc(validitySize);
7764+ rc = sqlite3_blob_read(*blobChunksValidity, (void *)*bufferChunksValidity,
7765+ validitySize, 0);
7766+ if (rc != SQLITE_OK) {
7767+ vtab_set_error(&p->base,
7768+ VEC_INTERAL_ERROR
7769+ "could not read validity blob newly created chunk "
7770+ "%s.%s.%lld",
7771+ p->schemaName, p->shadowChunksName, *chunk_rowid);
7772+ goto cleanup;
7773+ }
7774+ }
7775+
7776+ rc = SQLITE_OK;
7777+
7778+cleanup:
7779+ return rc;
7780+}
7781+
7782+/**
7783+ * @brief Write the vector data into the provided vector blob at the given
7784+ * offset
7785+ *
7786+ * @param blobVectors SQLite BLOB to write to
7787+ * @param chunk_offset the "offset" (ie validity bitmap position) to write the
7788+ * vector to
7789+ * @param bVector pointer to the vector containing data
7790+ * @param dimensions how many dimensions the vector has
7791+ * @param element_type the vector type
7792+ * @return result of sqlite3_blob_write, SQLITE_OK on success, otherwise failure
7793+ */
7794+static int
7795+vec0_write_vector_to_vector_blob(sqlite3_blob *blobVectors, i64 chunk_offset,
7796+ const void *bVector, size_t dimensions,
7797+ enum VectorElementType element_type) {
7798+ int n;
7799+ int offset;
7800+
7801+ switch (element_type) {
7802+ case SQLITE_VEC_ELEMENT_TYPE_FLOAT32:
7803+ n = dimensions * sizeof(f32);
7804+ offset = chunk_offset * dimensions * sizeof(f32);
7805+ break;
7806+ case SQLITE_VEC_ELEMENT_TYPE_INT8:
7807+ n = dimensions * sizeof(i8);
7808+ offset = chunk_offset * dimensions * sizeof(i8);
7809+ break;
7810+ case SQLITE_VEC_ELEMENT_TYPE_BIT:
7811+ n = dimensions / CHAR_BIT;
7812+ offset = chunk_offset * dimensions / CHAR_BIT;
7813+ break;
7814+ }
7815+
7816+ return sqlite3_blob_write(blobVectors, bVector, n, offset);
7817+}
7818+
7819+/**
7820+ * @brief
7821+ *
7822+ * @param p vec0 virtual table
7823+ * @param chunk_rowid: which chunk to write to
7824+ * @param chunk_offset: the offset inside the chunk to write the vector to.
7825+ * @param rowid: the rowid of the inserting row
7826+ * @param vectorDatas: array of the vector data to insert
7827+ * @param blobValidity: writeable validity blob of the row's assigned chunk.
7828+ * @param validity: snapshot buffer of the valdity column from the row's
7829+ * assigned chunk.
7830+ * @return int SQLITE_OK on success, error code on failure
7831+ */
7832+int vec0Update_InsertWriteFinalStep(vec0_vtab *p, i64 chunk_rowid,
7833+ i64 chunk_offset, i64 rowid,
7834+ void *vectorDatas[],
7835+ sqlite3_blob *blobChunksValidity,
7836+ const unsigned char *bufferChunksValidity) {
7837+ int rc, brc;
7838+ sqlite3_blob *blobChunksRowids = NULL;
7839+
7840+ // mark the validity bit for this row in the chunk's validity bitmap
7841+ // Get the byte offset of the bitmap
7842+ char unsigned bx = bufferChunksValidity[chunk_offset / CHAR_BIT];
7843+ // set the bit at the chunk_offset position inside that byte
7844+ bx = bx | (1 << (chunk_offset % CHAR_BIT));
7845+ // write that 1 byte
7846+ rc = sqlite3_blob_write(blobChunksValidity, &bx, 1, chunk_offset / CHAR_BIT);
7847+ if (rc != SQLITE_OK) {
7848+ vtab_set_error(&p->base, VEC_INTERAL_ERROR "could not mark validity bit ");
7849+ return rc;
7850+ }
7851+
7852+ // Go insert the vector data into the vector chunk shadow tables
7853+ for (int i = 0; i < p->numVectorColumns; i++) {
7854+ sqlite3_blob *blobVectors;
7855+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowVectorChunksNames[i],
7856+ "vectors", chunk_rowid, 1, &blobVectors);
7857+ if (rc != SQLITE_OK) {
7858+ vtab_set_error(&p->base, "Error opening vector blob at %s.%s.%lld",
7859+ p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid);
7860+ goto cleanup;
7861+ }
7862+
7863+ i64 expected =
7864+ p->chunk_size * vector_column_byte_size(p->vector_columns[i]);
7865+ i64 actual = sqlite3_blob_bytes(blobVectors);
7866+
7867+ if (actual != expected) {
7868+ // IMP: V16386_00456
7869+ vtab_set_error(
7870+ &p->base,
7871+ VEC_INTERAL_ERROR
7872+ "vector blob size mismatch on %s.%s.%lld. Expected %lld, actual %lld",
7873+ p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid, expected,
7874+ actual);
7875+ rc = SQLITE_ERROR;
7876+ // already error, can ignore result code
7877+ sqlite3_blob_close(blobVectors);
7878+ goto cleanup;
7879+ };
7880+
7881+ rc = vec0_write_vector_to_vector_blob(
7882+ blobVectors, chunk_offset, vectorDatas[i],
7883+ p->vector_columns[i].dimensions, p->vector_columns[i].element_type);
7884+ if (rc != SQLITE_OK) {
7885+ vtab_set_error(&p->base,
7886+ VEC_INTERAL_ERROR
7887+ "could not write vector blob on %s.%s.%lld",
7888+ p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid);
7889+ rc = SQLITE_ERROR;
7890+ // already error, can ignore result code
7891+ sqlite3_blob_close(blobVectors);
7892+ goto cleanup;
7893+ }
7894+ rc = sqlite3_blob_close(blobVectors);
7895+ if (rc != SQLITE_OK) {
7896+ vtab_set_error(&p->base,
7897+ VEC_INTERAL_ERROR
7898+ "could not close vector blob on %s.%s.%lld",
7899+ p->schemaName, p->shadowVectorChunksNames[i], chunk_rowid);
7900+ rc = SQLITE_ERROR;
7901+ goto cleanup;
7902+ }
7903+ }
7904+
7905+ // write the new rowid to the rowids column of the _chunks table
7906+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "rowids",
7907+ chunk_rowid, 1, &blobChunksRowids);
7908+ if (rc != SQLITE_OK) {
7909+ // IMP: V09221_26060
7910+ vtab_set_error(&p->base,
7911+ VEC_INTERAL_ERROR "could not open rowids blob on %s.%s.%lld",
7912+ p->schemaName, p->shadowChunksName, chunk_rowid);
7913+ goto cleanup;
7914+ }
7915+ i64 expected = p->chunk_size * sizeof(i64);
7916+ i64 actual = sqlite3_blob_bytes(blobChunksRowids);
7917+ if (expected != actual) {
7918+ // IMP: V12779_29618
7919+ vtab_set_error(
7920+ &p->base,
7921+ VEC_INTERAL_ERROR
7922+ "rowids blob size mismatch on %s.%s.%lld. Expected %lld, actual %lld",
7923+ p->schemaName, p->shadowChunksName, chunk_rowid, expected, actual);
7924+ rc = SQLITE_ERROR;
7925+ goto cleanup;
7926+ }
7927+ rc = sqlite3_blob_write(blobChunksRowids, &rowid, sizeof(i64),
7928+ chunk_offset * sizeof(i64));
7929+ if (rc != SQLITE_OK) {
7930+ vtab_set_error(
7931+ &p->base, VEC_INTERAL_ERROR "could not write rowids blob on %s.%s.%lld",
7932+ p->schemaName, p->shadowChunksName, chunk_rowid);
7933+ rc = SQLITE_ERROR;
7934+ goto cleanup;
7935+ }
7936+
7937+ // Now with all the vectors inserted, go back and update the _rowids table
7938+ // with the new chunk_rowid/chunk_offset values
7939+ rc = vec0_rowids_update_position(p, rowid, chunk_rowid, chunk_offset);
7940+
7941+cleanup:
7942+ brc = sqlite3_blob_close(blobChunksRowids);
7943+ if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) {
7944+ vtab_set_error(
7945+ &p->base, VEC_INTERAL_ERROR "could not close rowids blob on %s.%s.%lld",
7946+ p->schemaName, p->shadowChunksName, chunk_rowid);
7947+ return brc;
7948+ }
7949+ return rc;
7950+}
7951+
7952+int vec0_write_metadata_value(vec0_vtab *p, int metadata_column_idx, i64 rowid, i64 chunk_id, i64 chunk_offset, sqlite3_value * v, int isupdate) {
7953+ int rc;
7954+ struct Vec0MetadataColumnDefinition * metadata_column = &p->metadata_columns[metadata_column_idx];
7955+ vec0_metadata_column_kind kind = metadata_column->kind;
7956+
7957+ // verify input value matches column type
7958+ switch(kind) {
7959+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
7960+ if(sqlite3_value_type(v) != SQLITE_INTEGER || ((sqlite3_value_int(v) != 0) && (sqlite3_value_int(v) != 1))) {
7961+ rc = SQLITE_ERROR;
7962+ vtab_set_error(&p->base, "Expected 0 or 1 for BOOLEAN metadata column %.*s", metadata_column->name_length, metadata_column->name);
7963+ goto done;
7964+ }
7965+ break;
7966+ }
7967+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
7968+ if(sqlite3_value_type(v) != SQLITE_INTEGER) {
7969+ rc = SQLITE_ERROR;
7970+ vtab_set_error(&p->base, "Expected integer for INTEGER metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v)));
7971+ goto done;
7972+ }
7973+ break;
7974+ }
7975+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
7976+ if(sqlite3_value_type(v) != SQLITE_FLOAT) {
7977+ rc = SQLITE_ERROR;
7978+ vtab_set_error(&p->base, "Expected float for FLOAT metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v)));
7979+ goto done;
7980+ }
7981+ break;
7982+ }
7983+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
7984+ if(sqlite3_value_type(v) != SQLITE_TEXT) {
7985+ rc = SQLITE_ERROR;
7986+ vtab_set_error(&p->base, "Expected text for TEXT metadata column %.*s, received %s", metadata_column->name_length, metadata_column->name, type_name(sqlite3_value_type(v)));
7987+ goto done;
7988+ }
7989+ break;
7990+ }
7991+ }
7992+
7993+ sqlite3_blob * blobValue = NULL;
7994+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_column_idx], "data", chunk_id, 1, &blobValue);
7995+ if(rc != SQLITE_OK) {
7996+ goto done;
7997+ }
7998+
7999+ switch(kind) {
8000+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
8001+ u8 block;
8002+ int value = sqlite3_value_int(v);
8003+ rc = sqlite3_blob_read(blobValue, &block, sizeof(u8), (int) (chunk_offset / CHAR_BIT));
8004+ if(rc != SQLITE_OK) {
8005+ goto done;
8006+ }
8007+
8008+ if (value) {
8009+ block |= 1 << (chunk_offset % CHAR_BIT);
8010+ } else {
8011+ block &= ~(1 << (chunk_offset % CHAR_BIT));
8012+ }
8013+
8014+ rc = sqlite3_blob_write(blobValue, &block, sizeof(u8), chunk_offset / CHAR_BIT);
8015+ break;
8016+ }
8017+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
8018+ i64 value = sqlite3_value_int64(v);
8019+ rc = sqlite3_blob_write(blobValue, &value, sizeof(value), chunk_offset * sizeof(i64));
8020+ break;
8021+ }
8022+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
8023+ double value = sqlite3_value_double(v);
8024+ rc = sqlite3_blob_write(blobValue, &value, sizeof(value), chunk_offset * sizeof(double));
8025+ break;
8026+ }
8027+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
8028+ int prev_n;
8029+ rc = sqlite3_blob_read(blobValue, &prev_n, sizeof(int), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8030+ if(rc != SQLITE_OK) {
8031+ goto done;
8032+ }
8033+
8034+ const char * s = (const char *) sqlite3_value_text(v);
8035+ int n = sqlite3_value_bytes(v);
8036+ u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
8037+ memset(view, 0, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8038+ memcpy(view, &n, sizeof(int));
8039+ memcpy(view+4, s, min(n, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH-4));
8040+
8041+ rc = sqlite3_blob_write(blobValue, &view, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH, chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8042+ if(n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
8043+ const char * zSql;
8044+
8045+ if(isupdate && (prev_n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH)) {
8046+ zSql = sqlite3_mprintf("UPDATE " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " SET data = ?2 WHERE rowid = ?1", p->schemaName, p->tableName, metadata_column_idx);
8047+ }else {
8048+ zSql = sqlite3_mprintf("INSERT INTO " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " (rowid, data) VALUES (?1, ?2)", p->schemaName, p->tableName, metadata_column_idx);
8049+ }
8050+ if(!zSql) {
8051+ rc = SQLITE_NOMEM;
8052+ goto done;
8053+ }
8054+ sqlite3_stmt * stmt;
8055+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8056+ if(rc != SQLITE_OK) {
8057+ goto done;
8058+ }
8059+ sqlite3_bind_int64(stmt, 1, rowid);
8060+ sqlite3_bind_text(stmt, 2, s, n, SQLITE_STATIC);
8061+ rc = sqlite3_step(stmt);
8062+ sqlite3_finalize(stmt);
8063+
8064+ if(rc != SQLITE_DONE) {
8065+ rc = SQLITE_ERROR;
8066+ goto done;
8067+ }
8068+ }
8069+ else if(prev_n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
8070+ const char * zSql = sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_column_idx);
8071+ if(!zSql) {
8072+ rc = SQLITE_NOMEM;
8073+ goto done;
8074+ }
8075+ sqlite3_stmt * stmt;
8076+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8077+ if(rc != SQLITE_OK) {
8078+ goto done;
8079+ }
8080+ sqlite3_bind_int64(stmt, 1, rowid);
8081+ rc = sqlite3_step(stmt);
8082+ sqlite3_finalize(stmt);
8083+
8084+ if(rc != SQLITE_DONE) {
8085+ rc = SQLITE_ERROR;
8086+ goto done;
8087+ }
8088+ }
8089+ break;
8090+ }
8091+ }
8092+
8093+ if(rc != SQLITE_OK) {
8094+
8095+ }
8096+ rc = sqlite3_blob_close(blobValue);
8097+ if(rc != SQLITE_OK) {
8098+ goto done;
8099+ }
8100+
8101+ done:
8102+ return rc;
8103+}
8104+
8105+
8106+/**
8107+ * @brief Handles INSERT INTO operations on a vec0 table.
8108+ *
8109+ * @return int SQLITE_OK on success, otherwise error code on failure
8110+ */
8111+int vec0Update_Insert(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv,
8112+ sqlite_int64 *pRowid) {
8113+ UNUSED_PARAMETER(argc);
8114+ vec0_vtab *p = (vec0_vtab *)pVTab;
8115+ int rc;
8116+ // Rowid for the inserted row, deterimined by the inserted ID + _rowids shadow
8117+ // table
8118+ i64 rowid;
8119+
8120+ // Array to hold the vector data of the inserted row. Individual elements will
8121+ // have a lifetime bound to the argv[..] values.
8122+ void *vectorDatas[VEC0_MAX_VECTOR_COLUMNS];
8123+ // Array to hold cleanup functions for vectorDatas[]
8124+ vector_cleanup cleanups[VEC0_MAX_VECTOR_COLUMNS];
8125+
8126+ sqlite3_value * partitionKeyValues[VEC0_MAX_PARTITION_COLUMNS];
8127+
8128+ // Rowid of the chunk in the _chunks shadow table that the row will be a part
8129+ // of.
8130+ i64 chunk_rowid;
8131+ // offset within the chunk where the rowid belongs
8132+ i64 chunk_offset;
8133+
8134+ // a write-able blob of the validity column for the given chunk. Used to mark
8135+ // validity bit
8136+ sqlite3_blob *blobChunksValidity = NULL;
8137+ // buffer for the valididty column for the given chunk. Maybe not needed here?
8138+ const unsigned char *bufferChunksValidity = NULL;
8139+ int numReadVectors = 0;
8140+
8141+ // Read all provided partition key values into partitionKeyValues
8142+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8143+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_PARTITION) {
8144+ continue;
8145+ }
8146+ int partition_key_idx = p->user_column_idxs[i];
8147+ partitionKeyValues[partition_key_idx] = argv[2+VEC0_COLUMN_USERN_START + i];
8148+
8149+ int new_value_type = sqlite3_value_type(partitionKeyValues[partition_key_idx]);
8150+ if((new_value_type != SQLITE_NULL) && (new_value_type != p->paritition_columns[partition_key_idx].type)) {
8151+ // IMP: V11454_28292
8152+ vtab_set_error(
8153+ pVTab,
8154+ "Parition key type mismatch: The partition key column %.*s has type %s, but %s was provided.",
8155+ p->paritition_columns[partition_key_idx].name_length,
8156+ p->paritition_columns[partition_key_idx].name,
8157+ type_name(p->paritition_columns[partition_key_idx].type),
8158+ type_name(new_value_type)
8159+ );
8160+ rc = SQLITE_ERROR;
8161+ goto cleanup;
8162+ }
8163+ }
8164+
8165+ // read all the inserted vectors into vectorDatas, validate their lengths.
8166+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8167+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) {
8168+ continue;
8169+ }
8170+ int vector_column_idx = p->user_column_idxs[i];
8171+ sqlite3_value *valueVector = argv[2 + VEC0_COLUMN_USERN_START + i];
8172+ size_t dimensions;
8173+
8174+ char *pzError;
8175+ enum VectorElementType elementType;
8176+ rc = vector_from_value(valueVector, &vectorDatas[vector_column_idx], &dimensions,
8177+ &elementType, &cleanups[vector_column_idx], &pzError);
8178+ if (rc != SQLITE_OK) {
8179+ // IMP: V06519_23358
8180+ vtab_set_error(
8181+ pVTab, "Inserted vector for the \"%.*s\" column is invalid: %z",
8182+ p->vector_columns[vector_column_idx].name_length, p->vector_columns[vector_column_idx].name, pzError);
8183+ rc = SQLITE_ERROR;
8184+ goto cleanup;
8185+ }
8186+
8187+ numReadVectors++;
8188+ if (elementType != p->vector_columns[vector_column_idx].element_type) {
8189+ // IMP: V08221_25059
8190+ vtab_set_error(
8191+ pVTab,
8192+ "Inserted vector for the \"%.*s\" column is expected to be of type "
8193+ "%s, but a %s vector was provided.",
8194+ p->vector_columns[i].name_length, p->vector_columns[i].name,
8195+ vector_subtype_name(p->vector_columns[i].element_type),
8196+ vector_subtype_name(elementType));
8197+ rc = SQLITE_ERROR;
8198+ goto cleanup;
8199+ }
8200+
8201+ if (dimensions != p->vector_columns[vector_column_idx].dimensions) {
8202+ // IMP: V01145_17984
8203+ vtab_set_error(
8204+ pVTab,
8205+ "Dimension mismatch for inserted vector for the \"%.*s\" column. "
8206+ "Expected %d dimensions but received %d.",
8207+ p->vector_columns[vector_column_idx].name_length, p->vector_columns[vector_column_idx].name,
8208+ p->vector_columns[vector_column_idx].dimensions, dimensions);
8209+ rc = SQLITE_ERROR;
8210+ goto cleanup;
8211+ }
8212+ }
8213+
8214+ // Cannot insert a value in the hidden "distance" column
8215+ if (sqlite3_value_type(argv[2 + vec0_column_distance_idx(p)]) !=
8216+ SQLITE_NULL) {
8217+ // IMP: V24228_08298
8218+ vtab_set_error(pVTab,
8219+ "A value was provided for the hidden \"distance\" column.");
8220+ rc = SQLITE_ERROR;
8221+ goto cleanup;
8222+ }
8223+ // Cannot insert a value in the hidden "k" column
8224+ if (sqlite3_value_type(argv[2 + vec0_column_k_idx(p)]) != SQLITE_NULL) {
8225+ // IMP: V11875_28713
8226+ vtab_set_error(pVTab, "A value was provided for the hidden \"k\" column.");
8227+ rc = SQLITE_ERROR;
8228+ goto cleanup;
8229+ }
8230+
8231+ // Step #1: Insert/get a rowid for this row, from the _rowids table.
8232+ rc = vec0Update_InsertRowidStep(p, argv[2 + VEC0_COLUMN_ID], &rowid);
8233+ if (rc != SQLITE_OK) {
8234+ goto cleanup;
8235+ }
8236+
8237+ // Step #2: Find the next "available" position in the _chunks table for this
8238+ // row.
8239+ rc = vec0Update_InsertNextAvailableStep(p, partitionKeyValues,
8240+ &chunk_rowid, &chunk_offset,
8241+ &blobChunksValidity,
8242+ &bufferChunksValidity);
8243+ if (rc != SQLITE_OK) {
8244+ goto cleanup;
8245+ }
8246+
8247+ // Step #3: With the next available chunk position, write out all the vectors
8248+ // to their specified location.
8249+ rc = vec0Update_InsertWriteFinalStep(p, chunk_rowid, chunk_offset, rowid,
8250+ vectorDatas, blobChunksValidity,
8251+ bufferChunksValidity);
8252+ if (rc != SQLITE_OK) {
8253+ goto cleanup;
8254+ }
8255+
8256+ if(p->numAuxiliaryColumns > 0) {
8257+ sqlite3_stmt *stmt;
8258+ sqlite3_str * s = sqlite3_str_new(NULL);
8259+ sqlite3_str_appendf(s, "INSERT INTO " VEC0_SHADOW_AUXILIARY_NAME "(rowid ", p->schemaName, p->tableName);
8260+ for(int i = 0; i < p->numAuxiliaryColumns; i++) {
8261+ sqlite3_str_appendf(s, ", value%02d", i);
8262+ }
8263+ sqlite3_str_appendall(s, ") VALUES (? ");
8264+ for(int i = 0; i < p->numAuxiliaryColumns; i++) {
8265+ sqlite3_str_appendall(s, ", ?");
8266+ }
8267+ sqlite3_str_appendall(s, ")");
8268+ char * zSql = sqlite3_str_finish(s);
8269+ // TODO double check error handling ehre
8270+ if(!zSql) {
8271+ rc = SQLITE_NOMEM;
8272+ goto cleanup;
8273+ }
8274+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8275+ if(rc != SQLITE_OK) {
8276+ goto cleanup;
8277+ }
8278+ sqlite3_bind_int64(stmt, 1, rowid);
8279+
8280+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8281+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY) {
8282+ continue;
8283+ }
8284+ int auxiliary_key_idx = p->user_column_idxs[i];
8285+ sqlite3_value * v = argv[2+VEC0_COLUMN_USERN_START + i];
8286+ int v_type = sqlite3_value_type(v);
8287+ if(v_type != SQLITE_NULL && (v_type != p->auxiliary_columns[auxiliary_key_idx].type)) {
8288+ sqlite3_finalize(stmt);
8289+ rc = SQLITE_CONSTRAINT;
8290+ vtab_set_error(
8291+ pVTab,
8292+ "Auxiliary column type mismatch: The auxiliary column %.*s has type %s, but %s was provided.",
8293+ p->auxiliary_columns[auxiliary_key_idx].name_length,
8294+ p->auxiliary_columns[auxiliary_key_idx].name,
8295+ type_name(p->auxiliary_columns[auxiliary_key_idx].type),
8296+ type_name(v_type)
8297+ );
8298+ goto cleanup;
8299+ }
8300+ // first 1 is for 1-based indexing on sqlite3_bind_*, second 1 is to account for initial rowid parameter
8301+ sqlite3_bind_value(stmt, 1 + 1 + auxiliary_key_idx, v);
8302+ }
8303+
8304+ rc = sqlite3_step(stmt);
8305+ if(rc != SQLITE_DONE) {
8306+ sqlite3_finalize(stmt);
8307+ rc = SQLITE_ERROR;
8308+ goto cleanup;
8309+ }
8310+ sqlite3_finalize(stmt);
8311+ }
8312+
8313+
8314+ for(int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8315+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) {
8316+ continue;
8317+ }
8318+ int metadata_idx = p->user_column_idxs[i];
8319+ sqlite3_value *v = argv[2 + VEC0_COLUMN_USERN_START + i];
8320+ rc = vec0_write_metadata_value(p, metadata_idx, rowid, chunk_rowid, chunk_offset, v, 0);
8321+ if(rc != SQLITE_OK) {
8322+ goto cleanup;
8323+ }
8324+ }
8325+
8326+ *pRowid = rowid;
8327+ rc = SQLITE_OK;
8328+
8329+cleanup:
8330+ for (int i = 0; i < numReadVectors; i++) {
8331+ cleanups[i](vectorDatas[i]);
8332+ }
8333+ sqlite3_free((void *)bufferChunksValidity);
8334+ int brc = sqlite3_blob_close(blobChunksValidity);
8335+ if ((rc == SQLITE_OK) && (brc != SQLITE_OK)) {
8336+ vtab_set_error(&p->base,
8337+ VEC_INTERAL_ERROR "unknown error, blobChunksValidity could "
8338+ "not be closed, please file an issue");
8339+ return brc;
8340+ }
8341+ return rc;
8342+}
8343+
8344+int vec0Update_Delete_ClearValidity(vec0_vtab *p, i64 chunk_id,
8345+ u64 chunk_offset) {
8346+ int rc, brc;
8347+ sqlite3_blob *blobChunksValidity = NULL;
8348+ char unsigned bx;
8349+ int validityOffset = chunk_offset / CHAR_BIT;
8350+
8351+ // 2. ensure chunks.validity bit is 1, then set to 0
8352+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "validity",
8353+ chunk_id, 1, &blobChunksValidity);
8354+ if (rc != SQLITE_OK) {
8355+ // IMP: V26002_10073
8356+ vtab_set_error(&p->base, "could not open validity blob for %s.%s.%lld",
8357+ p->schemaName, p->shadowChunksName, chunk_id);
8358+ return SQLITE_ERROR;
8359+ }
8360+ // will skip the sqlite3_blob_bytes(blobChunksValidity) check for now,
8361+ // the read below would catch it
8362+
8363+ rc = sqlite3_blob_read(blobChunksValidity, &bx, sizeof(bx), validityOffset);
8364+ if (rc != SQLITE_OK) {
8365+ // IMP: V21193_05263
8366+ vtab_set_error(
8367+ &p->base, "could not read validity blob for %s.%s.%lld at %d",
8368+ p->schemaName, p->shadowChunksName, chunk_id, validityOffset);
8369+ goto cleanup;
8370+ }
8371+ if (!(bx >> (chunk_offset % CHAR_BIT))) {
8372+ // IMP: V21193_05263
8373+ rc = SQLITE_ERROR;
8374+ vtab_set_error(
8375+ &p->base,
8376+ "vec0 deletion error: validity bit is not set for %s.%s.%lld at %d",
8377+ p->schemaName, p->shadowChunksName, chunk_id, validityOffset);
8378+ goto cleanup;
8379+ }
8380+ char unsigned mask = ~(1 << (chunk_offset % CHAR_BIT));
8381+ char result = bx & mask;
8382+ rc = sqlite3_blob_write(blobChunksValidity, &result, sizeof(bx),
8383+ validityOffset);
8384+ if (rc != SQLITE_OK) {
8385+ vtab_set_error(
8386+ &p->base, "could not write to validity blob for %s.%s.%lld at %d",
8387+ p->schemaName, p->shadowChunksName, chunk_id, validityOffset);
8388+ goto cleanup;
8389+ }
8390+
8391+cleanup:
8392+
8393+ brc = sqlite3_blob_close(blobChunksValidity);
8394+ if (rc != SQLITE_OK)
8395+ return rc;
8396+ if (brc != SQLITE_OK) {
8397+ vtab_set_error(&p->base,
8398+ "vec0 deletion error: Error commiting validity blob "
8399+ "transaction on %s.%s.%lld at %d",
8400+ p->schemaName, p->shadowChunksName, chunk_id,
8401+ validityOffset);
8402+ return brc;
8403+ }
8404+ return SQLITE_OK;
8405+}
8406+
8407+int vec0Update_Delete_DeleteRowids(vec0_vtab *p, i64 rowid) {
8408+ int rc;
8409+ sqlite3_stmt *stmt = NULL;
8410+
8411+ char *zSql =
8412+ sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_ROWIDS_NAME " WHERE rowid = ?",
8413+ p->schemaName, p->tableName);
8414+ if (!zSql) {
8415+ return SQLITE_NOMEM;
8416+ }
8417+
8418+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8419+ sqlite3_free(zSql);
8420+ if (rc != SQLITE_OK) {
8421+ goto cleanup;
8422+ }
8423+ sqlite3_bind_int64(stmt, 1, rowid);
8424+ rc = sqlite3_step(stmt);
8425+ if (rc != SQLITE_DONE) {
8426+ goto cleanup;
8427+ }
8428+ rc = SQLITE_OK;
8429+
8430+cleanup:
8431+ sqlite3_finalize(stmt);
8432+ return rc;
8433+}
8434+
8435+int vec0Update_Delete_DeleteAux(vec0_vtab *p, i64 rowid) {
8436+ int rc;
8437+ sqlite3_stmt *stmt = NULL;
8438+
8439+ char *zSql =
8440+ sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_AUXILIARY_NAME " WHERE rowid = ?",
8441+ p->schemaName, p->tableName);
8442+ if (!zSql) {
8443+ return SQLITE_NOMEM;
8444+ }
8445+
8446+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8447+ sqlite3_free(zSql);
8448+ if (rc != SQLITE_OK) {
8449+ goto cleanup;
8450+ }
8451+ sqlite3_bind_int64(stmt, 1, rowid);
8452+ rc = sqlite3_step(stmt);
8453+ if (rc != SQLITE_DONE) {
8454+ goto cleanup;
8455+ }
8456+ rc = SQLITE_OK;
8457+
8458+cleanup:
8459+ sqlite3_finalize(stmt);
8460+ return rc;
8461+}
8462+
8463+int vec0Update_Delete_ClearMetadata(vec0_vtab *p, int metadata_idx, i64 rowid, i64 chunk_id,
8464+ u64 chunk_offset) {
8465+ int rc;
8466+ sqlite3_blob * blobValue;
8467+ vec0_metadata_column_kind kind = p->metadata_columns[metadata_idx].kind;
8468+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowMetadataChunksNames[metadata_idx], "data", chunk_id, 1, &blobValue);
8469+ if(rc != SQLITE_OK) {
8470+ return rc;
8471+ }
8472+
8473+ switch(kind) {
8474+ case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
8475+ u8 block;
8476+ rc = sqlite3_blob_read(blobValue, &block, sizeof(u8), (int) (chunk_offset / CHAR_BIT));
8477+ if(rc != SQLITE_OK) {
8478+ goto done;
8479+ }
8480+
8481+ block &= ~(1 << (chunk_offset % CHAR_BIT));
8482+ rc = sqlite3_blob_write(blobValue, &block, sizeof(u8), chunk_offset / CHAR_BIT);
8483+ break;
8484+ }
8485+ case VEC0_METADATA_COLUMN_KIND_INTEGER: {
8486+ i64 v = 0;
8487+ rc = sqlite3_blob_write(blobValue, &v, sizeof(v), chunk_offset * sizeof(i64));
8488+ break;
8489+ }
8490+ case VEC0_METADATA_COLUMN_KIND_FLOAT: {
8491+ double v = 0;
8492+ rc = sqlite3_blob_write(blobValue, &v, sizeof(v), chunk_offset * sizeof(double));
8493+ break;
8494+ }
8495+ case VEC0_METADATA_COLUMN_KIND_TEXT: {
8496+ int n;
8497+ rc = sqlite3_blob_read(blobValue, &n, sizeof(int), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8498+ if(rc != SQLITE_OK) {
8499+ goto done;
8500+ }
8501+
8502+ u8 view[VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH];
8503+ memset(view, 0, VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8504+ rc = sqlite3_blob_write(blobValue, &view, sizeof(view), chunk_offset * VEC0_METADATA_TEXT_VIEW_BUFFER_LENGTH);
8505+ if(rc != SQLITE_OK) {
8506+ goto done;
8507+ }
8508+
8509+ if(n > VEC0_METADATA_TEXT_VIEW_DATA_LENGTH) {
8510+ const char * zSql = sqlite3_mprintf("DELETE FROM " VEC0_SHADOW_METADATA_TEXT_DATA_NAME " WHERE rowid = ?", p->schemaName, p->tableName, metadata_idx);
8511+ if(!zSql) {
8512+ rc = SQLITE_NOMEM;
8513+ goto done;
8514+ }
8515+ sqlite3_stmt * stmt;
8516+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8517+ if(rc != SQLITE_OK) {
8518+ goto done;
8519+ }
8520+ sqlite3_bind_int64(stmt, 1, rowid);
8521+ rc = sqlite3_step(stmt);
8522+ if(rc != SQLITE_DONE) {
8523+ rc = SQLITE_ERROR;
8524+ goto done;
8525+ }
8526+ sqlite3_finalize(stmt);
8527+ }
8528+ break;
8529+ }
8530+ }
8531+ int rc2;
8532+ done:
8533+ rc2 = sqlite3_blob_close(blobValue);
8534+ if(rc == SQLITE_OK) {
8535+ return rc2;
8536+ }
8537+ return rc;
8538+}
8539+
8540+int vec0Update_Delete(sqlite3_vtab *pVTab, sqlite3_value *idValue) {
8541+ vec0_vtab *p = (vec0_vtab *)pVTab;
8542+ int rc;
8543+ i64 rowid;
8544+ i64 chunk_id;
8545+ i64 chunk_offset;
8546+
8547+ if (p->pkIsText) {
8548+ rc = vec0_rowid_from_id(p, idValue, &rowid);
8549+ if (rc != SQLITE_OK) {
8550+ return rc;
8551+ }
8552+ } else {
8553+ rowid = sqlite3_value_int64(idValue);
8554+ }
8555+
8556+ // 1. Find chunk position for given rowid
8557+ // 2. Ensure that validity bit for position is 1, then set to 0
8558+ // 3. Zero out rowid in chunks.rowid
8559+ // 4. Zero out vector data in all vector column chunks
8560+ // 5. Delete value in _rowids table
8561+
8562+ // 1. get chunk_id and chunk_offset from _rowids
8563+ rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset);
8564+ if (rc != SQLITE_OK) {
8565+ return rc;
8566+ }
8567+
8568+ rc = vec0Update_Delete_ClearValidity(p, chunk_id, chunk_offset);
8569+ if (rc != SQLITE_OK) {
8570+ return rc;
8571+ }
8572+
8573+ // 3. zero out rowid in chunks.rowids
8574+ // https://github.com/asg017/sqlite-vec/issues/54
8575+
8576+ // 4. zero out any data in vector chunks tables
8577+ // https://github.com/asg017/sqlite-vec/issues/54
8578+
8579+ // 5. delete from _rowids table
8580+ rc = vec0Update_Delete_DeleteRowids(p, rowid);
8581+ if (rc != SQLITE_OK) {
8582+ return rc;
8583+ }
8584+
8585+ // 6. delete any auxiliary rows
8586+ if(p->numAuxiliaryColumns > 0) {
8587+ rc = vec0Update_Delete_DeleteAux(p, rowid);
8588+ if (rc != SQLITE_OK) {
8589+ return rc;
8590+ }
8591+ }
8592+
8593+ // 6. delete metadata
8594+ for(int i = 0; i < p->numMetadataColumns; i++) {
8595+ rc = vec0Update_Delete_ClearMetadata(p, i, rowid, chunk_id, chunk_offset);
8596+ }
8597+
8598+ return SQLITE_OK;
8599+}
8600+
8601+int vec0Update_UpdateAuxColumn(vec0_vtab *p, int auxiliary_column_idx, sqlite3_value * value, i64 rowid) {
8602+ int rc;
8603+ sqlite3_stmt *stmt;
8604+ const char * zSql = sqlite3_mprintf("UPDATE " VEC0_SHADOW_AUXILIARY_NAME " SET value%02d = ? WHERE rowid = ?", p->schemaName, p->tableName, auxiliary_column_idx);
8605+ if(!zSql) {
8606+ return SQLITE_NOMEM;
8607+ }
8608+ rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
8609+ if(rc != SQLITE_OK) {
8610+ return rc;
8611+ }
8612+ sqlite3_bind_value(stmt, 1, value);
8613+ sqlite3_bind_int64(stmt, 2, rowid);
8614+ rc = sqlite3_step(stmt);
8615+ if(rc != SQLITE_DONE) {
8616+ sqlite3_finalize(stmt);
8617+ return SQLITE_ERROR;
8618+ }
8619+ sqlite3_finalize(stmt);
8620+ return SQLITE_OK;
8621+}
8622+
8623+int vec0Update_UpdateVectorColumn(vec0_vtab *p, i64 chunk_id, i64 chunk_offset,
8624+ int i, sqlite3_value *valueVector) {
8625+ int rc;
8626+
8627+ sqlite3_blob *blobVectors = NULL;
8628+
8629+ char *pzError;
8630+ size_t dimensions;
8631+ enum VectorElementType elementType;
8632+ void *vector;
8633+ vector_cleanup cleanup = vector_cleanup_noop;
8634+ // https://github.com/asg017/sqlite-vec/issues/53
8635+ rc = vector_from_value(valueVector, &vector, &dimensions, &elementType,
8636+ &cleanup, &pzError);
8637+ if (rc != SQLITE_OK) {
8638+ // IMP: V15203_32042
8639+ vtab_set_error(
8640+ &p->base, "Updated vector for the \"%.*s\" column is invalid: %z",
8641+ p->vector_columns[i].name_length, p->vector_columns[i].name, pzError);
8642+ rc = SQLITE_ERROR;
8643+ goto cleanup;
8644+ }
8645+ if (elementType != p->vector_columns[i].element_type) {
8646+ // IMP: V03643_20481
8647+ vtab_set_error(
8648+ &p->base,
8649+ "Updated vector for the \"%.*s\" column is expected to be of type "
8650+ "%s, but a %s vector was provided.",
8651+ p->vector_columns[i].name_length, p->vector_columns[i].name,
8652+ vector_subtype_name(p->vector_columns[i].element_type),
8653+ vector_subtype_name(elementType));
8654+ rc = SQLITE_ERROR;
8655+ goto cleanup;
8656+ }
8657+ if (dimensions != p->vector_columns[i].dimensions) {
8658+ // IMP: V25739_09810
8659+ vtab_set_error(
8660+ &p->base,
8661+ "Dimension mismatch for new updated vector for the \"%.*s\" column. "
8662+ "Expected %d dimensions but received %d.",
8663+ p->vector_columns[i].name_length, p->vector_columns[i].name,
8664+ p->vector_columns[i].dimensions, dimensions);
8665+ rc = SQLITE_ERROR;
8666+ goto cleanup;
8667+ }
8668+
8669+ rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowVectorChunksNames[i],
8670+ "vectors", chunk_id, 1, &blobVectors);
8671+ if (rc != SQLITE_OK) {
8672+ vtab_set_error(&p->base, "Could not open vectors blob for %s.%s.%lld",
8673+ p->schemaName, p->shadowVectorChunksNames[i], chunk_id);
8674+ goto cleanup;
8675+ }
8676+ rc = vec0_write_vector_to_vector_blob(blobVectors, chunk_offset, vector,
8677+ p->vector_columns[i].dimensions,
8678+ p->vector_columns[i].element_type);
8679+ if (rc != SQLITE_OK) {
8680+ vtab_set_error(&p->base, "Could not write to vectors blob for %s.%s.%lld",
8681+ p->schemaName, p->shadowVectorChunksNames[i], chunk_id);
8682+ goto cleanup;
8683+ }
8684+
8685+cleanup:
8686+ cleanup(vector);
8687+ int brc = sqlite3_blob_close(blobVectors);
8688+ if (rc != SQLITE_OK) {
8689+ return rc;
8690+ }
8691+ if (brc != SQLITE_OK) {
8692+ vtab_set_error(
8693+ &p->base,
8694+ "Could not commit blob transaction for vectors blob for %s.%s.%lld",
8695+ p->schemaName, p->shadowVectorChunksNames[i], chunk_id);
8696+ return brc;
8697+ }
8698+ return SQLITE_OK;
8699+}
8700+
8701+int vec0Update_Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv) {
8702+ UNUSED_PARAMETER(argc);
8703+ vec0_vtab *p = (vec0_vtab *)pVTab;
8704+ int rc;
8705+ i64 chunk_id;
8706+ i64 chunk_offset;
8707+
8708+ i64 rowid;
8709+ if (p->pkIsText) {
8710+ const char *a = (const char *)sqlite3_value_text(argv[0]);
8711+ const char *b = (const char *)sqlite3_value_text(argv[1]);
8712+ // IMP: V08886_25725
8713+ if ((sqlite3_value_bytes(argv[0]) != sqlite3_value_bytes(argv[1])) ||
8714+ strncmp(a, b, sqlite3_value_bytes(argv[0])) != 0) {
8715+ vtab_set_error(pVTab,
8716+ "UPDATEs on vec0 primary key values are not allowed.");
8717+ return SQLITE_ERROR;
8718+ }
8719+ rc = vec0_rowid_from_id(p, argv[0], &rowid);
8720+ if (rc != SQLITE_OK) {
8721+ return rc;
8722+ }
8723+ } else {
8724+ rowid = sqlite3_value_int64(argv[0]);
8725+ }
8726+
8727+ // 1) get chunk_id and chunk_offset from _rowids
8728+ rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset);
8729+ if (rc != SQLITE_OK) {
8730+ return rc;
8731+ }
8732+
8733+ // 2) update any partition key values
8734+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8735+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_PARTITION) {
8736+ continue;
8737+ }
8738+ sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i];
8739+ if(sqlite3_value_nochange(value)) {
8740+ continue;
8741+ }
8742+ vtab_set_error(pVTab, "UPDATE on partition key columns are not supported yet. ");
8743+ return SQLITE_ERROR;
8744+ }
8745+
8746+ // 3) handle auxiliary column updates
8747+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8748+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_AUXILIARY) {
8749+ continue;
8750+ }
8751+ int auxiliary_column_idx = p->user_column_idxs[i];
8752+ sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i];
8753+ if(sqlite3_value_nochange(value)) {
8754+ continue;
8755+ }
8756+ rc = vec0Update_UpdateAuxColumn(p, auxiliary_column_idx, value, rowid);
8757+ if(rc != SQLITE_OK) {
8758+ return SQLITE_ERROR;
8759+ }
8760+ }
8761+
8762+ // 4) handle metadata column updates
8763+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8764+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_METADATA) {
8765+ continue;
8766+ }
8767+ int metadata_column_idx = p->user_column_idxs[i];
8768+ sqlite3_value * value = argv[2+VEC0_COLUMN_USERN_START + i];
8769+ if(sqlite3_value_nochange(value)) {
8770+ continue;
8771+ }
8772+ rc = vec0_write_metadata_value(p, metadata_column_idx, rowid, chunk_id, chunk_offset, value, 1);
8773+ if(rc != SQLITE_OK) {
8774+ return rc;
8775+ }
8776+ }
8777+
8778+ // 5) iterate over all new vectors, update the vectors
8779+ for (int i = 0; i < vec0_num_defined_user_columns(p); i++) {
8780+ if(p->user_column_kinds[i] != SQLITE_VEC0_USER_COLUMN_KIND_VECTOR) {
8781+ continue;
8782+ }
8783+ int vector_idx = p->user_column_idxs[i];
8784+ sqlite3_value *valueVector = argv[2 + VEC0_COLUMN_USERN_START + i];
8785+ // in vec0Column, we check sqlite3_vtab_nochange() on vector columns.
8786+ // If the vector column isn't being changed, we return NULL;
8787+ // That's not great, that means vector columns can never be NULLABLE
8788+ // (bc we cant distinguish if an updated vector is truly NULL or nochange).
8789+ // Also it means that if someone tries to run `UPDATE v SET X = NULL`,
8790+ // we can't effectively detect and raise an error.
8791+ // A better solution would be to use a custom result_type for "empty",
8792+ // but subtypes don't appear to survive xColumn -> xUpdate, it's always 0.
8793+ // So for now, we'll just use NULL and warn people to not SET X = NULL
8794+ // in the docs.
8795+ if (sqlite3_value_type(valueVector) == SQLITE_NULL) {
8796+ continue;
8797+ }
8798+
8799+ rc = vec0Update_UpdateVectorColumn(p, chunk_id, chunk_offset, vector_idx,
8800+ valueVector);
8801+ if (rc != SQLITE_OK) {
8802+ return SQLITE_ERROR;
8803+ }
8804+ }
8805+
8806+ return SQLITE_OK;
8807+}
8808+
8809+static int vec0Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv,
8810+ sqlite_int64 *pRowid) {
8811+ // DELETE operation
8812+ if (argc == 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
8813+ return vec0Update_Delete(pVTab, argv[0]);
8814+ }
8815+ // INSERT operation
8816+ else if (argc > 1 && sqlite3_value_type(argv[0]) == SQLITE_NULL) {
8817+ return vec0Update_Insert(pVTab, argc, argv, pRowid);
8818+ }
8819+ // UPDATE operation
8820+ else if (argc > 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
8821+ return vec0Update_Update(pVTab, argc, argv);
8822+ } else {
8823+ vtab_set_error(pVTab, "Unrecognized xUpdate operation provided for vec0.");
8824+ return SQLITE_ERROR;
8825+ }
8826+}
8827+
8828+static int vec0ShadowName(const char *zName) {
8829+ static const char *azName[] = {
8830+ "rowids", "chunks", "auxiliary", "info",
8831+
8832+ // Up to VEC0_MAX_METADATA_COLUMNS
8833+ // TODO be smarter about this man
8834+ "metadatachunks00",
8835+ "metadatachunks01",
8836+ "metadatachunks02",
8837+ "metadatachunks03",
8838+ "metadatachunks04",
8839+ "metadatachunks05",
8840+ "metadatachunks06",
8841+ "metadatachunks07",
8842+ "metadatachunks08",
8843+ "metadatachunks09",
8844+ "metadatachunks10",
8845+ "metadatachunks11",
8846+ "metadatachunks12",
8847+ "metadatachunks13",
8848+ "metadatachunks14",
8849+ "metadatachunks15",
8850+
8851+ // Up to
8852+ "metadatatext00",
8853+ "metadatatext01",
8854+ "metadatatext02",
8855+ "metadatatext03",
8856+ "metadatatext04",
8857+ "metadatatext05",
8858+ "metadatatext06",
8859+ "metadatatext07",
8860+ "metadatatext08",
8861+ "metadatatext09",
8862+ "metadatatext10",
8863+ "metadatatext11",
8864+ "metadatatext12",
8865+ "metadatatext13",
8866+ "metadatatext14",
8867+ "metadatatext15",
8868+ };
8869+
8870+ for (size_t i = 0; i < sizeof(azName) / sizeof(azName[0]); i++) {
8871+ if (sqlite3_stricmp(zName, azName[i]) == 0)
8872+ return 1;
8873+ }
8874+ //for(size_t i = 0; i < )"vector_chunks", "metadatachunks"
8875+ return 0;
8876+}
8877+
8878+static int vec0Begin(sqlite3_vtab *pVTab) {
8879+ UNUSED_PARAMETER(pVTab);
8880+ return SQLITE_OK;
8881+}
8882+static int vec0Sync(sqlite3_vtab *pVTab) {
8883+ UNUSED_PARAMETER(pVTab);
8884+ vec0_vtab *p = (vec0_vtab *)pVTab;
8885+ if (p->stmtLatestChunk) {
8886+ sqlite3_finalize(p->stmtLatestChunk);
8887+ p->stmtLatestChunk = NULL;
8888+ }
8889+ if (p->stmtRowidsInsertRowid) {
8890+ sqlite3_finalize(p->stmtRowidsInsertRowid);
8891+ p->stmtRowidsInsertRowid = NULL;
8892+ }
8893+ if (p->stmtRowidsInsertId) {
8894+ sqlite3_finalize(p->stmtRowidsInsertId);
8895+ p->stmtRowidsInsertId = NULL;
8896+ }
8897+ if (p->stmtRowidsUpdatePosition) {
8898+ sqlite3_finalize(p->stmtRowidsUpdatePosition);
8899+ p->stmtRowidsUpdatePosition = NULL;
8900+ }
8901+ if (p->stmtRowidsGetChunkPosition) {
8902+ sqlite3_finalize(p->stmtRowidsGetChunkPosition);
8903+ p->stmtRowidsGetChunkPosition = NULL;
8904+ }
8905+ return SQLITE_OK;
8906+}
8907+static int vec0Commit(sqlite3_vtab *pVTab) {
8908+ UNUSED_PARAMETER(pVTab);
8909+ return SQLITE_OK;
8910+}
8911+static int vec0Rollback(sqlite3_vtab *pVTab) {
8912+ UNUSED_PARAMETER(pVTab);
8913+ return SQLITE_OK;
8914+}
8915+
8916+static sqlite3_module vec0Module = {
8917+ /* iVersion */ 3,
8918+ /* xCreate */ vec0Create,
8919+ /* xConnect */ vec0Connect,
8920+ /* xBestIndex */ vec0BestIndex,
8921+ /* xDisconnect */ vec0Disconnect,
8922+ /* xDestroy */ vec0Destroy,
8923+ /* xOpen */ vec0Open,
8924+ /* xClose */ vec0Close,
8925+ /* xFilter */ vec0Filter,
8926+ /* xNext */ vec0Next,
8927+ /* xEof */ vec0Eof,
8928+ /* xColumn */ vec0Column,
8929+ /* xRowid */ vec0Rowid,
8930+ /* xUpdate */ vec0Update,
8931+ /* xBegin */ vec0Begin,
8932+ /* xSync */ vec0Sync,
8933+ /* xCommit */ vec0Commit,
8934+ /* xRollback */ vec0Rollback,
8935+ /* xFindFunction */ 0,
8936+ /* xRename */ 0, // https://github.com/asg017/sqlite-vec/issues/43
8937+ /* xSavepoint */ 0,
8938+ /* xRelease */ 0,
8939+ /* xRollbackTo */ 0,
8940+ /* xShadowName */ vec0ShadowName,
8941+#if SQLITE_VERSION_NUMBER >= 3044000
8942+ /* xIntegrity */ 0, // https://github.com/asg017/sqlite-vec/issues/44
8943+#endif
8944+};
8945+#pragma endregion
8946+
8947+static char *POINTER_NAME_STATIC_BLOB_DEF = "vec0-static_blob_def";
8948+struct static_blob_definition {
8949+ void *p;
8950+ size_t dimensions;
8951+ size_t nvectors;
8952+ enum VectorElementType element_type;
8953+};
8954+static void vec_static_blob_from_raw(sqlite3_context *context, int argc,
8955+ sqlite3_value **argv) {
8956+
8957+ assert(argc == 4);
8958+ struct static_blob_definition *p;
8959+ p = sqlite3_malloc(sizeof(*p));
8960+ if (!p) {
8961+ sqlite3_result_error_nomem(context);
8962+ return;
8963+ }
8964+ memset(p, 0, sizeof(*p));
8965+ p->p = (void *)sqlite3_value_int64(argv[0]);
8966+ p->element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
8967+ p->dimensions = sqlite3_value_int64(argv[2]);
8968+ p->nvectors = sqlite3_value_int64(argv[3]);
8969+ sqlite3_result_pointer(context, p, POINTER_NAME_STATIC_BLOB_DEF,
8970+ sqlite3_free);
8971+}
8972+#pragma region vec_static_blobs() table function
8973+
8974+#define MAX_STATIC_BLOBS 16
8975+
8976+typedef struct static_blob static_blob;
8977+struct static_blob {
8978+ char *name;
8979+ void *p;
8980+ size_t dimensions;
8981+ size_t nvectors;
8982+ enum VectorElementType element_type;
8983+};
8984+
8985+typedef struct vec_static_blob_data vec_static_blob_data;
8986+struct vec_static_blob_data {
8987+ static_blob static_blobs[MAX_STATIC_BLOBS];
8988+};
8989+
8990+typedef struct vec_static_blobs_vtab vec_static_blobs_vtab;
8991+struct vec_static_blobs_vtab {
8992+ sqlite3_vtab base;
8993+ vec_static_blob_data *data;
8994+};
8995+
8996+typedef struct vec_static_blobs_cursor vec_static_blobs_cursor;
8997+struct vec_static_blobs_cursor {
8998+ sqlite3_vtab_cursor base;
8999+ sqlite3_int64 iRowid;
9000+};
9001+
9002+static int vec_static_blobsConnect(sqlite3 *db, void *pAux, int argc,
9003+ const char *const *argv,
9004+ sqlite3_vtab **ppVtab, char **pzErr) {
9005+ UNUSED_PARAMETER(argc);
9006+ UNUSED_PARAMETER(argv);
9007+ UNUSED_PARAMETER(pzErr);
9008+
9009+ vec_static_blobs_vtab *pNew;
9010+#define VEC_STATIC_BLOBS_NAME 0
9011+#define VEC_STATIC_BLOBS_DATA 1
9012+#define VEC_STATIC_BLOBS_DIMENSIONS 2
9013+#define VEC_STATIC_BLOBS_COUNT 3
9014+ int rc = sqlite3_declare_vtab(
9015+ db, "CREATE TABLE x(name, data, dimensions hidden, count hidden)");
9016+ if (rc == SQLITE_OK) {
9017+ pNew = sqlite3_malloc(sizeof(*pNew));
9018+ *ppVtab = (sqlite3_vtab *)pNew;
9019+ if (pNew == 0)
9020+ return SQLITE_NOMEM;
9021+ memset(pNew, 0, sizeof(*pNew));
9022+ pNew->data = pAux;
9023+ }
9024+ return rc;
9025+}
9026+
9027+static int vec_static_blobsDisconnect(sqlite3_vtab *pVtab) {
9028+ vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pVtab;
9029+ sqlite3_free(p);
9030+ return SQLITE_OK;
9031+}
9032+
9033+static int vec_static_blobsUpdate(sqlite3_vtab *pVTab, int argc,
9034+ sqlite3_value **argv, sqlite_int64 *pRowid) {
9035+ UNUSED_PARAMETER(pRowid);
9036+ vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pVTab;
9037+ // DELETE operation
9038+ if (argc == 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
9039+ return SQLITE_ERROR;
9040+ }
9041+ // INSERT operation
9042+ else if (argc > 1 && sqlite3_value_type(argv[0]) == SQLITE_NULL) {
9043+ const char *key =
9044+ (const char *)sqlite3_value_text(argv[2 + VEC_STATIC_BLOBS_NAME]);
9045+ int idx = -1;
9046+ for (int i = 0; i < MAX_STATIC_BLOBS; i++) {
9047+ if (!p->data->static_blobs[i].name) {
9048+ p->data->static_blobs[i].name = sqlite3_mprintf("%s", key);
9049+ idx = i;
9050+ break;
9051+ }
9052+ }
9053+ if (idx < 0)
9054+ abort();
9055+ struct static_blob_definition *def = sqlite3_value_pointer(
9056+ argv[2 + VEC_STATIC_BLOBS_DATA], POINTER_NAME_STATIC_BLOB_DEF);
9057+ p->data->static_blobs[idx].p = def->p;
9058+ p->data->static_blobs[idx].dimensions = def->dimensions;
9059+ p->data->static_blobs[idx].nvectors = def->nvectors;
9060+ p->data->static_blobs[idx].element_type = def->element_type;
9061+
9062+ return SQLITE_OK;
9063+ }
9064+ // UPDATE operation
9065+ else if (argc > 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
9066+ return SQLITE_ERROR;
9067+ }
9068+ return SQLITE_ERROR;
9069+}
9070+
9071+static int vec_static_blobsOpen(sqlite3_vtab *p,
9072+ sqlite3_vtab_cursor **ppCursor) {
9073+ UNUSED_PARAMETER(p);
9074+ vec_static_blobs_cursor *pCur;
9075+ pCur = sqlite3_malloc(sizeof(*pCur));
9076+ if (pCur == 0)
9077+ return SQLITE_NOMEM;
9078+ memset(pCur, 0, sizeof(*pCur));
9079+ *ppCursor = &pCur->base;
9080+ return SQLITE_OK;
9081+}
9082+
9083+static int vec_static_blobsClose(sqlite3_vtab_cursor *cur) {
9084+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9085+ sqlite3_free(pCur);
9086+ return SQLITE_OK;
9087+}
9088+
9089+static int vec_static_blobsBestIndex(sqlite3_vtab *pVTab,
9090+ sqlite3_index_info *pIdxInfo) {
9091+ UNUSED_PARAMETER(pVTab);
9092+ pIdxInfo->idxNum = 1;
9093+ pIdxInfo->estimatedCost = (double)10;
9094+ pIdxInfo->estimatedRows = 10;
9095+ return SQLITE_OK;
9096+}
9097+
9098+static int vec_static_blobsNext(sqlite3_vtab_cursor *cur);
9099+static int vec_static_blobsFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
9100+ const char *idxStr, int argc,
9101+ sqlite3_value **argv) {
9102+ UNUSED_PARAMETER(idxNum);
9103+ UNUSED_PARAMETER(idxStr);
9104+ UNUSED_PARAMETER(argc);
9105+ UNUSED_PARAMETER(argv);
9106+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)pVtabCursor;
9107+ pCur->iRowid = -1;
9108+ vec_static_blobsNext(pVtabCursor);
9109+ return SQLITE_OK;
9110+}
9111+
9112+static int vec_static_blobsRowid(sqlite3_vtab_cursor *cur,
9113+ sqlite_int64 *pRowid) {
9114+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9115+ *pRowid = pCur->iRowid;
9116+ return SQLITE_OK;
9117+}
9118+
9119+static int vec_static_blobsNext(sqlite3_vtab_cursor *cur) {
9120+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9121+ vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pCur->base.pVtab;
9122+ pCur->iRowid++;
9123+ while (pCur->iRowid < MAX_STATIC_BLOBS) {
9124+ if (p->data->static_blobs[pCur->iRowid].name) {
9125+ return SQLITE_OK;
9126+ }
9127+ pCur->iRowid++;
9128+ }
9129+ return SQLITE_OK;
9130+}
9131+
9132+static int vec_static_blobsEof(sqlite3_vtab_cursor *cur) {
9133+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9134+ return pCur->iRowid >= MAX_STATIC_BLOBS;
9135+}
9136+
9137+static int vec_static_blobsColumn(sqlite3_vtab_cursor *cur,
9138+ sqlite3_context *context, int i) {
9139+ vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
9140+ vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)cur->pVtab;
9141+ switch (i) {
9142+ case VEC_STATIC_BLOBS_NAME:
9143+ sqlite3_result_text(context, p->data->static_blobs[pCur->iRowid].name, -1,
9144+ SQLITE_TRANSIENT);
9145+ break;
9146+ case VEC_STATIC_BLOBS_DATA:
9147+ sqlite3_result_null(context);
9148+ break;
9149+ case VEC_STATIC_BLOBS_DIMENSIONS:
9150+ sqlite3_result_int64(context,
9151+ p->data->static_blobs[pCur->iRowid].dimensions);
9152+ break;
9153+ case VEC_STATIC_BLOBS_COUNT:
9154+ sqlite3_result_int64(context, p->data->static_blobs[pCur->iRowid].nvectors);
9155+ break;
9156+ }
9157+ return SQLITE_OK;
9158+}
9159+
9160+static sqlite3_module vec_static_blobsModule = {
9161+ /* iVersion */ 3,
9162+ /* xCreate */ 0,
9163+ /* xConnect */ vec_static_blobsConnect,
9164+ /* xBestIndex */ vec_static_blobsBestIndex,
9165+ /* xDisconnect */ vec_static_blobsDisconnect,
9166+ /* xDestroy */ 0,
9167+ /* xOpen */ vec_static_blobsOpen,
9168+ /* xClose */ vec_static_blobsClose,
9169+ /* xFilter */ vec_static_blobsFilter,
9170+ /* xNext */ vec_static_blobsNext,
9171+ /* xEof */ vec_static_blobsEof,
9172+ /* xColumn */ vec_static_blobsColumn,
9173+ /* xRowid */ vec_static_blobsRowid,
9174+ /* xUpdate */ vec_static_blobsUpdate,
9175+ /* xBegin */ 0,
9176+ /* xSync */ 0,
9177+ /* xCommit */ 0,
9178+ /* xRollback */ 0,
9179+ /* xFindMethod */ 0,
9180+ /* xRename */ 0,
9181+ /* xSavepoint */ 0,
9182+ /* xRelease */ 0,
9183+ /* xRollbackTo */ 0,
9184+ /* xShadowName */ 0,
9185+#if SQLITE_VERSION_NUMBER >= 3044000
9186+ /* xIntegrity */ 0
9187+#endif
9188+};
9189+#pragma endregion
9190+
9191+#pragma region vec_static_blob_entries() table function
9192+
9193+typedef struct vec_static_blob_entries_vtab vec_static_blob_entries_vtab;
9194+struct vec_static_blob_entries_vtab {
9195+ sqlite3_vtab base;
9196+ static_blob *blob;
9197+};
9198+typedef enum {
9199+ VEC_SBE__QUERYPLAN_FULLSCAN = 1,
9200+ VEC_SBE__QUERYPLAN_KNN = 2
9201+} vec_sbe_query_plan;
9202+
9203+struct sbe_query_knn_data {
9204+ i64 k;
9205+ i64 k_used;
9206+ // Array of rowids of size k. Must be freed with sqlite3_free().
9207+ i32 *rowids;
9208+ // Array of distances of size k. Must be freed with sqlite3_free().
9209+ f32 *distances;
9210+ i64 current_idx;
9211+};
9212+void sbe_query_knn_data_clear(struct sbe_query_knn_data *knn_data) {
9213+ if (!knn_data)
9214+ return;
9215+
9216+ if (knn_data->rowids) {
9217+ sqlite3_free(knn_data->rowids);
9218+ knn_data->rowids = NULL;
9219+ }
9220+ if (knn_data->distances) {
9221+ sqlite3_free(knn_data->distances);
9222+ knn_data->distances = NULL;
9223+ }
9224+}
9225+
9226+typedef struct vec_static_blob_entries_cursor vec_static_blob_entries_cursor;
9227+struct vec_static_blob_entries_cursor {
9228+ sqlite3_vtab_cursor base;
9229+ sqlite3_int64 iRowid;
9230+ vec_sbe_query_plan query_plan;
9231+ struct sbe_query_knn_data *knn_data;
9232+};
9233+
9234+static int vec_static_blob_entriesConnect(sqlite3 *db, void *pAux, int argc,
9235+ const char *const *argv,
9236+ sqlite3_vtab **ppVtab, char **pzErr) {
9237+ UNUSED_PARAMETER(argc);
9238+ UNUSED_PARAMETER(argv);
9239+ UNUSED_PARAMETER(pzErr);
9240+ vec_static_blob_data *blob_data = pAux;
9241+ int idx = -1;
9242+ for (int i = 0; i < MAX_STATIC_BLOBS; i++) {
9243+ if (!blob_data->static_blobs[i].name)
9244+ continue;
9245+ if (strncmp(blob_data->static_blobs[i].name, argv[3],
9246+ strlen(blob_data->static_blobs[i].name)) == 0) {
9247+ idx = i;
9248+ break;
9249+ }
9250+ }
9251+ if (idx < 0)
9252+ abort();
9253+ vec_static_blob_entries_vtab *pNew;
9254+#define VEC_STATIC_BLOB_ENTRIES_VECTOR 0
9255+#define VEC_STATIC_BLOB_ENTRIES_DISTANCE 1
9256+#define VEC_STATIC_BLOB_ENTRIES_K 2
9257+ int rc = sqlite3_declare_vtab(
9258+ db, "CREATE TABLE x(vector, distance hidden, k hidden)");
9259+ if (rc == SQLITE_OK) {
9260+ pNew = sqlite3_malloc(sizeof(*pNew));
9261+ *ppVtab = (sqlite3_vtab *)pNew;
9262+ if (pNew == 0)
9263+ return SQLITE_NOMEM;
9264+ memset(pNew, 0, sizeof(*pNew));
9265+ pNew->blob = &blob_data->static_blobs[idx];
9266+ }
9267+ return rc;
9268+}
9269+
9270+static int vec_static_blob_entriesCreate(sqlite3 *db, void *pAux, int argc,
9271+ const char *const *argv,
9272+ sqlite3_vtab **ppVtab, char **pzErr) {
9273+ return vec_static_blob_entriesConnect(db, pAux, argc, argv, ppVtab, pzErr);
9274+}
9275+
9276+static int vec_static_blob_entriesDisconnect(sqlite3_vtab *pVtab) {
9277+ vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVtab;
9278+ sqlite3_free(p);
9279+ return SQLITE_OK;
9280+}
9281+
9282+static int vec_static_blob_entriesOpen(sqlite3_vtab *p,
9283+ sqlite3_vtab_cursor **ppCursor) {
9284+ UNUSED_PARAMETER(p);
9285+ vec_static_blob_entries_cursor *pCur;
9286+ pCur = sqlite3_malloc(sizeof(*pCur));
9287+ if (pCur == 0)
9288+ return SQLITE_NOMEM;
9289+ memset(pCur, 0, sizeof(*pCur));
9290+ *ppCursor = &pCur->base;
9291+ return SQLITE_OK;
9292+}
9293+
9294+static int vec_static_blob_entriesClose(sqlite3_vtab_cursor *cur) {
9295+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9296+ sqlite3_free(pCur->knn_data);
9297+ sqlite3_free(pCur);
9298+ return SQLITE_OK;
9299+}
9300+
9301+static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab,
9302+ sqlite3_index_info *pIdxInfo) {
9303+ vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVTab;
9304+ int iMatchTerm = -1;
9305+ int iLimitTerm = -1;
9306+ // int iRowidTerm = -1; // https://github.com/asg017/sqlite-vec/issues/47
9307+ int iKTerm = -1;
9308+
9309+ for (int i = 0; i < pIdxInfo->nConstraint; i++) {
9310+ if (!pIdxInfo->aConstraint[i].usable)
9311+ continue;
9312+
9313+ int iColumn = pIdxInfo->aConstraint[i].iColumn;
9314+ int op = pIdxInfo->aConstraint[i].op;
9315+ if (op == SQLITE_INDEX_CONSTRAINT_MATCH &&
9316+ iColumn == VEC_STATIC_BLOB_ENTRIES_VECTOR) {
9317+ if (iMatchTerm > -1) {
9318+ // https://github.com/asg017/sqlite-vec/issues/51
9319+ return SQLITE_ERROR;
9320+ }
9321+ iMatchTerm = i;
9322+ }
9323+ if (op == SQLITE_INDEX_CONSTRAINT_LIMIT) {
9324+ iLimitTerm = i;
9325+ }
9326+ if (op == SQLITE_INDEX_CONSTRAINT_EQ &&
9327+ iColumn == VEC_STATIC_BLOB_ENTRIES_K) {
9328+ iKTerm = i;
9329+ }
9330+ }
9331+ if (iMatchTerm >= 0) {
9332+ if (iLimitTerm < 0 && iKTerm < 0) {
9333+ // https://github.com/asg017/sqlite-vec/issues/51
9334+ return SQLITE_ERROR;
9335+ }
9336+ if (iLimitTerm >= 0 && iKTerm >= 0) {
9337+ return SQLITE_ERROR; // limit or k, not both
9338+ }
9339+ if (pIdxInfo->nOrderBy < 1) {
9340+ vtab_set_error(pVTab, "ORDER BY distance required");
9341+ return SQLITE_CONSTRAINT;
9342+ }
9343+ if (pIdxInfo->nOrderBy > 1) {
9344+ // https://github.com/asg017/sqlite-vec/issues/51
9345+ vtab_set_error(pVTab, "more than 1 ORDER BY clause provided");
9346+ return SQLITE_CONSTRAINT;
9347+ }
9348+ if (pIdxInfo->aOrderBy[0].iColumn != VEC_STATIC_BLOB_ENTRIES_DISTANCE) {
9349+ vtab_set_error(pVTab, "ORDER BY must be on the distance column");
9350+ return SQLITE_CONSTRAINT;
9351+ }
9352+ if (pIdxInfo->aOrderBy[0].desc) {
9353+ vtab_set_error(pVTab,
9354+ "Only ascending in ORDER BY distance clause is supported, "
9355+ "DESC is not supported yet.");
9356+ return SQLITE_CONSTRAINT;
9357+ }
9358+
9359+ pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_KNN;
9360+ pIdxInfo->estimatedCost = (double)10;
9361+ pIdxInfo->estimatedRows = 10;
9362+
9363+ pIdxInfo->orderByConsumed = 1;
9364+ pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = 1;
9365+ pIdxInfo->aConstraintUsage[iMatchTerm].omit = 1;
9366+ if (iLimitTerm >= 0) {
9367+ pIdxInfo->aConstraintUsage[iLimitTerm].argvIndex = 2;
9368+ pIdxInfo->aConstraintUsage[iLimitTerm].omit = 1;
9369+ } else {
9370+ pIdxInfo->aConstraintUsage[iKTerm].argvIndex = 2;
9371+ pIdxInfo->aConstraintUsage[iKTerm].omit = 1;
9372+ }
9373+
9374+ } else {
9375+ pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_FULLSCAN;
9376+ pIdxInfo->estimatedCost = (double)p->blob->nvectors;
9377+ pIdxInfo->estimatedRows = p->blob->nvectors;
9378+ }
9379+ return SQLITE_OK;
9380+}
9381+
9382+static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor,
9383+ int idxNum, const char *idxStr,
9384+ int argc, sqlite3_value **argv) {
9385+ UNUSED_PARAMETER(idxStr);
9386+ assert(argc >= 0 && argc <= 3);
9387+ vec_static_blob_entries_cursor *pCur =
9388+ (vec_static_blob_entries_cursor *)pVtabCursor;
9389+ vec_static_blob_entries_vtab *p =
9390+ (vec_static_blob_entries_vtab *)pCur->base.pVtab;
9391+
9392+ if (idxNum == VEC_SBE__QUERYPLAN_KNN) {
9393+ assert(argc == 2);
9394+ pCur->query_plan = VEC_SBE__QUERYPLAN_KNN;
9395+ struct sbe_query_knn_data *knn_data;
9396+ knn_data = sqlite3_malloc(sizeof(*knn_data));
9397+ if (!knn_data) {
9398+ return SQLITE_NOMEM;
9399+ }
9400+ memset(knn_data, 0, sizeof(*knn_data));
9401+
9402+ void *queryVector;
9403+ size_t dimensions;
9404+ enum VectorElementType elementType;
9405+ vector_cleanup cleanup;
9406+ char *err;
9407+ int rc = vector_from_value(argv[0], &queryVector, &dimensions, &elementType,
9408+ &cleanup, &err);
9409+ if (rc != SQLITE_OK) {
9410+ return SQLITE_ERROR;
9411+ }
9412+ if (elementType != p->blob->element_type) {
9413+ return SQLITE_ERROR;
9414+ }
9415+ if (dimensions != p->blob->dimensions) {
9416+ return SQLITE_ERROR;
9417+ }
9418+
9419+ i64 k = min(sqlite3_value_int64(argv[1]), (i64)p->blob->nvectors);
9420+ if (k < 0) {
9421+ // HANDLE https://github.com/asg017/sqlite-vec/issues/55
9422+ return SQLITE_ERROR;
9423+ }
9424+ if (k == 0) {
9425+ knn_data->k = 0;
9426+ pCur->knn_data = knn_data;
9427+ return SQLITE_OK;
9428+ }
9429+
9430+ size_t bsize = (p->blob->nvectors + 7) & ~7;
9431+
9432+ i32 *topk_rowids = sqlite3_malloc(k * sizeof(i32));
9433+ if (!topk_rowids) {
9434+ // HANDLE https://github.com/asg017/sqlite-vec/issues/55
9435+ return SQLITE_ERROR;
9436+ }
9437+ f32 *distances = sqlite3_malloc(bsize * sizeof(f32));
9438+ if (!distances) {
9439+ // HANDLE https://github.com/asg017/sqlite-vec/issues/55
9440+ return SQLITE_ERROR;
9441+ }
9442+
9443+ for (size_t i = 0; i < p->blob->nvectors; i++) {
9444+ // https://github.com/asg017/sqlite-vec/issues/52
9445+ float *v = ((float *)p->blob->p) + (i * p->blob->dimensions);
9446+ distances[i] =
9447+ distance_l2_sqr_float(v, (float *)queryVector, &p->blob->dimensions);
9448+ }
9449+ u8 *candidates = bitmap_new(bsize);
9450+ assert(candidates);
9451+
9452+ u8 *taken = bitmap_new(bsize);
9453+ assert(taken);
9454+
9455+ bitmap_fill(candidates, bsize);
9456+ for (size_t i = bsize; i >= p->blob->nvectors; i--) {
9457+ bitmap_set(candidates, i, 0);
9458+ }
9459+ i32 k_used = 0;
9460+ min_idx(distances, bsize, candidates, topk_rowids, k, taken, &k_used);
9461+ knn_data->current_idx = 0;
9462+ knn_data->distances = distances;
9463+ knn_data->k = k;
9464+ knn_data->rowids = topk_rowids;
9465+
9466+ pCur->knn_data = knn_data;
9467+ } else {
9468+ pCur->query_plan = VEC_SBE__QUERYPLAN_FULLSCAN;
9469+ pCur->iRowid = 0;
9470+ }
9471+
9472+ return SQLITE_OK;
9473+}
9474+
9475+static int vec_static_blob_entriesRowid(sqlite3_vtab_cursor *cur,
9476+ sqlite_int64 *pRowid) {
9477+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9478+ switch (pCur->query_plan) {
9479+ case VEC_SBE__QUERYPLAN_FULLSCAN: {
9480+ *pRowid = pCur->iRowid;
9481+ return SQLITE_OK;
9482+ }
9483+ case VEC_SBE__QUERYPLAN_KNN: {
9484+ i32 rowid = ((i32 *)pCur->knn_data->rowids)[pCur->knn_data->current_idx];
9485+ *pRowid = (sqlite3_int64)rowid;
9486+ return SQLITE_OK;
9487+ }
9488+ }
9489+ return SQLITE_ERROR;
9490+}
9491+
9492+static int vec_static_blob_entriesNext(sqlite3_vtab_cursor *cur) {
9493+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9494+ switch (pCur->query_plan) {
9495+ case VEC_SBE__QUERYPLAN_FULLSCAN: {
9496+ pCur->iRowid++;
9497+ return SQLITE_OK;
9498+ }
9499+ case VEC_SBE__QUERYPLAN_KNN: {
9500+ pCur->knn_data->current_idx++;
9501+ return SQLITE_OK;
9502+ }
9503+ }
9504+ return SQLITE_ERROR;
9505+}
9506+
9507+static int vec_static_blob_entriesEof(sqlite3_vtab_cursor *cur) {
9508+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9509+ vec_static_blob_entries_vtab *p =
9510+ (vec_static_blob_entries_vtab *)pCur->base.pVtab;
9511+ switch (pCur->query_plan) {
9512+ case VEC_SBE__QUERYPLAN_FULLSCAN: {
9513+ return (size_t)pCur->iRowid >= p->blob->nvectors;
9514+ }
9515+ case VEC_SBE__QUERYPLAN_KNN: {
9516+ return pCur->knn_data->current_idx >= pCur->knn_data->k;
9517+ }
9518+ }
9519+ return SQLITE_ERROR;
9520+}
9521+
9522+static int vec_static_blob_entriesColumn(sqlite3_vtab_cursor *cur,
9523+ sqlite3_context *context, int i) {
9524+ vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
9525+ vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)cur->pVtab;
9526+
9527+ switch (pCur->query_plan) {
9528+ case VEC_SBE__QUERYPLAN_FULLSCAN: {
9529+ switch (i) {
9530+ case VEC_STATIC_BLOB_ENTRIES_VECTOR:
9531+
9532+ sqlite3_result_blob(
9533+ context,
9534+ ((unsigned char *)p->blob->p) +
9535+ (pCur->iRowid * p->blob->dimensions * sizeof(float)),
9536+ p->blob->dimensions * sizeof(float), SQLITE_TRANSIENT);
9537+ sqlite3_result_subtype(context, p->blob->element_type);
9538+ break;
9539+ }
9540+ return SQLITE_OK;
9541+ }
9542+ case VEC_SBE__QUERYPLAN_KNN: {
9543+ switch (i) {
9544+ case VEC_STATIC_BLOB_ENTRIES_VECTOR: {
9545+ i32 rowid = ((i32 *)pCur->knn_data->rowids)[pCur->knn_data->current_idx];
9546+ sqlite3_result_blob(context,
9547+ ((unsigned char *)p->blob->p) +
9548+ (rowid * p->blob->dimensions * sizeof(float)),
9549+ p->blob->dimensions * sizeof(float),
9550+ SQLITE_TRANSIENT);
9551+ sqlite3_result_subtype(context, p->blob->element_type);
9552+ break;
9553+ }
9554+ }
9555+ return SQLITE_OK;
9556+ }
9557+ }
9558+ return SQLITE_ERROR;
9559+}
9560+
9561+static sqlite3_module vec_static_blob_entriesModule = {
9562+ /* iVersion */ 3,
9563+ /* xCreate */
9564+ vec_static_blob_entriesCreate, // handle rm?
9565+ // https://github.com/asg017/sqlite-vec/issues/55
9566+ /* xConnect */ vec_static_blob_entriesConnect,
9567+ /* xBestIndex */ vec_static_blob_entriesBestIndex,
9568+ /* xDisconnect */ vec_static_blob_entriesDisconnect,
9569+ /* xDestroy */ vec_static_blob_entriesDisconnect,
9570+ /* xOpen */ vec_static_blob_entriesOpen,
9571+ /* xClose */ vec_static_blob_entriesClose,
9572+ /* xFilter */ vec_static_blob_entriesFilter,
9573+ /* xNext */ vec_static_blob_entriesNext,
9574+ /* xEof */ vec_static_blob_entriesEof,
9575+ /* xColumn */ vec_static_blob_entriesColumn,
9576+ /* xRowid */ vec_static_blob_entriesRowid,
9577+ /* xUpdate */ 0,
9578+ /* xBegin */ 0,
9579+ /* xSync */ 0,
9580+ /* xCommit */ 0,
9581+ /* xRollback */ 0,
9582+ /* xFindMethod */ 0,
9583+ /* xRename */ 0,
9584+ /* xSavepoint */ 0,
9585+ /* xRelease */ 0,
9586+ /* xRollbackTo */ 0,
9587+ /* xShadowName */ 0,
9588+#if SQLITE_VERSION_NUMBER >= 3044000
9589+ /* xIntegrity */ 0
9590+#endif
9591+};
9592+#pragma endregion
9593+
9594+#ifdef SQLITE_VEC_ENABLE_AVX
9595+#define SQLITE_VEC_DEBUG_BUILD_AVX "avx"
9596+#else
9597+#define SQLITE_VEC_DEBUG_BUILD_AVX ""
9598+#endif
9599+#ifdef SQLITE_VEC_ENABLE_NEON
9600+#define SQLITE_VEC_DEBUG_BUILD_NEON "neon"
9601+#else
9602+#define SQLITE_VEC_DEBUG_BUILD_NEON ""
9603+#endif
9604+
9605+#define SQLITE_VEC_DEBUG_BUILD \
9606+ SQLITE_VEC_DEBUG_BUILD_AVX " " SQLITE_VEC_DEBUG_BUILD_NEON
9607+
9608+#define SQLITE_VEC_DEBUG_STRING \
9609+ "Version: " SQLITE_VEC_VERSION "\n" \
9610+ "Date: " SQLITE_VEC_DATE "\n" \
9611+ "Commit: " SQLITE_VEC_SOURCE "\n" \
9612+ "Build flags: " SQLITE_VEC_DEBUG_BUILD
9613+
9614+SQLITE_VEC_API int sqlite3_vec_init(sqlite3 *db, char **pzErrMsg,
9615+ const sqlite3_api_routines *pApi) {
9616+#ifndef SQLITE_CORE
9617+ SQLITE_EXTENSION_INIT2(pApi);
9618+#endif
9619+ int rc = SQLITE_OK;
9620+
9621+#define DEFAULT_FLAGS (SQLITE_UTF8 | SQLITE_INNOCUOUS | SQLITE_DETERMINISTIC)
9622+
9623+ rc = sqlite3_create_function_v2(db, "vec_version", 0, DEFAULT_FLAGS,
9624+ SQLITE_VEC_VERSION, _static_text_func, NULL,
9625+ NULL, NULL);
9626+ if (rc != SQLITE_OK) {
9627+ return rc;
9628+ }
9629+ rc = sqlite3_create_function_v2(db, "vec_debug", 0, DEFAULT_FLAGS,
9630+ SQLITE_VEC_DEBUG_STRING, _static_text_func,
9631+ NULL, NULL, NULL);
9632+ if (rc != SQLITE_OK) {
9633+ return rc;
9634+ }
9635+ static struct {
9636+ const char *zFName;
9637+ void (*xFunc)(sqlite3_context *, int, sqlite3_value **);
9638+ int nArg;
9639+ int flags;
9640+ } aFunc[] = {
9641+ // clang-format off
9642+ //{"vec_version", _static_text_func, 0, DEFAULT_FLAGS, (void *) SQLITE_VEC_VERSION },
9643+ //{"vec_debug", _static_text_func, 0, DEFAULT_FLAGS, (void *) SQLITE_VEC_DEBUG_STRING },
9644+ {"vec_distance_l2", vec_distance_l2, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9645+ {"vec_distance_l1", vec_distance_l1, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9646+ {"vec_distance_hamming",vec_distance_hamming, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9647+ {"vec_distance_cosine", vec_distance_cosine, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9648+ {"vec_length", vec_length, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE, },
9649+ {"vec_type", vec_type, 1, DEFAULT_FLAGS, },
9650+ {"vec_to_json", vec_to_json, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9651+ {"vec_add", vec_add, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9652+ {"vec_sub", vec_sub, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9653+ {"vec_slice", vec_slice, 3, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9654+ {"vec_normalize", vec_normalize, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9655+ {"vec_f32", vec_f32, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9656+ {"vec_bit", vec_bit, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9657+ {"vec_int8", vec_int8, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9658+ {"vec_quantize_int8", vec_quantize_int8, 2, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9659+ {"vec_quantize_binary", vec_quantize_binary, 1, DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, },
9660+ // clang-format on
9661+ };
9662+
9663+ static struct {
9664+ char *name;
9665+ const sqlite3_module *module;
9666+ void *p;
9667+ void (*xDestroy)(void *);
9668+ } aMod[] = {
9669+ // clang-format off
9670+ {"vec0", &vec0Module, NULL, NULL},
9671+ {"vec_each", &vec_eachModule, NULL, NULL},
9672+ // clang-format on
9673+ };
9674+
9675+ for (unsigned long i = 0; i < countof(aFunc) && rc == SQLITE_OK; i++) {
9676+ rc = sqlite3_create_function_v2(db, aFunc[i].zFName, aFunc[i].nArg,
9677+ aFunc[i].flags, NULL, aFunc[i].xFunc, NULL,
9678+ NULL, NULL);
9679+ if (rc != SQLITE_OK) {
9680+ *pzErrMsg = sqlite3_mprintf("Error creating function %s: %s",
9681+ aFunc[i].zFName, sqlite3_errmsg(db));
9682+ return rc;
9683+ }
9684+ }
9685+
9686+ for (unsigned long i = 0; i < countof(aMod) && rc == SQLITE_OK; i++) {
9687+ rc = sqlite3_create_module_v2(db, aMod[i].name, aMod[i].module, NULL, NULL);
9688+ if (rc != SQLITE_OK) {
9689+ *pzErrMsg = sqlite3_mprintf("Error creating module %s: %s", aMod[i].name,
9690+ sqlite3_errmsg(db));
9691+ return rc;
9692+ }
9693+ }
9694+
9695+ return SQLITE_OK;
9696+}
9697+
9698+#ifndef SQLITE_VEC_OMIT_FS
9699+SQLITE_VEC_API int sqlite3_vec_numpy_init(sqlite3 *db, char **pzErrMsg,
9700+ const sqlite3_api_routines *pApi) {
9701+ UNUSED_PARAMETER(pzErrMsg);
9702+#ifndef SQLITE_CORE
9703+ SQLITE_EXTENSION_INIT2(pApi);
9704+#endif
9705+ int rc = SQLITE_OK;
9706+ rc = sqlite3_create_function_v2(db, "vec_npy_file", 1, SQLITE_RESULT_SUBTYPE,
9707+ NULL, vec_npy_file, NULL, NULL, NULL);
9708+ if(rc != SQLITE_OK) {
9709+ return rc;
9710+ }
9711+ rc = sqlite3_create_module_v2(db, "vec_npy_each", &vec_npy_eachModule, NULL, NULL);
9712+ return rc;
9713+}
9714+#endif
9715+
9716+SQLITE_VEC_API int
9717+sqlite3_vec_static_blobs_init(sqlite3 *db, char **pzErrMsg,
9718+ const sqlite3_api_routines *pApi) {
9719+ UNUSED_PARAMETER(pzErrMsg);
9720+#ifndef SQLITE_CORE
9721+ SQLITE_EXTENSION_INIT2(pApi);
9722+#endif
9723+
9724+ int rc = SQLITE_OK;
9725+ vec_static_blob_data *static_blob_data;
9726+ static_blob_data = sqlite3_malloc(sizeof(*static_blob_data));
9727+ if (!static_blob_data) {
9728+ return SQLITE_NOMEM;
9729+ }
9730+ memset(static_blob_data, 0, sizeof(*static_blob_data));
9731+
9732+ rc = sqlite3_create_function_v2(
9733+ db, "vec_static_blob_from_raw", 4,
9734+ DEFAULT_FLAGS | SQLITE_SUBTYPE | SQLITE_RESULT_SUBTYPE, NULL,
9735+ vec_static_blob_from_raw, NULL, NULL, NULL);
9736+ if (rc != SQLITE_OK)
9737+ return rc;
9738+
9739+ rc = sqlite3_create_module_v2(db, "vec_static_blobs", &vec_static_blobsModule,
9740+ static_blob_data, sqlite3_free);
9741+ if (rc != SQLITE_OK)
9742+ return rc;
9743+ rc = sqlite3_create_module_v2(db, "vec_static_blob_entries",
9744+ &vec_static_blob_entriesModule,
9745+ static_blob_data, NULL);
9746+ if (rc != SQLITE_OK)
9747+ return rc;
9748+ return rc;
9749+}
added spike/nim/vendor/sqlite-vec.h +41 -0
new file mode 100644
@@ -0,0 +1,41 @@
1+#ifndef SQLITE_VEC_H
2+#define SQLITE_VEC_H
3+
4+#ifndef SQLITE_CORE
5+#include "sqlite3ext.h"
6+#else
7+#include "sqlite3.h"
8+#endif
9+
10+#ifdef SQLITE_VEC_STATIC
11+ #define SQLITE_VEC_API
12+#else
13+ #ifdef _WIN32
14+ #define SQLITE_VEC_API __declspec(dllexport)
15+ #else
16+ #define SQLITE_VEC_API
17+ #endif
18+#endif
19+
20+#define SQLITE_VEC_VERSION "v0.1.6"
21+// TODO rm
22+#define SQLITE_VEC_DATE "2024-11-20T16:38:29Z+0000"
23+#define SQLITE_VEC_SOURCE "639fca5739fe056fdc98f3d539c4cd79328d7dc7"
24+
25+
26+#define SQLITE_VEC_VERSION_MAJOR 0
27+#define SQLITE_VEC_VERSION_MINOR 1
28+#define SQLITE_VEC_VERSION_PATCH 6
29+
30+#ifdef __cplusplus
31+extern "C" {
32+#endif
33+
34+SQLITE_VEC_API int sqlite3_vec_init(sqlite3 *db, char **pzErrMsg,
35+ const sqlite3_api_routines *pApi);
36+
37+#ifdef __cplusplus
38+} /* end of the 'extern "C"' block */
39+#endif
40+
41+#endif /* ifndef SQLITE_VEC_H */
new file mode 100644
@@ -0,0 +1,41 @@
1+#ifndef SQLITE_VEC_H
2+#define SQLITE_VEC_H
3+
4+#ifndef SQLITE_CORE
5+#include "sqlite3ext.h"
6+#else
7+#include "sqlite3.h"
8+#endif
9+
10+#ifdef SQLITE_VEC_STATIC
11+ #define SQLITE_VEC_API
12+#else
13+ #ifdef _WIN32
14+ #define SQLITE_VEC_API __declspec(dllexport)
15+ #else
16+ #define SQLITE_VEC_API
17+ #endif
18+#endif
19+
20+#define SQLITE_VEC_VERSION "v0.1.6"
21+// TODO rm
22+#define SQLITE_VEC_DATE "2024-11-20T16:38:29Z+0000"
23+#define SQLITE_VEC_SOURCE "639fca5739fe056fdc98f3d539c4cd79328d7dc7"
24+
25+
26+#define SQLITE_VEC_VERSION_MAJOR 0
27+#define SQLITE_VEC_VERSION_MINOR 1
28+#define SQLITE_VEC_VERSION_PATCH 6
29+
30+#ifdef __cplusplus
31+extern "C" {
32+#endif
33+
34+SQLITE_VEC_API int sqlite3_vec_init(sqlite3 *db, char **pzErrMsg,
35+ const sqlite3_api_routines *pApi);
36+
37+#ifdef __cplusplus
38+} /* end of the 'extern "C"' block */
39+#endif
40+
41+#endif /* ifndef SQLITE_VEC_H */
added spike/nim/verify/go.mod +3 -0
new file mode 100644
@@ -0,0 +1,3 @@
1+module glean.spike/verify
2+
3+go 1.26.2
new file mode 100644
@@ -0,0 +1,3 @@
1+module glean.spike/verify
2+
3+go 1.26.2
added spike/nim/verify/verify.go +178 -0
new file mode 100644
@@ -0,0 +1,178 @@
1+// Independent check on the Nim DPoP spike.
2+//
3+// Self-verification in the same library would only prove internal
4+// consistency. This re-implements the verifier side from Go's standard
5+// library -- the same stack the current server runs on -- so agreement here
6+// means the Nim output is genuinely interoperable, not merely self-coherent.
7+//
8+// ./spike/nim/dpop_probe | go run ./verify
9+package main
10+
11+import (
12+ "crypto/ecdsa"
13+ "crypto/elliptic"
14+ "crypto/sha256"
15+ "encoding/base64"
16+ "encoding/json"
17+ "fmt"
18+ "io"
19+ "math/big"
20+ "os"
21+ "strings"
22+)
23+
24+type jwk struct {
25+ Crv string `json:"crv"`
26+ Kty string `json:"kty"`
27+ X string `json:"x"`
28+ Y string `json:"y"`
29+}
30+
31+type spikeOutput struct {
32+ JWT string `json:"jwt"`
33+ JWK jwk `json:"jwk"`
34+ Thumbprint string `json:"thumbprint"`
35+}
36+
37+var failures int
38+
39+func check(name string, ok bool, detail string) {
40+ label := "PASS"
41+ if !ok {
42+ label = "FAIL"
43+ failures++
44+ }
45+ if detail != "" {
46+ fmt.Printf(" %s %s -- %s\n", label, name, detail)
47+ return
48+ }
49+ fmt.Printf(" %s %s\n", label, name)
50+}
51+
52+func main() {
53+ raw, err := io.ReadAll(os.Stdin)
54+ if err != nil {
55+ fmt.Fprintln(os.Stderr, "read stdin:", err)
56+ os.Exit(2)
57+ }
58+ var out spikeOutput
59+ if err := json.Unmarshal(raw, &out); err != nil {
60+ fmt.Fprintln(os.Stderr, "parse spike output:", err)
61+ os.Exit(2)
62+ }
63+
64+ fmt.Println("Go cross-verification of the Nim DPoP proof")
65+
66+ parts := strings.Split(out.JWT, ".")
67+ check("JWT has three parts", len(parts) == 3, fmt.Sprintf("%d", len(parts)))
68+ if len(parts) != 3 {
69+ os.Exit(1)
70+ }
71+
72+ // Header: ATProto requires typ=dpop+jwt and the embedded public key.
73+ headerBytes, err := base64.RawURLEncoding.DecodeString(parts[0])
74+ check("header is unpadded base64url", err == nil, errText(err))
75+ var header struct {
76+ Typ string `json:"typ"`
77+ Alg string `json:"alg"`
78+ JWK jwk `json:"jwk"`
79+ }
80+ if err == nil {
81+ err = json.Unmarshal(headerBytes, &header)
82+ }
83+ check("header parses", err == nil, errText(err))
84+ check("typ is dpop+jwt", header.Typ == "dpop+jwt", header.Typ)
85+ check("alg is ES256", header.Alg == "ES256", header.Alg)
86+ check("header jwk matches the reported key",
87+ header.JWK == out.JWK, "")
88+
89+ payloadBytes, err := base64.RawURLEncoding.DecodeString(parts[1])
90+ check("payload is unpadded base64url", err == nil, errText(err))
91+ var payload map[string]any
92+ if err == nil {
93+ err = json.Unmarshal(payloadBytes, &payload)
94+ }
95+ check("payload parses", err == nil, errText(err))
96+ for _, claim := range []string{"jti", "htm", "htu", "iat"} {
97+ _, ok := payload[claim]
98+ check("payload has "+claim, ok, "")
99+ }
100+
101+ // Signature: JOSE requires raw R||S, not DER. A DER signature here would
102+ // be the single easiest mistake to make and would fail every PDS.
103+ sig, err := base64.RawURLEncoding.DecodeString(parts[2])
104+ check("signature decodes", err == nil, errText(err))
105+ check("signature is 64 raw bytes (R||S, not DER)", len(sig) == 64,
106+ fmt.Sprintf("%d bytes", len(sig)))
107+ if len(sig) != 64 {
108+ os.Exit(1)
109+ }
110+ check("signature is not DER-wrapped", sig[0] != 0x30, fmt.Sprintf("0x%02x", sig[0]))
111+
112+ pub, err := publicKey(out.JWK)
113+ check("JWK converts to a P-256 public key", err == nil, errText(err))
114+ if err != nil {
115+ os.Exit(1)
116+ }
117+
118+ signingInput := parts[0] + "." + parts[1]
119+ digest := sha256.Sum256([]byte(signingInput))
120+ r := new(big.Int).SetBytes(sig[:32])
121+ s := new(big.Int).SetBytes(sig[32:])
122+ check("ECDSA signature verifies", ecdsa.Verify(pub, digest[:], r, s), "")
123+
124+ // Negative control: a verifier that accepts anything proves nothing.
125+ tampered := []byte(signingInput)
126+ tampered[len(tampered)-1] ^= 0x01
127+ badDigest := sha256.Sum256(tampered)
128+ check("a tampered payload is rejected",
129+ !ecdsa.Verify(pub, badDigest[:], r, s), "")
130+
131+ // RFC 7638 thumbprint, recomputed independently.
132+ canonical := fmt.Sprintf(`{"crv":"%s","kty":"%s","x":"%s","y":"%s"}`,
133+ out.JWK.Crv, out.JWK.Kty, out.JWK.X, out.JWK.Y)
134+ sum := sha256.Sum256([]byte(canonical))
135+ want := base64.RawURLEncoding.EncodeToString(sum[:])
136+ check("RFC 7638 thumbprint matches", want == out.Thumbprint, out.Thumbprint)
137+
138+ fmt.Println()
139+ if failures == 0 {
140+ fmt.Println("RESULT: the Nim proof is valid and interoperable.")
141+ return
142+ }
143+ fmt.Printf("RESULT: %d check(s) failed.\n", failures)
144+ os.Exit(1)
145+}
146+
147+func publicKey(k jwk) (*ecdsa.PublicKey, error) {
148+ if k.Kty != "EC" || k.Crv != "P-256" {
149+ return nil, fmt.Errorf("unexpected key type %s/%s", k.Kty, k.Crv)
150+ }
151+ x, err := base64.RawURLEncoding.DecodeString(k.X)
152+ if err != nil {
153+ return nil, fmt.Errorf("x: %w", err)
154+ }
155+ y, err := base64.RawURLEncoding.DecodeString(k.Y)
156+ if err != nil {
157+ return nil, fmt.Errorf("y: %w", err)
158+ }
159+ if len(x) != 32 || len(y) != 32 {
160+ return nil, fmt.Errorf("coordinates must be 32 bytes, got %d/%d", len(x), len(y))
161+ }
162+ pub := &ecdsa.PublicKey{
163+ Curve: elliptic.P256(),
164+ X: new(big.Int).SetBytes(x),
165+ Y: new(big.Int).SetBytes(y),
166+ }
167+ if !pub.Curve.IsOnCurve(pub.X, pub.Y) {
168+ return nil, fmt.Errorf("point is not on P-256")
169+ }
170+ return pub, nil
171+}
172+
173+func errText(err error) string {
174+ if err == nil {
175+ return ""
176+ }
177+ return err.Error()
178+}
new file mode 100644
@@ -0,0 +1,178 @@
1+// Independent check on the Nim DPoP spike.
2+//
3+// Self-verification in the same library would only prove internal
4+// consistency. This re-implements the verifier side from Go's standard
5+// library -- the same stack the current server runs on -- so agreement here
6+// means the Nim output is genuinely interoperable, not merely self-coherent.
7+//
8+// ./spike/nim/dpop_probe | go run ./verify
9+package main
10+
11+import (
12+ "crypto/ecdsa"
13+ "crypto/elliptic"
14+ "crypto/sha256"
15+ "encoding/base64"
16+ "encoding/json"
17+ "fmt"
18+ "io"
19+ "math/big"
20+ "os"
21+ "strings"
22+)
23+
24+type jwk struct {
25+ Crv string `json:"crv"`
26+ Kty string `json:"kty"`
27+ X string `json:"x"`
28+ Y string `json:"y"`
29+}
30+
31+type spikeOutput struct {
32+ JWT string `json:"jwt"`
33+ JWK jwk `json:"jwk"`
34+ Thumbprint string `json:"thumbprint"`
35+}
36+
37+var failures int
38+
39+func check(name string, ok bool, detail string) {
40+ label := "PASS"
41+ if !ok {
42+ label = "FAIL"
43+ failures++
44+ }
45+ if detail != "" {
46+ fmt.Printf(" %s %s -- %s\n", label, name, detail)
47+ return
48+ }
49+ fmt.Printf(" %s %s\n", label, name)
50+}
51+
52+func main() {
53+ raw, err := io.ReadAll(os.Stdin)
54+ if err != nil {
55+ fmt.Fprintln(os.Stderr, "read stdin:", err)
56+ os.Exit(2)
57+ }
58+ var out spikeOutput
59+ if err := json.Unmarshal(raw, &out); err != nil {
60+ fmt.Fprintln(os.Stderr, "parse spike output:", err)
61+ os.Exit(2)
62+ }
63+
64+ fmt.Println("Go cross-verification of the Nim DPoP proof")
65+
66+ parts := strings.Split(out.JWT, ".")
67+ check("JWT has three parts", len(parts) == 3, fmt.Sprintf("%d", len(parts)))
68+ if len(parts) != 3 {
69+ os.Exit(1)
70+ }
71+
72+ // Header: ATProto requires typ=dpop+jwt and the embedded public key.
73+ headerBytes, err := base64.RawURLEncoding.DecodeString(parts[0])
74+ check("header is unpadded base64url", err == nil, errText(err))
75+ var header struct {
76+ Typ string `json:"typ"`
77+ Alg string `json:"alg"`
78+ JWK jwk `json:"jwk"`
79+ }
80+ if err == nil {
81+ err = json.Unmarshal(headerBytes, &header)
82+ }
83+ check("header parses", err == nil, errText(err))
84+ check("typ is dpop+jwt", header.Typ == "dpop+jwt", header.Typ)
85+ check("alg is ES256", header.Alg == "ES256", header.Alg)
86+ check("header jwk matches the reported key",
87+ header.JWK == out.JWK, "")
88+
89+ payloadBytes, err := base64.RawURLEncoding.DecodeString(parts[1])
90+ check("payload is unpadded base64url", err == nil, errText(err))
91+ var payload map[string]any
92+ if err == nil {
93+ err = json.Unmarshal(payloadBytes, &payload)
94+ }
95+ check("payload parses", err == nil, errText(err))
96+ for _, claim := range []string{"jti", "htm", "htu", "iat"} {
97+ _, ok := payload[claim]
98+ check("payload has "+claim, ok, "")
99+ }
100+
101+ // Signature: JOSE requires raw R||S, not DER. A DER signature here would
102+ // be the single easiest mistake to make and would fail every PDS.
103+ sig, err := base64.RawURLEncoding.DecodeString(parts[2])
104+ check("signature decodes", err == nil, errText(err))
105+ check("signature is 64 raw bytes (R||S, not DER)", len(sig) == 64,
106+ fmt.Sprintf("%d bytes", len(sig)))
107+ if len(sig) != 64 {
108+ os.Exit(1)
109+ }
110+ check("signature is not DER-wrapped", sig[0] != 0x30, fmt.Sprintf("0x%02x", sig[0]))
111+
112+ pub, err := publicKey(out.JWK)
113+ check("JWK converts to a P-256 public key", err == nil, errText(err))
114+ if err != nil {
115+ os.Exit(1)
116+ }
117+
118+ signingInput := parts[0] + "." + parts[1]
119+ digest := sha256.Sum256([]byte(signingInput))
120+ r := new(big.Int).SetBytes(sig[:32])
121+ s := new(big.Int).SetBytes(sig[32:])
122+ check("ECDSA signature verifies", ecdsa.Verify(pub, digest[:], r, s), "")
123+
124+ // Negative control: a verifier that accepts anything proves nothing.
125+ tampered := []byte(signingInput)
126+ tampered[len(tampered)-1] ^= 0x01
127+ badDigest := sha256.Sum256(tampered)
128+ check("a tampered payload is rejected",
129+ !ecdsa.Verify(pub, badDigest[:], r, s), "")
130+
131+ // RFC 7638 thumbprint, recomputed independently.
132+ canonical := fmt.Sprintf(`{"crv":"%s","kty":"%s","x":"%s","y":"%s"}`,
133+ out.JWK.Crv, out.JWK.Kty, out.JWK.X, out.JWK.Y)
134+ sum := sha256.Sum256([]byte(canonical))
135+ want := base64.RawURLEncoding.EncodeToString(sum[:])
136+ check("RFC 7638 thumbprint matches", want == out.Thumbprint, out.Thumbprint)
137+
138+ fmt.Println()
139+ if failures == 0 {
140+ fmt.Println("RESULT: the Nim proof is valid and interoperable.")
141+ return
142+ }
143+ fmt.Printf("RESULT: %d check(s) failed.\n", failures)
144+ os.Exit(1)
145+}
146+
147+func publicKey(k jwk) (*ecdsa.PublicKey, error) {
148+ if k.Kty != "EC" || k.Crv != "P-256" {
149+ return nil, fmt.Errorf("unexpected key type %s/%s", k.Kty, k.Crv)
150+ }
151+ x, err := base64.RawURLEncoding.DecodeString(k.X)
152+ if err != nil {
153+ return nil, fmt.Errorf("x: %w", err)
154+ }
155+ y, err := base64.RawURLEncoding.DecodeString(k.Y)
156+ if err != nil {
157+ return nil, fmt.Errorf("y: %w", err)
158+ }
159+ if len(x) != 32 || len(y) != 32 {
160+ return nil, fmt.Errorf("coordinates must be 32 bytes, got %d/%d", len(x), len(y))
161+ }
162+ pub := &ecdsa.PublicKey{
163+ Curve: elliptic.P256(),
164+ X: new(big.Int).SetBytes(x),
165+ Y: new(big.Int).SetBytes(y),
166+ }
167+ if !pub.Curve.IsOnCurve(pub.X, pub.Y) {
168+ return nil, fmt.Errorf("point is not on P-256")
169+ }
170+ return pub, nil
171+}
172+
173+func errText(err error) string {
174+ if err == nil {
175+ return ""
176+ }
177+ return err.Error()
178+}