rickub
nandi/gleanpublic Fork
⑂ 775d793
Branches
claude/modal-container-min-76527a 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.

Download ZIP

Delete the ATProto half

Ten modules and eight probes go: OAuth, DPoP, XRPC, identity resolution,
Jetstream and its supervisor, the session store, the social records, and the
clustering engine. None of it is reachable from a reader that has no
accounts.

That cascades further than the file list. The schema loses sixteen tables --
oauth, follows, likes, annotations, settings, and every derived
recommendation table -- which in turn removes the whole `recs` database, so
two SQLite files instead of three. sqlite-vec goes with the embeddings it
existed for, taking 305KB of vendored C and the SQLITE_CORE define that had
to be got exactly right. exp() and log() were registered only for the
clustering SQL, so they go too; this time the claim that nothing calls them
is backed by having deleted every caller.

articlestore no longer counts likes, so Article loses likeCount and
hasLiked, and each query that selected them loses a bind parameter.

`atproto/` is renamed `reader/`, since nothing in it is ATProto any more and
the old name would actively mislead.

Trimming the schema by filtering the generated statements needed two passes:
the first dropped the users table, because its follows_dirty column matched
a naive "follows" marker, and the second had to catch indexes whose ON
clause names an unqualified table the CREATE INDEX prefixes with a schema.

Verified after each step: both suites pass, and a clean run subscribes,
ingests, searches, scrapes and marks read.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
nandithebull committed 2026-09-21T19:47:21-07:00 Browse files
775d793 parent: 4cd4748
deleted spike/nim/atproto/cluster.nim +0 -276
deleted file mode 100644
@@ -1,276 +0,0 @@
1-## The similarity computations behind recommendations, ported from
2-## internal/cluster/jaccard.go and the pure helpers of scoring.go.
3-##
4-## Two kinds of similarity are precomputed on a schedule, because both are
5-## O(pairs) and neither needs to be fresh to the second:
6-##
7-## feed_similarity which feeds are subscribed to by the same people
8-## user_similarity which people subscribe to the same feeds
9-##
10-## Both are time-decayed. An overlap from three years ago says much less
11-## about what someone wants now than one from last week, and without decay a
12-## long-lived account's early subscriptions dominate its recommendations
13-## forever.
14-##
15-## The decay constant is `exp(-0.023 * days)`, a half-life of about a month
16-## (ln 2 / 0.023 ≈ 30). It is written inline in the SQL rather than named,
17-## the same way the Go implementation writes it.
18-
19-import std/[algorithm, math, options, random]
20-import ./gleandb
21-import ./sqlite
22-
23-type
24- Config* = object
25- followBoost*: float
26- likesWeight*: float
27- tagsWeight*: float
28- descriptionWeight*: float
29-
30- SignalWeights* = object
31- wSub*: float
32- wLike*: float
33- wTag*: float
34- wSocial*: float
35- wPop*: float
36- wCategory*: float
37- wContent*: float
38-
39- FeedRecommendation* = object
40- feedUrl*: string
41- title*: string
42- siteUrl*: string
43- description*: string
44- subscriberCount*: int
45- faviconUrl*: string
46- score*: float
47-
48- PersonRecommendation* = object
49- did*: string
50- handle*: string
51- displayName*: string
52- avatarUrl*: string
53- commonFeeds*: int
54- commonLikes*: int
55- commonTags*: int
56- isFollowed*: bool
57- score*: float
58-
59-const
60- ## Half-life of roughly 30 days: ln(2) / 0.023.
61- DecayPerDay* = 0.023
62-
63-proc defaultConfig*(): Config =
64- Config(followBoost: 0.5, likesWeight: 0.3, tagsWeight: 0.2,
65- descriptionWeight: 0.15)
66-
67-proc defaultWeights*(): SignalWeights =
68- SignalWeights(wSub: 1.0, wLike: 0.5, wTag: 0.3, wSocial: 0.7,
69- wPop: 0.2, wCategory: 0.4, wContent: 0.4)
70-
71-proc decay*(ageDays: float): float =
72- ## The weight an event of this age still carries.
73- ##
74- ## Exposed so the decay can be checked directly; the queries compute it in
75- ## SQL because pulling every like into the application to weight it would
76- ## defeat the point of precomputing.
77- exp(-DecayPerDay * ageDays)
78-
79-# --- similarity ------------------------------------------------------------
80-
81-## Both computations stage into a temp table and swap in a second
82-## transaction. The first pass is the expensive one and holds no lock on the
83-## live table; readers keep seeing the previous generation until the swap,
84-## rather than an empty table for the length of the computation.
85-
86-proc computeFeedSimilarity*(g: GleanDb) =
87- ## For each pair of feeds, how strongly their subscriber sets overlap.
88- ##
89- ## The numerator is a decayed count of shared subscribers and the
90- ## denominator is the union of the two subscriber counts, which makes this
91- ## a Jaccard index with recency folded into the intersection.
92- g.db.transaction:
93- g.db.exec """CREATE TEMP TABLE IF NOT EXISTS _feed_sim_staging (
94- feed_a TEXT NOT NULL, feed_b TEXT NOT NULL, jaccard REAL NOT NULL,
95- PRIMARY KEY (feed_a, feed_b))"""
96- g.db.exec "DELETE FROM _feed_sim_staging"
97- g.db.exec """
98- INSERT INTO _feed_sim_staging (feed_a, feed_b, jaccard)
99- SELECT feed_a, feed_b, jaccard FROM (
100- SELECT
101- s1.feed_url AS feed_a,
102- s2.feed_url AS feed_b,
103- SUM(EXP(-0.023 * CAST(julianday('now') -
104- julianday(MIN(s1.added_at, s2.added_at)) AS REAL)))
105- / NULLIF(f1.subscriber_count + f2.subscriber_count
106- - CAST(COUNT(*) AS REAL), 0) AS jaccard
107- FROM articles.subscriptions s1
108- -- feed_url < feed_url keeps one row per unordered pair, and drops
109- -- the self-pair without a separate predicate.
110- JOIN articles.subscriptions s2
111- ON s1.user_did = s2.user_did AND s1.feed_url < s2.feed_url
112- JOIN articles.feeds f1 ON f1.feed_url = s1.feed_url
113- JOIN articles.feeds f2 ON f2.feed_url = s2.feed_url
114- WHERE s1.added_at IS NOT NULL AND s2.added_at IS NOT NULL
115- GROUP BY s1.feed_url, s2.feed_url
116- ) WHERE jaccard IS NOT NULL"""
117-
118- g.db.transaction:
119- g.db.exec "DELETE FROM recs.feed_similarity"
120- g.db.exec """INSERT INTO recs.feed_similarity (feed_a, feed_b, jaccard)
121- SELECT feed_a, feed_b, jaccard FROM _feed_sim_staging"""
122-
123-proc computeUserSimilarity*(g: GleanDb) =
124- ## For each pair of people, how strongly their subscriptions overlap.
125- ##
126- ## Undecayed, unlike feed similarity: a subscription is a standing choice
127- ## rather than an event, so its age says little. Shared *likes* are decayed
128- ## separately below, because those are events.
129- g.db.transaction:
130- g.db.exec """CREATE TEMP TABLE IF NOT EXISTS _user_sim_staging (
131- user_a TEXT NOT NULL, user_b TEXT NOT NULL, jaccard REAL NOT NULL,
132- common_feeds INT NOT NULL DEFAULT 0,
133- common_likes INT NOT NULL DEFAULT 0,
134- common_tags INT NOT NULL DEFAULT 0,
135- PRIMARY KEY (user_a, user_b))"""
136- g.db.exec "DELETE FROM _user_sim_staging"
137- g.db.exec """
138- INSERT INTO _user_sim_staging
139- (user_a, user_b, jaccard, common_feeds, common_likes, common_tags)
140- SELECT
141- s1.user_did, s2.user_did,
142- CAST(COUNT(*) AS REAL) / (
143- (SELECT COUNT(*) FROM articles.subscriptions WHERE user_did = s1.user_did) +
144- (SELECT COUNT(*) FROM articles.subscriptions WHERE user_did = s2.user_did) -
145- CAST(COUNT(*) AS REAL)
146- ),
147- COUNT(*), 0, 0
148- FROM articles.subscriptions s1
149- JOIN articles.subscriptions s2
150- ON s1.feed_url = s2.feed_url AND s1.user_did < s2.user_did
151- GROUP BY s1.user_did, s2.user_did"""
152-
153- # Shared likes, decayed on both sides: two people who liked the same
154- # thing years apart have less in common than two who liked it the same
155- # week, and the product of the two decays expresses that.
156- #
157- # ROUND, not CAST AS INTEGER. The Go implementation truncates, and since
158- # each decay factor is at most 1 their product is always just under it --
159- # so a pair sharing one recent like sums to 0.9999 and stores 0. That
160- # number is shown to readers as the reason for a recommendation ("N
161- # shared likes"), so truncation makes a real overlap read as none.
162- g.db.exec """CREATE TEMP TABLE IF NOT EXISTS _likes_overlap (
163- user_a TEXT, user_b TEXT, common INT, PRIMARY KEY(user_a, user_b))"""
164- g.db.exec "DELETE FROM _likes_overlap"
165- g.db.exec """
166- INSERT INTO _likes_overlap (user_a, user_b, common)
167- SELECT l1.author_did, l2.author_did,
168- CAST(ROUND(SUM(
169- EXP(-0.023 * CAST(julianday('now') - julianday(l1.created_at) AS REAL))
170- * EXP(-0.023 * CAST(julianday('now') - julianday(l2.created_at) AS REAL))
171- )) AS INTEGER)
172- FROM articles.likes l1
173- JOIN articles.likes l2
174- ON l1.feed_url = l2.feed_url AND l1.article_url = l2.article_url
175- AND l1.author_did < l2.author_did
176- WHERE l1.created_at IS NOT NULL AND l2.created_at IS NOT NULL
177- GROUP BY l1.author_did, l2.author_did"""
178-
179- # A pair that shares likes but no feeds has no staging row yet, so this
180- # inserts before it updates.
181- g.db.exec """
182- INSERT INTO _user_sim_staging (user_a, user_b, jaccard, common_feeds,
183- common_likes, common_tags)
184- SELECT o.user_a, o.user_b, 0, 0, o.common, 0
185- FROM _likes_overlap o
186- WHERE NOT EXISTS (SELECT 1 FROM _user_sim_staging s
187- WHERE s.user_a = o.user_a AND s.user_b = o.user_b)"""
188- g.db.exec """
189- UPDATE _user_sim_staging SET common_likes = (
190- SELECT o.common FROM _likes_overlap o
191- WHERE o.user_a = _user_sim_staging.user_a
192- AND o.user_b = _user_sim_staging.user_b)
193- WHERE EXISTS (SELECT 1 FROM _likes_overlap o
194- WHERE o.user_a = _user_sim_staging.user_a
195- AND o.user_b = _user_sim_staging.user_b)"""
196-
197- g.db.transaction:
198- g.db.exec "DELETE FROM recs.user_similarity"
199- g.db.exec """INSERT INTO recs.user_similarity
200- (user_a, user_b, jaccard, common_feeds, common_likes, common_tags)
201- SELECT user_a, user_b, jaccard, common_feeds, common_likes,
202- common_tags
203- FROM _user_sim_staging"""
204-
205-proc feedSimilarity*(g: GleanDb, feedA, feedB: string): float =
206- ## Stored pairs are ordered, so look up whichever way round the caller asks.
207- let (a, b) = if feedA < feedB: (feedA, feedB) else: (feedB, feedA)
208- let row = g.db.queryFirst(
209- "SELECT jaccard FROM recs.feed_similarity WHERE feed_a = ? AND feed_b = ?",
210- [p(a), p(b)], proc(r: Row): float = r.f(0))
211- if row.isSome: row.get else: 0.0
212-
213-proc userSimilarity*(g: GleanDb, userA, userB: string): float =
214- let (a, b) = if userA < userB: (userA, userB) else: (userB, userA)
215- let row = g.db.queryFirst(
216- "SELECT jaccard FROM recs.user_similarity WHERE user_a = ? AND user_b = ?",
217- [p(a), p(b)], proc(r: Row): float = r.f(0))
218- if row.isSome: row.get else: 0.0
219-
220-# --- scoring helpers -------------------------------------------------------
221-
222-proc normalizeScores*(scores: var openArray[float]) =
223- ## Rescale to 0..1.
224- ##
225- ## Scores from different signals are not comparable in magnitude -- a
226- ## like-overlap sum and a log-popularity term live on unrelated scales --
227- ## so they are normalised before being blended or ranked against each
228- ## other. A single item, or a set where every score is equal, is left
229- ## alone: there is no spread to stretch, and dividing by zero would turn
230- ## the whole list into NaN.
231- if scores.len < 2: return
232- var lo = scores[0]
233- var hi = scores[0]
234- for s in scores:
235- lo = min(lo, s)
236- hi = max(hi, s)
237- if hi == lo: return
238- let span = hi - lo
239- for i in 0 ..< scores.len:
240- scores[i] = (scores[i] - lo) / span
241-
242-proc normalizeFeedScores*(recs: var seq[FeedRecommendation]) =
243- var s = newSeq[float](recs.len)
244- for i, r in recs: s[i] = r.score
245- normalizeScores(s)
246- for i in 0 ..< recs.len: recs[i].score = s[i]
247-
248-proc normalizePersonScores*(recs: var seq[PersonRecommendation]) =
249- var s = newSeq[float](recs.len)
250- for i, r in recs: s[i] = r.score
251- normalizeScores(s)
252- for i in 0 ..< recs.len: recs[i].score = s[i]
253-
254-proc samplePeople*(pool: seq[PersonRecommendation], half: int,
255- rng: var Rand): seq[PersonRecommendation] =
256- ## Up to `half` from inside the reader's network and `half` from outside.
257- ##
258- ## The split is the point: ranking people purely by similarity returns the
259- ## reader's existing network, which they already have. Reserving half the
260- ## slots for strangers is what makes this discovery rather than a mirror.
261- ##
262- ## Shuffled rather than taken in order so the same faces do not appear
263- ## every time -- the pool changes far more slowly than the reader revisits.
264- var inNet, outNet: seq[PersonRecommendation]
265- for p in pool:
266- if p.isFollowed: inNet.add p else: outNet.add p
267-
268- rng.shuffle(inNet)
269- rng.shuffle(outNet)
270- if inNet.len > half: inNet.setLen(half)
271- if outNet.len > half: outNet.setLen(half)
272- inNet & outNet
273-
274-proc rankByScore*[T](items: var seq[T]) =
275- ## Highest first, which is what every recommendation listing wants.
276- items.sort(proc(a, b: T): int = cmp(b.score, a.score))
deleted file mode 100644
@@ -1,276 +0,0 @@
1-## The similarity computations behind recommendations, ported from
2-## internal/cluster/jaccard.go and the pure helpers of scoring.go.
3-##
4-## Two kinds of similarity are precomputed on a schedule, because both are
5-## O(pairs) and neither needs to be fresh to the second:
6-##
7-## feed_similarity which feeds are subscribed to by the same people
8-## user_similarity which people subscribe to the same feeds
9-##
10-## Both are time-decayed. An overlap from three years ago says much less
11-## about what someone wants now than one from last week, and without decay a
12-## long-lived account's early subscriptions dominate its recommendations
13-## forever.
14-##
15-## The decay constant is `exp(-0.023 * days)`, a half-life of about a month
16-## (ln 2 / 0.023 ≈ 30). It is written inline in the SQL rather than named,
17-## the same way the Go implementation writes it.
18-
19-import std/[algorithm, math, options, random]
20-import ./gleandb
21-import ./sqlite
22-
23-type
24- Config* = object
25- followBoost*: float
26- likesWeight*: float
27- tagsWeight*: float
28- descriptionWeight*: float
29-
30- SignalWeights* = object
31- wSub*: float
32- wLike*: float
33- wTag*: float
34- wSocial*: float
35- wPop*: float
36- wCategory*: float
37- wContent*: float
38-
39- FeedRecommendation* = object
40- feedUrl*: string
41- title*: string
42- siteUrl*: string
43- description*: string
44- subscriberCount*: int
45- faviconUrl*: string
46- score*: float
47-
48- PersonRecommendation* = object
49- did*: string
50- handle*: string
51- displayName*: string
52- avatarUrl*: string
53- commonFeeds*: int
54- commonLikes*: int
55- commonTags*: int
56- isFollowed*: bool
57- score*: float
58-
59-const
60- ## Half-life of roughly 30 days: ln(2) / 0.023.
61- DecayPerDay* = 0.023
62-
63-proc defaultConfig*(): Config =
64- Config(followBoost: 0.5, likesWeight: 0.3, tagsWeight: 0.2,
65- descriptionWeight: 0.15)
66-
67-proc defaultWeights*(): SignalWeights =
68- SignalWeights(wSub: 1.0, wLike: 0.5, wTag: 0.3, wSocial: 0.7,
69- wPop: 0.2, wCategory: 0.4, wContent: 0.4)
70-
71-proc decay*(ageDays: float): float =
72- ## The weight an event of this age still carries.
73- ##
74- ## Exposed so the decay can be checked directly; the queries compute it in
75- ## SQL because pulling every like into the application to weight it would
76- ## defeat the point of precomputing.
77- exp(-DecayPerDay * ageDays)
78-
79-# --- similarity ------------------------------------------------------------
80-
81-## Both computations stage into a temp table and swap in a second
82-## transaction. The first pass is the expensive one and holds no lock on the
83-## live table; readers keep seeing the previous generation until the swap,
84-## rather than an empty table for the length of the computation.
85-
86-proc computeFeedSimilarity*(g: GleanDb) =
87- ## For each pair of feeds, how strongly their subscriber sets overlap.
88- ##
89- ## The numerator is a decayed count of shared subscribers and the
90- ## denominator is the union of the two subscriber counts, which makes this
91- ## a Jaccard index with recency folded into the intersection.
92- g.db.transaction:
93- g.db.exec """CREATE TEMP TABLE IF NOT EXISTS _feed_sim_staging (
94- feed_a TEXT NOT NULL, feed_b TEXT NOT NULL, jaccard REAL NOT NULL,
95- PRIMARY KEY (feed_a, feed_b))"""
96- g.db.exec "DELETE FROM _feed_sim_staging"
97- g.db.exec """
98- INSERT INTO _feed_sim_staging (feed_a, feed_b, jaccard)
99- SELECT feed_a, feed_b, jaccard FROM (
100- SELECT
101- s1.feed_url AS feed_a,
102- s2.feed_url AS feed_b,
103- SUM(EXP(-0.023 * CAST(julianday('now') -
104- julianday(MIN(s1.added_at, s2.added_at)) AS REAL)))
105- / NULLIF(f1.subscriber_count + f2.subscriber_count
106- - CAST(COUNT(*) AS REAL), 0) AS jaccard
107- FROM articles.subscriptions s1
108- -- feed_url < feed_url keeps one row per unordered pair, and drops
109- -- the self-pair without a separate predicate.
110- JOIN articles.subscriptions s2
111- ON s1.user_did = s2.user_did AND s1.feed_url < s2.feed_url
112- JOIN articles.feeds f1 ON f1.feed_url = s1.feed_url
113- JOIN articles.feeds f2 ON f2.feed_url = s2.feed_url
114- WHERE s1.added_at IS NOT NULL AND s2.added_at IS NOT NULL
115- GROUP BY s1.feed_url, s2.feed_url
116- ) WHERE jaccard IS NOT NULL"""
117-
118- g.db.transaction:
119- g.db.exec "DELETE FROM recs.feed_similarity"
120- g.db.exec """INSERT INTO recs.feed_similarity (feed_a, feed_b, jaccard)
121- SELECT feed_a, feed_b, jaccard FROM _feed_sim_staging"""
122-
123-proc computeUserSimilarity*(g: GleanDb) =
124- ## For each pair of people, how strongly their subscriptions overlap.
125- ##
126- ## Undecayed, unlike feed similarity: a subscription is a standing choice
127- ## rather than an event, so its age says little. Shared *likes* are decayed
128- ## separately below, because those are events.
129- g.db.transaction:
130- g.db.exec """CREATE TEMP TABLE IF NOT EXISTS _user_sim_staging (
131- user_a TEXT NOT NULL, user_b TEXT NOT NULL, jaccard REAL NOT NULL,
132- common_feeds INT NOT NULL DEFAULT 0,
133- common_likes INT NOT NULL DEFAULT 0,
134- common_tags INT NOT NULL DEFAULT 0,
135- PRIMARY KEY (user_a, user_b))"""
136- g.db.exec "DELETE FROM _user_sim_staging"
137- g.db.exec """
138- INSERT INTO _user_sim_staging
139- (user_a, user_b, jaccard, common_feeds, common_likes, common_tags)
140- SELECT
141- s1.user_did, s2.user_did,
142- CAST(COUNT(*) AS REAL) / (
143- (SELECT COUNT(*) FROM articles.subscriptions WHERE user_did = s1.user_did) +
144- (SELECT COUNT(*) FROM articles.subscriptions WHERE user_did = s2.user_did) -
145- CAST(COUNT(*) AS REAL)
146- ),
147- COUNT(*), 0, 0
148- FROM articles.subscriptions s1
149- JOIN articles.subscriptions s2
150- ON s1.feed_url = s2.feed_url AND s1.user_did < s2.user_did
151- GROUP BY s1.user_did, s2.user_did"""
152-
153- # Shared likes, decayed on both sides: two people who liked the same
154- # thing years apart have less in common than two who liked it the same
155- # week, and the product of the two decays expresses that.
156- #
157- # ROUND, not CAST AS INTEGER. The Go implementation truncates, and since
158- # each decay factor is at most 1 their product is always just under it --
159- # so a pair sharing one recent like sums to 0.9999 and stores 0. That
160- # number is shown to readers as the reason for a recommendation ("N
161- # shared likes"), so truncation makes a real overlap read as none.
162- g.db.exec """CREATE TEMP TABLE IF NOT EXISTS _likes_overlap (
163- user_a TEXT, user_b TEXT, common INT, PRIMARY KEY(user_a, user_b))"""
164- g.db.exec "DELETE FROM _likes_overlap"
165- g.db.exec """
166- INSERT INTO _likes_overlap (user_a, user_b, common)
167- SELECT l1.author_did, l2.author_did,
168- CAST(ROUND(SUM(
169- EXP(-0.023 * CAST(julianday('now') - julianday(l1.created_at) AS REAL))
170- * EXP(-0.023 * CAST(julianday('now') - julianday(l2.created_at) AS REAL))
171- )) AS INTEGER)
172- FROM articles.likes l1
173- JOIN articles.likes l2
174- ON l1.feed_url = l2.feed_url AND l1.article_url = l2.article_url
175- AND l1.author_did < l2.author_did
176- WHERE l1.created_at IS NOT NULL AND l2.created_at IS NOT NULL
177- GROUP BY l1.author_did, l2.author_did"""
178-
179- # A pair that shares likes but no feeds has no staging row yet, so this
180- # inserts before it updates.
181- g.db.exec """
182- INSERT INTO _user_sim_staging (user_a, user_b, jaccard, common_feeds,
183- common_likes, common_tags)
184- SELECT o.user_a, o.user_b, 0, 0, o.common, 0
185- FROM _likes_overlap o
186- WHERE NOT EXISTS (SELECT 1 FROM _user_sim_staging s
187- WHERE s.user_a = o.user_a AND s.user_b = o.user_b)"""
188- g.db.exec """
189- UPDATE _user_sim_staging SET common_likes = (
190- SELECT o.common FROM _likes_overlap o
191- WHERE o.user_a = _user_sim_staging.user_a
192- AND o.user_b = _user_sim_staging.user_b)
193- WHERE EXISTS (SELECT 1 FROM _likes_overlap o
194- WHERE o.user_a = _user_sim_staging.user_a
195- AND o.user_b = _user_sim_staging.user_b)"""
196-
197- g.db.transaction:
198- g.db.exec "DELETE FROM recs.user_similarity"
199- g.db.exec """INSERT INTO recs.user_similarity
200- (user_a, user_b, jaccard, common_feeds, common_likes, common_tags)
201- SELECT user_a, user_b, jaccard, common_feeds, common_likes,
202- common_tags
203- FROM _user_sim_staging"""
204-
205-proc feedSimilarity*(g: GleanDb, feedA, feedB: string): float =
206- ## Stored pairs are ordered, so look up whichever way round the caller asks.
207- let (a, b) = if feedA < feedB: (feedA, feedB) else: (feedB, feedA)
208- let row = g.db.queryFirst(
209- "SELECT jaccard FROM recs.feed_similarity WHERE feed_a = ? AND feed_b = ?",
210- [p(a), p(b)], proc(r: Row): float = r.f(0))
211- if row.isSome: row.get else: 0.0
212-
213-proc userSimilarity*(g: GleanDb, userA, userB: string): float =
214- let (a, b) = if userA < userB: (userA, userB) else: (userB, userA)
215- let row = g.db.queryFirst(
216- "SELECT jaccard FROM recs.user_similarity WHERE user_a = ? AND user_b = ?",
217- [p(a), p(b)], proc(r: Row): float = r.f(0))
218- if row.isSome: row.get else: 0.0
219-
220-# --- scoring helpers -------------------------------------------------------
221-
222-proc normalizeScores*(scores: var openArray[float]) =
223- ## Rescale to 0..1.
224- ##
225- ## Scores from different signals are not comparable in magnitude -- a
226- ## like-overlap sum and a log-popularity term live on unrelated scales --
227- ## so they are normalised before being blended or ranked against each
228- ## other. A single item, or a set where every score is equal, is left
229- ## alone: there is no spread to stretch, and dividing by zero would turn
230- ## the whole list into NaN.
231- if scores.len < 2: return
232- var lo = scores[0]
233- var hi = scores[0]
234- for s in scores:
235- lo = min(lo, s)
236- hi = max(hi, s)
237- if hi == lo: return
238- let span = hi - lo
239- for i in 0 ..< scores.len:
240- scores[i] = (scores[i] - lo) / span
241-
242-proc normalizeFeedScores*(recs: var seq[FeedRecommendation]) =
243- var s = newSeq[float](recs.len)
244- for i, r in recs: s[i] = r.score
245- normalizeScores(s)
246- for i in 0 ..< recs.len: recs[i].score = s[i]
247-
248-proc normalizePersonScores*(recs: var seq[PersonRecommendation]) =
249- var s = newSeq[float](recs.len)
250- for i, r in recs: s[i] = r.score
251- normalizeScores(s)
252- for i in 0 ..< recs.len: recs[i].score = s[i]
253-
254-proc samplePeople*(pool: seq[PersonRecommendation], half: int,
255- rng: var Rand): seq[PersonRecommendation] =
256- ## Up to `half` from inside the reader's network and `half` from outside.
257- ##
258- ## The split is the point: ranking people purely by similarity returns the
259- ## reader's existing network, which they already have. Reserving half the
260- ## slots for strangers is what makes this discovery rather than a mirror.
261- ##
262- ## Shuffled rather than taken in order so the same faces do not appear
263- ## every time -- the pool changes far more slowly than the reader revisits.
264- var inNet, outNet: seq[PersonRecommendation]
265- for p in pool:
266- if p.isFollowed: inNet.add p else: outNet.add p
267-
268- rng.shuffle(inNet)
269- rng.shuffle(outNet)
270- if inNet.len > half: inNet.setLen(half)
271- if outNet.len > half: outNet.setLen(half)
272- inNet & outNet
273-
274-proc rankByScore*[T](items: var seq[T]) =
275- ## Highest first, which is what every recommendation listing wants.
276- items.sort(proc(a, b: T): int = cmp(b.score, a.score))
deleted spike/nim/atproto/cursor_store.nim +0 -59
deleted file mode 100644
@@ -1,59 +0,0 @@
1-## A Jetstream cursor store backed by SQLite, with batched writes.
2-##
3-## Bridges `store` and `jetstream_run`. Separate from both so the supervisor
4-## does not drag a database dependency behind it, and so an application can
5-## substitute its own persistence.
6-##
7-## Writes are batched because the cursor advances on *every* event -- often
8-## hundreds a second -- and a synchronous write per event would make the
9-## database the bottleneck for a firehose. The cost is that the stored cursor
10-## trails the live one, which is exactly the gap the supervisor's rewind
11-## exists to cover.
12-
13-import std/times
14-import ./jetstream_run
15-import ./sqlite
16-import ./store
17-
18-type DbCursorStore* = ref object of CursorStore
19- db*: Db
20- flushEvery*: int64 ## seconds
21- lastFlush: int64
22- pending: int64
23- dirty: bool
24-
25-proc newDbCursorStore*(db: Db, flushEvery = 5'i64): DbCursorStore =
26- result = DbCursorStore(db: db, flushEvery: flushEvery,
27- lastFlush: getTime().toUnix)
28- result.cursor = loadCursor(db)
29- result.pending = result.cursor
30-
31-method load*(s: DbCursorStore): int64 {.gcsafe, raises: [].} =
32- s.cursor
33-
34-proc flushNow*(s: DbCursorStore) =
35- ## Write the newest cursor out. Safe to call at any time; a failure is
36- ## swallowed, because losing a cursor update costs a replay and raising
37- ## here would cost the subscription.
38- if not s.dirty: return
39- try:
40- saveCursor(s.db, s.pending)
41- s.cursor = s.pending
42- s.dirty = false
43- s.lastFlush = getTime().toUnix
44- except CatchableError:
45- discard
46-
47-method save*(s: DbCursorStore, cursor: int64) {.gcsafe, raises: [].} =
48- ## Called per event. Records in memory and writes at most once per
49- ## flushEvery seconds.
50- if cursor <= s.pending: return
51- s.pending = cursor
52- s.dirty = true
53- let now = try: getTime().toUnix except CatchableError: return
54- if now - s.lastFlush >= s.flushEvery:
55- s.flushNow()
56-
57-proc close*(s: DbCursorStore) =
58- ## Final flush, so a clean shutdown does not throw away the last batch.
59- s.flushNow()
deleted file mode 100644
@@ -1,59 +0,0 @@
1-## A Jetstream cursor store backed by SQLite, with batched writes.
2-##
3-## Bridges `store` and `jetstream_run`. Separate from both so the supervisor
4-## does not drag a database dependency behind it, and so an application can
5-## substitute its own persistence.
6-##
7-## Writes are batched because the cursor advances on *every* event -- often
8-## hundreds a second -- and a synchronous write per event would make the
9-## database the bottleneck for a firehose. The cost is that the stored cursor
10-## trails the live one, which is exactly the gap the supervisor's rewind
11-## exists to cover.
12-
13-import std/times
14-import ./jetstream_run
15-import ./sqlite
16-import ./store
17-
18-type DbCursorStore* = ref object of CursorStore
19- db*: Db
20- flushEvery*: int64 ## seconds
21- lastFlush: int64
22- pending: int64
23- dirty: bool
24-
25-proc newDbCursorStore*(db: Db, flushEvery = 5'i64): DbCursorStore =
26- result = DbCursorStore(db: db, flushEvery: flushEvery,
27- lastFlush: getTime().toUnix)
28- result.cursor = loadCursor(db)
29- result.pending = result.cursor
30-
31-method load*(s: DbCursorStore): int64 {.gcsafe, raises: [].} =
32- s.cursor
33-
34-proc flushNow*(s: DbCursorStore) =
35- ## Write the newest cursor out. Safe to call at any time; a failure is
36- ## swallowed, because losing a cursor update costs a replay and raising
37- ## here would cost the subscription.
38- if not s.dirty: return
39- try:
40- saveCursor(s.db, s.pending)
41- s.cursor = s.pending
42- s.dirty = false
43- s.lastFlush = getTime().toUnix
44- except CatchableError:
45- discard
46-
47-method save*(s: DbCursorStore, cursor: int64) {.gcsafe, raises: [].} =
48- ## Called per event. Records in memory and writes at most once per
49- ## flushEvery seconds.
50- if cursor <= s.pending: return
51- s.pending = cursor
52- s.dirty = true
53- let now = try: getTime().toUnix except CatchableError: return
54- if now - s.lastFlush >= s.flushEvery:
55- s.flushNow()
56-
57-proc close*(s: DbCursorStore) =
58- ## Final flush, so a clean shutdown does not throw away the last batch.
59- s.flushNow()
deleted spike/nim/atproto/dpop.nim +0 -153
deleted file mode 100644
@@ -1,153 +0,0 @@
1-## ES256 keys and DPoP proofs (RFC 9449) for ATProto OAuth.
2-##
3-## What has to be exactly right:
4-## * ES256 signatures in JOSE form -- raw R||S, 64 bytes, NOT the DER that
5-## most crypto libraries hand you.
6-## * An RFC 7638 JWK thumbprint, which is a SHA-256 over canonical JSON:
7-## exactly the members crv/kty/x/y, lexicographically ordered, no spaces.
8-## * base64url without padding, everywhere.
9-##
10-## bearssl is used because br_ecdsa_sign_raw already emits R||S, so there is no
11-## DER unwrapping step to get wrong. dpop_probe.nim cross-checks the output
12-## against a Go verifier.
13-
14-import std/[base64, json, strutils, times, strformat]
15-import bearssl/[ec, rand]
16-import bearssl/abi/[bearssl_ec, bearssl_rand, bearssl_hash]
17-import nimcrypto/[sha2, hash]
18-
19-type P256Key* = object
20- priv*: array[32, byte]
21- pubX*: array[32, byte]
22- pubY*: array[32, byte]
23-
24-proc toSeq(a: openArray[byte]): seq[byte] =
25- result = newSeq[byte](a.len)
26- for i, b in a: result[i] = b
27-
28-proc b64u*(data: openArray[byte]): string =
29- ## base64url, unpadded (RFC 7515 §2).
30- base64.encode(data, safe = true).strip(chars = {'='})
31-
32-proc b64u*(s: string): string =
33- b64u(s.toOpenArrayByte(0, s.high).toSeq)
34-
35-proc generateKey*(rng: var HmacDrbgContext): P256Key =
36- ## P-256 keypair. bearssl returns the public point uncompressed:
37- ## 0x04 || X(32) || Y(32).
38- var
39- skBuf: array[EC_KBUF_PRIV_MAX_SIZE, byte]
40- pkBuf: array[EC_KBUF_PUB_MAX_SIZE, byte]
41- sk: EcPrivateKey
42- pk: EcPublicKey
43-
44- let n = ecKeygen(PrngClassPointerConst(addr rng.vtable), addr ecPrimeI31, addr sk,
45- addr skBuf[0], EC_secp256r1.cint)
46- doAssert n > 0, "ecKeygen failed"
47- let m = ecComputePub(addr ecPrimeI31, addr pk, addr pkBuf[0], addr sk)
48- doAssert m > 0, "ecComputePub failed"
49-
50- doAssert sk.xlen == 32, &"unexpected private key length {sk.xlen}"
51- doAssert pk.qlen == 65 and cast[ptr UncheckedArray[byte]](pk.q)[0] == 0x04'u8,
52- "expected an uncompressed public point"
53-
54- let
55- skBytes = cast[ptr UncheckedArray[byte]](sk.x)
56- pkBytes = cast[ptr UncheckedArray[byte]](pk.q)
57- for i in 0 ..< 32:
58- result.priv[i] = skBytes[i]
59- result.pubX[i] = pkBytes[1 + i]
60- result.pubY[i] = pkBytes[33 + i]
61-
62-proc publicJwk*(k: P256Key): JsonNode =
63- %*{"crv": "P-256", "kty": "EC", "x": b64u(k.pubX), "y": b64u(k.pubY)}
64-
65-proc thumbprint*(k: P256Key): string =
66- ## RFC 7638: SHA-256 over the canonical JWK. The member set and ordering are
67- ## fixed by the spec, so this is built by hand rather than serialised from a
68- ## JsonNode whose key order is incidental.
69- let canonical = &"""{{"crv":"P-256","kty":"EC","x":"{b64u(k.pubX)}","y":"{b64u(k.pubY)}"}}"""
70- b64u(sha256.digest(canonical).data)
71-
72-proc signEs256*(k: P256Key, rng: var HmacDrbgContext, signingInput: string): seq[byte] =
73- var
74- sk: EcPrivateKey
75- priv = k.priv
76- sk.curve = EC_secp256r1.cint
77- sk.x = cast[ptr byte](addr priv[0])
78- sk.xlen = 32
79-
80- let digest = sha256.digest(signingInput).data
81- var sig: array[64, byte]
82- let n = ecdsaSignRawGetDefault()(
83- addr ecPrimeI31, addr sha256Vtable, unsafeAddr digest[0], addr sk, addr sig[0])
84- doAssert n == 64, &"expected a 64-byte raw signature, got {n}"
85- result = sig.toSeq
86-
87-proc b64uDecode*(s: string): string =
88- ## Inverse of b64u: restore the padding base64url drops, then decode.
89- var t = s.replace('-', '+').replace('_', '/')
90- while t.len mod 4 != 0: t.add '='
91- base64.decode(t)
92-
93-proc accessTokenHash*(accessToken: string): string =
94- ## The `ath` claim: base64url(SHA-256(access token)).
95- ##
96- ## This is what binds a proof to one specific token. A resource-server proof
97- ## without it is rejected, and a proof carrying the wrong one lets a stolen
98- ## token be replayed with a fresh proof -- which is the whole thing DPoP
99- ## exists to prevent.
100- b64u(sha256.digest(accessToken).data)
101-
102-proc verifyEs256*(k: P256Key, signingInput: string,
103- sig: openArray[byte]): bool =
104- ## Verify a raw R||S signature against this key's public half.
105- ##
106- ## Mostly useful for proving a key survived a round trip: ECDSA is
107- ## randomised, so a reloaded key cannot be checked by re-signing and
108- ## comparing bytes. Signing with the reloaded key and verifying against the
109- ## original public point does establish the private scalar came back intact.
110- if sig.len != 64: return false
111- var q: array[65, byte]
112- q[0] = 0x04
113- for i in 0 ..< 32:
114- q[1 + i] = k.pubX[i]
115- q[33 + i] = k.pubY[i]
116-
117- var pk: EcPublicKey
118- pk.curve = EC_secp256r1.cint
119- pk.q = cast[ptr byte](addr q[0])
120- pk.qlen = 65
121-
122- let digest = sha256.digest(signingInput).data
123- ecdsaVrfyRawGetDefault()(
124- addr ecPrimeI31, unsafeAddr digest[0], digest.len.uint,
125- addr pk, unsafeAddr sig[0], sig.len.uint) == 1
126-
127-proc dpopProof*(k: P256Key, rng: var HmacDrbgContext, htm, htu: string,
128- nonce = "", accessToken = "", issuer = ""): string =
129- ## A DPoP proof JWT (RFC 9449).
130- ##
131- ## Auth-server and resource-server proofs differ, even when the same host
132- ## fills both roles: a request to a PDS carries `ath` (bound to the access
133- ## token) and `iss` (the auth server that issued it), while the PAR and
134- ## token requests that mint that token carry neither.
135- let header = %*{"typ": "dpop+jwt", "alg": "ES256", "jwk": k.publicJwk}
136- let issued = now().utc.toTime.toUnix
137- var payload = %*{
138- "jti": b64u(sha256.digest(&"{htm}{htu}{epochTime()}").data)[0 ..< 16],
139- "htm": htm,
140- "htu": htu,
141- "iat": issued,
142- "exp": issued + 300,
143- }
144- if nonce.len > 0:
145- payload["nonce"] = %nonce
146- if accessToken.len > 0:
147- payload["ath"] = %accessTokenHash(accessToken)
148- if issuer.len > 0:
149- payload["iss"] = %issuer
150-
151- let signingInput = b64u($header) & "." & b64u($payload)
152- signingInput & "." & b64u(k.signEs256(rng, signingInput))
153-
deleted file mode 100644
@@ -1,153 +0,0 @@
1-## ES256 keys and DPoP proofs (RFC 9449) for ATProto OAuth.
2-##
3-## What has to be exactly right:
4-## * ES256 signatures in JOSE form -- raw R||S, 64 bytes, NOT the DER that
5-## most crypto libraries hand you.
6-## * An RFC 7638 JWK thumbprint, which is a SHA-256 over canonical JSON:
7-## exactly the members crv/kty/x/y, lexicographically ordered, no spaces.
8-## * base64url without padding, everywhere.
9-##
10-## bearssl is used because br_ecdsa_sign_raw already emits R||S, so there is no
11-## DER unwrapping step to get wrong. dpop_probe.nim cross-checks the output
12-## against a Go verifier.
13-
14-import std/[base64, json, strutils, times, strformat]
15-import bearssl/[ec, rand]
16-import bearssl/abi/[bearssl_ec, bearssl_rand, bearssl_hash]
17-import nimcrypto/[sha2, hash]
18-
19-type P256Key* = object
20- priv*: array[32, byte]
21- pubX*: array[32, byte]
22- pubY*: array[32, byte]
23-
24-proc toSeq(a: openArray[byte]): seq[byte] =
25- result = newSeq[byte](a.len)
26- for i, b in a: result[i] = b
27-
28-proc b64u*(data: openArray[byte]): string =
29- ## base64url, unpadded (RFC 7515 §2).
30- base64.encode(data, safe = true).strip(chars = {'='})
31-
32-proc b64u*(s: string): string =
33- b64u(s.toOpenArrayByte(0, s.high).toSeq)
34-
35-proc generateKey*(rng: var HmacDrbgContext): P256Key =
36- ## P-256 keypair. bearssl returns the public point uncompressed:
37- ## 0x04 || X(32) || Y(32).
38- var
39- skBuf: array[EC_KBUF_PRIV_MAX_SIZE, byte]
40- pkBuf: array[EC_KBUF_PUB_MAX_SIZE, byte]
41- sk: EcPrivateKey
42- pk: EcPublicKey
43-
44- let n = ecKeygen(PrngClassPointerConst(addr rng.vtable), addr ecPrimeI31, addr sk,
45- addr skBuf[0], EC_secp256r1.cint)
46- doAssert n > 0, "ecKeygen failed"
47- let m = ecComputePub(addr ecPrimeI31, addr pk, addr pkBuf[0], addr sk)
48- doAssert m > 0, "ecComputePub failed"
49-
50- doAssert sk.xlen == 32, &"unexpected private key length {sk.xlen}"
51- doAssert pk.qlen == 65 and cast[ptr UncheckedArray[byte]](pk.q)[0] == 0x04'u8,
52- "expected an uncompressed public point"
53-
54- let
55- skBytes = cast[ptr UncheckedArray[byte]](sk.x)
56- pkBytes = cast[ptr UncheckedArray[byte]](pk.q)
57- for i in 0 ..< 32:
58- result.priv[i] = skBytes[i]
59- result.pubX[i] = pkBytes[1 + i]
60- result.pubY[i] = pkBytes[33 + i]
61-
62-proc publicJwk*(k: P256Key): JsonNode =
63- %*{"crv": "P-256", "kty": "EC", "x": b64u(k.pubX), "y": b64u(k.pubY)}
64-
65-proc thumbprint*(k: P256Key): string =
66- ## RFC 7638: SHA-256 over the canonical JWK. The member set and ordering are
67- ## fixed by the spec, so this is built by hand rather than serialised from a
68- ## JsonNode whose key order is incidental.
69- let canonical = &"""{{"crv":"P-256","kty":"EC","x":"{b64u(k.pubX)}","y":"{b64u(k.pubY)}"}}"""
70- b64u(sha256.digest(canonical).data)
71-
72-proc signEs256*(k: P256Key, rng: var HmacDrbgContext, signingInput: string): seq[byte] =
73- var
74- sk: EcPrivateKey
75- priv = k.priv
76- sk.curve = EC_secp256r1.cint
77- sk.x = cast[ptr byte](addr priv[0])
78- sk.xlen = 32
79-
80- let digest = sha256.digest(signingInput).data
81- var sig: array[64, byte]
82- let n = ecdsaSignRawGetDefault()(
83- addr ecPrimeI31, addr sha256Vtable, unsafeAddr digest[0], addr sk, addr sig[0])
84- doAssert n == 64, &"expected a 64-byte raw signature, got {n}"
85- result = sig.toSeq
86-
87-proc b64uDecode*(s: string): string =
88- ## Inverse of b64u: restore the padding base64url drops, then decode.
89- var t = s.replace('-', '+').replace('_', '/')
90- while t.len mod 4 != 0: t.add '='
91- base64.decode(t)
92-
93-proc accessTokenHash*(accessToken: string): string =
94- ## The `ath` claim: base64url(SHA-256(access token)).
95- ##
96- ## This is what binds a proof to one specific token. A resource-server proof
97- ## without it is rejected, and a proof carrying the wrong one lets a stolen
98- ## token be replayed with a fresh proof -- which is the whole thing DPoP
99- ## exists to prevent.
100- b64u(sha256.digest(accessToken).data)
101-
102-proc verifyEs256*(k: P256Key, signingInput: string,
103- sig: openArray[byte]): bool =
104- ## Verify a raw R||S signature against this key's public half.
105- ##
106- ## Mostly useful for proving a key survived a round trip: ECDSA is
107- ## randomised, so a reloaded key cannot be checked by re-signing and
108- ## comparing bytes. Signing with the reloaded key and verifying against the
109- ## original public point does establish the private scalar came back intact.
110- if sig.len != 64: return false
111- var q: array[65, byte]
112- q[0] = 0x04
113- for i in 0 ..< 32:
114- q[1 + i] = k.pubX[i]
115- q[33 + i] = k.pubY[i]
116-
117- var pk: EcPublicKey
118- pk.curve = EC_secp256r1.cint
119- pk.q = cast[ptr byte](addr q[0])
120- pk.qlen = 65
121-
122- let digest = sha256.digest(signingInput).data
123- ecdsaVrfyRawGetDefault()(
124- addr ecPrimeI31, unsafeAddr digest[0], digest.len.uint,
125- addr pk, unsafeAddr sig[0], sig.len.uint) == 1
126-
127-proc dpopProof*(k: P256Key, rng: var HmacDrbgContext, htm, htu: string,
128- nonce = "", accessToken = "", issuer = ""): string =
129- ## A DPoP proof JWT (RFC 9449).
130- ##
131- ## Auth-server and resource-server proofs differ, even when the same host
132- ## fills both roles: a request to a PDS carries `ath` (bound to the access
133- ## token) and `iss` (the auth server that issued it), while the PAR and
134- ## token requests that mint that token carry neither.
135- let header = %*{"typ": "dpop+jwt", "alg": "ES256", "jwk": k.publicJwk}
136- let issued = now().utc.toTime.toUnix
137- var payload = %*{
138- "jti": b64u(sha256.digest(&"{htm}{htu}{epochTime()}").data)[0 ..< 16],
139- "htm": htm,
140- "htu": htu,
141- "iat": issued,
142- "exp": issued + 300,
143- }
144- if nonce.len > 0:
145- payload["nonce"] = %nonce
146- if accessToken.len > 0:
147- payload["ath"] = %accessTokenHash(accessToken)
148- if issuer.len > 0:
149- payload["iss"] = %issuer
150-
151- let signingInput = b64u($header) & "." & b64u($payload)
152- signingInput & "." & b64u(k.signEs256(rng, signingInput))
153-
deleted spike/nim/atproto/gleandb.nim +0 -150
deleted file mode 100644
@@ -1,150 +0,0 @@
1-## Glean's database: three SQLite files behind one connection.
2-##
3-## `main` holds users and auth, `articles` the feeds and their contents, and
4-## `recs` the derived recommendation tables. They are separate files attached
5-## together, so a query can join across them while each still checkpoints and
6-## vacuums on its own -- and, more practically, so the derived tables can be
7-## deleted and rebuilt without touching anything a user wrote.
8-##
9-## The pragmas are copied from internal/db/db.go. They are not decoration:
10-## WAL is what lets readers run during a write, and the per-database repeats
11-## matter because PRAGMA applies to one schema at a time -- setting WAL on
12-## `main` leaves `articles` journalling the slow way.
13-
14-import std/[math, options, strformat, strutils]
15-import ./schema
16-import ./sqlite
17-
18-type
19- GleanDb* = object
20- db*: Db
21- basePath*: string
22-
23-{.push importc, cdecl.}
24-proc sqlite3_create_function_v2(db: pointer, name: cstring, nArg: cint,
25- eTextRep: cint, pApp: pointer,
26- xFunc, xStep, xFinal, xDestroy: pointer): cint
27-proc sqlite3_value_double(v: pointer): cdouble
28-proc sqlite3_result_double(ctx: pointer, v: cdouble)
29-proc sqlite3_result_null(ctx: pointer)
30-{.pop.}
31-
32-const
33- SQLITE_UTF8 = 1.cint
34- SQLITE_DETERMINISTIC = 0x000000800.cint
35-
36-# SQLite ships maths functions only when built with
37-# SQLITE_ENABLE_MATH_FUNCTIONS, which the library here is not, so these are
38-# registered exactly as the Go connection registers them. They are not
39-# optional: the clustering SQL calls EXP and LOG in seven places for time
40-# decay and popularity normalisation, and without them every recommendation
41-# query fails at runtime rather than at startup.
42-proc sqlExp(ctx: pointer, n: cint, args: ptr UncheckedArray[pointer])
43- {.cdecl.} =
44- if n != 1:
45- sqlite3_result_null(ctx)
46- return
47- sqlite3_result_double(ctx, exp(sqlite3_value_double(args[0])).cdouble)
48-
49-proc sqlLog(ctx: pointer, n: cint, args: ptr UncheckedArray[pointer])
50- {.cdecl.} =
51- if n != 1:
52- sqlite3_result_null(ctx)
53- return
54- let v = sqlite3_value_double(args[0]).float
55- # log(0) is -inf and log(negative) is NaN; SQL callers want neither, and a
56- # NULL propagates through an ORDER BY far more gracefully.
57- if v <= 0:
58- sqlite3_result_null(ctx)
59- else:
60- sqlite3_result_double(ctx, ln(v).cdouble)
61-
62-proc registerFunctions*(db: Db) =
63- for (name, fn) in [("exp", sqlExp), ("log", sqlLog)]:
64- let rc = sqlite3_create_function_v2(
65- db.handle, name.cstring, 1,
66- SQLITE_UTF8 or SQLITE_DETERMINISTIC,
67- nil, cast[pointer](fn), nil, nil, nil)
68- if rc != SQLITE_OK:
69- raise newException(SqliteError, "registering SQL function " & name)
70-
71-const
72- MainPragmas = [
73- "journal_mode=WAL",
74- "wal_autocheckpoint=1000",
75- "busy_timeout=30000",
76- "synchronous=NORMAL",
77- "temp_store=FILE",
78- "mmap_size=268435456",
79- "auto_vacuum=INCREMENTAL",
80- ]
81- AttachedPragmas = [
82- "journal_mode=WAL",
83- "wal_autocheckpoint=1000",
84- "busy_timeout=30000",
85- "synchronous=NORMAL",
86- "temp_store=FILE",
87- "mmap_size=268435456",
88- ]
89-
90-proc applyPragmas(db: Db, schemaName: string, pragmas: openArray[string]) =
91- for pragma in pragmas:
92- # A pragma that will not take is not worth failing a boot over -- some are
93- # advisory, and mmap_size in particular is refused on filesystems that
94- # cannot map. The ones that matter are asserted after opening instead.
95- try:
96- if schemaName.len == 0:
97- db.exec &"PRAGMA {pragma}"
98- else:
99- db.exec &"PRAGMA {schemaName}.{pragma}"
100- except SqliteError:
101- discard
102-
103-proc open*(basePath: string): GleanDb =
104- ## Open `<basePath>_users` and attach `_articles` and `_recs`.
105- ##
106- ## The suffixes match the Go layout, so a database written by either side is
107- ## readable by the other.
108- let
109- usersPath = basePath & "_users"
110- articlesPath = basePath & "_articles"
111- recsPath = basePath & "_recs"
112-
113- var db = sqlite.open(usersPath)
114- registerFunctions(db)
115- applyPragmas(db, "", MainPragmas)
116-
117- # Quoted and single-quote-escaped: a path is not necessarily tame, and this
118- # one can come from configuration.
119- for (name, path) in [("articles", articlesPath), ("recs", recsPath)]:
120- db.exec &"ATTACH DATABASE '{path.replace(\"'\", \"''\")}' AS {name}"
121- applyPragmas(db, name, AttachedPragmas)
122-
123- GleanDb(db: db, basePath: basePath)
124-
125-proc migrate*(g: GleanDb) =
126- ## Apply the schema. Every statement is IF NOT EXISTS, so this runs on every
127- ## boot rather than being gated on a version.
128- for stmt in usersSchema: g.db.exec stmt
129- for stmt in articlesSchema: g.db.exec stmt
130- for stmt in recsSchema: g.db.exec stmt
131-
132-proc createEmbeddingTables*(g: GleanDb, dimension: int) =
133- ## The vec0 tables carry their dimension in the DDL, so they cannot live in
134- ## the static schema -- it comes from GLEAN_EMBED_DIMENSION.
135- g.db.exec &"""CREATE VIRTUAL TABLE IF NOT EXISTS recs.feed_embeddings
136- USING vec0(feed_url TEXT PRIMARY KEY, embedding float[{dimension}])"""
137- g.db.exec &"""CREATE VIRTUAL TABLE IF NOT EXISTS recs.article_embeddings
138- USING vec0(article_id INTEGER PRIMARY KEY, embedding float[{dimension}])"""
139-
140-proc close*(g: var GleanDb) =
141- g.db.close()
142-
143-proc journalMode*(g: GleanDb, schemaName = ""): string =
144- let sql = if schemaName.len == 0: "PRAGMA journal_mode"
145- else: &"PRAGMA {schemaName}.journal_mode"
146- g.db.queryText(sql).get("")
147-
148-proc attachedSchemas*(g: GleanDb): seq[string] =
149- for row in g.db.rows("PRAGMA database_list"):
150- result.add row[1]
deleted file mode 100644
@@ -1,150 +0,0 @@
1-## Glean's database: three SQLite files behind one connection.
2-##
3-## `main` holds users and auth, `articles` the feeds and their contents, and
4-## `recs` the derived recommendation tables. They are separate files attached
5-## together, so a query can join across them while each still checkpoints and
6-## vacuums on its own -- and, more practically, so the derived tables can be
7-## deleted and rebuilt without touching anything a user wrote.
8-##
9-## The pragmas are copied from internal/db/db.go. They are not decoration:
10-## WAL is what lets readers run during a write, and the per-database repeats
11-## matter because PRAGMA applies to one schema at a time -- setting WAL on
12-## `main` leaves `articles` journalling the slow way.
13-
14-import std/[math, options, strformat, strutils]
15-import ./schema
16-import ./sqlite
17-
18-type
19- GleanDb* = object
20- db*: Db
21- basePath*: string
22-
23-{.push importc, cdecl.}
24-proc sqlite3_create_function_v2(db: pointer, name: cstring, nArg: cint,
25- eTextRep: cint, pApp: pointer,
26- xFunc, xStep, xFinal, xDestroy: pointer): cint
27-proc sqlite3_value_double(v: pointer): cdouble
28-proc sqlite3_result_double(ctx: pointer, v: cdouble)
29-proc sqlite3_result_null(ctx: pointer)
30-{.pop.}
31-
32-const
33- SQLITE_UTF8 = 1.cint
34- SQLITE_DETERMINISTIC = 0x000000800.cint
35-
36-# SQLite ships maths functions only when built with
37-# SQLITE_ENABLE_MATH_FUNCTIONS, which the library here is not, so these are
38-# registered exactly as the Go connection registers them. They are not
39-# optional: the clustering SQL calls EXP and LOG in seven places for time
40-# decay and popularity normalisation, and without them every recommendation
41-# query fails at runtime rather than at startup.
42-proc sqlExp(ctx: pointer, n: cint, args: ptr UncheckedArray[pointer])
43- {.cdecl.} =
44- if n != 1:
45- sqlite3_result_null(ctx)
46- return
47- sqlite3_result_double(ctx, exp(sqlite3_value_double(args[0])).cdouble)
48-
49-proc sqlLog(ctx: pointer, n: cint, args: ptr UncheckedArray[pointer])
50- {.cdecl.} =
51- if n != 1:
52- sqlite3_result_null(ctx)
53- return
54- let v = sqlite3_value_double(args[0]).float
55- # log(0) is -inf and log(negative) is NaN; SQL callers want neither, and a
56- # NULL propagates through an ORDER BY far more gracefully.
57- if v <= 0:
58- sqlite3_result_null(ctx)
59- else:
60- sqlite3_result_double(ctx, ln(v).cdouble)
61-
62-proc registerFunctions*(db: Db) =
63- for (name, fn) in [("exp", sqlExp), ("log", sqlLog)]:
64- let rc = sqlite3_create_function_v2(
65- db.handle, name.cstring, 1,
66- SQLITE_UTF8 or SQLITE_DETERMINISTIC,
67- nil, cast[pointer](fn), nil, nil, nil)
68- if rc != SQLITE_OK:
69- raise newException(SqliteError, "registering SQL function " & name)
70-
71-const
72- MainPragmas = [
73- "journal_mode=WAL",
74- "wal_autocheckpoint=1000",
75- "busy_timeout=30000",
76- "synchronous=NORMAL",
77- "temp_store=FILE",
78- "mmap_size=268435456",
79- "auto_vacuum=INCREMENTAL",
80- ]
81- AttachedPragmas = [
82- "journal_mode=WAL",
83- "wal_autocheckpoint=1000",
84- "busy_timeout=30000",
85- "synchronous=NORMAL",
86- "temp_store=FILE",
87- "mmap_size=268435456",
88- ]
89-
90-proc applyPragmas(db: Db, schemaName: string, pragmas: openArray[string]) =
91- for pragma in pragmas:
92- # A pragma that will not take is not worth failing a boot over -- some are
93- # advisory, and mmap_size in particular is refused on filesystems that
94- # cannot map. The ones that matter are asserted after opening instead.
95- try:
96- if schemaName.len == 0:
97- db.exec &"PRAGMA {pragma}"
98- else:
99- db.exec &"PRAGMA {schemaName}.{pragma}"
100- except SqliteError:
101- discard
102-
103-proc open*(basePath: string): GleanDb =
104- ## Open `<basePath>_users` and attach `_articles` and `_recs`.
105- ##
106- ## The suffixes match the Go layout, so a database written by either side is
107- ## readable by the other.
108- let
109- usersPath = basePath & "_users"
110- articlesPath = basePath & "_articles"
111- recsPath = basePath & "_recs"
112-
113- var db = sqlite.open(usersPath)
114- registerFunctions(db)
115- applyPragmas(db, "", MainPragmas)
116-
117- # Quoted and single-quote-escaped: a path is not necessarily tame, and this
118- # one can come from configuration.
119- for (name, path) in [("articles", articlesPath), ("recs", recsPath)]:
120- db.exec &"ATTACH DATABASE '{path.replace(\"'\", \"''\")}' AS {name}"
121- applyPragmas(db, name, AttachedPragmas)
122-
123- GleanDb(db: db, basePath: basePath)
124-
125-proc migrate*(g: GleanDb) =
126- ## Apply the schema. Every statement is IF NOT EXISTS, so this runs on every
127- ## boot rather than being gated on a version.
128- for stmt in usersSchema: g.db.exec stmt
129- for stmt in articlesSchema: g.db.exec stmt
130- for stmt in recsSchema: g.db.exec stmt
131-
132-proc createEmbeddingTables*(g: GleanDb, dimension: int) =
133- ## The vec0 tables carry their dimension in the DDL, so they cannot live in
134- ## the static schema -- it comes from GLEAN_EMBED_DIMENSION.
135- g.db.exec &"""CREATE VIRTUAL TABLE IF NOT EXISTS recs.feed_embeddings
136- USING vec0(feed_url TEXT PRIMARY KEY, embedding float[{dimension}])"""
137- g.db.exec &"""CREATE VIRTUAL TABLE IF NOT EXISTS recs.article_embeddings
138- USING vec0(article_id INTEGER PRIMARY KEY, embedding float[{dimension}])"""
139-
140-proc close*(g: var GleanDb) =
141- g.db.close()
142-
143-proc journalMode*(g: GleanDb, schemaName = ""): string =
144- let sql = if schemaName.len == 0: "PRAGMA journal_mode"
145- else: &"PRAGMA {schemaName}.journal_mode"
146- g.db.queryText(sql).get("")
147-
148-proc attachedSchemas*(g: GleanDb): seq[string] =
149- for row in g.db.rows("PRAGMA database_list"):
150- result.add row[1]
deleted spike/nim/atproto/identity.nim +0 -149
deleted file mode 100644
@@ -1,149 +0,0 @@
1-## Handle and DID resolution: the chain that turns "alice.bsky.social" into
2-## the PDS host an OAuth flow has to talk to.
3-##
4-## The steps, matching what indigo's identity directory does:
5-## handle -> DID via DNS TXT _atproto.<handle>, else
6-## https://<handle>/.well-known/atproto-did
7-## DID -> doc via the PLC directory for did:plc, or
8-## https://<host>/.well-known/did.json for did:web
9-## doc -> PDS the service entry with id "#atproto_pds"
10-
11-import std/[httpclient, json, options, os, osproc, strutils, uri]
12-
13-const
14- DefaultPlcUrl* = "https://plc.directory"
15- UserAgent* = "glean-nim/0.1"
16- HttpTimeoutMs = 10_000
17-
18-type
19- IdentityError* = object of CatchableError
20-
21- Identity* = object
22- did*: string
23- handle*: string
24- pdsEndpoint*: string
25-
26-proc newClient(): HttpClient =
27- newHttpClient(userAgent = UserAgent, timeout = HttpTimeoutMs)
28-
29-proc isValidHandle*(handle: string): bool =
30- ## Deliberately loose: a dotted, non-empty ASCII domain. The authoritative
31- ## check is whether it resolves.
32- if handle.len == 0 or handle.len > 253 or '.' notin handle:
33- return false
34- for label in handle.split('.'):
35- if label.len == 0: return false
36- for ch in label:
37- if ch notin {'a'..'z', 'A'..'Z', '0'..'9', '-'}: return false
38- true
39-
40-proc isDid*(s: string): bool =
41- s.startsWith("did:plc:") or s.startsWith("did:web:")
42-
43-proc resolveHandleDns(handle: string): Option[string] =
44- ## _atproto.<handle> TXT record containing "did=did:plc:...".
45- ##
46- ## Nim has no resolver in the stdlib, so this shells out to `dig`. A real
47- ## port should bind a resolver rather than depend on a binary being present;
48- ## a missing `dig` is treated as "no DNS answer" so the HTTP fallback runs.
49- let dig = findExe("dig")
50- if dig.len == 0:
51- return none(string)
52- let (output, code) = execCmdEx(dig & " +short +time=3 +tries=1 TXT _atproto." &
53- quoteShell(handle))
54- if code != 0:
55- return none(string)
56- for rawLine in output.splitLines:
57- let line = rawLine.strip(chars = {'"', ' ', '\t'})
58- if line.startsWith("did="):
59- let did = line[4 .. ^1].strip(chars = {'"'})
60- if did.isDid:
61- return some(did)
62- none(string)
63-
64-proc resolveHandleHttp(handle: string): Option[string] =
65- let client = newClient()
66- defer: client.close()
67- try:
68- let res = client.get("https://" & handle & "/.well-known/atproto-did")
69- if res.code.int != 200:
70- return none(string)
71- let did = res.body.strip()
72- if did.isDid: some(did) else: none(string)
73- except CatchableError:
74- none(string)
75-
76-proc resolveHandle*(handle: string): string =
77- ## DNS first, then the well-known endpoint, as the spec prefers.
78- if not handle.isValidHandle:
79- raise newException(IdentityError, "not a valid handle: " & handle)
80- let viaDns = resolveHandleDns(handle)
81- if viaDns.isSome:
82- return viaDns.get
83- let viaHttp = resolveHandleHttp(handle)
84- if viaHttp.isSome:
85- return viaHttp.get
86- raise newException(IdentityError, "could not resolve handle: " & handle)
87-
88-proc resolveDidDoc*(did: string, plcUrl = DefaultPlcUrl): JsonNode =
89- if not did.isDid:
90- raise newException(IdentityError, "unsupported DID method: " & did)
91-
92- let url =
93- if did.startsWith("did:plc:"):
94- plcUrl & "/" & did
95- else:
96- # did:web:example.com -> https://example.com/.well-known/did.json
97- let host = did["did:web:".len .. ^1].replace("%3A", ":")
98- "https://" & host & "/.well-known/did.json"
99-
100- let client = newClient()
101- defer: client.close()
102- let res = client.get(url)
103- if res.code.int != 200:
104- raise newException(IdentityError,
105- "DID document lookup failed (HTTP " & $res.code.int & "): " & did)
106- try:
107- result = parseJson(res.body)
108- except JsonParsingError:
109- raise newException(IdentityError, "DID document is not valid JSON: " & did)
110-
111-proc pdsEndpoint*(doc: JsonNode): string =
112- ## The service entry whose id ends in "#atproto_pds".
113- if doc.kind != JObject or "service" notin doc:
114- return ""
115- for svc in doc["service"]:
116- let id = svc{"id"}.getStr
117- if id.endsWith("#atproto_pds"):
118- return svc{"serviceEndpoint"}.getStr.strip(chars = {'/'})
119- ""
120-
121-proc handleFromDoc(doc: JsonNode): string =
122- for aka in doc{"alsoKnownAs"}:
123- let v = aka.getStr
124- if v.startsWith("at://"):
125- return v["at://".len .. ^1]
126- ""
127-
128-proc lookup*(identifier: string, plcUrl = DefaultPlcUrl): Identity =
129- ## Resolve a handle or DID all the way to a PDS endpoint.
130- let did = if identifier.isDid: identifier else: resolveHandle(identifier)
131- let doc = resolveDidDoc(did, plcUrl)
132- result = Identity(
133- did: did,
134- handle: if identifier.isDid: handleFromDoc(doc) else: identifier,
135- pdsEndpoint: pdsEndpoint(doc),
136- )
137- if result.pdsEndpoint.len == 0:
138- raise newException(IdentityError, "identity has no atproto PDS: " & did)
139-
140-proc hostOf*(url: string): string =
141- ## Host (with port, if given) from a URL.
142- ##
143- ## parseUri treats a scheme-less string as a path, so "example.com" would
144- ## otherwise yield an empty host and build URLs like "https:///xrpc/...".
145- ## Callers legitimately pass bare hostnames, so assume https when no scheme
146- ## is present.
147- let withScheme = if "://" in url: url else: "https://" & url
148- let u = parseUri(withScheme)
149- if u.port.len > 0: u.hostname & ":" & u.port else: u.hostname
deleted file mode 100644
@@ -1,149 +0,0 @@
1-## Handle and DID resolution: the chain that turns "alice.bsky.social" into
2-## the PDS host an OAuth flow has to talk to.
3-##
4-## The steps, matching what indigo's identity directory does:
5-## handle -> DID via DNS TXT _atproto.<handle>, else
6-## https://<handle>/.well-known/atproto-did
7-## DID -> doc via the PLC directory for did:plc, or
8-## https://<host>/.well-known/did.json for did:web
9-## doc -> PDS the service entry with id "#atproto_pds"
10-
11-import std/[httpclient, json, options, os, osproc, strutils, uri]
12-
13-const
14- DefaultPlcUrl* = "https://plc.directory"
15- UserAgent* = "glean-nim/0.1"
16- HttpTimeoutMs = 10_000
17-
18-type
19- IdentityError* = object of CatchableError
20-
21- Identity* = object
22- did*: string
23- handle*: string
24- pdsEndpoint*: string
25-
26-proc newClient(): HttpClient =
27- newHttpClient(userAgent = UserAgent, timeout = HttpTimeoutMs)
28-
29-proc isValidHandle*(handle: string): bool =
30- ## Deliberately loose: a dotted, non-empty ASCII domain. The authoritative
31- ## check is whether it resolves.
32- if handle.len == 0 or handle.len > 253 or '.' notin handle:
33- return false
34- for label in handle.split('.'):
35- if label.len == 0: return false
36- for ch in label:
37- if ch notin {'a'..'z', 'A'..'Z', '0'..'9', '-'}: return false
38- true
39-
40-proc isDid*(s: string): bool =
41- s.startsWith("did:plc:") or s.startsWith("did:web:")
42-
43-proc resolveHandleDns(handle: string): Option[string] =
44- ## _atproto.<handle> TXT record containing "did=did:plc:...".
45- ##
46- ## Nim has no resolver in the stdlib, so this shells out to `dig`. A real
47- ## port should bind a resolver rather than depend on a binary being present;
48- ## a missing `dig` is treated as "no DNS answer" so the HTTP fallback runs.
49- let dig = findExe("dig")
50- if dig.len == 0:
51- return none(string)
52- let (output, code) = execCmdEx(dig & " +short +time=3 +tries=1 TXT _atproto." &
53- quoteShell(handle))
54- if code != 0:
55- return none(string)
56- for rawLine in output.splitLines:
57- let line = rawLine.strip(chars = {'"', ' ', '\t'})
58- if line.startsWith("did="):
59- let did = line[4 .. ^1].strip(chars = {'"'})
60- if did.isDid:
61- return some(did)
62- none(string)
63-
64-proc resolveHandleHttp(handle: string): Option[string] =
65- let client = newClient()
66- defer: client.close()
67- try:
68- let res = client.get("https://" & handle & "/.well-known/atproto-did")
69- if res.code.int != 200:
70- return none(string)
71- let did = res.body.strip()
72- if did.isDid: some(did) else: none(string)
73- except CatchableError:
74- none(string)
75-
76-proc resolveHandle*(handle: string): string =
77- ## DNS first, then the well-known endpoint, as the spec prefers.
78- if not handle.isValidHandle:
79- raise newException(IdentityError, "not a valid handle: " & handle)
80- let viaDns = resolveHandleDns(handle)
81- if viaDns.isSome:
82- return viaDns.get
83- let viaHttp = resolveHandleHttp(handle)
84- if viaHttp.isSome:
85- return viaHttp.get
86- raise newException(IdentityError, "could not resolve handle: " & handle)
87-
88-proc resolveDidDoc*(did: string, plcUrl = DefaultPlcUrl): JsonNode =
89- if not did.isDid:
90- raise newException(IdentityError, "unsupported DID method: " & did)
91-
92- let url =
93- if did.startsWith("did:plc:"):
94- plcUrl & "/" & did
95- else:
96- # did:web:example.com -> https://example.com/.well-known/did.json
97- let host = did["did:web:".len .. ^1].replace("%3A", ":")
98- "https://" & host & "/.well-known/did.json"
99-
100- let client = newClient()
101- defer: client.close()
102- let res = client.get(url)
103- if res.code.int != 200:
104- raise newException(IdentityError,
105- "DID document lookup failed (HTTP " & $res.code.int & "): " & did)
106- try:
107- result = parseJson(res.body)
108- except JsonParsingError:
109- raise newException(IdentityError, "DID document is not valid JSON: " & did)
110-
111-proc pdsEndpoint*(doc: JsonNode): string =
112- ## The service entry whose id ends in "#atproto_pds".
113- if doc.kind != JObject or "service" notin doc:
114- return ""
115- for svc in doc["service"]:
116- let id = svc{"id"}.getStr
117- if id.endsWith("#atproto_pds"):
118- return svc{"serviceEndpoint"}.getStr.strip(chars = {'/'})
119- ""
120-
121-proc handleFromDoc(doc: JsonNode): string =
122- for aka in doc{"alsoKnownAs"}:
123- let v = aka.getStr
124- if v.startsWith("at://"):
125- return v["at://".len .. ^1]
126- ""
127-
128-proc lookup*(identifier: string, plcUrl = DefaultPlcUrl): Identity =
129- ## Resolve a handle or DID all the way to a PDS endpoint.
130- let did = if identifier.isDid: identifier else: resolveHandle(identifier)
131- let doc = resolveDidDoc(did, plcUrl)
132- result = Identity(
133- did: did,
134- handle: if identifier.isDid: handleFromDoc(doc) else: identifier,
135- pdsEndpoint: pdsEndpoint(doc),
136- )
137- if result.pdsEndpoint.len == 0:
138- raise newException(IdentityError, "identity has no atproto PDS: " & did)
139-
140-proc hostOf*(url: string): string =
141- ## Host (with port, if given) from a URL.
142- ##
143- ## parseUri treats a scheme-less string as a path, so "example.com" would
144- ## otherwise yield an empty host and build URLs like "https:///xrpc/...".
145- ## Callers legitimately pass bare hostnames, so assume https when no scheme
146- ## is present.
147- let withScheme = if "://" in url: url else: "https://" & url
148- let u = parseUri(withScheme)
149- if u.port.len > 0: u.hostname & ":" & u.port else: u.hostname
deleted spike/nim/atproto/jetstream.nim +0 -184
deleted file mode 100644
@@ -1,184 +0,0 @@
1-## Jetstream consumer: the ATProto firehose as JSON over a websocket.
2-##
3-## Jetstream is a public, unauthenticated feed of repository commits, so this
4-## is the one part of the port that can be verified end to end without a
5-## consented session.
6-##
7-## Filtering happens server-side via query parameters (wantedCollections,
8-## wantedDids), which matters: the unfiltered firehose is the entire network.
9-## A subscription that silently loses its filters still delivers events, so a
10-## consumer that looks healthy can be reading the wrong stream entirely --
11-## see the note on websock's Uri overload in `connect`.
12-
13-import std/[json, strutils, uri]
14-import pkg/chronicles
15-import pkg/chronos
16-import pkg/websock/websock
17-
18-type
19- JetstreamError* = object of CatchableError
20-
21- EventKind* = enum
22- ekCommit = "commit"
23- ekIdentity = "identity"
24- ekAccount = "account"
25- ekUnknown = "unknown"
26-
27- CommitOperation* = enum
28- coCreate = "create"
29- coUpdate = "update"
30- coDelete = "delete"
31-
32- Commit* = object
33- rev*: string
34- operation*: CommitOperation
35- collection*: string
36- rkey*: string
37- cid*: string
38- record*: JsonNode ## nil for a delete
39-
40- Event* = object
41- did*: string
42- timeUs*: int64 ## microseconds; also the replay cursor
43- kind*: EventKind
44- commit*: Commit ## meaningful only when kind is ekCommit
45- raw*: JsonNode
46-
47- JetstreamConfig* = object
48- url*: string
49- wantedCollections*: seq[string]
50- wantedDids*: seq[string]
51- cursor*: int64 ## 0 to start live
52- userAgent*: string
53-
54- EventHandler* = proc(ev: Event) {.gcsafe, raises: [].}
55-
56-const
57- DefaultJetstream* = "wss://jetstream1.us-east.bsky.network"
58- DefaultUserAgent* = "glean-nim/0.1"
59-
60-proc bytesToString(bs: openArray[byte]): string =
61- ## websock hands back seq[byte]; Jetstream frames are UTF-8 JSON text.
62- result = newString(bs.len)
63- if bs.len > 0:
64- copyMem(result[0].addr, bs[0].unsafeAddr, bs.len)
65-
66-proc parseEventKind(s: string): EventKind =
67- case s
68- of "commit": ekCommit
69- of "identity": ekIdentity
70- of "account": ekAccount
71- else: ekUnknown
72-
73-proc parseOperation(s: string): CommitOperation =
74- case s
75- of "create": coCreate
76- of "update": coUpdate
77- of "delete": coDelete
78- else: coCreate
79-
80-proc parseEvent*(raw: string): Event =
81- ## Decode one Jetstream frame.
82- let j =
83- try: parseJson(raw)
84- except CatchableError:
85- raise newException(JetstreamError, "frame is not valid JSON")
86-
87- result = Event(
88- did: j{"did"}.getStr,
89- timeUs: j{"time_us"}.getBiggestInt.int64,
90- kind: parseEventKind(j{"kind"}.getStr),
91- raw: j,
92- )
93-
94- if result.kind == ekCommit and j.hasKey("commit"):
95- let c = j["commit"]
96- result.commit = Commit(
97- rev: c{"rev"}.getStr,
98- operation: parseOperation(c{"operation"}.getStr),
99- collection: c{"collection"}.getStr,
100- rkey: c{"rkey"}.getStr,
101- cid: c{"cid"}.getStr,
102- record: c{"record"},
103- )
104-
105-proc subscribeUrl*(cfg: JetstreamConfig): string =
106- ## Build the /subscribe URL, including filters.
107- var base = cfg.url.strip(chars = {'/'})
108- if not base.endsWith("/subscribe"):
109- base &= "/subscribe"
110-
111- var params: seq[(string, string)]
112- # Repeated keys, not a comma-joined list: Jetstream reads these as a
113- # multi-valued parameter.
114- for c in cfg.wantedCollections:
115- params.add ("wantedCollections", c)
116- for d in cfg.wantedDids:
117- params.add ("wantedDids", d)
118- if cfg.cursor > 0:
119- params.add ("cursor", $cfg.cursor)
120-
121- if params.len > 0: base & "?" & params.encodeQuery else: base
122-
123-proc connect*(cfg: JetstreamConfig): Future[WSSession] {.async.} =
124- ## Open the websocket.
125- ##
126- ## Deliberately not websock's `Uri` overload: it forwards only `uri.path`
127- ## and drops `uri.query`, which for Jetstream means losing every filter and
128- ## silently subscribing to the entire network firehose instead of the
129- ## handful of collections asked for.
130- let full = subscribeUrl(cfg)
131- let u = parseUri(full)
132- if u.scheme notin ["ws", "wss"]:
133- raise newException(JetstreamError, "not a websocket URL: " & full)
134-
135- let
136- secure = u.scheme == "wss"
137- port = if u.port.len > 0: u.port else: (if secure: "443" else: "80")
138- pathAndQuery = (if u.path.len > 0: u.path else: "/") &
139- (if u.query.len > 0: "?" & u.query else: "")
140-
141- await WebSocket.connect(
142- host = u.hostname & ":" & port,
143- path = pathAndQuery,
144- hostName = u.hostname,
145- secure = secure,
146- )
147-
148-proc readEvents*(ws: WSSession, handler: EventHandler,
149- maxEvents = 0) {.async.} =
150- ## Pump frames until the socket closes, or until `maxEvents` have been
151- ## handled when that is non-zero.
152- var seen = 0
153- while ws.readyState != ReadyState.Closed:
154- let frame = await ws.recvMsg()
155- if frame.len == 0:
156- # A zero-length read means the peer closed; recvMsg does not raise.
157- break
158- var ev: Event
159- try:
160- ev = parseEvent(bytesToString(frame))
161- except JetstreamError:
162- # One malformed frame should not end a long-lived subscription.
163- continue
164- handler(ev)
165- inc seen
166- if maxEvents > 0 and seen >= maxEvents:
167- break
168-
169-proc collectEvents*(cfg: JetstreamConfig, count: int,
170- timeout = 30.seconds): Future[seq[Event]] {.async.} =
171- ## Connect, gather `count` events, and close. For probes and tests.
172- var got: seq[Event]
173- let ws = await connect(cfg)
174- defer:
175- try: await ws.close()
176- except CatchableError: discard
177-
178- proc onEvent(ev: Event) {.gcsafe, raises: [].} =
179- got.add ev
180-
181- let pump = readEvents(ws, onEvent, maxEvents = count)
182- if not await pump.withTimeout(timeout):
183- await pump.cancelAndWait()
184- got
deleted file mode 100644
@@ -1,184 +0,0 @@
1-## Jetstream consumer: the ATProto firehose as JSON over a websocket.
2-##
3-## Jetstream is a public, unauthenticated feed of repository commits, so this
4-## is the one part of the port that can be verified end to end without a
5-## consented session.
6-##
7-## Filtering happens server-side via query parameters (wantedCollections,
8-## wantedDids), which matters: the unfiltered firehose is the entire network.
9-## A subscription that silently loses its filters still delivers events, so a
10-## consumer that looks healthy can be reading the wrong stream entirely --
11-## see the note on websock's Uri overload in `connect`.
12-
13-import std/[json, strutils, uri]
14-import pkg/chronicles
15-import pkg/chronos
16-import pkg/websock/websock
17-
18-type
19- JetstreamError* = object of CatchableError
20-
21- EventKind* = enum
22- ekCommit = "commit"
23- ekIdentity = "identity"
24- ekAccount = "account"
25- ekUnknown = "unknown"
26-
27- CommitOperation* = enum
28- coCreate = "create"
29- coUpdate = "update"
30- coDelete = "delete"
31-
32- Commit* = object
33- rev*: string
34- operation*: CommitOperation
35- collection*: string
36- rkey*: string
37- cid*: string
38- record*: JsonNode ## nil for a delete
39-
40- Event* = object
41- did*: string
42- timeUs*: int64 ## microseconds; also the replay cursor
43- kind*: EventKind
44- commit*: Commit ## meaningful only when kind is ekCommit
45- raw*: JsonNode
46-
47- JetstreamConfig* = object
48- url*: string
49- wantedCollections*: seq[string]
50- wantedDids*: seq[string]
51- cursor*: int64 ## 0 to start live
52- userAgent*: string
53-
54- EventHandler* = proc(ev: Event) {.gcsafe, raises: [].}
55-
56-const
57- DefaultJetstream* = "wss://jetstream1.us-east.bsky.network"
58- DefaultUserAgent* = "glean-nim/0.1"
59-
60-proc bytesToString(bs: openArray[byte]): string =
61- ## websock hands back seq[byte]; Jetstream frames are UTF-8 JSON text.
62- result = newString(bs.len)
63- if bs.len > 0:
64- copyMem(result[0].addr, bs[0].unsafeAddr, bs.len)
65-
66-proc parseEventKind(s: string): EventKind =
67- case s
68- of "commit": ekCommit
69- of "identity": ekIdentity
70- of "account": ekAccount
71- else: ekUnknown
72-
73-proc parseOperation(s: string): CommitOperation =
74- case s
75- of "create": coCreate
76- of "update": coUpdate
77- of "delete": coDelete
78- else: coCreate
79-
80-proc parseEvent*(raw: string): Event =
81- ## Decode one Jetstream frame.
82- let j =
83- try: parseJson(raw)
84- except CatchableError:
85- raise newException(JetstreamError, "frame is not valid JSON")
86-
87- result = Event(
88- did: j{"did"}.getStr,
89- timeUs: j{"time_us"}.getBiggestInt.int64,
90- kind: parseEventKind(j{"kind"}.getStr),
91- raw: j,
92- )
93-
94- if result.kind == ekCommit and j.hasKey("commit"):
95- let c = j["commit"]
96- result.commit = Commit(
97- rev: c{"rev"}.getStr,
98- operation: parseOperation(c{"operation"}.getStr),
99- collection: c{"collection"}.getStr,
100- rkey: c{"rkey"}.getStr,
101- cid: c{"cid"}.getStr,
102- record: c{"record"},
103- )
104-
105-proc subscribeUrl*(cfg: JetstreamConfig): string =
106- ## Build the /subscribe URL, including filters.
107- var base = cfg.url.strip(chars = {'/'})
108- if not base.endsWith("/subscribe"):
109- base &= "/subscribe"
110-
111- var params: seq[(string, string)]
112- # Repeated keys, not a comma-joined list: Jetstream reads these as a
113- # multi-valued parameter.
114- for c in cfg.wantedCollections:
115- params.add ("wantedCollections", c)
116- for d in cfg.wantedDids:
117- params.add ("wantedDids", d)
118- if cfg.cursor > 0:
119- params.add ("cursor", $cfg.cursor)
120-
121- if params.len > 0: base & "?" & params.encodeQuery else: base
122-
123-proc connect*(cfg: JetstreamConfig): Future[WSSession] {.async.} =
124- ## Open the websocket.
125- ##
126- ## Deliberately not websock's `Uri` overload: it forwards only `uri.path`
127- ## and drops `uri.query`, which for Jetstream means losing every filter and
128- ## silently subscribing to the entire network firehose instead of the
129- ## handful of collections asked for.
130- let full = subscribeUrl(cfg)
131- let u = parseUri(full)
132- if u.scheme notin ["ws", "wss"]:
133- raise newException(JetstreamError, "not a websocket URL: " & full)
134-
135- let
136- secure = u.scheme == "wss"
137- port = if u.port.len > 0: u.port else: (if secure: "443" else: "80")
138- pathAndQuery = (if u.path.len > 0: u.path else: "/") &
139- (if u.query.len > 0: "?" & u.query else: "")
140-
141- await WebSocket.connect(
142- host = u.hostname & ":" & port,
143- path = pathAndQuery,
144- hostName = u.hostname,
145- secure = secure,
146- )
147-
148-proc readEvents*(ws: WSSession, handler: EventHandler,
149- maxEvents = 0) {.async.} =
150- ## Pump frames until the socket closes, or until `maxEvents` have been
151- ## handled when that is non-zero.
152- var seen = 0
153- while ws.readyState != ReadyState.Closed:
154- let frame = await ws.recvMsg()
155- if frame.len == 0:
156- # A zero-length read means the peer closed; recvMsg does not raise.
157- break
158- var ev: Event
159- try:
160- ev = parseEvent(bytesToString(frame))
161- except JetstreamError:
162- # One malformed frame should not end a long-lived subscription.
163- continue
164- handler(ev)
165- inc seen
166- if maxEvents > 0 and seen >= maxEvents:
167- break
168-
169-proc collectEvents*(cfg: JetstreamConfig, count: int,
170- timeout = 30.seconds): Future[seq[Event]] {.async.} =
171- ## Connect, gather `count` events, and close. For probes and tests.
172- var got: seq[Event]
173- let ws = await connect(cfg)
174- defer:
175- try: await ws.close()
176- except CatchableError: discard
177-
178- proc onEvent(ev: Event) {.gcsafe, raises: [].} =
179- got.add ev
180-
181- let pump = readEvents(ws, onEvent, maxEvents = count)
182- if not await pump.withTimeout(timeout):
183- await pump.cancelAndWait()
184- got
deleted spike/nim/atproto/jetstream_run.nim +0 -167
deleted file mode 100644
@@ -1,167 +0,0 @@
1-## Keeping a Jetstream subscription alive.
2-##
3-## `jetstream.nim` reads a stream; this survives it ending. Three concerns,
4-## all of which are about not losing events rather than about uptime:
5-##
6-## 1. **Resume from a cursor, rewound.** Cursors are persisted in batches, so
7-## the last one written is always slightly behind the last one *seen*. On
8-## reconnect the consumer rewinds past that gap and accepts re-delivery --
9-## duplicate events are cheap when handling is idempotent, and a hole in
10-## the record is not recoverable.
11-##
12-## 2. **Rotate connections on a timer.** Server-side filters are fixed for the
13-## life of a connection, so a subscription that never reconnects never
14-## learns about newly-known DIDs. Glean's Go consumer rotates every 15
15-## minutes for exactly this reason.
16-##
17-## 3. **Back off, with jitter.** The Go consumer waits a flat 5s between
18-## attempts, which is fine against one healthy server and unkind to one
19-## that is down. This grows the delay to a cap, and jitters it so a fleet
20-## of consumers does not reconnect in lockstep.
21-
22-import std/[math, random]
23-import pkg/chronicles
24-import pkg/chronos
25-import pkg/websock/websock
26-import ./jetstream
27-
28-type
29- BackoffConfig* = object
30- initial*: Duration
31- maxDelay*: Duration
32- factor*: float
33- jitter*: float ## 0.0-1.0 fraction of the delay to randomise
34-
35- SupervisorConfig* = object
36- jetstream*: JetstreamConfig
37- backoff*: BackoffConfig
38- rewind*: Duration
39- ## How far behind the persisted cursor to resume. Covers the window
40- ## between the last cursor flush and the disconnection.
41- rotateEvery*: Duration
42- ## Force a reconnect this often so filter changes take effect. Zero
43- ## disables rotation.
44-
45- CursorStore* = ref object of RootObj ## Override to persist across restarts.
46- cursor*: int64
47-
48- Supervisor* = ref object
49- cfg*: SupervisorConfig
50- store*: CursorStore
51- handler*: EventHandler
52- running*: bool
53- reconnects*: int ## observable for tests and metrics
54-
55-# Both are `raises: []` on purpose. Saving a cursor happens on the event path,
56-# and a persistent store that throws -- a locked database, a full disk -- must
57-# not tear down the subscription. An implementation that can fail is expected
58-# to swallow and log, and to lose at most the rewind window.
59-method load*(s: CursorStore): int64 {.base, gcsafe, raises: [].} =
60- s.cursor
61-method save*(s: CursorStore, cursor: int64) {.base, gcsafe, raises: [].} =
62- s.cursor = cursor
63-
64-const
65- DefaultBackoff* = BackoffConfig(
66- initial: 1.seconds, maxDelay: 60.seconds, factor: 2.0, jitter: 0.2)
67- DefaultRewind* = 5.seconds
68- DefaultRotateEvery* = 15.minutes
69-
70-proc newSupervisorConfig*(js: JetstreamConfig): SupervisorConfig =
71- SupervisorConfig(
72- jetstream: js,
73- backoff: DefaultBackoff,
74- rewind: DefaultRewind,
75- rotateEvery: DefaultRotateEvery,
76- )
77-
78-proc delayFor*(b: BackoffConfig, attempt: int, rand01: float = -1.0): Duration =
79- ## Delay before attempt `attempt` (0 = first retry).
80- ##
81- ## `rand01` is injectable so the jitter can be tested deterministically;
82- ## left at -1 it draws from the global RNG.
83- ##
84- ## The first retry is jittered like every other. Short-circuiting it to the
85- ## bare initial delay would leave the one moment a herd is most likely --
86- ## every consumer dropped at once by the same server restart -- as the one
87- ## moment they all wake together.
88- let
89- exp = max(attempt, 0).float
90- grown = b.initial.milliseconds.float * pow(b.factor, exp)
91- capped = min(grown, b.maxDelay.milliseconds.float)
92- if b.jitter <= 0:
93- return milliseconds(capped.int64)
94- # Jitter downward only, so the cap stays a real ceiling.
95- let r = if rand01 >= 0: rand01 else: rand(1.0)
96- let jittered = capped * (1.0 - b.jitter * r)
97- milliseconds(max(jittered.int64, 1'i64))
98-
99-proc rewoundCursor*(store: CursorStore, rewind: Duration): int64 =
100- ## The persisted cursor, moved back to cover the unflushed window. Zero
101- ## (start live) stays zero -- there is nothing to rewind to.
102- let c = store.load()
103- if c <= 0: return 0
104- max(c - rewind.microseconds, 0'i64)
105-
106-proc newSupervisor*(cfg: SupervisorConfig, handler: EventHandler,
107- store: CursorStore = nil): Supervisor =
108- Supervisor(
109- cfg: cfg,
110- store: if store != nil: store else: CursorStore(),
111- handler: handler,
112- )
113-
114-proc runOnce(sup: Supervisor) {.async.} =
115- ## One connection, read until it ends or the rotation timer fires.
116- var cfg = sup.cfg.jetstream
117- cfg.cursor = rewoundCursor(sup.store, sup.cfg.rewind)
118-
119- let ws = await connect(cfg)
120- defer:
121- try: await ws.close()
122- except CatchableError: discard
123-
124- # Record every event's cursor, so a drop resumes near where it stopped.
125- let store = sup.store
126- let inner = sup.handler
127- proc onEvent(ev: Event) {.gcsafe, raises: [].} =
128- if ev.timeUs > 0:
129- store.save(ev.timeUs)
130- inner(ev)
131-
132- let pump = readEvents(ws, onEvent)
133- if sup.cfg.rotateEvery > ZeroDuration:
134- if not await pump.withTimeout(sup.cfg.rotateEvery):
135- # Rotation is a normal end to a connection, not a failure.
136- await pump.cancelAndWait()
137- else:
138- await pump
139-
140-proc run*(sup: Supervisor) {.async.} =
141- ## Reconnect forever. Stop by setting `running` to false.
142- sup.running = true
143- var attempt = 0
144- while sup.running:
145- try:
146- await sup.runOnce()
147- # A clean end (rotation, or the server closing) is not a failure, so
148- # the next attempt starts from the bottom of the backoff curve.
149- attempt = 0
150- except CancelledError as e:
151- raise e
152- except CatchableError as e:
153- inc sup.reconnects
154- warn "jetstream connection lost", err = e.msg, attempt = attempt
155- let delay = delayFor(sup.cfg.backoff, attempt)
156- inc attempt
157- await sleepAsync(delay)
158- continue
159-
160- if not sup.running:
161- break
162- # Even a clean rotation waits a moment, so a server that closes
163- # immediately cannot spin this loop.
164- await sleepAsync(sup.cfg.backoff.initial)
165-
166-proc stop*(sup: Supervisor) =
167- sup.running = false
deleted file mode 100644
@@ -1,167 +0,0 @@
1-## Keeping a Jetstream subscription alive.
2-##
3-## `jetstream.nim` reads a stream; this survives it ending. Three concerns,
4-## all of which are about not losing events rather than about uptime:
5-##
6-## 1. **Resume from a cursor, rewound.** Cursors are persisted in batches, so
7-## the last one written is always slightly behind the last one *seen*. On
8-## reconnect the consumer rewinds past that gap and accepts re-delivery --
9-## duplicate events are cheap when handling is idempotent, and a hole in
10-## the record is not recoverable.
11-##
12-## 2. **Rotate connections on a timer.** Server-side filters are fixed for the
13-## life of a connection, so a subscription that never reconnects never
14-## learns about newly-known DIDs. Glean's Go consumer rotates every 15
15-## minutes for exactly this reason.
16-##
17-## 3. **Back off, with jitter.** The Go consumer waits a flat 5s between
18-## attempts, which is fine against one healthy server and unkind to one
19-## that is down. This grows the delay to a cap, and jitters it so a fleet
20-## of consumers does not reconnect in lockstep.
21-
22-import std/[math, random]
23-import pkg/chronicles
24-import pkg/chronos
25-import pkg/websock/websock
26-import ./jetstream
27-
28-type
29- BackoffConfig* = object
30- initial*: Duration
31- maxDelay*: Duration
32- factor*: float
33- jitter*: float ## 0.0-1.0 fraction of the delay to randomise
34-
35- SupervisorConfig* = object
36- jetstream*: JetstreamConfig
37- backoff*: BackoffConfig
38- rewind*: Duration
39- ## How far behind the persisted cursor to resume. Covers the window
40- ## between the last cursor flush and the disconnection.
41- rotateEvery*: Duration
42- ## Force a reconnect this often so filter changes take effect. Zero
43- ## disables rotation.
44-
45- CursorStore* = ref object of RootObj ## Override to persist across restarts.
46- cursor*: int64
47-
48- Supervisor* = ref object
49- cfg*: SupervisorConfig
50- store*: CursorStore
51- handler*: EventHandler
52- running*: bool
53- reconnects*: int ## observable for tests and metrics
54-
55-# Both are `raises: []` on purpose. Saving a cursor happens on the event path,
56-# and a persistent store that throws -- a locked database, a full disk -- must
57-# not tear down the subscription. An implementation that can fail is expected
58-# to swallow and log, and to lose at most the rewind window.
59-method load*(s: CursorStore): int64 {.base, gcsafe, raises: [].} =
60- s.cursor
61-method save*(s: CursorStore, cursor: int64) {.base, gcsafe, raises: [].} =
62- s.cursor = cursor
63-
64-const
65- DefaultBackoff* = BackoffConfig(
66- initial: 1.seconds, maxDelay: 60.seconds, factor: 2.0, jitter: 0.2)
67- DefaultRewind* = 5.seconds
68- DefaultRotateEvery* = 15.minutes
69-
70-proc newSupervisorConfig*(js: JetstreamConfig): SupervisorConfig =
71- SupervisorConfig(
72- jetstream: js,
73- backoff: DefaultBackoff,
74- rewind: DefaultRewind,
75- rotateEvery: DefaultRotateEvery,
76- )
77-
78-proc delayFor*(b: BackoffConfig, attempt: int, rand01: float = -1.0): Duration =
79- ## Delay before attempt `attempt` (0 = first retry).
80- ##
81- ## `rand01` is injectable so the jitter can be tested deterministically;
82- ## left at -1 it draws from the global RNG.
83- ##
84- ## The first retry is jittered like every other. Short-circuiting it to the
85- ## bare initial delay would leave the one moment a herd is most likely --
86- ## every consumer dropped at once by the same server restart -- as the one
87- ## moment they all wake together.
88- let
89- exp = max(attempt, 0).float
90- grown = b.initial.milliseconds.float * pow(b.factor, exp)
91- capped = min(grown, b.maxDelay.milliseconds.float)
92- if b.jitter <= 0:
93- return milliseconds(capped.int64)
94- # Jitter downward only, so the cap stays a real ceiling.
95- let r = if rand01 >= 0: rand01 else: rand(1.0)
96- let jittered = capped * (1.0 - b.jitter * r)
97- milliseconds(max(jittered.int64, 1'i64))
98-
99-proc rewoundCursor*(store: CursorStore, rewind: Duration): int64 =
100- ## The persisted cursor, moved back to cover the unflushed window. Zero
101- ## (start live) stays zero -- there is nothing to rewind to.
102- let c = store.load()
103- if c <= 0: return 0
104- max(c - rewind.microseconds, 0'i64)
105-
106-proc newSupervisor*(cfg: SupervisorConfig, handler: EventHandler,
107- store: CursorStore = nil): Supervisor =
108- Supervisor(
109- cfg: cfg,
110- store: if store != nil: store else: CursorStore(),
111- handler: handler,
112- )
113-
114-proc runOnce(sup: Supervisor) {.async.} =
115- ## One connection, read until it ends or the rotation timer fires.
116- var cfg = sup.cfg.jetstream
117- cfg.cursor = rewoundCursor(sup.store, sup.cfg.rewind)
118-
119- let ws = await connect(cfg)
120- defer:
121- try: await ws.close()
122- except CatchableError: discard
123-
124- # Record every event's cursor, so a drop resumes near where it stopped.
125- let store = sup.store
126- let inner = sup.handler
127- proc onEvent(ev: Event) {.gcsafe, raises: [].} =
128- if ev.timeUs > 0:
129- store.save(ev.timeUs)
130- inner(ev)
131-
132- let pump = readEvents(ws, onEvent)
133- if sup.cfg.rotateEvery > ZeroDuration:
134- if not await pump.withTimeout(sup.cfg.rotateEvery):
135- # Rotation is a normal end to a connection, not a failure.
136- await pump.cancelAndWait()
137- else:
138- await pump
139-
140-proc run*(sup: Supervisor) {.async.} =
141- ## Reconnect forever. Stop by setting `running` to false.
142- sup.running = true
143- var attempt = 0
144- while sup.running:
145- try:
146- await sup.runOnce()
147- # A clean end (rotation, or the server closing) is not a failure, so
148- # the next attempt starts from the bottom of the backoff curve.
149- attempt = 0
150- except CancelledError as e:
151- raise e
152- except CatchableError as e:
153- inc sup.reconnects
154- warn "jetstream connection lost", err = e.msg, attempt = attempt
155- let delay = delayFor(sup.cfg.backoff, attempt)
156- inc attempt
157- await sleepAsync(delay)
158- continue
159-
160- if not sup.running:
161- break
162- # Even a clean rotation waits a moment, so a server that closes
163- # immediately cannot spin this loop.
164- await sleepAsync(sup.cfg.backoff.initial)
165-
166-proc stop*(sup: Supervisor) =
167- sup.running = false
deleted spike/nim/atproto/oauth.nim +0 -329
deleted file mode 100644
@@ -1,329 +0,0 @@
1-## ATProto OAuth client: metadata discovery, PAR, authorization URL, and the
2-## token exchange, ported from the shape of bluesky-social/indigo's
3-## atproto/auth/oauth package.
4-##
5-## Two details drive most of the design:
6-##
7-## 1. Every request to the auth server carries a DPoP proof bound to a
8-## per-session P-256 key, and the server will reject the first one with
9-## `use_dpop_nonce` plus a DPoP-Nonce header. That is normal, not an error:
10-## the request must be retried once with the supplied nonce. Note the
11-## status is 400 on the auth server (not the 401 a resource server uses).
12-##
13-## 2. PKCE is mandatory. The verifier is kept alongside the request so the
14-## callback can present it.
15-
16-import std/[httpclient, json, strutils, strformat, sysrand, uri]
17-import nimcrypto/[sha2, hash]
18-import bearssl/rand
19-import ./dpop
20-import ./identity
21-
22-type
23- OAuthError* = object of CatchableError
24-
25- ClientConfig* = object
26- clientId*: string
27- callbackUrl*: string
28- scopes*: seq[string]
29-
30- AuthServerMetadata* = object
31- issuer*: string
32- authorizationEndpoint*: string
33- tokenEndpoint*: string
34- parEndpoint*: string
35- revocationEndpoint*: string
36-
37- AuthRequestData* = object ## Persisted between the redirect and the callback.
38- state*: string
39- authServerUrl*: string
40- scopes*: seq[string]
41- pkceVerifier*: string
42- requestUri*: string
43- tokenEndpoint*: string
44- revocationEndpoint*: string
45- dpopNonce*: string
46- dpopKey*: P256Key
47- accountDid*: string
48-
49- TokenResponse* = object
50- accessToken*: string
51- refreshToken*: string
52- tokenType*: string
53- expiresIn*: int
54- scope*: string
55- sub*: string
56-
57-const HttpTimeoutMs = 15_000
58-
59-proc newClient(): HttpClient =
60- newHttpClient(userAgent = UserAgent, timeout = HttpTimeoutMs)
61-
62-proc scopeStr*(scopes: seq[string]): string = scopes.join(" ")
63-
64-proc newPublicConfig*(clientId, callbackUrl: string,
65- scopes = @["atproto", "transition:generic"]): ClientConfig =
66- ClientConfig(clientId: clientId, callbackUrl: callbackUrl, scopes: scopes)
67-
68-proc newLocalhostConfig*(callbackUrl: string,
69- scopes = @["atproto", "transition:generic"]): ClientConfig =
70- ## Localhost development clients have no hosted metadata document; the
71- ## client_id encodes the redirect and scopes instead.
72- var params = @[("redirect_uri", callbackUrl), ("scope", scopeStr(scopes))]
73- ClientConfig(
74- clientId: "http://localhost" & "?" & params.encodeQuery,
75- callbackUrl: callbackUrl,
76- scopes: scopes,
77- )
78-
79-proc clientMetadata*(cfg: ClientConfig): JsonNode =
80- ## Served at the client_id URL for a public (non-confidential) client.
81- %*{
82- "client_id": cfg.clientId,
83- "application_type": "web",
84- "client_name": "glean",
85- "grant_types": ["authorization_code", "refresh_token"],
86- "scope": scopeStr(cfg.scopes),
87- "response_types": ["code"],
88- "redirect_uris": [cfg.callbackUrl],
89- "token_endpoint_auth_method": "none",
90- "dpop_bound_access_tokens": true,
91- }
92-
93-proc secureRandomB64*(n: int): string =
94- var buf = newSeq[byte](n)
95- doAssert urandom(buf), "system randomness unavailable"
96- b64u(buf)
97-
98-proc s256Challenge*(verifier: string): string =
99- b64u(sha256.digest(verifier).data)
100-
101-# --- discovery -------------------------------------------------------------
102-
103-proc resolveAuthServerUrl*(pdsUrl: string): string =
104- ## The PDS advertises which authorization server governs it.
105- let url = &"https://{hostOf(pdsUrl)}/.well-known/oauth-protected-resource"
106- let client = newClient()
107- defer: client.close()
108- let res = client.get(url)
109- if res.code.int != 200:
110- raise newException(OAuthError,
111- &"protected resource metadata failed (HTTP {res.code.int}) at {url}")
112- let body = parseJson(res.body)
113- let servers = body{"authorization_servers"}
114- if servers == nil or servers.kind != JArray or servers.len == 0:
115- raise newException(OAuthError, "PDS lists no authorization servers: " & pdsUrl)
116- servers[0].getStr.strip(chars = {'/'})
117-
118-proc resolveAuthServerMetadata*(authServerUrl: string): AuthServerMetadata =
119- let url = &"https://{hostOf(authServerUrl)}/.well-known/oauth-authorization-server"
120- let client = newClient()
121- defer: client.close()
122- let res = client.get(url)
123- if res.code.int != 200:
124- raise newException(OAuthError,
125- &"auth server metadata failed (HTTP {res.code.int}) at {url}")
126- let b = parseJson(res.body)
127-
128- result = AuthServerMetadata(
129- issuer: b{"issuer"}.getStr,
130- authorizationEndpoint: b{"authorization_endpoint"}.getStr,
131- tokenEndpoint: b{"token_endpoint"}.getStr,
132- parEndpoint: b{"pushed_authorization_request_endpoint"}.getStr,
133- revocationEndpoint: b{"revocation_endpoint"}.getStr,
134- )
135-
136- # An auth server missing any of these cannot complete an ATProto flow, and
137- # failing here beats a confusing error three requests later.
138- for (name, value) in {
139- "issuer": result.issuer,
140- "authorization_endpoint": result.authorizationEndpoint,
141- "token_endpoint": result.tokenEndpoint,
142- "pushed_authorization_request_endpoint": result.parEndpoint,
143- }:
144- if value.len == 0:
145- raise newException(OAuthError, "auth server metadata is missing " & name)
146-
147- if result.issuer.strip(chars = {'/'}) != authServerUrl.strip(chars = {'/'}):
148- raise newException(OAuthError,
149- &"auth server issuer mismatch: metadata says {result.issuer}, fetched from {authServerUrl}")
150-
151-# --- requests --------------------------------------------------------------
152-
153-proc errorReason(body: string): string =
154- try:
155- let j = parseJson(body)
156- result = j{"error"}.getStr
157- let desc = j{"error_description"}.getStr
158- if desc.len > 0:
159- result &= ": " & desc
160- except CatchableError:
161- result = "unknown"
162- if result.len == 0:
163- result = "unknown"
164-
165-proc errorCode(body: string): string =
166- try: parseJson(body){"error"}.getStr
167- except CatchableError: ""
168-
169-proc postWithDpop(url, body: string, key: P256Key, rng: var HmacDrbgContext,
170- nonce: var string): Response =
171- ## POST with a DPoP proof, retrying once when the server demands a nonce.
172- ##
173- ## The first request of a session always lacks a nonce, so the retry is the
174- ## expected path rather than an error case. The auth server signals this
175- ## with 400 + error=use_dpop_nonce and supplies DPoP-Nonce.
176- let client = newClient()
177- defer: client.close()
178-
179- for attempt in 0 .. 1:
180- let proof = dpopProof(key, rng, "POST", url, nonce = nonce)
181- client.headers = newHttpHeaders({
182- "Content-Type": "application/x-www-form-urlencoded",
183- "DPoP": proof,
184- })
185- result = client.post(url, body = body)
186-
187- let supplied = result.headers.getOrDefault("DPoP-Nonce").string
188- if supplied.len > 0:
189- nonce = supplied
190-
191- if attempt == 0 and result.code.int == 400 and supplied.len > 0:
192- if errorCode(result.body) == "use_dpop_nonce":
193- continue
194- return result
195-
196-proc sendAuthRequest*(cfg: ClientConfig, meta: AuthServerMetadata,
197- rng: var HmacDrbgContext, loginHint = ""): AuthRequestData =
198- ## Pushed Authorization Request: hand the parameters to the auth server up
199- ## front and get back a short-lived request_uri to redirect the user with.
200- let
201- state = secureRandomB64(16)
202- verifier = secureRandomB64(48)
203- key = generateKey(rng)
204-
205- var params = @[
206- ("client_id", cfg.clientId),
207- ("state", state),
208- ("redirect_uri", cfg.callbackUrl),
209- ("scope", scopeStr(cfg.scopes)),
210- ("response_type", "code"),
211- ("code_challenge", s256Challenge(verifier)),
212- ("code_challenge_method", "S256"),
213- ]
214- if loginHint.len > 0:
215- params.add ("login_hint", loginHint)
216-
217- var nonce = ""
218- let res = postWithDpop(meta.parEndpoint, params.encodeQuery, key, rng, nonce)
219- if res.code.int notin [200, 201]:
220- raise newException(OAuthError,
221- &"PAR failed (HTTP {res.code.int}): {errorReason(res.body)}")
222-
223- let requestUri = parseJson(res.body){"request_uri"}.getStr
224- if requestUri.len == 0:
225- raise newException(OAuthError, "PAR response has no request_uri")
226-
227- AuthRequestData(
228- state: state,
229- authServerUrl: meta.issuer,
230- scopes: cfg.scopes,
231- pkceVerifier: verifier,
232- requestUri: requestUri,
233- tokenEndpoint: meta.tokenEndpoint,
234- revocationEndpoint: meta.revocationEndpoint,
235- dpopNonce: nonce,
236- dpopKey: key,
237- )
238-
239-proc authorizeUrl*(cfg: ClientConfig, meta: AuthServerMetadata,
240- info: AuthRequestData): string =
241- let params = @[("client_id", cfg.clientId), ("request_uri", info.requestUri)]
242- meta.authorizationEndpoint & "?" & params.encodeQuery
243-
244-proc startAuthFlow*(cfg: ClientConfig, identifier: string,
245- rng: var HmacDrbgContext,
246- plcUrl = DefaultPlcUrl): (string, AuthRequestData) =
247- ## Resolve an account (or take an auth server URL directly), push the
248- ## request, and return the URL to send the user to.
249- var
250- authServerUrl: string
251- accountDid: string
252- loginHint: string
253-
254- if identifier.startsWith("https://"):
255- authServerUrl = identifier.strip(chars = {'/'})
256- else:
257- let ident = lookup(identifier, plcUrl)
258- accountDid = ident.did
259- loginHint = identifier
260- authServerUrl = resolveAuthServerUrl(ident.pdsEndpoint)
261-
262- let meta = resolveAuthServerMetadata(authServerUrl)
263- var info = sendAuthRequest(cfg, meta, rng, loginHint)
264- info.accountDid = accountDid
265- (authorizeUrl(cfg, meta, info), info)
266-
267-proc exchangeCode*(cfg: ClientConfig, info: var AuthRequestData,
268- code: string, rng: var HmacDrbgContext): TokenResponse =
269- ## Trade the authorization code for tokens, proving possession of both the
270- ## PKCE verifier and the session's DPoP key.
271- let params = @[
272- ("client_id", cfg.clientId),
273- ("redirect_uri", cfg.callbackUrl),
274- ("grant_type", "authorization_code"),
275- ("code", code),
276- ("code_verifier", info.pkceVerifier),
277- ]
278- var nonce = info.dpopNonce
279- let res = postWithDpop(info.tokenEndpoint, params.encodeQuery,
280- info.dpopKey, rng, nonce)
281- info.dpopNonce = nonce
282- if res.code.int != 200:
283- raise newException(OAuthError,
284- &"token request failed (HTTP {res.code.int}): {errorReason(res.body)}")
285-
286- let b = parseJson(res.body)
287- result = TokenResponse(
288- accessToken: b{"access_token"}.getStr,
289- refreshToken: b{"refresh_token"}.getStr,
290- tokenType: b{"token_type"}.getStr,
291- expiresIn: b{"expires_in"}.getInt,
292- scope: b{"scope"}.getStr,
293- sub: b{"sub"}.getStr,
294- )
295-
296- # DPoP-bound tokens are useless as bearer tokens; a server answering
297- # token_type=Bearer means the binding silently did not happen.
298- if not result.tokenType.toLowerAscii.startsWith("dpop"):
299- raise newException(OAuthError,
300- "expected a DPoP-bound token, got token_type=" & result.tokenType)
301- if result.sub.len == 0:
302- raise newException(OAuthError, "token response has no subject DID")
303- if info.accountDid.len > 0 and result.sub != info.accountDid:
304- raise newException(OAuthError,
305- &"token subject {result.sub} does not match requested account {info.accountDid}")
306-
307-proc refresh*(cfg: ClientConfig, info: var AuthRequestData,
308- refreshToken: string, rng: var HmacDrbgContext): TokenResponse =
309- let params = @[
310- ("client_id", cfg.clientId),
311- ("grant_type", "refresh_token"),
312- ("refresh_token", refreshToken),
313- ]
314- var nonce = info.dpopNonce
315- let res = postWithDpop(info.tokenEndpoint, params.encodeQuery,
316- info.dpopKey, rng, nonce)
317- info.dpopNonce = nonce
318- if res.code.int != 200:
319- raise newException(OAuthError,
320- &"refresh failed (HTTP {res.code.int}): {errorReason(res.body)}")
321- let b = parseJson(res.body)
322- TokenResponse(
323- accessToken: b{"access_token"}.getStr,
324- refreshToken: b{"refresh_token"}.getStr,
325- tokenType: b{"token_type"}.getStr,
326- expiresIn: b{"expires_in"}.getInt,
327- scope: b{"scope"}.getStr,
328- sub: b{"sub"}.getStr,
329- )
deleted file mode 100644
@@ -1,329 +0,0 @@
1-## ATProto OAuth client: metadata discovery, PAR, authorization URL, and the
2-## token exchange, ported from the shape of bluesky-social/indigo's
3-## atproto/auth/oauth package.
4-##
5-## Two details drive most of the design:
6-##
7-## 1. Every request to the auth server carries a DPoP proof bound to a
8-## per-session P-256 key, and the server will reject the first one with
9-## `use_dpop_nonce` plus a DPoP-Nonce header. That is normal, not an error:
10-## the request must be retried once with the supplied nonce. Note the
11-## status is 400 on the auth server (not the 401 a resource server uses).
12-##
13-## 2. PKCE is mandatory. The verifier is kept alongside the request so the
14-## callback can present it.
15-
16-import std/[httpclient, json, strutils, strformat, sysrand, uri]
17-import nimcrypto/[sha2, hash]
18-import bearssl/rand
19-import ./dpop
20-import ./identity
21-
22-type
23- OAuthError* = object of CatchableError
24-
25- ClientConfig* = object
26- clientId*: string
27- callbackUrl*: string
28- scopes*: seq[string]
29-
30- AuthServerMetadata* = object
31- issuer*: string
32- authorizationEndpoint*: string
33- tokenEndpoint*: string
34- parEndpoint*: string
35- revocationEndpoint*: string
36-
37- AuthRequestData* = object ## Persisted between the redirect and the callback.
38- state*: string
39- authServerUrl*: string
40- scopes*: seq[string]
41- pkceVerifier*: string
42- requestUri*: string
43- tokenEndpoint*: string
44- revocationEndpoint*: string
45- dpopNonce*: string
46- dpopKey*: P256Key
47- accountDid*: string
48-
49- TokenResponse* = object
50- accessToken*: string
51- refreshToken*: string
52- tokenType*: string
53- expiresIn*: int
54- scope*: string
55- sub*: string
56-
57-const HttpTimeoutMs = 15_000
58-
59-proc newClient(): HttpClient =
60- newHttpClient(userAgent = UserAgent, timeout = HttpTimeoutMs)
61-
62-proc scopeStr*(scopes: seq[string]): string = scopes.join(" ")
63-
64-proc newPublicConfig*(clientId, callbackUrl: string,
65- scopes = @["atproto", "transition:generic"]): ClientConfig =
66- ClientConfig(clientId: clientId, callbackUrl: callbackUrl, scopes: scopes)
67-
68-proc newLocalhostConfig*(callbackUrl: string,
69- scopes = @["atproto", "transition:generic"]): ClientConfig =
70- ## Localhost development clients have no hosted metadata document; the
71- ## client_id encodes the redirect and scopes instead.
72- var params = @[("redirect_uri", callbackUrl), ("scope", scopeStr(scopes))]
73- ClientConfig(
74- clientId: "http://localhost" & "?" & params.encodeQuery,
75- callbackUrl: callbackUrl,
76- scopes: scopes,
77- )
78-
79-proc clientMetadata*(cfg: ClientConfig): JsonNode =
80- ## Served at the client_id URL for a public (non-confidential) client.
81- %*{
82- "client_id": cfg.clientId,
83- "application_type": "web",
84- "client_name": "glean",
85- "grant_types": ["authorization_code", "refresh_token"],
86- "scope": scopeStr(cfg.scopes),
87- "response_types": ["code"],
88- "redirect_uris": [cfg.callbackUrl],
89- "token_endpoint_auth_method": "none",
90- "dpop_bound_access_tokens": true,
91- }
92-
93-proc secureRandomB64*(n: int): string =
94- var buf = newSeq[byte](n)
95- doAssert urandom(buf), "system randomness unavailable"
96- b64u(buf)
97-
98-proc s256Challenge*(verifier: string): string =
99- b64u(sha256.digest(verifier).data)
100-
101-# --- discovery -------------------------------------------------------------
102-
103-proc resolveAuthServerUrl*(pdsUrl: string): string =
104- ## The PDS advertises which authorization server governs it.
105- let url = &"https://{hostOf(pdsUrl)}/.well-known/oauth-protected-resource"
106- let client = newClient()
107- defer: client.close()
108- let res = client.get(url)
109- if res.code.int != 200:
110- raise newException(OAuthError,
111- &"protected resource metadata failed (HTTP {res.code.int}) at {url}")
112- let body = parseJson(res.body)
113- let servers = body{"authorization_servers"}
114- if servers == nil or servers.kind != JArray or servers.len == 0:
115- raise newException(OAuthError, "PDS lists no authorization servers: " & pdsUrl)
116- servers[0].getStr.strip(chars = {'/'})
117-
118-proc resolveAuthServerMetadata*(authServerUrl: string): AuthServerMetadata =
119- let url = &"https://{hostOf(authServerUrl)}/.well-known/oauth-authorization-server"
120- let client = newClient()
121- defer: client.close()
122- let res = client.get(url)
123- if res.code.int != 200:
124- raise newException(OAuthError,
125- &"auth server metadata failed (HTTP {res.code.int}) at {url}")
126- let b = parseJson(res.body)
127-
128- result = AuthServerMetadata(
129- issuer: b{"issuer"}.getStr,
130- authorizationEndpoint: b{"authorization_endpoint"}.getStr,
131- tokenEndpoint: b{"token_endpoint"}.getStr,
132- parEndpoint: b{"pushed_authorization_request_endpoint"}.getStr,
133- revocationEndpoint: b{"revocation_endpoint"}.getStr,
134- )
135-
136- # An auth server missing any of these cannot complete an ATProto flow, and
137- # failing here beats a confusing error three requests later.
138- for (name, value) in {
139- "issuer": result.issuer,
140- "authorization_endpoint": result.authorizationEndpoint,
141- "token_endpoint": result.tokenEndpoint,
142- "pushed_authorization_request_endpoint": result.parEndpoint,
143- }:
144- if value.len == 0:
145- raise newException(OAuthError, "auth server metadata is missing " & name)
146-
147- if result.issuer.strip(chars = {'/'}) != authServerUrl.strip(chars = {'/'}):
148- raise newException(OAuthError,
149- &"auth server issuer mismatch: metadata says {result.issuer}, fetched from {authServerUrl}")
150-
151-# --- requests --------------------------------------------------------------
152-
153-proc errorReason(body: string): string =
154- try:
155- let j = parseJson(body)
156- result = j{"error"}.getStr
157- let desc = j{"error_description"}.getStr
158- if desc.len > 0:
159- result &= ": " & desc
160- except CatchableError:
161- result = "unknown"
162- if result.len == 0:
163- result = "unknown"
164-
165-proc errorCode(body: string): string =
166- try: parseJson(body){"error"}.getStr
167- except CatchableError: ""
168-
169-proc postWithDpop(url, body: string, key: P256Key, rng: var HmacDrbgContext,
170- nonce: var string): Response =
171- ## POST with a DPoP proof, retrying once when the server demands a nonce.
172- ##
173- ## The first request of a session always lacks a nonce, so the retry is the
174- ## expected path rather than an error case. The auth server signals this
175- ## with 400 + error=use_dpop_nonce and supplies DPoP-Nonce.
176- let client = newClient()
177- defer: client.close()
178-
179- for attempt in 0 .. 1:
180- let proof = dpopProof(key, rng, "POST", url, nonce = nonce)
181- client.headers = newHttpHeaders({
182- "Content-Type": "application/x-www-form-urlencoded",
183- "DPoP": proof,
184- })
185- result = client.post(url, body = body)
186-
187- let supplied = result.headers.getOrDefault("DPoP-Nonce").string
188- if supplied.len > 0:
189- nonce = supplied
190-
191- if attempt == 0 and result.code.int == 400 and supplied.len > 0:
192- if errorCode(result.body) == "use_dpop_nonce":
193- continue
194- return result
195-
196-proc sendAuthRequest*(cfg: ClientConfig, meta: AuthServerMetadata,
197- rng: var HmacDrbgContext, loginHint = ""): AuthRequestData =
198- ## Pushed Authorization Request: hand the parameters to the auth server up
199- ## front and get back a short-lived request_uri to redirect the user with.
200- let
201- state = secureRandomB64(16)
202- verifier = secureRandomB64(48)
203- key = generateKey(rng)
204-
205- var params = @[
206- ("client_id", cfg.clientId),
207- ("state", state),
208- ("redirect_uri", cfg.callbackUrl),
209- ("scope", scopeStr(cfg.scopes)),
210- ("response_type", "code"),
211- ("code_challenge", s256Challenge(verifier)),
212- ("code_challenge_method", "S256"),
213- ]
214- if loginHint.len > 0:
215- params.add ("login_hint", loginHint)
216-
217- var nonce = ""
218- let res = postWithDpop(meta.parEndpoint, params.encodeQuery, key, rng, nonce)
219- if res.code.int notin [200, 201]:
220- raise newException(OAuthError,
221- &"PAR failed (HTTP {res.code.int}): {errorReason(res.body)}")
222-
223- let requestUri = parseJson(res.body){"request_uri"}.getStr
224- if requestUri.len == 0:
225- raise newException(OAuthError, "PAR response has no request_uri")
226-
227- AuthRequestData(
228- state: state,
229- authServerUrl: meta.issuer,
230- scopes: cfg.scopes,
231- pkceVerifier: verifier,
232- requestUri: requestUri,
233- tokenEndpoint: meta.tokenEndpoint,
234- revocationEndpoint: meta.revocationEndpoint,
235- dpopNonce: nonce,
236- dpopKey: key,
237- )
238-
239-proc authorizeUrl*(cfg: ClientConfig, meta: AuthServerMetadata,
240- info: AuthRequestData): string =
241- let params = @[("client_id", cfg.clientId), ("request_uri", info.requestUri)]
242- meta.authorizationEndpoint & "?" & params.encodeQuery
243-
244-proc startAuthFlow*(cfg: ClientConfig, identifier: string,
245- rng: var HmacDrbgContext,
246- plcUrl = DefaultPlcUrl): (string, AuthRequestData) =
247- ## Resolve an account (or take an auth server URL directly), push the
248- ## request, and return the URL to send the user to.
249- var
250- authServerUrl: string
251- accountDid: string
252- loginHint: string
253-
254- if identifier.startsWith("https://"):
255- authServerUrl = identifier.strip(chars = {'/'})
256- else:
257- let ident = lookup(identifier, plcUrl)
258- accountDid = ident.did
259- loginHint = identifier
260- authServerUrl = resolveAuthServerUrl(ident.pdsEndpoint)
261-
262- let meta = resolveAuthServerMetadata(authServerUrl)
263- var info = sendAuthRequest(cfg, meta, rng, loginHint)
264- info.accountDid = accountDid
265- (authorizeUrl(cfg, meta, info), info)
266-
267-proc exchangeCode*(cfg: ClientConfig, info: var AuthRequestData,
268- code: string, rng: var HmacDrbgContext): TokenResponse =
269- ## Trade the authorization code for tokens, proving possession of both the
270- ## PKCE verifier and the session's DPoP key.
271- let params = @[
272- ("client_id", cfg.clientId),
273- ("redirect_uri", cfg.callbackUrl),
274- ("grant_type", "authorization_code"),
275- ("code", code),
276- ("code_verifier", info.pkceVerifier),
277- ]
278- var nonce = info.dpopNonce
279- let res = postWithDpop(info.tokenEndpoint, params.encodeQuery,
280- info.dpopKey, rng, nonce)
281- info.dpopNonce = nonce
282- if res.code.int != 200:
283- raise newException(OAuthError,
284- &"token request failed (HTTP {res.code.int}): {errorReason(res.body)}")
285-
286- let b = parseJson(res.body)
287- result = TokenResponse(
288- accessToken: b{"access_token"}.getStr,
289- refreshToken: b{"refresh_token"}.getStr,
290- tokenType: b{"token_type"}.getStr,
291- expiresIn: b{"expires_in"}.getInt,
292- scope: b{"scope"}.getStr,
293- sub: b{"sub"}.getStr,
294- )
295-
296- # DPoP-bound tokens are useless as bearer tokens; a server answering
297- # token_type=Bearer means the binding silently did not happen.
298- if not result.tokenType.toLowerAscii.startsWith("dpop"):
299- raise newException(OAuthError,
300- "expected a DPoP-bound token, got token_type=" & result.tokenType)
301- if result.sub.len == 0:
302- raise newException(OAuthError, "token response has no subject DID")
303- if info.accountDid.len > 0 and result.sub != info.accountDid:
304- raise newException(OAuthError,
305- &"token subject {result.sub} does not match requested account {info.accountDid}")
306-
307-proc refresh*(cfg: ClientConfig, info: var AuthRequestData,
308- refreshToken: string, rng: var HmacDrbgContext): TokenResponse =
309- let params = @[
310- ("client_id", cfg.clientId),
311- ("grant_type", "refresh_token"),
312- ("refresh_token", refreshToken),
313- ]
314- var nonce = info.dpopNonce
315- let res = postWithDpop(info.tokenEndpoint, params.encodeQuery,
316- info.dpopKey, rng, nonce)
317- info.dpopNonce = nonce
318- if res.code.int != 200:
319- raise newException(OAuthError,
320- &"refresh failed (HTTP {res.code.int}): {errorReason(res.body)}")
321- let b = parseJson(res.body)
322- TokenResponse(
323- accessToken: b{"access_token"}.getStr,
324- refreshToken: b{"refresh_token"}.getStr,
325- tokenType: b{"token_type"}.getStr,
326- expiresIn: b{"expires_in"}.getInt,
327- scope: b{"scope"}.getStr,
328- sub: b{"sub"}.getStr,
329- )
deleted spike/nim/atproto/schema.nim +0 -232
deleted file mode 100644
@@ -1,232 +0,0 @@
1-## Glean's schema, in the three databases it uses.
2-##
3-## Extracted mechanically from internal/db/db.go rather than retyped: 44 DDL
4-## statements transcribed by hand would drift from the Go definition, and a
5-## column that differs by a default or a CHECK is the kind of thing that only
6-## shows up as a constraint violation months later.
7-##
8-## Regenerate with tools/extract_schema.py after changing the Go schema.
9-
10-const usersSchema* = [
11- """CREATE TABLE IF NOT EXISTS users (
12- did TEXT PRIMARY KEY,
13- indexed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
14- updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
15- follows_dirty BOOLEAN NOT NULL DEFAULT 1
16- )""",
17- """CREATE TABLE IF NOT EXISTS follows (
18- user_did TEXT NOT NULL,
19- target_did TEXT NOT NULL,
20- uri TEXT,
21- cid TEXT,
22- followed_at DATETIME,
23- PRIMARY KEY (user_did, target_did)
24- )""",
25- """CREATE TABLE IF NOT EXISTS oauth_auth_requests (
26- state TEXT PRIMARY KEY,
27- data TEXT NOT NULL
28- )""",
29- """CREATE TABLE IF NOT EXISTS oauth_sessions (
30- account_did TEXT NOT NULL,
31- session_id TEXT NOT NULL,
32- data TEXT NOT NULL,
33- PRIMARY KEY (account_did, session_id)
34- )""",
35- """CREATE INDEX IF NOT EXISTS idx_follows_target ON follows(target_did)""",
36- """CREATE INDEX IF NOT EXISTS idx_follows_uri ON follows(uri)""",
37- """CREATE TABLE IF NOT EXISTS user_settings (
38- did TEXT PRIMARY KEY,
39- languages TEXT,
40- expanded_view BOOLEAN NOT NULL DEFAULT 0,
41- digest_enabled BOOLEAN NOT NULL DEFAULT 0,
42- updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
43- )""",
44- """CREATE TABLE IF NOT EXISTS dismissed_recommendations (
45- user_did TEXT NOT NULL,
46- target_type TEXT NOT NULL CHECK(target_type IN ('feed', 'article', 'person')),
47- target_id TEXT NOT NULL,
48- reason TEXT,
49- dismissed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
50- PRIMARY KEY (user_did, target_type, target_id)
51- )""",
52- """CREATE TABLE IF NOT EXISTS recommendation_impressions (
53- user_did TEXT NOT NULL,
54- target_type TEXT NOT NULL CHECK(target_type IN ('feed', 'article', 'person')),
55- target_id TEXT NOT NULL,
56- first_shown_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
57- last_shown_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
58- shown_count INTEGER NOT NULL DEFAULT 1,
59- acted BOOLEAN NOT NULL DEFAULT 0,
60- PRIMARY KEY (user_did, target_type, target_id)
61- )""",
62- """CREATE INDEX IF NOT EXISTS idx_impressions_user_unacted ON recommendation_impressions(user_did, acted, shown_count)""",
63- """CREATE TABLE IF NOT EXISTS jetstream_cursor (
64- id INTEGER PRIMARY KEY CHECK(id = 1),
65- cursor_us INTEGER NOT NULL
66- )""",
67-]
68-
69-const articlesSchema* = [
70- """CREATE TABLE IF NOT EXISTS articles.feeds (
71- feed_url TEXT PRIMARY KEY,
72- title TEXT,
73- site_url TEXT,
74- description TEXT,
75- feed_type TEXT CHECK(feed_type IN ('rss', 'atom', 'json', 'atproto')),
76- last_fetched_at DATETIME,
77- last_error TEXT,
78- subscriber_count INTEGER NOT NULL DEFAULT 0,
79- consecutive_empty_fetches INTEGER NOT NULL DEFAULT 0,
80- error_count INTEGER NOT NULL DEFAULT 0,
81- favicon_url TEXT
82- )""",
83- """CREATE TABLE IF NOT EXISTS articles.subscriptions (
84- id INTEGER PRIMARY KEY AUTOINCREMENT,
85- user_did TEXT NOT NULL,
86- feed_url TEXT NOT NULL,
87- title TEXT,
88- category TEXT,
89- added_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
90- uri TEXT,
91- cid TEXT,
92- UNIQUE(user_did, feed_url)
93- )""",
94- """CREATE TABLE IF NOT EXISTS articles.articles (
95- id INTEGER PRIMARY KEY AUTOINCREMENT,
96- feed_url TEXT NOT NULL,
97- guid TEXT NOT NULL,
98- title TEXT NOT NULL DEFAULT '',
99- url TEXT,
100- author TEXT,
101- summary TEXT,
102- content TEXT,
103- full_content TEXT,
104- published DATETIME,
105- updated DATETIME,
106- fetched_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
107- language TEXT NOT NULL DEFAULT '',
108- UNIQUE(feed_url, guid)
109- )""",
110- """CREATE TABLE IF NOT EXISTS articles.read_state (
111- user_did TEXT NOT NULL,
112- article_id INTEGER NOT NULL,
113- is_read BOOLEAN NOT NULL DEFAULT 0,
114- read_at DATETIME,
115- PRIMARY KEY (user_did, article_id)
116- )""",
117- """CREATE TABLE IF NOT EXISTS articles.read_state_history (
118- user_did TEXT NOT NULL,
119- feed_url TEXT NOT NULL,
120- guid TEXT NOT NULL,
121- is_read BOOLEAN NOT NULL DEFAULT 0,
122- read_at DATETIME,
123- PRIMARY KEY (user_did, feed_url, guid)
124- )""",
125- """CREATE TABLE IF NOT EXISTS articles.annotations (
126- id INTEGER PRIMARY KEY AUTOINCREMENT,
127- uri TEXT NOT NULL UNIQUE,
128- author_did TEXT NOT NULL,
129- feed_url TEXT NOT NULL,
130- article_url TEXT NOT NULL,
131- quote TEXT,
132- note TEXT,
133- tags TEXT,
134- rating INTEGER,
135- created_at DATETIME NOT NULL,
136- cid TEXT
137- )""",
138- """CREATE TABLE IF NOT EXISTS articles.likes (
139- id INTEGER PRIMARY KEY AUTOINCREMENT,
140- uri TEXT NOT NULL UNIQUE,
141- author_did TEXT NOT NULL,
142- feed_url TEXT NOT NULL,
143- article_url TEXT NOT NULL,
144- created_at DATETIME NOT NULL,
145- cid TEXT,
146- UNIQUE(author_did, feed_url, article_url)
147- )""",
148- """CREATE INDEX IF NOT EXISTS articles.idx_subscriptions_feed_user ON subscriptions(feed_url, user_did)""",
149- """CREATE INDEX IF NOT EXISTS articles.idx_subscriptions_uri ON subscriptions(uri)""",
150- """CREATE INDEX IF NOT EXISTS articles.idx_likes_author_feed ON likes(author_did, feed_url, created_at)""",
151- """CREATE INDEX IF NOT EXISTS articles.idx_articles_url ON articles(url)""",
152- """CREATE INDEX IF NOT EXISTS articles.idx_read_state_unread ON read_state(user_did, is_read) WHERE is_read = 0""",
153- """CREATE INDEX IF NOT EXISTS articles.idx_annotations_article ON annotations(article_url)""",
154- """CREATE INDEX IF NOT EXISTS articles.idx_annotations_author ON annotations(author_did)""",
155- """CREATE INDEX IF NOT EXISTS articles.idx_annotations_created_at ON annotations(created_at DESC)""",
156- """CREATE INDEX IF NOT EXISTS articles.idx_likes_created_at ON likes(created_at DESC)""",
157- """CREATE INDEX IF NOT EXISTS articles.idx_articles_language ON articles(language)""",
158- """CREATE INDEX IF NOT EXISTS articles.idx_likes_article_created ON likes(feed_url, article_url, created_at)""",
159- """CREATE INDEX IF NOT EXISTS articles.idx_annotations_feed_article_created ON annotations(feed_url, article_url, created_at)""",
160- """CREATE INDEX IF NOT EXISTS articles.idx_likes_created_author ON likes(created_at, author_did)""",
161- """CREATE VIRTUAL TABLE IF NOT EXISTS articles.articles_fts USING fts5(title, summary, content, author, content=articles, content_rowid=id)""",
162- """CREATE TRIGGER IF NOT EXISTS articles.articles_ai AFTER INSERT ON articles BEGIN
163- INSERT INTO articles_fts(rowid, title, summary, content, author) VALUES (new.id, new.title, new.summary, new.content, new.author);
164- END""",
165- """CREATE TRIGGER IF NOT EXISTS articles.articles_ad AFTER DELETE ON articles BEGIN
166- INSERT INTO articles_fts(articles_fts, rowid, title, summary, content, author) VALUES('delete', old.id, old.title, old.summary, old.content, old.author);
167- END""",
168- """CREATE TRIGGER IF NOT EXISTS articles.articles_au AFTER UPDATE ON articles BEGIN
169- INSERT INTO articles_fts(articles_fts, rowid, title, summary, content, author) VALUES('delete', old.id, old.title, old.summary, old.content, old.author);
170- INSERT INTO articles_fts(rowid, title, summary, content, author) VALUES (new.id, new.title, new.summary, new.content, new.author);
171- END""",
172-]
173-
174-const recsSchema* = [
175- """CREATE TABLE IF NOT EXISTS recs.feed_similarity (
176- feed_a TEXT NOT NULL,
177- feed_b TEXT NOT NULL,
178- jaccard REAL NOT NULL,
179- computed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
180- PRIMARY KEY (feed_a, feed_b),
181- CHECK(feed_a < feed_b)
182- )""",
183- """CREATE TABLE IF NOT EXISTS recs.user_similarity (
184- user_a TEXT NOT NULL,
185- user_b TEXT NOT NULL,
186- jaccard REAL NOT NULL,
187- common_feeds INTEGER NOT NULL,
188- common_likes INTEGER NOT NULL DEFAULT 0,
189- common_tags INTEGER NOT NULL DEFAULT 0,
190- computed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
191- PRIMARY KEY (user_a, user_b),
192- CHECK(user_a < user_b)
193- )""",
194- """CREATE TABLE IF NOT EXISTS recs.follow_distances (
195- user_a TEXT NOT NULL,
196- user_b TEXT NOT NULL,
197- distance INTEGER NOT NULL CHECK(distance IN (1, 2, 3)),
198- PRIMARY KEY (user_a, user_b)
199- )""",
200- """CREATE TABLE IF NOT EXISTS recs.user_signal_weights (
201- user_did TEXT PRIMARY KEY,
202- w_sub REAL NOT NULL DEFAULT 1.0,
203- w_like REAL NOT NULL DEFAULT 0.5,
204- w_tag REAL NOT NULL DEFAULT 0.3,
205- w_social REAL NOT NULL DEFAULT 0.7,
206- w_pop REAL NOT NULL DEFAULT 0.2,
207- w_category REAL NOT NULL DEFAULT 0.4,
208- w_content REAL NOT NULL DEFAULT 0.4,
209- updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
210- )""",
211- """CREATE TABLE IF NOT EXISTS recs.user_signal_profiles (
212- user_did TEXT PRIMARY KEY,
213- total_likes INTEGER NOT NULL DEFAULT 0,
214- total_tags INTEGER NOT NULL DEFAULT 0,
215- top_categories TEXT,
216- updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
217- )""",
218- """CREATE TABLE IF NOT EXISTS recs.feed_embedding_meta (
219- feed_url TEXT PRIMARY KEY,
220- source_text TEXT NOT NULL DEFAULT ''
221- )""",
222- """CREATE INDEX IF NOT EXISTS recs.idx_follow_distances_a_dist ON follow_distances(user_a, distance)""",
223- """CREATE INDEX IF NOT EXISTS recs.idx_user_similarity_b ON user_similarity(user_b)""",
224- """CREATE TABLE IF NOT EXISTS recs.precomputed_recommendations (
225- user_did TEXT NOT NULL,
226- rec_type TEXT NOT NULL CHECK(rec_type IN ('feed', 'article', 'person')),
227- data TEXT NOT NULL,
228- computed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
229- PRIMARY KEY (user_did, rec_type)
230- )""",
231-]
232-
deleted file mode 100644
@@ -1,232 +0,0 @@
1-## Glean's schema, in the three databases it uses.
2-##
3-## Extracted mechanically from internal/db/db.go rather than retyped: 44 DDL
4-## statements transcribed by hand would drift from the Go definition, and a
5-## column that differs by a default or a CHECK is the kind of thing that only
6-## shows up as a constraint violation months later.
7-##
8-## Regenerate with tools/extract_schema.py after changing the Go schema.
9-
10-const usersSchema* = [
11- """CREATE TABLE IF NOT EXISTS users (
12- did TEXT PRIMARY KEY,
13- indexed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
14- updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
15- follows_dirty BOOLEAN NOT NULL DEFAULT 1
16- )""",
17- """CREATE TABLE IF NOT EXISTS follows (
18- user_did TEXT NOT NULL,
19- target_did TEXT NOT NULL,
20- uri TEXT,
21- cid TEXT,
22- followed_at DATETIME,
23- PRIMARY KEY (user_did, target_did)
24- )""",
25- """CREATE TABLE IF NOT EXISTS oauth_auth_requests (
26- state TEXT PRIMARY KEY,
27- data TEXT NOT NULL
28- )""",
29- """CREATE TABLE IF NOT EXISTS oauth_sessions (
30- account_did TEXT NOT NULL,
31- session_id TEXT NOT NULL,
32- data TEXT NOT NULL,
33- PRIMARY KEY (account_did, session_id)
34- )""",
35- """CREATE INDEX IF NOT EXISTS idx_follows_target ON follows(target_did)""",
36- """CREATE INDEX IF NOT EXISTS idx_follows_uri ON follows(uri)""",
37- """CREATE TABLE IF NOT EXISTS user_settings (
38- did TEXT PRIMARY KEY,
39- languages TEXT,
40- expanded_view BOOLEAN NOT NULL DEFAULT 0,
41- digest_enabled BOOLEAN NOT NULL DEFAULT 0,
42- updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
43- )""",
44- """CREATE TABLE IF NOT EXISTS dismissed_recommendations (
45- user_did TEXT NOT NULL,
46- target_type TEXT NOT NULL CHECK(target_type IN ('feed', 'article', 'person')),
47- target_id TEXT NOT NULL,
48- reason TEXT,
49- dismissed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
50- PRIMARY KEY (user_did, target_type, target_id)
51- )""",
52- """CREATE TABLE IF NOT EXISTS recommendation_impressions (
53- user_did TEXT NOT NULL,
54- target_type TEXT NOT NULL CHECK(target_type IN ('feed', 'article', 'person')),
55- target_id TEXT NOT NULL,
56- first_shown_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
57- last_shown_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
58- shown_count INTEGER NOT NULL DEFAULT 1,
59- acted BOOLEAN NOT NULL DEFAULT 0,
60- PRIMARY KEY (user_did, target_type, target_id)
61- )""",
62- """CREATE INDEX IF NOT EXISTS idx_impressions_user_unacted ON recommendation_impressions(user_did, acted, shown_count)""",
63- """CREATE TABLE IF NOT EXISTS jetstream_cursor (
64- id INTEGER PRIMARY KEY CHECK(id = 1),
65- cursor_us INTEGER NOT NULL
66- )""",
67-]
68-
69-const articlesSchema* = [
70- """CREATE TABLE IF NOT EXISTS articles.feeds (
71- feed_url TEXT PRIMARY KEY,
72- title TEXT,
73- site_url TEXT,
74- description TEXT,
75- feed_type TEXT CHECK(feed_type IN ('rss', 'atom', 'json', 'atproto')),
76- last_fetched_at DATETIME,
77- last_error TEXT,
78- subscriber_count INTEGER NOT NULL DEFAULT 0,
79- consecutive_empty_fetches INTEGER NOT NULL DEFAULT 0,
80- error_count INTEGER NOT NULL DEFAULT 0,
81- favicon_url TEXT
82- )""",
83- """CREATE TABLE IF NOT EXISTS articles.subscriptions (
84- id INTEGER PRIMARY KEY AUTOINCREMENT,
85- user_did TEXT NOT NULL,
86- feed_url TEXT NOT NULL,
87- title TEXT,
88- category TEXT,
89- added_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
90- uri TEXT,
91- cid TEXT,
92- UNIQUE(user_did, feed_url)
93- )""",
94- """CREATE TABLE IF NOT EXISTS articles.articles (
95- id INTEGER PRIMARY KEY AUTOINCREMENT,
96- feed_url TEXT NOT NULL,
97- guid TEXT NOT NULL,
98- title TEXT NOT NULL DEFAULT '',
99- url TEXT,
100- author TEXT,
101- summary TEXT,
102- content TEXT,
103- full_content TEXT,
104- published DATETIME,
105- updated DATETIME,
106- fetched_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
107- language TEXT NOT NULL DEFAULT '',
108- UNIQUE(feed_url, guid)
109- )""",
110- """CREATE TABLE IF NOT EXISTS articles.read_state (
111- user_did TEXT NOT NULL,
112- article_id INTEGER NOT NULL,
113- is_read BOOLEAN NOT NULL DEFAULT 0,
114- read_at DATETIME,
115- PRIMARY KEY (user_did, article_id)
116- )""",
117- """CREATE TABLE IF NOT EXISTS articles.read_state_history (
118- user_did TEXT NOT NULL,
119- feed_url TEXT NOT NULL,
120- guid TEXT NOT NULL,
121- is_read BOOLEAN NOT NULL DEFAULT 0,
122- read_at DATETIME,
123- PRIMARY KEY (user_did, feed_url, guid)
124- )""",
125- """CREATE TABLE IF NOT EXISTS articles.annotations (
126- id INTEGER PRIMARY KEY AUTOINCREMENT,
127- uri TEXT NOT NULL UNIQUE,
128- author_did TEXT NOT NULL,
129- feed_url TEXT NOT NULL,
130- article_url TEXT NOT NULL,
131- quote TEXT,
132- note TEXT,
133- tags TEXT,
134- rating INTEGER,
135- created_at DATETIME NOT NULL,
136- cid TEXT
137- )""",
138- """CREATE TABLE IF NOT EXISTS articles.likes (
139- id INTEGER PRIMARY KEY AUTOINCREMENT,
140- uri TEXT NOT NULL UNIQUE,
141- author_did TEXT NOT NULL,
142- feed_url TEXT NOT NULL,
143- article_url TEXT NOT NULL,
144- created_at DATETIME NOT NULL,
145- cid TEXT,
146- UNIQUE(author_did, feed_url, article_url)
147- )""",
148- """CREATE INDEX IF NOT EXISTS articles.idx_subscriptions_feed_user ON subscriptions(feed_url, user_did)""",
149- """CREATE INDEX IF NOT EXISTS articles.idx_subscriptions_uri ON subscriptions(uri)""",
150- """CREATE INDEX IF NOT EXISTS articles.idx_likes_author_feed ON likes(author_did, feed_url, created_at)""",
151- """CREATE INDEX IF NOT EXISTS articles.idx_articles_url ON articles(url)""",
152- """CREATE INDEX IF NOT EXISTS articles.idx_read_state_unread ON read_state(user_did, is_read) WHERE is_read = 0""",
153- """CREATE INDEX IF NOT EXISTS articles.idx_annotations_article ON annotations(article_url)""",
154- """CREATE INDEX IF NOT EXISTS articles.idx_annotations_author ON annotations(author_did)""",
155- """CREATE INDEX IF NOT EXISTS articles.idx_annotations_created_at ON annotations(created_at DESC)""",
156- """CREATE INDEX IF NOT EXISTS articles.idx_likes_created_at ON likes(created_at DESC)""",
157- """CREATE INDEX IF NOT EXISTS articles.idx_articles_language ON articles(language)""",
158- """CREATE INDEX IF NOT EXISTS articles.idx_likes_article_created ON likes(feed_url, article_url, created_at)""",
159- """CREATE INDEX IF NOT EXISTS articles.idx_annotations_feed_article_created ON annotations(feed_url, article_url, created_at)""",
160- """CREATE INDEX IF NOT EXISTS articles.idx_likes_created_author ON likes(created_at, author_did)""",
161- """CREATE VIRTUAL TABLE IF NOT EXISTS articles.articles_fts USING fts5(title, summary, content, author, content=articles, content_rowid=id)""",
162- """CREATE TRIGGER IF NOT EXISTS articles.articles_ai AFTER INSERT ON articles BEGIN
163- INSERT INTO articles_fts(rowid, title, summary, content, author) VALUES (new.id, new.title, new.summary, new.content, new.author);
164- END""",
165- """CREATE TRIGGER IF NOT EXISTS articles.articles_ad AFTER DELETE ON articles BEGIN
166- INSERT INTO articles_fts(articles_fts, rowid, title, summary, content, author) VALUES('delete', old.id, old.title, old.summary, old.content, old.author);
167- END""",
168- """CREATE TRIGGER IF NOT EXISTS articles.articles_au AFTER UPDATE ON articles BEGIN
169- INSERT INTO articles_fts(articles_fts, rowid, title, summary, content, author) VALUES('delete', old.id, old.title, old.summary, old.content, old.author);
170- INSERT INTO articles_fts(rowid, title, summary, content, author) VALUES (new.id, new.title, new.summary, new.content, new.author);
171- END""",
172-]
173-
174-const recsSchema* = [
175- """CREATE TABLE IF NOT EXISTS recs.feed_similarity (
176- feed_a TEXT NOT NULL,
177- feed_b TEXT NOT NULL,
178- jaccard REAL NOT NULL,
179- computed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
180- PRIMARY KEY (feed_a, feed_b),
181- CHECK(feed_a < feed_b)
182- )""",
183- """CREATE TABLE IF NOT EXISTS recs.user_similarity (
184- user_a TEXT NOT NULL,
185- user_b TEXT NOT NULL,
186- jaccard REAL NOT NULL,
187- common_feeds INTEGER NOT NULL,
188- common_likes INTEGER NOT NULL DEFAULT 0,
189- common_tags INTEGER NOT NULL DEFAULT 0,
190- computed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
191- PRIMARY KEY (user_a, user_b),
192- CHECK(user_a < user_b)
193- )""",
194- """CREATE TABLE IF NOT EXISTS recs.follow_distances (
195- user_a TEXT NOT NULL,
196- user_b TEXT NOT NULL,
197- distance INTEGER NOT NULL CHECK(distance IN (1, 2, 3)),
198- PRIMARY KEY (user_a, user_b)
199- )""",
200- """CREATE TABLE IF NOT EXISTS recs.user_signal_weights (
201- user_did TEXT PRIMARY KEY,
202- w_sub REAL NOT NULL DEFAULT 1.0,
203- w_like REAL NOT NULL DEFAULT 0.5,
204- w_tag REAL NOT NULL DEFAULT 0.3,
205- w_social REAL NOT NULL DEFAULT 0.7,
206- w_pop REAL NOT NULL DEFAULT 0.2,
207- w_category REAL NOT NULL DEFAULT 0.4,
208- w_content REAL NOT NULL DEFAULT 0.4,
209- updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
210- )""",
211- """CREATE TABLE IF NOT EXISTS recs.user_signal_profiles (
212- user_did TEXT PRIMARY KEY,
213- total_likes INTEGER NOT NULL DEFAULT 0,
214- total_tags INTEGER NOT NULL DEFAULT 0,
215- top_categories TEXT,
216- updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
217- )""",
218- """CREATE TABLE IF NOT EXISTS recs.feed_embedding_meta (
219- feed_url TEXT PRIMARY KEY,
220- source_text TEXT NOT NULL DEFAULT ''
221- )""",
222- """CREATE INDEX IF NOT EXISTS recs.idx_follow_distances_a_dist ON follow_distances(user_a, distance)""",
223- """CREATE INDEX IF NOT EXISTS recs.idx_user_similarity_b ON user_similarity(user_b)""",
224- """CREATE TABLE IF NOT EXISTS recs.precomputed_recommendations (
225- user_did TEXT NOT NULL,
226- rec_type TEXT NOT NULL CHECK(rec_type IN ('feed', 'article', 'person')),
227- data TEXT NOT NULL,
228- computed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
229- PRIMARY KEY (user_did, rec_type)
230- )""",
231-]
232-
deleted spike/nim/atproto/socialstore.nim +0 -388
deleted file mode 100644
@@ -1,388 +0,0 @@
1-## Likes, annotations and the social reads built on them, ported from
2-## internal/db/social.go.
3-##
4-## Likes and annotations are ATProto records first and rows second. They are
5-## keyed by `uri` (the at:// record URI) because that is the identity the PDS
6-## owns, and every one of them can arrive twice -- once from the Jetstream
7-## firehose and once from a PDS sync -- so every write here is idempotent by
8-## construction rather than by checking first.
9-
10-import std/[options, strutils, times]
11-import ./articlestore
12-import ./feedstore
13-import ./gleandb
14-import ./sqlite
15-
16-type
17- Annotation* = object
18- id*: int64
19- uri*: string
20- authorDid*: string
21- feedUrl*: string
22- articleUrl*: string
23- quote*: string
24- note*: string
25- tags*: seq[string]
26- rating*: Option[int]
27- createdAt*: Option[DateTime]
28- cid*: string
29-
30- Like* = object
31- id*: int64
32- uri*: string
33- authorDid*: string
34- feedUrl*: string
35- articleUrl*: string
36- createdAt*: Option[DateTime]
37- cid*: string
38-
39- TrendingItem* = object
40- articleId*: int64
41- title*: string
42- url*: string
43- author*: string
44- summary*: string
45- feedUrl*: string
46- feedTitle*: string
47- faviconUrl*: string
48- likeCount*: int
49- annotationCount*: int
50- hasLiked*: bool
51-
52-# --- helpers ---------------------------------------------------------------
53-
54-proc splitTags(s: string): seq[string] =
55- for t in s.split(','):
56- let v = t.strip()
57- if v.len > 0: result.add v
58-
59-proc joinTags(tags: seq[string]): string = tags.join(",")
60-
61-proc langFilter*(languages: seq[string], prefix: string): (string, seq[Param]) =
62- ## Restrict to the reader's languages, always keeping unknown ones.
63- ##
64- ## `language = ''` means detection has not run yet, not "no language". It is
65- ## included deliberately: excluding it would make new articles vanish for
66- ## anyone with a filter set, and reappear minutes later once classified.
67- if languages.len == 0:
68- return ("", @[])
69- var holes: seq[string]
70- var params: seq[Param]
71- for l in languages:
72- holes.add "?"
73- params.add p(l)
74- (" AND (" & prefix & "language IN (" & holes.join(",") & ") OR " &
75- prefix & "language = '')", params)
76-
77-# --- annotations -----------------------------------------------------------
78-
79-proc readAnnotation(r: Row): Annotation =
80- Annotation(
81- id: r.i64(0),
82- uri: r.str(1),
83- authorDid: r.str(2),
84- feedUrl: r.str(3),
85- articleUrl: r.str(4),
86- quote: r.str(5),
87- note: r.str(6),
88- tags: splitTags(r.str(7)),
89- rating: (if r.isNull(8): none(int) else: some(r.i(8))),
90- createdAt: parseSqliteTime(r.str(9)),
91- cid: r.str(10),
92- )
93-
94-const annotationColumns = """an.id, an.uri, an.author_did, an.feed_url,
95- an.article_url, an.quote, an.note, an.tags, an.rating, an.created_at, an.cid"""
96-
97-proc createAnnotation*(g: GleanDb, a: Annotation) =
98- ## Upsert on `uri`: the same record can arrive from Jetstream and from a
99- ## sync, and an edit keeps the URI while changing the content.
100- g.db.run("""INSERT INTO articles.annotations
101- (uri, author_did, feed_url, article_url, quote, note, tags,
102- rating, created_at, cid)
103- VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
104- ON CONFLICT(uri) DO UPDATE SET
105- quote = excluded.quote,
106- note = excluded.note,
107- tags = excluded.tags,
108- rating = excluded.rating,
109- cid = excluded.cid""",
110- p(a.uri), p(a.authorDid), p(a.feedUrl), p(a.articleUrl),
111- pOrNull(a.quote), pOrNull(a.note), pOrNull(joinTags(a.tags)),
112- (if a.rating.isSome: p(a.rating.get.int64) else: nullParam()),
113- p(formatSqliteTime(a.createdAt.get(now().utc))),
114- pOrNull(a.cid))
115-
116-proc getAnnotation*(g: GleanDb, id: int64): Option[Annotation] =
117- g.db.queryFirst("SELECT " & annotationColumns &
118- " FROM articles.annotations an WHERE an.id = ?",
119- [p(id)], readAnnotation)
120-
121-proc annotationExists*(g: GleanDb, uri: string): bool =
122- g.db.queryInt("SELECT 1 FROM articles.annotations WHERE uri = ?",
123- p(uri)).isSome
124-
125-proc deleteAnnotation*(g: GleanDb, uri: string) =
126- g.db.run("DELETE FROM articles.annotations WHERE uri = ?", p(uri))
127-
128-proc listAnnotations*(g: GleanDb, feedUrl = "", articleUrl = "",
129- authorDid = "", limit = 50, offset = 0): seq[Annotation] =
130- var sql = "SELECT " & annotationColumns &
131- " FROM articles.annotations an WHERE 1=1"
132- var params: seq[Param]
133- if feedUrl.len > 0:
134- sql &= " AND an.feed_url = ?"
135- params.add p(feedUrl)
136- if articleUrl.len > 0:
137- sql &= " AND an.article_url = ?"
138- params.add p(articleUrl)
139- if authorDid.len > 0:
140- sql &= " AND an.author_did = ?"
141- params.add p(authorDid)
142- sql &= " ORDER BY an.created_at DESC LIMIT ? OFFSET ?"
143- params.add p(limit.int64)
144- params.add p(offset.int64)
145- g.db.queryAll(sql, params, readAnnotation)
146-
147-proc batchCreateAnnotations*(g: GleanDb, annotations: seq[Annotation]) =
148- if annotations.len == 0: return
149- g.db.transaction:
150- var ps = g.db.prepared("""
151- INSERT INTO articles.annotations
152- (uri, author_did, feed_url, article_url, quote, note, tags, rating,
153- created_at, cid)
154- VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
155- ON CONFLICT(uri) DO NOTHING""")
156- defer: ps.finalize()
157- for a in annotations:
158- ps.exec(p(a.uri), p(a.authorDid), p(a.feedUrl), p(a.articleUrl),
159- pOrNull(a.quote), pOrNull(a.note), pOrNull(joinTags(a.tags)),
160- (if a.rating.isSome: p(a.rating.get.int64) else: nullParam()),
161- p(formatSqliteTime(a.createdAt.get(now().utc))), pOrNull(a.cid))
162-
163-# --- likes -----------------------------------------------------------------
164-
165-proc readLike(r: Row): Like =
166- Like(
167- id: r.i64(0),
168- uri: r.str(1),
169- authorDid: r.str(2),
170- feedUrl: r.str(3),
171- articleUrl: r.str(4),
172- createdAt: parseSqliteTime(r.str(5)),
173- cid: r.str(6),
174- )
175-
176-const likeColumns = """l.id, l.uri, l.author_did, l.feed_url, l.article_url,
177- l.created_at, l.cid"""
178-
179-proc createLike*(g: GleanDb, l: Like) =
180- g.db.run("""INSERT INTO articles.likes
181- (uri, author_did, feed_url, article_url, created_at, cid)
182- VALUES (?, ?, ?, ?, ?, ?)
183- ON CONFLICT(uri) DO NOTHING""",
184- p(l.uri), p(l.authorDid), p(l.feedUrl), p(l.articleUrl),
185- p(formatSqliteTime(l.createdAt.get(now().utc))), pOrNull(l.cid))
186-
187-proc batchCreateLikes*(g: GleanDb, likes: seq[Like]) =
188- if likes.len == 0: return
189- g.db.transaction:
190- var ps = g.db.prepared("""
191- INSERT INTO articles.likes
192- (uri, author_did, feed_url, article_url, created_at, cid)
193- VALUES (?, ?, ?, ?, ?, ?)
194- ON CONFLICT(uri) DO NOTHING""")
195- defer: ps.finalize()
196- for l in likes:
197- ps.exec(p(l.uri), p(l.authorDid), p(l.feedUrl), p(l.articleUrl),
198- p(formatSqliteTime(l.createdAt.get(now().utc))), pOrNull(l.cid))
199-
200-proc deleteLike*(g: GleanDb, uri: string) =
201- g.db.run("DELETE FROM articles.likes WHERE uri = ?", p(uri))
202-
203-proc deleteLikeByUserArticle*(g: GleanDb, authorDid, feedUrl, articleUrl: string) =
204- g.db.run("""DELETE FROM articles.likes
205- WHERE author_did = ? AND feed_url = ? AND article_url = ?""",
206- p(authorDid), p(feedUrl), p(articleUrl))
207-
208-proc getLike*(g: GleanDb, authorDid, feedUrl, articleUrl: string): Option[Like] =
209- g.db.queryFirst("SELECT " & likeColumns & """
210- FROM articles.likes l
211- WHERE l.author_did = ? AND l.feed_url = ? AND l.article_url = ?""",
212- [p(authorDid), p(feedUrl), p(articleUrl)], readLike)
213-
214-proc hasLiked*(g: GleanDb, authorDid, feedUrl, articleUrl: string): bool =
215- getLike(g, authorDid, feedUrl, articleUrl).isSome
216-
217-proc getLikeCount*(g: GleanDb, feedUrl, articleUrl: string): int =
218- g.db.queryInt("""SELECT count(*) FROM articles.likes
219- WHERE feed_url = ? AND article_url = ?""",
220- p(feedUrl), p(articleUrl)).get(0).int
221-
222-proc listLikes*(g: GleanDb, authorDid: string; feedUrl = "",
223- limit = 50, offset = 0): seq[Like] =
224- var sql = "SELECT " & likeColumns &
225- " FROM articles.likes l WHERE l.author_did = ?"
226- var params = @[p(authorDid)]
227- if feedUrl.len > 0:
228- sql &= " AND l.feed_url = ?"
229- params.add p(feedUrl)
230- sql &= " ORDER BY l.created_at DESC LIMIT ? OFFSET ?"
231- params.add p(limit.int64)
232- params.add p(offset.int64)
233- g.db.queryAll(sql, params, readLike)
234-
235-# --- reconciliation --------------------------------------------------------
236-
237-proc deleteOrphanedLikes*(g: GleanDb, userDid: string, activeUris: seq[string]) =
238- ## Remove local likes the user's PDS no longer has.
239- ##
240- ## This is the destructive half of a sync, so an empty active set is treated
241- ## as "delete everything this user liked" only because that is what an empty
242- ## PDS collection means. A caller that failed to *fetch* the collection must
243- ## not reach here -- the difference between "you unliked everything" and
244- ## "the request failed" is not visible from this side.
245- if activeUris.len == 0:
246- g.db.run("DELETE FROM articles.likes WHERE author_did = ?", p(userDid))
247- return
248- g.db.transaction:
249- g.db.run("CREATE TEMP TABLE IF NOT EXISTS active_uris (uri TEXT PRIMARY KEY)")
250- g.db.run("DELETE FROM active_uris")
251- var ps = g.db.prepared(
252- "INSERT OR IGNORE INTO active_uris (uri) VALUES (?)")
253- defer: ps.finalize()
254- for uri in activeUris:
255- ps.exec(p(uri))
256- g.db.run("""DELETE FROM articles.likes
257- WHERE author_did = ?
258- AND uri NOT IN (SELECT uri FROM active_uris)""", p(userDid))
259-
260-proc deleteOrphanedAnnotations*(g: GleanDb, userDid: string,
261- activeUris: seq[string]) =
262- if activeUris.len == 0:
263- g.db.run("DELETE FROM articles.annotations WHERE author_did = ?", p(userDid))
264- return
265- g.db.transaction:
266- g.db.run("CREATE TEMP TABLE IF NOT EXISTS active_uris (uri TEXT PRIMARY KEY)")
267- g.db.run("DELETE FROM active_uris")
268- var ps = g.db.prepared(
269- "INSERT OR IGNORE INTO active_uris (uri) VALUES (?)")
270- defer: ps.finalize()
271- for uri in activeUris:
272- ps.exec(p(uri))
273- g.db.run("""DELETE FROM articles.annotations
274- WHERE author_did = ?
275- AND uri NOT IN (SELECT uri FROM active_uris)""", p(userDid))
276-
277-# --- follows ---------------------------------------------------------------
278-
279-proc upsertFollow*(g: GleanDb, userDid, targetDid, uri, cid: string) =
280- g.db.run("""INSERT INTO follows (user_did, target_did, uri, cid, followed_at)
281- VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
282- ON CONFLICT(user_did, target_did) DO UPDATE SET
283- uri = COALESCE(excluded.uri, follows.uri),
284- cid = COALESCE(excluded.cid, follows.cid)""",
285- p(userDid), p(targetDid), pOrNull(uri), pOrNull(cid))
286-
287-proc deleteFollow*(g: GleanDb, userDid, targetDid: string) =
288- g.db.run("DELETE FROM follows WHERE user_did = ? AND target_did = ?",
289- p(userDid), p(targetDid))
290-
291-proc listFollows*(g: GleanDb, userDid: string): seq[string] =
292- for r in g.db.query(
293- "SELECT target_did FROM follows WHERE user_did = ? ORDER BY target_did",
294- p(userDid)):
295- result.add r.str(0)
296-
297-# --- trending --------------------------------------------------------------
298-
299-proc readTrending(r: Row): TrendingItem =
300- TrendingItem(
301- articleId: r.i64(0),
302- title: r.str(1),
303- url: r.str(2),
304- author: r.str(3),
305- summary: r.str(4),
306- feedUrl: r.str(5),
307- feedTitle: r.str(6),
308- faviconUrl: r.str(7),
309- likeCount: r.i(8),
310- annotationCount: r.i(9),
311- hasLiked: r.b(10),
312- )
313-
314-const trendingSelect = """
315- SELECT ar.id, ar.title, COALESCE(ar.url, ''), COALESCE(ar.author, ''),
316- COALESCE(ar.summary, ''), l.feed_url, COALESCE(f.title, ''),
317- COALESCE(f.favicon_url, ''),
318- COUNT(DISTINCT l.id) AS like_count,
319- COUNT(DISTINCT an.id) AS annotation_count,
320- COALESCE(MAX(CASE WHEN ul.id IS NOT NULL THEN 1 ELSE 0 END), 0)
321- FROM articles.likes l
322- JOIN articles.articles ar
323- ON ar.url = l.article_url AND ar.feed_url = l.feed_url
324- LEFT JOIN articles.feeds f ON f.feed_url = l.feed_url
325- LEFT JOIN articles.annotations an
326- ON an.feed_url = l.feed_url AND an.article_url = l.article_url
327- AND an.created_at >= ?
328- LEFT JOIN articles.likes ul
329- ON ul.feed_url = l.feed_url AND ul.article_url = l.article_url
330- AND ul.author_did = ?
331- WHERE l.created_at >= ?"""
332-
333-## Future-dated articles sort last: a feed that publishes with a scheduled
334-## timestamp would otherwise pin itself to the top of trending indefinitely.
335-const trendingOrder = """
336- GROUP BY ar.id
337- ORDER BY like_count DESC, annotation_count DESC,
338- (CASE WHEN ar.published > datetime('now') THEN 1 ELSE 0 END),
339- ar.published DESC
340- LIMIT ? OFFSET ?"""
341-
342-proc listTrendingArticles*(g: GleanDb, userDid, since: string,
343- limit = 25, offset = 0): seq[TrendingItem] =
344- ## Global trending: everyone's likes.
345- g.db.queryAll(trendingSelect & trendingOrder,
346- [p(since), p(userDid), p(since), p(limit.int64), p(offset.int64)],
347- readTrending)
348-
349-proc listTrendingArticlesForUser*(g: GleanDb, userDid, since: string,
350- languages: seq[string] = @[],
351- limit = 25, offset = 0): seq[TrendingItem] =
352- ## Trending within the reader's own graph: people they follow, people the
353- ## similarity table pairs them with, and themselves.
354- let (filter, langParams) = langFilter(languages, "ar.")
355- var params = @[p(since), p(userDid), p(since),
356- p(userDid), p(userDid), p(userDid), p(userDid), p(userDid)]
357- params.add langParams
358- params.add p(limit.int64)
359- params.add p(offset.int64)
360- g.db.queryAll(trendingSelect & """
361- AND l.author_did IN (
362- SELECT CASE WHEN us.user_a = ? THEN us.user_b ELSE us.user_a END
363- FROM recs.user_similarity us
364- WHERE us.user_a = ? OR us.user_b = ?
365- UNION SELECT ?
366- UNION SELECT fo.target_did FROM follows fo WHERE fo.user_did = ?
367- )""" & filter & trendingOrder, params, readTrending)
368-
369-proc listLikedArticles*(g: GleanDb, userDid: string,
370- limit = 25, offset = 0): seq[Article] =
371- ## The reader's own likes, as articles.
372- g.db.queryAll("""
373- SELECT a.id, a.feed_url, COALESCE(f.title, a.feed_url), f.favicon_url,
374- a.guid, a.title, a.url, a.author, a.summary, a.content,
375- a.full_content, a.published, a.updated,
376- COALESCE(r.is_read, 0),
377- (SELECT count(*) FROM articles.likes l2 WHERE l2.article_url = a.url),
378- 1
379- FROM articles.likes l
380- JOIN articles.articles a
381- ON a.url = l.article_url AND a.feed_url = l.feed_url
382- LEFT JOIN articles.feeds f ON f.feed_url = a.feed_url
383- LEFT JOIN articles.read_state r ON r.article_id = a.id AND r.user_did = ?
384- WHERE l.author_did = ?
385- ORDER BY l.created_at DESC
386- LIMIT ? OFFSET ?""",
387- [p(userDid), p(userDid), p(limit.int64), p(offset.int64)],
388- proc(r: Row): Article = readArticleRow(r))
deleted file mode 100644
@@ -1,388 +0,0 @@
1-## Likes, annotations and the social reads built on them, ported from
2-## internal/db/social.go.
3-##
4-## Likes and annotations are ATProto records first and rows second. They are
5-## keyed by `uri` (the at:// record URI) because that is the identity the PDS
6-## owns, and every one of them can arrive twice -- once from the Jetstream
7-## firehose and once from a PDS sync -- so every write here is idempotent by
8-## construction rather than by checking first.
9-
10-import std/[options, strutils, times]
11-import ./articlestore
12-import ./feedstore
13-import ./gleandb
14-import ./sqlite
15-
16-type
17- Annotation* = object
18- id*: int64
19- uri*: string
20- authorDid*: string
21- feedUrl*: string
22- articleUrl*: string
23- quote*: string
24- note*: string
25- tags*: seq[string]
26- rating*: Option[int]
27- createdAt*: Option[DateTime]
28- cid*: string
29-
30- Like* = object
31- id*: int64
32- uri*: string
33- authorDid*: string
34- feedUrl*: string
35- articleUrl*: string
36- createdAt*: Option[DateTime]
37- cid*: string
38-
39- TrendingItem* = object
40- articleId*: int64
41- title*: string
42- url*: string
43- author*: string
44- summary*: string
45- feedUrl*: string
46- feedTitle*: string
47- faviconUrl*: string
48- likeCount*: int
49- annotationCount*: int
50- hasLiked*: bool
51-
52-# --- helpers ---------------------------------------------------------------
53-
54-proc splitTags(s: string): seq[string] =
55- for t in s.split(','):
56- let v = t.strip()
57- if v.len > 0: result.add v
58-
59-proc joinTags(tags: seq[string]): string = tags.join(",")
60-
61-proc langFilter*(languages: seq[string], prefix: string): (string, seq[Param]) =
62- ## Restrict to the reader's languages, always keeping unknown ones.
63- ##
64- ## `language = ''` means detection has not run yet, not "no language". It is
65- ## included deliberately: excluding it would make new articles vanish for
66- ## anyone with a filter set, and reappear minutes later once classified.
67- if languages.len == 0:
68- return ("", @[])
69- var holes: seq[string]
70- var params: seq[Param]
71- for l in languages:
72- holes.add "?"
73- params.add p(l)
74- (" AND (" & prefix & "language IN (" & holes.join(",") & ") OR " &
75- prefix & "language = '')", params)
76-
77-# --- annotations -----------------------------------------------------------
78-
79-proc readAnnotation(r: Row): Annotation =
80- Annotation(
81- id: r.i64(0),
82- uri: r.str(1),
83- authorDid: r.str(2),
84- feedUrl: r.str(3),
85- articleUrl: r.str(4),
86- quote: r.str(5),
87- note: r.str(6),
88- tags: splitTags(r.str(7)),
89- rating: (if r.isNull(8): none(int) else: some(r.i(8))),
90- createdAt: parseSqliteTime(r.str(9)),
91- cid: r.str(10),
92- )
93-
94-const annotationColumns = """an.id, an.uri, an.author_did, an.feed_url,
95- an.article_url, an.quote, an.note, an.tags, an.rating, an.created_at, an.cid"""
96-
97-proc createAnnotation*(g: GleanDb, a: Annotation) =
98- ## Upsert on `uri`: the same record can arrive from Jetstream and from a
99- ## sync, and an edit keeps the URI while changing the content.
100- g.db.run("""INSERT INTO articles.annotations
101- (uri, author_did, feed_url, article_url, quote, note, tags,
102- rating, created_at, cid)
103- VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
104- ON CONFLICT(uri) DO UPDATE SET
105- quote = excluded.quote,
106- note = excluded.note,
107- tags = excluded.tags,
108- rating = excluded.rating,
109- cid = excluded.cid""",
110- p(a.uri), p(a.authorDid), p(a.feedUrl), p(a.articleUrl),
111- pOrNull(a.quote), pOrNull(a.note), pOrNull(joinTags(a.tags)),
112- (if a.rating.isSome: p(a.rating.get.int64) else: nullParam()),
113- p(formatSqliteTime(a.createdAt.get(now().utc))),
114- pOrNull(a.cid))
115-
116-proc getAnnotation*(g: GleanDb, id: int64): Option[Annotation] =
117- g.db.queryFirst("SELECT " & annotationColumns &
118- " FROM articles.annotations an WHERE an.id = ?",
119- [p(id)], readAnnotation)
120-
121-proc annotationExists*(g: GleanDb, uri: string): bool =
122- g.db.queryInt("SELECT 1 FROM articles.annotations WHERE uri = ?",
123- p(uri)).isSome
124-
125-proc deleteAnnotation*(g: GleanDb, uri: string) =
126- g.db.run("DELETE FROM articles.annotations WHERE uri = ?", p(uri))
127-
128-proc listAnnotations*(g: GleanDb, feedUrl = "", articleUrl = "",
129- authorDid = "", limit = 50, offset = 0): seq[Annotation] =
130- var sql = "SELECT " & annotationColumns &
131- " FROM articles.annotations an WHERE 1=1"
132- var params: seq[Param]
133- if feedUrl.len > 0:
134- sql &= " AND an.feed_url = ?"
135- params.add p(feedUrl)
136- if articleUrl.len > 0:
137- sql &= " AND an.article_url = ?"
138- params.add p(articleUrl)
139- if authorDid.len > 0:
140- sql &= " AND an.author_did = ?"
141- params.add p(authorDid)
142- sql &= " ORDER BY an.created_at DESC LIMIT ? OFFSET ?"
143- params.add p(limit.int64)
144- params.add p(offset.int64)
145- g.db.queryAll(sql, params, readAnnotation)
146-
147-proc batchCreateAnnotations*(g: GleanDb, annotations: seq[Annotation]) =
148- if annotations.len == 0: return
149- g.db.transaction:
150- var ps = g.db.prepared("""
151- INSERT INTO articles.annotations
152- (uri, author_did, feed_url, article_url, quote, note, tags, rating,
153- created_at, cid)
154- VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
155- ON CONFLICT(uri) DO NOTHING""")
156- defer: ps.finalize()
157- for a in annotations:
158- ps.exec(p(a.uri), p(a.authorDid), p(a.feedUrl), p(a.articleUrl),
159- pOrNull(a.quote), pOrNull(a.note), pOrNull(joinTags(a.tags)),
160- (if a.rating.isSome: p(a.rating.get.int64) else: nullParam()),
161- p(formatSqliteTime(a.createdAt.get(now().utc))), pOrNull(a.cid))
162-
163-# --- likes -----------------------------------------------------------------
164-
165-proc readLike(r: Row): Like =
166- Like(
167- id: r.i64(0),
168- uri: r.str(1),
169- authorDid: r.str(2),
170- feedUrl: r.str(3),
171- articleUrl: r.str(4),
172- createdAt: parseSqliteTime(r.str(5)),
173- cid: r.str(6),
174- )
175-
176-const likeColumns = """l.id, l.uri, l.author_did, l.feed_url, l.article_url,
177- l.created_at, l.cid"""
178-
179-proc createLike*(g: GleanDb, l: Like) =
180- g.db.run("""INSERT INTO articles.likes
181- (uri, author_did, feed_url, article_url, created_at, cid)
182- VALUES (?, ?, ?, ?, ?, ?)
183- ON CONFLICT(uri) DO NOTHING""",
184- p(l.uri), p(l.authorDid), p(l.feedUrl), p(l.articleUrl),
185- p(formatSqliteTime(l.createdAt.get(now().utc))), pOrNull(l.cid))
186-
187-proc batchCreateLikes*(g: GleanDb, likes: seq[Like]) =
188- if likes.len == 0: return
189- g.db.transaction:
190- var ps = g.db.prepared("""
191- INSERT INTO articles.likes
192- (uri, author_did, feed_url, article_url, created_at, cid)
193- VALUES (?, ?, ?, ?, ?, ?)
194- ON CONFLICT(uri) DO NOTHING""")
195- defer: ps.finalize()
196- for l in likes:
197- ps.exec(p(l.uri), p(l.authorDid), p(l.feedUrl), p(l.articleUrl),
198- p(formatSqliteTime(l.createdAt.get(now().utc))), pOrNull(l.cid))
199-
200-proc deleteLike*(g: GleanDb, uri: string) =
201- g.db.run("DELETE FROM articles.likes WHERE uri = ?", p(uri))
202-
203-proc deleteLikeByUserArticle*(g: GleanDb, authorDid, feedUrl, articleUrl: string) =
204- g.db.run("""DELETE FROM articles.likes
205- WHERE author_did = ? AND feed_url = ? AND article_url = ?""",
206- p(authorDid), p(feedUrl), p(articleUrl))
207-
208-proc getLike*(g: GleanDb, authorDid, feedUrl, articleUrl: string): Option[Like] =
209- g.db.queryFirst("SELECT " & likeColumns & """
210- FROM articles.likes l
211- WHERE l.author_did = ? AND l.feed_url = ? AND l.article_url = ?""",
212- [p(authorDid), p(feedUrl), p(articleUrl)], readLike)
213-
214-proc hasLiked*(g: GleanDb, authorDid, feedUrl, articleUrl: string): bool =
215- getLike(g, authorDid, feedUrl, articleUrl).isSome
216-
217-proc getLikeCount*(g: GleanDb, feedUrl, articleUrl: string): int =
218- g.db.queryInt("""SELECT count(*) FROM articles.likes
219- WHERE feed_url = ? AND article_url = ?""",
220- p(feedUrl), p(articleUrl)).get(0).int
221-
222-proc listLikes*(g: GleanDb, authorDid: string; feedUrl = "",
223- limit = 50, offset = 0): seq[Like] =
224- var sql = "SELECT " & likeColumns &
225- " FROM articles.likes l WHERE l.author_did = ?"
226- var params = @[p(authorDid)]
227- if feedUrl.len > 0:
228- sql &= " AND l.feed_url = ?"
229- params.add p(feedUrl)
230- sql &= " ORDER BY l.created_at DESC LIMIT ? OFFSET ?"
231- params.add p(limit.int64)
232- params.add p(offset.int64)
233- g.db.queryAll(sql, params, readLike)
234-
235-# --- reconciliation --------------------------------------------------------
236-
237-proc deleteOrphanedLikes*(g: GleanDb, userDid: string, activeUris: seq[string]) =
238- ## Remove local likes the user's PDS no longer has.
239- ##
240- ## This is the destructive half of a sync, so an empty active set is treated
241- ## as "delete everything this user liked" only because that is what an empty
242- ## PDS collection means. A caller that failed to *fetch* the collection must
243- ## not reach here -- the difference between "you unliked everything" and
244- ## "the request failed" is not visible from this side.
245- if activeUris.len == 0:
246- g.db.run("DELETE FROM articles.likes WHERE author_did = ?", p(userDid))
247- return
248- g.db.transaction:
249- g.db.run("CREATE TEMP TABLE IF NOT EXISTS active_uris (uri TEXT PRIMARY KEY)")
250- g.db.run("DELETE FROM active_uris")
251- var ps = g.db.prepared(
252- "INSERT OR IGNORE INTO active_uris (uri) VALUES (?)")
253- defer: ps.finalize()
254- for uri in activeUris:
255- ps.exec(p(uri))
256- g.db.run("""DELETE FROM articles.likes
257- WHERE author_did = ?
258- AND uri NOT IN (SELECT uri FROM active_uris)""", p(userDid))
259-
260-proc deleteOrphanedAnnotations*(g: GleanDb, userDid: string,
261- activeUris: seq[string]) =
262- if activeUris.len == 0:
263- g.db.run("DELETE FROM articles.annotations WHERE author_did = ?", p(userDid))
264- return
265- g.db.transaction:
266- g.db.run("CREATE TEMP TABLE IF NOT EXISTS active_uris (uri TEXT PRIMARY KEY)")
267- g.db.run("DELETE FROM active_uris")
268- var ps = g.db.prepared(
269- "INSERT OR IGNORE INTO active_uris (uri) VALUES (?)")
270- defer: ps.finalize()
271- for uri in activeUris:
272- ps.exec(p(uri))
273- g.db.run("""DELETE FROM articles.annotations
274- WHERE author_did = ?
275- AND uri NOT IN (SELECT uri FROM active_uris)""", p(userDid))
276-
277-# --- follows ---------------------------------------------------------------
278-
279-proc upsertFollow*(g: GleanDb, userDid, targetDid, uri, cid: string) =
280- g.db.run("""INSERT INTO follows (user_did, target_did, uri, cid, followed_at)
281- VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
282- ON CONFLICT(user_did, target_did) DO UPDATE SET
283- uri = COALESCE(excluded.uri, follows.uri),
284- cid = COALESCE(excluded.cid, follows.cid)""",
285- p(userDid), p(targetDid), pOrNull(uri), pOrNull(cid))
286-
287-proc deleteFollow*(g: GleanDb, userDid, targetDid: string) =
288- g.db.run("DELETE FROM follows WHERE user_did = ? AND target_did = ?",
289- p(userDid), p(targetDid))
290-
291-proc listFollows*(g: GleanDb, userDid: string): seq[string] =
292- for r in g.db.query(
293- "SELECT target_did FROM follows WHERE user_did = ? ORDER BY target_did",
294- p(userDid)):
295- result.add r.str(0)
296-
297-# --- trending --------------------------------------------------------------
298-
299-proc readTrending(r: Row): TrendingItem =
300- TrendingItem(
301- articleId: r.i64(0),
302- title: r.str(1),
303- url: r.str(2),
304- author: r.str(3),
305- summary: r.str(4),
306- feedUrl: r.str(5),
307- feedTitle: r.str(6),
308- faviconUrl: r.str(7),
309- likeCount: r.i(8),
310- annotationCount: r.i(9),
311- hasLiked: r.b(10),
312- )
313-
314-const trendingSelect = """
315- SELECT ar.id, ar.title, COALESCE(ar.url, ''), COALESCE(ar.author, ''),
316- COALESCE(ar.summary, ''), l.feed_url, COALESCE(f.title, ''),
317- COALESCE(f.favicon_url, ''),
318- COUNT(DISTINCT l.id) AS like_count,
319- COUNT(DISTINCT an.id) AS annotation_count,
320- COALESCE(MAX(CASE WHEN ul.id IS NOT NULL THEN 1 ELSE 0 END), 0)
321- FROM articles.likes l
322- JOIN articles.articles ar
323- ON ar.url = l.article_url AND ar.feed_url = l.feed_url
324- LEFT JOIN articles.feeds f ON f.feed_url = l.feed_url
325- LEFT JOIN articles.annotations an
326- ON an.feed_url = l.feed_url AND an.article_url = l.article_url
327- AND an.created_at >= ?
328- LEFT JOIN articles.likes ul
329- ON ul.feed_url = l.feed_url AND ul.article_url = l.article_url
330- AND ul.author_did = ?
331- WHERE l.created_at >= ?"""
332-
333-## Future-dated articles sort last: a feed that publishes with a scheduled
334-## timestamp would otherwise pin itself to the top of trending indefinitely.
335-const trendingOrder = """
336- GROUP BY ar.id
337- ORDER BY like_count DESC, annotation_count DESC,
338- (CASE WHEN ar.published > datetime('now') THEN 1 ELSE 0 END),
339- ar.published DESC
340- LIMIT ? OFFSET ?"""
341-
342-proc listTrendingArticles*(g: GleanDb, userDid, since: string,
343- limit = 25, offset = 0): seq[TrendingItem] =
344- ## Global trending: everyone's likes.
345- g.db.queryAll(trendingSelect & trendingOrder,
346- [p(since), p(userDid), p(since), p(limit.int64), p(offset.int64)],
347- readTrending)
348-
349-proc listTrendingArticlesForUser*(g: GleanDb, userDid, since: string,
350- languages: seq[string] = @[],
351- limit = 25, offset = 0): seq[TrendingItem] =
352- ## Trending within the reader's own graph: people they follow, people the
353- ## similarity table pairs them with, and themselves.
354- let (filter, langParams) = langFilter(languages, "ar.")
355- var params = @[p(since), p(userDid), p(since),
356- p(userDid), p(userDid), p(userDid), p(userDid), p(userDid)]
357- params.add langParams
358- params.add p(limit.int64)
359- params.add p(offset.int64)
360- g.db.queryAll(trendingSelect & """
361- AND l.author_did IN (
362- SELECT CASE WHEN us.user_a = ? THEN us.user_b ELSE us.user_a END
363- FROM recs.user_similarity us
364- WHERE us.user_a = ? OR us.user_b = ?
365- UNION SELECT ?
366- UNION SELECT fo.target_did FROM follows fo WHERE fo.user_did = ?
367- )""" & filter & trendingOrder, params, readTrending)
368-
369-proc listLikedArticles*(g: GleanDb, userDid: string,
370- limit = 25, offset = 0): seq[Article] =
371- ## The reader's own likes, as articles.
372- g.db.queryAll("""
373- SELECT a.id, a.feed_url, COALESCE(f.title, a.feed_url), f.favicon_url,
374- a.guid, a.title, a.url, a.author, a.summary, a.content,
375- a.full_content, a.published, a.updated,
376- COALESCE(r.is_read, 0),
377- (SELECT count(*) FROM articles.likes l2 WHERE l2.article_url = a.url),
378- 1
379- FROM articles.likes l
380- JOIN articles.articles a
381- ON a.url = l.article_url AND a.feed_url = l.feed_url
382- LEFT JOIN articles.feeds f ON f.feed_url = a.feed_url
383- LEFT JOIN articles.read_state r ON r.article_id = a.id AND r.user_did = ?
384- WHERE l.author_did = ?
385- ORDER BY l.created_at DESC
386- LIMIT ? OFFSET ?""",
387- [p(userDid), p(userDid), p(limit.int64), p(offset.int64)],
388- proc(r: Row): Article = readArticleRow(r))
deleted spike/nim/atproto/store.nim +0 -217
deleted file mode 100644
@@ -1,217 +0,0 @@
1-## Persisting what an OAuth session and a firehose subscription need to
2-## survive a restart.
3-##
4-## Three things are stored, matching the tables Glean's Go backend uses:
5-##
6-## oauth_auth_requests a flow in progress, keyed by `state`
7-## oauth_sessions a live session, keyed by (account_did, session_id)
8-## jetstream_cursor a single row: where the firehose got to
9-##
10-## The rows hold JSON, as the Go side does, so the schema does not have to
11-## change every time a field is added to a session.
12-##
13-## **This stores key material.** A session row contains the DPoP private key
14-## and the refresh token; together they are the account. The file wants the
15-## same protection as a password store: not world-readable, and not in a
16-## backup that travels somewhere less careful.
17-
18-import std/[json, options, sequtils, times]
19-import ./dpop
20-import ./oauth
21-import ./sqlite
22-
23-type
24- StoreError* = object of CatchableError
25-
26- SessionRecord* = object ## A usable OAuth session.
27- accountDid*: string
28- sessionId*: string
29- pdsUrl*: string
30- authServerUrl*: string
31- tokenEndpoint*: string
32- accessToken*: string
33- refreshToken*: string
34- dpopKey*: P256Key
35- hostNonce*: string
36- authNonce*: string
37- scopes*: seq[string]
38- createdAt*: int64
39-
40-const Schema = [
41- """CREATE TABLE IF NOT EXISTS oauth_auth_requests (
42- state TEXT PRIMARY KEY,
43- data TEXT NOT NULL,
44- created_at INTEGER NOT NULL
45- )""",
46- """CREATE TABLE IF NOT EXISTS oauth_sessions (
47- account_did TEXT NOT NULL,
48- session_id TEXT NOT NULL,
49- data TEXT NOT NULL,
50- PRIMARY KEY (account_did, session_id)
51- )""",
52- """CREATE TABLE IF NOT EXISTS jetstream_cursor (
53- id INTEGER PRIMARY KEY CHECK(id = 1),
54- cursor_us INTEGER NOT NULL
55- )""",
56-]
57-
58-proc migrate*(db: Db) =
59- for stmt in Schema:
60- db.exec stmt
61-
62-# --- key serialisation -----------------------------------------------------
63-
64-proc toJson*(k: P256Key): JsonNode =
65- ## The private scalar plus the public point.
66- ##
67- ## The public half is stored rather than re-derived on load: deriving it
68- ## needs a curve operation, and a stored key that cannot be read back
69- ## without one is a restart that silently drops every session.
70- %*{"d": b64u(k.priv), "x": b64u(k.pubX), "y": b64u(k.pubY)}
71-
72-proc keyFromJson*(j: JsonNode): P256Key =
73- proc part(name: string): seq[byte] =
74- let raw = b64uDecode(j{name}.getStr)
75- if raw.len != 32:
76- raise newException(StoreError, "stored key field " & name & " is not 32 bytes")
77- result = newSeq[byte](32)
78- copyMem(result[0].addr, raw[0].unsafeAddr, 32)
79-
80- let d = part("d")
81- let x = part("x")
82- let y = part("y")
83- for i in 0 ..< 32:
84- result.priv[i] = d[i]
85- result.pubX[i] = x[i]
86- result.pubY[i] = y[i]
87-
88-# --- auth requests ---------------------------------------------------------
89-
90-proc toJson(info: AuthRequestData): JsonNode =
91- %*{
92- "state": info.state,
93- "auth_server_url": info.authServerUrl,
94- "scopes": info.scopes,
95- "pkce_verifier": info.pkceVerifier,
96- "request_uri": info.requestUri,
97- "token_endpoint": info.tokenEndpoint,
98- "revocation_endpoint": info.revocationEndpoint,
99- "dpop_nonce": info.dpopNonce,
100- "dpop_key": info.dpopKey.toJson,
101- "account_did": info.accountDid,
102- }
103-
104-proc authRequestFromJson(j: JsonNode): AuthRequestData =
105- AuthRequestData(
106- state: j{"state"}.getStr,
107- authServerUrl: j{"auth_server_url"}.getStr,
108- scopes: j{"scopes"}.getElems.mapIt(it.getStr),
109- pkceVerifier: j{"pkce_verifier"}.getStr,
110- requestUri: j{"request_uri"}.getStr,
111- tokenEndpoint: j{"token_endpoint"}.getStr,
112- revocationEndpoint: j{"revocation_endpoint"}.getStr,
113- dpopNonce: j{"dpop_nonce"}.getStr,
114- dpopKey: keyFromJson(j{"dpop_key"}),
115- accountDid: j{"account_did"}.getStr,
116- )
117-
118-proc saveAuthRequest*(db: Db, info: AuthRequestData) =
119- db.run("""INSERT INTO oauth_auth_requests (state, data, created_at)
120- VALUES (?, ?, ?)
121- ON CONFLICT(state) DO UPDATE SET data = excluded.data""",
122- p(info.state), p($info.toJson), p(getTime().toUnix))
123-
124-proc takeAuthRequest*(db: Db, state: string): Option[AuthRequestData] =
125- ## Read **and delete**: an authorization code is single-use, so leaving the
126- ## request behind invites a replay of the callback.
127- let row = db.queryText(
128- "SELECT data FROM oauth_auth_requests WHERE state = ?", p(state))
129- if row.isNone:
130- return none(AuthRequestData)
131- db.run("DELETE FROM oauth_auth_requests WHERE state = ?", p(state))
132- try:
133- some(authRequestFromJson(parseJson(row.get)))
134- except JsonParsingError:
135- raise newException(StoreError, "stored auth request is not valid JSON")
136-
137-proc purgeAuthRequests*(db: Db, olderThan: int64) =
138- ## Flows that were started and never finished. Without this the table grows
139- ## forever, one row per abandoned sign-in.
140- db.run("DELETE FROM oauth_auth_requests WHERE created_at < ?",
141- p(getTime().toUnix - olderThan))
142-
143-# --- sessions --------------------------------------------------------------
144-
145-proc toJson(s: SessionRecord): JsonNode =
146- %*{
147- "account_did": s.accountDid,
148- "session_id": s.sessionId,
149- "pds_url": s.pdsUrl,
150- "auth_server_url": s.authServerUrl,
151- "token_endpoint": s.tokenEndpoint,
152- "access_token": s.accessToken,
153- "refresh_token": s.refreshToken,
154- "dpop_key": s.dpopKey.toJson,
155- "host_nonce": s.hostNonce,
156- "auth_nonce": s.authNonce,
157- "scopes": s.scopes,
158- "created_at": s.createdAt,
159- }
160-
161-proc sessionFromJson(j: JsonNode): SessionRecord =
162- SessionRecord(
163- accountDid: j{"account_did"}.getStr,
164- sessionId: j{"session_id"}.getStr,
165- pdsUrl: j{"pds_url"}.getStr,
166- authServerUrl: j{"auth_server_url"}.getStr,
167- tokenEndpoint: j{"token_endpoint"}.getStr,
168- accessToken: j{"access_token"}.getStr,
169- refreshToken: j{"refresh_token"}.getStr,
170- dpopKey: keyFromJson(j{"dpop_key"}),
171- hostNonce: j{"host_nonce"}.getStr,
172- authNonce: j{"auth_nonce"}.getStr,
173- scopes: j{"scopes"}.getElems.mapIt(it.getStr),
174- createdAt: j{"created_at"}.getBiggestInt.int64,
175- )
176-
177-proc saveSession*(db: Db, s: SessionRecord) =
178- if s.accountDid.len == 0 or s.sessionId.len == 0:
179- raise newException(StoreError, "a session needs both a DID and a session id")
180- db.run("""INSERT INTO oauth_sessions (account_did, session_id, data)
181- VALUES (?, ?, ?)
182- ON CONFLICT(account_did, session_id) DO UPDATE SET data = excluded.data""",
183- p(s.accountDid), p(s.sessionId), p($s.toJson))
184-
185-proc getSession*(db: Db, did, sessionId: string): Option[SessionRecord] =
186- let row = db.queryText(
187- "SELECT data FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
188- p(did), p(sessionId))
189- if row.isNone:
190- return none(SessionRecord)
191- try:
192- some(sessionFromJson(parseJson(row.get)))
193- except JsonParsingError:
194- raise newException(StoreError, "stored session is not valid JSON")
195-
196-proc deleteSession*(db: Db, did, sessionId: string) =
197- db.run("DELETE FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
198- p(did), p(sessionId))
199-
200-proc listSessions*(db: Db, did: string): seq[string] =
201- for row in db.rows(
202- "SELECT session_id FROM oauth_sessions WHERE account_did = ?", p(did)):
203- result.add row[0]
204-
205-proc listAccounts*(db: Db): seq[string] =
206- for row in db.rows("SELECT DISTINCT account_did FROM oauth_sessions"):
207- result.add row[0]
208-
209-# --- jetstream cursor ------------------------------------------------------
210-
211-proc saveCursor*(db: Db, cursorUs: int64) =
212- db.run("""INSERT INTO jetstream_cursor (id, cursor_us) VALUES (1, ?)
213- ON CONFLICT(id) DO UPDATE SET cursor_us = excluded.cursor_us""",
214- p(cursorUs))
215-
216-proc loadCursor*(db: Db): int64 =
217- db.queryInt("SELECT cursor_us FROM jetstream_cursor WHERE id = 1").get(0'i64)
deleted file mode 100644
@@ -1,217 +0,0 @@
1-## Persisting what an OAuth session and a firehose subscription need to
2-## survive a restart.
3-##
4-## Three things are stored, matching the tables Glean's Go backend uses:
5-##
6-## oauth_auth_requests a flow in progress, keyed by `state`
7-## oauth_sessions a live session, keyed by (account_did, session_id)
8-## jetstream_cursor a single row: where the firehose got to
9-##
10-## The rows hold JSON, as the Go side does, so the schema does not have to
11-## change every time a field is added to a session.
12-##
13-## **This stores key material.** A session row contains the DPoP private key
14-## and the refresh token; together they are the account. The file wants the
15-## same protection as a password store: not world-readable, and not in a
16-## backup that travels somewhere less careful.
17-
18-import std/[json, options, sequtils, times]
19-import ./dpop
20-import ./oauth
21-import ./sqlite
22-
23-type
24- StoreError* = object of CatchableError
25-
26- SessionRecord* = object ## A usable OAuth session.
27- accountDid*: string
28- sessionId*: string
29- pdsUrl*: string
30- authServerUrl*: string
31- tokenEndpoint*: string
32- accessToken*: string
33- refreshToken*: string
34- dpopKey*: P256Key
35- hostNonce*: string
36- authNonce*: string
37- scopes*: seq[string]
38- createdAt*: int64
39-
40-const Schema = [
41- """CREATE TABLE IF NOT EXISTS oauth_auth_requests (
42- state TEXT PRIMARY KEY,
43- data TEXT NOT NULL,
44- created_at INTEGER NOT NULL
45- )""",
46- """CREATE TABLE IF NOT EXISTS oauth_sessions (
47- account_did TEXT NOT NULL,
48- session_id TEXT NOT NULL,
49- data TEXT NOT NULL,
50- PRIMARY KEY (account_did, session_id)
51- )""",
52- """CREATE TABLE IF NOT EXISTS jetstream_cursor (
53- id INTEGER PRIMARY KEY CHECK(id = 1),
54- cursor_us INTEGER NOT NULL
55- )""",
56-]
57-
58-proc migrate*(db: Db) =
59- for stmt in Schema:
60- db.exec stmt
61-
62-# --- key serialisation -----------------------------------------------------
63-
64-proc toJson*(k: P256Key): JsonNode =
65- ## The private scalar plus the public point.
66- ##
67- ## The public half is stored rather than re-derived on load: deriving it
68- ## needs a curve operation, and a stored key that cannot be read back
69- ## without one is a restart that silently drops every session.
70- %*{"d": b64u(k.priv), "x": b64u(k.pubX), "y": b64u(k.pubY)}
71-
72-proc keyFromJson*(j: JsonNode): P256Key =
73- proc part(name: string): seq[byte] =
74- let raw = b64uDecode(j{name}.getStr)
75- if raw.len != 32:
76- raise newException(StoreError, "stored key field " & name & " is not 32 bytes")
77- result = newSeq[byte](32)
78- copyMem(result[0].addr, raw[0].unsafeAddr, 32)
79-
80- let d = part("d")
81- let x = part("x")
82- let y = part("y")
83- for i in 0 ..< 32:
84- result.priv[i] = d[i]
85- result.pubX[i] = x[i]
86- result.pubY[i] = y[i]
87-
88-# --- auth requests ---------------------------------------------------------
89-
90-proc toJson(info: AuthRequestData): JsonNode =
91- %*{
92- "state": info.state,
93- "auth_server_url": info.authServerUrl,
94- "scopes": info.scopes,
95- "pkce_verifier": info.pkceVerifier,
96- "request_uri": info.requestUri,
97- "token_endpoint": info.tokenEndpoint,
98- "revocation_endpoint": info.revocationEndpoint,
99- "dpop_nonce": info.dpopNonce,
100- "dpop_key": info.dpopKey.toJson,
101- "account_did": info.accountDid,
102- }
103-
104-proc authRequestFromJson(j: JsonNode): AuthRequestData =
105- AuthRequestData(
106- state: j{"state"}.getStr,
107- authServerUrl: j{"auth_server_url"}.getStr,
108- scopes: j{"scopes"}.getElems.mapIt(it.getStr),
109- pkceVerifier: j{"pkce_verifier"}.getStr,
110- requestUri: j{"request_uri"}.getStr,
111- tokenEndpoint: j{"token_endpoint"}.getStr,
112- revocationEndpoint: j{"revocation_endpoint"}.getStr,
113- dpopNonce: j{"dpop_nonce"}.getStr,
114- dpopKey: keyFromJson(j{"dpop_key"}),
115- accountDid: j{"account_did"}.getStr,
116- )
117-
118-proc saveAuthRequest*(db: Db, info: AuthRequestData) =
119- db.run("""INSERT INTO oauth_auth_requests (state, data, created_at)
120- VALUES (?, ?, ?)
121- ON CONFLICT(state) DO UPDATE SET data = excluded.data""",
122- p(info.state), p($info.toJson), p(getTime().toUnix))
123-
124-proc takeAuthRequest*(db: Db, state: string): Option[AuthRequestData] =
125- ## Read **and delete**: an authorization code is single-use, so leaving the
126- ## request behind invites a replay of the callback.
127- let row = db.queryText(
128- "SELECT data FROM oauth_auth_requests WHERE state = ?", p(state))
129- if row.isNone:
130- return none(AuthRequestData)
131- db.run("DELETE FROM oauth_auth_requests WHERE state = ?", p(state))
132- try:
133- some(authRequestFromJson(parseJson(row.get)))
134- except JsonParsingError:
135- raise newException(StoreError, "stored auth request is not valid JSON")
136-
137-proc purgeAuthRequests*(db: Db, olderThan: int64) =
138- ## Flows that were started and never finished. Without this the table grows
139- ## forever, one row per abandoned sign-in.
140- db.run("DELETE FROM oauth_auth_requests WHERE created_at < ?",
141- p(getTime().toUnix - olderThan))
142-
143-# --- sessions --------------------------------------------------------------
144-
145-proc toJson(s: SessionRecord): JsonNode =
146- %*{
147- "account_did": s.accountDid,
148- "session_id": s.sessionId,
149- "pds_url": s.pdsUrl,
150- "auth_server_url": s.authServerUrl,
151- "token_endpoint": s.tokenEndpoint,
152- "access_token": s.accessToken,
153- "refresh_token": s.refreshToken,
154- "dpop_key": s.dpopKey.toJson,
155- "host_nonce": s.hostNonce,
156- "auth_nonce": s.authNonce,
157- "scopes": s.scopes,
158- "created_at": s.createdAt,
159- }
160-
161-proc sessionFromJson(j: JsonNode): SessionRecord =
162- SessionRecord(
163- accountDid: j{"account_did"}.getStr,
164- sessionId: j{"session_id"}.getStr,
165- pdsUrl: j{"pds_url"}.getStr,
166- authServerUrl: j{"auth_server_url"}.getStr,
167- tokenEndpoint: j{"token_endpoint"}.getStr,
168- accessToken: j{"access_token"}.getStr,
169- refreshToken: j{"refresh_token"}.getStr,
170- dpopKey: keyFromJson(j{"dpop_key"}),
171- hostNonce: j{"host_nonce"}.getStr,
172- authNonce: j{"auth_nonce"}.getStr,
173- scopes: j{"scopes"}.getElems.mapIt(it.getStr),
174- createdAt: j{"created_at"}.getBiggestInt.int64,
175- )
176-
177-proc saveSession*(db: Db, s: SessionRecord) =
178- if s.accountDid.len == 0 or s.sessionId.len == 0:
179- raise newException(StoreError, "a session needs both a DID and a session id")
180- db.run("""INSERT INTO oauth_sessions (account_did, session_id, data)
181- VALUES (?, ?, ?)
182- ON CONFLICT(account_did, session_id) DO UPDATE SET data = excluded.data""",
183- p(s.accountDid), p(s.sessionId), p($s.toJson))
184-
185-proc getSession*(db: Db, did, sessionId: string): Option[SessionRecord] =
186- let row = db.queryText(
187- "SELECT data FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
188- p(did), p(sessionId))
189- if row.isNone:
190- return none(SessionRecord)
191- try:
192- some(sessionFromJson(parseJson(row.get)))
193- except JsonParsingError:
194- raise newException(StoreError, "stored session is not valid JSON")
195-
196-proc deleteSession*(db: Db, did, sessionId: string) =
197- db.run("DELETE FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
198- p(did), p(sessionId))
199-
200-proc listSessions*(db: Db, did: string): seq[string] =
201- for row in db.rows(
202- "SELECT session_id FROM oauth_sessions WHERE account_did = ?", p(did)):
203- result.add row[0]
204-
205-proc listAccounts*(db: Db): seq[string] =
206- for row in db.rows("SELECT DISTINCT account_did FROM oauth_sessions"):
207- result.add row[0]
208-
209-# --- jetstream cursor ------------------------------------------------------
210-
211-proc saveCursor*(db: Db, cursorUs: int64) =
212- db.run("""INSERT INTO jetstream_cursor (id, cursor_us) VALUES (1, ?)
213- ON CONFLICT(id) DO UPDATE SET cursor_us = excluded.cursor_us""",
214- p(cursorUs))
215-
216-proc loadCursor*(db: Db): int64 =
217- db.queryInt("SELECT cursor_us FROM jetstream_cursor WHERE id = 1").get(0'i64)
deleted spike/nim/atproto/xrpc.nim +0 -204
deleted file mode 100644
@@ -1,204 +0,0 @@
1-## Authenticated XRPC against a PDS, using a DPoP-bound OAuth access token.
2-##
3-## The important difference from the OAuth endpoints: a resource server
4-## signals a stale nonce with **401** and a `WWW-Authenticate: DPoP
5-## error="use_dpop_nonce"` header, where the auth server used 400 and a JSON
6-## body. Same protocol, different shape, so the retry logic cannot be shared
7-## by pattern-matching the status code alone.
8-##
9-## A session therefore tracks two nonces: one for the auth server (token
10-## refresh) and one for the PDS. They are issued independently and are not
11-## interchangeable.
12-
13-import std/[httpclient, json, strformat, strutils, uri]
14-import bearssl/rand
15-import ./dpop
16-import ./identity
17-import ./oauth
18-
19-type
20- XrpcError* = object of CatchableError
21- status*: int
22- xrpcError*: string ## the `error` field of an XRPC error body
23-
24- Session* = object ## Everything needed to make an authenticated call.
25- did*: string
26- pdsUrl*: string
27- authServerUrl*: string
28- tokenEndpoint*: string
29- accessToken*: string
30- refreshToken*: string
31- dpopKey*: P256Key
32- hostNonce*: string ## nonce issued by the PDS
33- authNonce*: string ## nonce issued by the auth server
34- clientId*: string
35-
36-const HttpTimeoutMs = 30_000
37-
38-proc sessionFromTokens*(info: AuthRequestData, tokens: TokenResponse,
39- pdsUrl, clientId: string): Session =
40- Session(
41- did: tokens.sub,
42- pdsUrl: pdsUrl.strip(chars = {'/'}),
43- authServerUrl: info.authServerUrl,
44- tokenEndpoint: info.tokenEndpoint,
45- accessToken: tokens.accessToken,
46- refreshToken: tokens.refreshToken,
47- dpopKey: info.dpopKey,
48- authNonce: info.dpopNonce,
49- clientId: clientId,
50- )
51-
52-proc dpopUrl*(url: string): string =
53- ## The `htu` claim is the request URL without query or fragment (RFC 9449
54- ## §4.2). Leaving the query on is a silent mismatch: the proof verifies
55- ## against a URL the server never compares it to, and every call 401s.
56- let u = parseUri(url)
57- var bare = u
58- bare.query = ""
59- bare.anchor = ""
60- $bare
61-
62-proc isNonceRetry(wwwAuth: string): bool =
63- ## WWW-Authenticate: DPoP error="use_dpop_nonce", error_description="..."
64- "use_dpop_nonce" in wwwAuth
65-
66-proc isExpiredToken(wwwAuth: string): bool =
67- "invalid_token" in wwwAuth
68-
69-proc parseXrpcError(body: string): (string, string) =
70- try:
71- let j = parseJson(body)
72- (j{"error"}.getStr, j{"message"}.getStr)
73- except CatchableError:
74- ("", "")
75-
76-proc raiseXrpc(status: int, body: string, what: string) =
77- let (err, msg) = parseXrpcError(body)
78- var e = newException(XrpcError,
79- &"{what} failed (HTTP {status})" &
80- (if err.len > 0: &": {err}" else: "") &
81- (if msg.len > 0: &" -- {msg}" else: ""))
82- e.status = status
83- e.xrpcError = err
84- raise e
85-
86-proc refreshSession*(sess: var Session, rng: var HmacDrbgContext) =
87- ## Exchange the refresh token for a new access token, reusing the session's
88- ## DPoP key -- the refresh token is bound to it just as the access token is.
89- if sess.refreshToken.len == 0:
90- raise newException(XrpcError, "session has no refresh token")
91- var info = AuthRequestData(
92- authServerUrl: sess.authServerUrl,
93- tokenEndpoint: sess.tokenEndpoint,
94- dpopKey: sess.dpopKey,
95- dpopNonce: sess.authNonce,
96- )
97- let cfg = ClientConfig(clientId: sess.clientId)
98- let tokens = refresh(cfg, info, sess.refreshToken, rng)
99- sess.accessToken = tokens.accessToken
100- if tokens.refreshToken.len > 0:
101- sess.refreshToken = tokens.refreshToken
102- sess.authNonce = info.dpopNonce
103-
104-proc request*(sess: var Session, rng: var HmacDrbgContext,
105- httpMethod: HttpMethod, url: string,
106- body = "", contentType = "application/json"): Response =
107- ## One authenticated request, handling both retry conditions.
108- ##
109- ## Up to three attempts: the original, one for a nonce the PDS only reveals
110- ## by rejecting us, and one after refreshing an expired access token.
111- let client = newHttpClient(userAgent = UserAgent, timeout = HttpTimeoutMs)
112- defer: client.close()
113-
114- let htu = dpopUrl(url)
115- # HttpMethod stringifies straight to the uppercase token the `htm` claim
116- # wants ("GET", "POST"), so no trimming.
117- let htm = $httpMethod
118- var refreshed = false
119- var noncedOnce = false
120-
121- for attempt in 0 .. 2:
122- let proof = dpopProof(sess.dpopKey, rng, htm, htu,
123- nonce = sess.hostNonce,
124- accessToken = sess.accessToken,
125- issuer = sess.authServerUrl)
126- var headers = @{
127- "Authorization": "DPoP " & sess.accessToken, # not Bearer
128- "DPoP": proof,
129- }
130- if body.len > 0:
131- headers.add ("Content-Type", contentType)
132- client.headers = newHttpHeaders(headers)
133-
134- result = client.request(url, httpMethod = httpMethod, body = body)
135-
136- # The PDS rotates nonces freely; adopt any it offers, success or not.
137- let supplied = result.headers.getOrDefault("DPoP-Nonce").string
138- if supplied.len > 0:
139- sess.hostNonce = supplied
140-
141- if result.code.int != 401:
142- return result
143-
144- let wwwAuth = result.headers.getOrDefault("WWW-Authenticate").string
145- if wwwAuth.len == 0:
146- return result
147-
148- if isNonceRetry(wwwAuth) and supplied.len > 0 and not noncedOnce:
149- noncedOnce = true
150- continue
151-
152- if isExpiredToken(wwwAuth) and not refreshed and sess.refreshToken.len > 0:
153- refreshed = true
154- refreshSession(sess, rng)
155- continue
156-
157- return result
158-
159- raise newException(XrpcError, "ran out of retries for " & url)
160-
161-proc query*(sess: var Session, rng: var HmacDrbgContext, nsid: string,
162- params: openArray[(string, string)] = []): JsonNode =
163- ## An XRPC query (GET), e.g. com.atproto.repo.listRecords.
164- var url = &"{sess.pdsUrl}/xrpc/{nsid}"
165- if params.len > 0:
166- url &= "?" & params.encodeQuery
167- let res = request(sess, rng, HttpGet, url)
168- if res.code.int != 200:
169- raiseXrpc(res.code.int, res.body, nsid)
170- try:
171- parseJson(res.body)
172- except JsonParsingError:
173- raise newException(XrpcError, nsid & " returned a non-JSON body")
174-
175-proc procedure*(sess: var Session, rng: var HmacDrbgContext, nsid: string,
176- body: JsonNode = nil): JsonNode =
177- ## An XRPC procedure (POST), e.g. com.atproto.repo.createRecord.
178- let url = &"{sess.pdsUrl}/xrpc/{nsid}"
179- let payload = if body == nil: "" else: $body
180- let res = request(sess, rng, HttpPost, url, body = payload)
181- if res.code.int notin [200, 201]:
182- raiseXrpc(res.code.int, res.body, nsid)
183- if res.body.strip.len == 0:
184- return newJObject()
185- try:
186- parseJson(res.body)
187- except JsonParsingError:
188- raise newException(XrpcError, nsid & " returned a non-JSON body")
189-
190-# --- unauthenticated helpers ----------------------------------------------
191-
192-proc publicQuery*(host, nsid: string,
193- params: openArray[(string, string)] = []): JsonNode =
194- ## Many XRPC endpoints need no auth. Useful on its own, and it isolates
195- ## "can we speak XRPC at all" from "is our DPoP signing right".
196- var url = &"https://{hostOf(host)}/xrpc/{nsid}"
197- if params.len > 0:
198- url &= "?" & params.encodeQuery
199- let client = newHttpClient(userAgent = UserAgent, timeout = HttpTimeoutMs)
200- defer: client.close()
201- let res = client.get(url)
202- if res.code.int != 200:
203- raiseXrpc(res.code.int, res.body, nsid)
204- parseJson(res.body)
deleted file mode 100644
@@ -1,204 +0,0 @@
1-## Authenticated XRPC against a PDS, using a DPoP-bound OAuth access token.
2-##
3-## The important difference from the OAuth endpoints: a resource server
4-## signals a stale nonce with **401** and a `WWW-Authenticate: DPoP
5-## error="use_dpop_nonce"` header, where the auth server used 400 and a JSON
6-## body. Same protocol, different shape, so the retry logic cannot be shared
7-## by pattern-matching the status code alone.
8-##
9-## A session therefore tracks two nonces: one for the auth server (token
10-## refresh) and one for the PDS. They are issued independently and are not
11-## interchangeable.
12-
13-import std/[httpclient, json, strformat, strutils, uri]
14-import bearssl/rand
15-import ./dpop
16-import ./identity
17-import ./oauth
18-
19-type
20- XrpcError* = object of CatchableError
21- status*: int
22- xrpcError*: string ## the `error` field of an XRPC error body
23-
24- Session* = object ## Everything needed to make an authenticated call.
25- did*: string
26- pdsUrl*: string
27- authServerUrl*: string
28- tokenEndpoint*: string
29- accessToken*: string
30- refreshToken*: string
31- dpopKey*: P256Key
32- hostNonce*: string ## nonce issued by the PDS
33- authNonce*: string ## nonce issued by the auth server
34- clientId*: string
35-
36-const HttpTimeoutMs = 30_000
37-
38-proc sessionFromTokens*(info: AuthRequestData, tokens: TokenResponse,
39- pdsUrl, clientId: string): Session =
40- Session(
41- did: tokens.sub,
42- pdsUrl: pdsUrl.strip(chars = {'/'}),
43- authServerUrl: info.authServerUrl,
44- tokenEndpoint: info.tokenEndpoint,
45- accessToken: tokens.accessToken,
46- refreshToken: tokens.refreshToken,
47- dpopKey: info.dpopKey,
48- authNonce: info.dpopNonce,
49- clientId: clientId,
50- )
51-
52-proc dpopUrl*(url: string): string =
53- ## The `htu` claim is the request URL without query or fragment (RFC 9449
54- ## §4.2). Leaving the query on is a silent mismatch: the proof verifies
55- ## against a URL the server never compares it to, and every call 401s.
56- let u = parseUri(url)
57- var bare = u
58- bare.query = ""
59- bare.anchor = ""
60- $bare
61-
62-proc isNonceRetry(wwwAuth: string): bool =
63- ## WWW-Authenticate: DPoP error="use_dpop_nonce", error_description="..."
64- "use_dpop_nonce" in wwwAuth
65-
66-proc isExpiredToken(wwwAuth: string): bool =
67- "invalid_token" in wwwAuth
68-
69-proc parseXrpcError(body: string): (string, string) =
70- try:
71- let j = parseJson(body)
72- (j{"error"}.getStr, j{"message"}.getStr)
73- except CatchableError:
74- ("", "")
75-
76-proc raiseXrpc(status: int, body: string, what: string) =
77- let (err, msg) = parseXrpcError(body)
78- var e = newException(XrpcError,
79- &"{what} failed (HTTP {status})" &
80- (if err.len > 0: &": {err}" else: "") &
81- (if msg.len > 0: &" -- {msg}" else: ""))
82- e.status = status
83- e.xrpcError = err
84- raise e
85-
86-proc refreshSession*(sess: var Session, rng: var HmacDrbgContext) =
87- ## Exchange the refresh token for a new access token, reusing the session's
88- ## DPoP key -- the refresh token is bound to it just as the access token is.
89- if sess.refreshToken.len == 0:
90- raise newException(XrpcError, "session has no refresh token")
91- var info = AuthRequestData(
92- authServerUrl: sess.authServerUrl,
93- tokenEndpoint: sess.tokenEndpoint,
94- dpopKey: sess.dpopKey,
95- dpopNonce: sess.authNonce,
96- )
97- let cfg = ClientConfig(clientId: sess.clientId)
98- let tokens = refresh(cfg, info, sess.refreshToken, rng)
99- sess.accessToken = tokens.accessToken
100- if tokens.refreshToken.len > 0:
101- sess.refreshToken = tokens.refreshToken
102- sess.authNonce = info.dpopNonce
103-
104-proc request*(sess: var Session, rng: var HmacDrbgContext,
105- httpMethod: HttpMethod, url: string,
106- body = "", contentType = "application/json"): Response =
107- ## One authenticated request, handling both retry conditions.
108- ##
109- ## Up to three attempts: the original, one for a nonce the PDS only reveals
110- ## by rejecting us, and one after refreshing an expired access token.
111- let client = newHttpClient(userAgent = UserAgent, timeout = HttpTimeoutMs)
112- defer: client.close()
113-
114- let htu = dpopUrl(url)
115- # HttpMethod stringifies straight to the uppercase token the `htm` claim
116- # wants ("GET", "POST"), so no trimming.
117- let htm = $httpMethod
118- var refreshed = false
119- var noncedOnce = false
120-
121- for attempt in 0 .. 2:
122- let proof = dpopProof(sess.dpopKey, rng, htm, htu,
123- nonce = sess.hostNonce,
124- accessToken = sess.accessToken,
125- issuer = sess.authServerUrl)
126- var headers = @{
127- "Authorization": "DPoP " & sess.accessToken, # not Bearer
128- "DPoP": proof,
129- }
130- if body.len > 0:
131- headers.add ("Content-Type", contentType)
132- client.headers = newHttpHeaders(headers)
133-
134- result = client.request(url, httpMethod = httpMethod, body = body)
135-
136- # The PDS rotates nonces freely; adopt any it offers, success or not.
137- let supplied = result.headers.getOrDefault("DPoP-Nonce").string
138- if supplied.len > 0:
139- sess.hostNonce = supplied
140-
141- if result.code.int != 401:
142- return result
143-
144- let wwwAuth = result.headers.getOrDefault("WWW-Authenticate").string
145- if wwwAuth.len == 0:
146- return result
147-
148- if isNonceRetry(wwwAuth) and supplied.len > 0 and not noncedOnce:
149- noncedOnce = true
150- continue
151-
152- if isExpiredToken(wwwAuth) and not refreshed and sess.refreshToken.len > 0:
153- refreshed = true
154- refreshSession(sess, rng)
155- continue
156-
157- return result
158-
159- raise newException(XrpcError, "ran out of retries for " & url)
160-
161-proc query*(sess: var Session, rng: var HmacDrbgContext, nsid: string,
162- params: openArray[(string, string)] = []): JsonNode =
163- ## An XRPC query (GET), e.g. com.atproto.repo.listRecords.
164- var url = &"{sess.pdsUrl}/xrpc/{nsid}"
165- if params.len > 0:
166- url &= "?" & params.encodeQuery
167- let res = request(sess, rng, HttpGet, url)
168- if res.code.int != 200:
169- raiseXrpc(res.code.int, res.body, nsid)
170- try:
171- parseJson(res.body)
172- except JsonParsingError:
173- raise newException(XrpcError, nsid & " returned a non-JSON body")
174-
175-proc procedure*(sess: var Session, rng: var HmacDrbgContext, nsid: string,
176- body: JsonNode = nil): JsonNode =
177- ## An XRPC procedure (POST), e.g. com.atproto.repo.createRecord.
178- let url = &"{sess.pdsUrl}/xrpc/{nsid}"
179- let payload = if body == nil: "" else: $body
180- let res = request(sess, rng, HttpPost, url, body = payload)
181- if res.code.int notin [200, 201]:
182- raiseXrpc(res.code.int, res.body, nsid)
183- if res.body.strip.len == 0:
184- return newJObject()
185- try:
186- parseJson(res.body)
187- except JsonParsingError:
188- raise newException(XrpcError, nsid & " returned a non-JSON body")
189-
190-# --- unauthenticated helpers ----------------------------------------------
191-
192-proc publicQuery*(host, nsid: string,
193- params: openArray[(string, string)] = []): JsonNode =
194- ## Many XRPC endpoints need no auth. Useful on its own, and it isolates
195- ## "can we speak XRPC at all" from "is our DPoP signing right".
196- var url = &"https://{hostOf(host)}/xrpc/{nsid}"
197- if params.len > 0:
198- url &= "?" & params.encodeQuery
199- let client = newHttpClient(userAgent = UserAgent, timeout = HttpTimeoutMs)
200- defer: client.close()
201- let res = client.get(url)
202- if res.code.int != 200:
203- raiseXrpc(res.code.int, res.body, nsid)
204- parseJson(res.body)
deleted spike/nim/cluster_probe.nim +0 -220
deleted file mode 100644
@@ -1,220 +0,0 @@
1-## Checks the similarity computations and scoring helpers.
2-##
3-## nim c -r -d:ssl cluster_probe.nim
4-##
5-## Recommendations are the one part of this system with no obviously correct
6-## answer, so the checks are about properties rather than values: that decay
7-## actually decays, that a pair is counted once, that normalisation does not
8-## produce NaN, that discovery is not just a mirror of who you already follow.
9-
10-import std/[math, options, os, random, sequtils, strformat, strutils, times]
11-import atproto/[articlestore, cluster, feedstore, gleandb, socialstore, sqlite]
12-
13-var failures = 0
14-
15-proc report(name: string, ok: bool, detail = "") =
16- if not ok: inc failures
17- let label = if ok: "PASS" else: "FAIL"
18- if detail.len > 0: echo &" {label} {name} -- {detail}"
19- else: echo &" {label} {name}"
20-
21-proc cleanup(base: string) =
22- for suffix in ["_users", "_articles", "_recs"]:
23- for ext in ["", "-wal", "-shm"]:
24- removeFile(base & suffix & ext)
25-
26-const
27- alice = "did:plc:alice"
28- bob = "did:plc:bob"
29- carol = "did:plc:carol"
30-
31-proc main() =
32- echo "Clustering probe"
33- echo ""
34-
35- # --- pure helpers -------------------------------------------------------
36- block:
37- report("decay at zero days is 1", abs(decay(0.0) - 1.0) < 1e-12)
38- # ln(2)/0.023 is about 30 days, so a month-old event is worth half.
39- report("decay halves in about a month",
40- abs(decay(30.1) - 0.5) < 0.01, &"{decay(30.1):.4f}")
41- report("decay keeps shrinking", decay(365.0) < decay(30.0))
42- report("decay never reaches zero", decay(3650.0) > 0.0,
43- &"{decay(3650.0):.3e}")
44-
45- block:
46- var s = @[0.0, 5.0, 10.0]
47- normalizeScores(s)
48- report("normalisation maps to 0..1", s == @[0.0, 0.5, 1.0], $s)
49-
50- # The case that would produce NaN: no spread to divide by.
51- var flat = @[3.0, 3.0, 3.0]
52- normalizeScores(flat)
53- report("an all-equal set is left alone", flat == @[3.0, 3.0, 3.0], $flat)
54- for v in flat:
55- report("and produces no NaN", not v.isNaN)
56-
57- var one = @[7.0]
58- normalizeScores(one)
59- report("a single score is left alone", one == @[7.0])
60-
61- var empty: seq[float] = @[]
62- normalizeScores(empty)
63- report("an empty set does not crash", empty.len == 0)
64-
65- var negatives = @[-10.0, -5.0, 0.0]
66- normalizeScores(negatives)
67- report("negative scores normalise too", negatives == @[0.0, 0.5, 1.0],
68- $negatives)
69-
70- block:
71- var rng = initRand(42)
72- var pool: seq[PersonRecommendation]
73- for i in 0 ..< 10:
74- pool.add PersonRecommendation(did: &"did:plc:in{i}", isFollowed: true)
75- for i in 0 ..< 10:
76- pool.add PersonRecommendation(did: &"did:plc:out{i}", isFollowed: false)
77-
78- let sampled = samplePeople(pool, 3, rng)
79- let followed = sampled.countIt(it.isFollowed)
80- let strangers = sampled.len - followed
81- report("sampling caps each side at the half", sampled.len == 6,
82- $sampled.len)
83- # Ranking purely by similarity returns the network the reader already
84- # has; reserving half the slots is what makes this discovery.
85- report("half the slots go to people not followed",
86- followed == 3 and strangers == 3,
87- &"{followed} followed, {strangers} strangers")
88-
89- # Fewer than `half` on one side must not pad from the other.
90- var lopsided: seq[PersonRecommendation]
91- lopsided.add PersonRecommendation(did: "did:plc:only", isFollowed: true)
92- for i in 0 ..< 10:
93- lopsided.add PersonRecommendation(did: &"did:plc:s{i}", isFollowed: false)
94- let small = samplePeople(lopsided, 3, rng)
95- report("a thin side is not padded from the other",
96- small.countIt(it.isFollowed) == 1 and small.len == 4, $small.len)
97-
98- # Shuffled, so the same faces do not recur every visit.
99- var rngA = initRand(1)
100- var rngB = initRand(2)
101- let a = samplePeople(pool, 3, rngA).mapIt(it.did).join(",")
102- let b = samplePeople(pool, 3, rngB).mapIt(it.did).join(",")
103- report("sampling varies between runs", a != b)
104-
105- block:
106- var recs = @[
107- FeedRecommendation(feedUrl: "c", score: 0.1),
108- FeedRecommendation(feedUrl: "a", score: 0.9),
109- FeedRecommendation(feedUrl: "b", score: 0.5),
110- ]
111- rankByScore(recs)
112- report("ranking puts the highest score first",
113- recs.mapIt(it.feedUrl) == @["a", "b", "c"],
114- recs.mapIt(it.feedUrl).join(","))
115-
116- # --- similarity against a real database ---------------------------------
117- let base = getTempDir() / "glean_cluster_probe"
118- cleanup(base)
119- defer: cleanup(base)
120-
121- var g = gleandb.open(base)
122- defer: g.close()
123- g.migrate()
124-
125- for did in [alice, bob, carol]:
126- g.db.run("INSERT INTO users (did) VALUES (?)", p(did))
127-
128- # alice and bob share two feeds; carol shares one with each.
129- for i in 1 .. 3:
130- g.upsertFeed(Feed(feedUrl: &"https://f{i}.test/feed", title: &"Feed {i}",
131- feedType: "rss"))
132- g.createSubscription(alice, "https://f1.test/feed", "", "", "", "")
133- g.createSubscription(alice, "https://f2.test/feed", "", "", "", "")
134- g.createSubscription(bob, "https://f1.test/feed", "", "", "", "")
135- g.createSubscription(bob, "https://f2.test/feed", "", "", "", "")
136- g.createSubscription(carol, "https://f1.test/feed", "", "", "", "")
137- g.createSubscription(carol, "https://f3.test/feed", "", "", "", "")
138-
139- block:
140- g.computeUserSimilarity()
141- let ab = g.userSimilarity(alice, bob)
142- let ac = g.userSimilarity(alice, carol)
143- # alice/bob share 2 of 2; alice/carol share 1 of 3.
144- report("identical subscribers score 1.0", abs(ab - 1.0) < 1e-9, &"{ab:.4f}")
145- report("partial overlap scores lower",
146- abs(ac - (1.0 / 3.0)) < 1e-9, &"{ac:.4f}")
147- report("similarity is symmetric on lookup",
148- abs(g.userSimilarity(bob, alice) - ab) < 1e-12)
149- report("an unrelated pair scores zero",
150- g.userSimilarity(alice, "did:plc:nobody") == 0.0)
151-
152- # One row per unordered pair, or every score would be counted twice
153- # downstream.
154- let rows = g.db.queryInt("SELECT count(*) FROM recs.user_similarity").get(0)
155- report("each pair is stored once", rows == 3, &"{rows} rows")
156-
157- block:
158- # Recomputing must replace rather than accumulate.
159- g.computeUserSimilarity()
160- let rows = g.db.queryInt("SELECT count(*) FROM recs.user_similarity").get(0)
161- report("recomputing replaces the previous generation", rows == 3,
162- &"{rows} rows")
163-
164- block:
165- g.computeFeedSimilarity()
166- let f12 = g.feedSimilarity("https://f1.test/feed", "https://f2.test/feed")
167- let f13 = g.feedSimilarity("https://f1.test/feed", "https://f3.test/feed")
168- # f1 and f2 share alice and bob; f1 and f3 share only carol.
169- report("feeds with more shared subscribers score higher", f12 > f13,
170- &"f1/f2={f12:.4f} f1/f3={f13:.4f}")
171- let reversed = g.feedSimilarity("https://f2.test/feed", "https://f1.test/feed")
172- report("feed similarity is symmetric on lookup", abs(reversed - f12) < 1e-12)
173- report("an unshared pair is absent",
174- g.feedSimilarity("https://f2.test/feed", "https://f3.test/feed") == 0.0)
175-
176- # --- decay in the query, not just in Nim --------------------------------
177- block:
178- # Two pairs of people who liked the same articles, one recently and one
179- # long ago. The old overlap must weigh less -- that is the whole reason
180- # EXP is registered as a SQL function.
181- g.upsertFeed(Feed(feedUrl: "https://n.test/feed", title: "News"))
182- let fresh = now().utc
183- let stale = now().utc - initDuration(days = 400)
184- for (who, at, n) in [(alice, fresh, "r"), (bob, fresh, "r"),
185- (carol, stale, "s")]:
186- g.createLike(Like(uri: &"at://{who}/like/{n}", authorDid: who,
187- feedUrl: "https://n.test/feed",
188- articleUrl: &"https://n.test/{n}",
189- createdAt: some(at)))
190- # carol needs a partner on the stale article for a pair to exist.
191- g.createLike(Like(uri: "at://dave/like/s", authorDid: "did:plc:dave",
192- feedUrl: "https://n.test/feed",
193- articleUrl: "https://n.test/s",
194- createdAt: some(stale)))
195-
196- g.computeUserSimilarity()
197- proc overlapOf(a, b: string): float =
198- let (x, y) = if a < b: (a, b) else: (b, a)
199- g.db.queryFirst(
200- "SELECT common_likes FROM recs.user_similarity WHERE user_a = ? AND user_b = ?",
201- [p(x), p(y)], proc(r: Row): float = r.f(0)).get(-1.0)
202-
203- let recent = overlapOf(alice, bob)
204- let old = overlapOf(carol, "did:plc:dave")
205- # The check that was vacuous before: both sides were 0 because the Go
206- # query truncates a decayed sum that is always just under 1.
207- report("a recent shared like counts as one", recent == 1.0, &"{recent}")
208- report("a 400-day-old shared like decays to nothing", old == 0.0, &"{old}")
209- report("decay makes old overlap weigh less than new", recent > old,
210- &"recent={recent} old={old}")
211-
212- echo ""
213- if failures == 0:
214- echo "RESULT: similarity and scoring behave."
215- else:
216- echo &"RESULT: {failures} check(s) failed."
217- quit 1
218-
219-when isMainModule:
220- main()
deleted file mode 100644
@@ -1,220 +0,0 @@
1-## Checks the similarity computations and scoring helpers.
2-##
3-## nim c -r -d:ssl cluster_probe.nim
4-##
5-## Recommendations are the one part of this system with no obviously correct
6-## answer, so the checks are about properties rather than values: that decay
7-## actually decays, that a pair is counted once, that normalisation does not
8-## produce NaN, that discovery is not just a mirror of who you already follow.
9-
10-import std/[math, options, os, random, sequtils, strformat, strutils, times]
11-import atproto/[articlestore, cluster, feedstore, gleandb, socialstore, sqlite]
12-
13-var failures = 0
14-
15-proc report(name: string, ok: bool, detail = "") =
16- if not ok: inc failures
17- let label = if ok: "PASS" else: "FAIL"
18- if detail.len > 0: echo &" {label} {name} -- {detail}"
19- else: echo &" {label} {name}"
20-
21-proc cleanup(base: string) =
22- for suffix in ["_users", "_articles", "_recs"]:
23- for ext in ["", "-wal", "-shm"]:
24- removeFile(base & suffix & ext)
25-
26-const
27- alice = "did:plc:alice"
28- bob = "did:plc:bob"
29- carol = "did:plc:carol"
30-
31-proc main() =
32- echo "Clustering probe"
33- echo ""
34-
35- # --- pure helpers -------------------------------------------------------
36- block:
37- report("decay at zero days is 1", abs(decay(0.0) - 1.0) < 1e-12)
38- # ln(2)/0.023 is about 30 days, so a month-old event is worth half.
39- report("decay halves in about a month",
40- abs(decay(30.1) - 0.5) < 0.01, &"{decay(30.1):.4f}")
41- report("decay keeps shrinking", decay(365.0) < decay(30.0))
42- report("decay never reaches zero", decay(3650.0) > 0.0,
43- &"{decay(3650.0):.3e}")
44-
45- block:
46- var s = @[0.0, 5.0, 10.0]
47- normalizeScores(s)
48- report("normalisation maps to 0..1", s == @[0.0, 0.5, 1.0], $s)
49-
50- # The case that would produce NaN: no spread to divide by.
51- var flat = @[3.0, 3.0, 3.0]
52- normalizeScores(flat)
53- report("an all-equal set is left alone", flat == @[3.0, 3.0, 3.0], $flat)
54- for v in flat:
55- report("and produces no NaN", not v.isNaN)
56-
57- var one = @[7.0]
58- normalizeScores(one)
59- report("a single score is left alone", one == @[7.0])
60-
61- var empty: seq[float] = @[]
62- normalizeScores(empty)
63- report("an empty set does not crash", empty.len == 0)
64-
65- var negatives = @[-10.0, -5.0, 0.0]
66- normalizeScores(negatives)
67- report("negative scores normalise too", negatives == @[0.0, 0.5, 1.0],
68- $negatives)
69-
70- block:
71- var rng = initRand(42)
72- var pool: seq[PersonRecommendation]
73- for i in 0 ..< 10:
74- pool.add PersonRecommendation(did: &"did:plc:in{i}", isFollowed: true)
75- for i in 0 ..< 10:
76- pool.add PersonRecommendation(did: &"did:plc:out{i}", isFollowed: false)
77-
78- let sampled = samplePeople(pool, 3, rng)
79- let followed = sampled.countIt(it.isFollowed)
80- let strangers = sampled.len - followed
81- report("sampling caps each side at the half", sampled.len == 6,
82- $sampled.len)
83- # Ranking purely by similarity returns the network the reader already
84- # has; reserving half the slots is what makes this discovery.
85- report("half the slots go to people not followed",
86- followed == 3 and strangers == 3,
87- &"{followed} followed, {strangers} strangers")
88-
89- # Fewer than `half` on one side must not pad from the other.
90- var lopsided: seq[PersonRecommendation]
91- lopsided.add PersonRecommendation(did: "did:plc:only", isFollowed: true)
92- for i in 0 ..< 10:
93- lopsided.add PersonRecommendation(did: &"did:plc:s{i}", isFollowed: false)
94- let small = samplePeople(lopsided, 3, rng)
95- report("a thin side is not padded from the other",
96- small.countIt(it.isFollowed) == 1 and small.len == 4, $small.len)
97-
98- # Shuffled, so the same faces do not recur every visit.
99- var rngA = initRand(1)
100- var rngB = initRand(2)
101- let a = samplePeople(pool, 3, rngA).mapIt(it.did).join(",")
102- let b = samplePeople(pool, 3, rngB).mapIt(it.did).join(",")
103- report("sampling varies between runs", a != b)
104-
105- block:
106- var recs = @[
107- FeedRecommendation(feedUrl: "c", score: 0.1),
108- FeedRecommendation(feedUrl: "a", score: 0.9),
109- FeedRecommendation(feedUrl: "b", score: 0.5),
110- ]
111- rankByScore(recs)
112- report("ranking puts the highest score first",
113- recs.mapIt(it.feedUrl) == @["a", "b", "c"],
114- recs.mapIt(it.feedUrl).join(","))
115-
116- # --- similarity against a real database ---------------------------------
117- let base = getTempDir() / "glean_cluster_probe"
118- cleanup(base)
119- defer: cleanup(base)
120-
121- var g = gleandb.open(base)
122- defer: g.close()
123- g.migrate()
124-
125- for did in [alice, bob, carol]:
126- g.db.run("INSERT INTO users (did) VALUES (?)", p(did))
127-
128- # alice and bob share two feeds; carol shares one with each.
129- for i in 1 .. 3:
130- g.upsertFeed(Feed(feedUrl: &"https://f{i}.test/feed", title: &"Feed {i}",
131- feedType: "rss"))
132- g.createSubscription(alice, "https://f1.test/feed", "", "", "", "")
133- g.createSubscription(alice, "https://f2.test/feed", "", "", "", "")
134- g.createSubscription(bob, "https://f1.test/feed", "", "", "", "")
135- g.createSubscription(bob, "https://f2.test/feed", "", "", "", "")
136- g.createSubscription(carol, "https://f1.test/feed", "", "", "", "")
137- g.createSubscription(carol, "https://f3.test/feed", "", "", "", "")
138-
139- block:
140- g.computeUserSimilarity()
141- let ab = g.userSimilarity(alice, bob)
142- let ac = g.userSimilarity(alice, carol)
143- # alice/bob share 2 of 2; alice/carol share 1 of 3.
144- report("identical subscribers score 1.0", abs(ab - 1.0) < 1e-9, &"{ab:.4f}")
145- report("partial overlap scores lower",
146- abs(ac - (1.0 / 3.0)) < 1e-9, &"{ac:.4f}")
147- report("similarity is symmetric on lookup",
148- abs(g.userSimilarity(bob, alice) - ab) < 1e-12)
149- report("an unrelated pair scores zero",
150- g.userSimilarity(alice, "did:plc:nobody") == 0.0)
151-
152- # One row per unordered pair, or every score would be counted twice
153- # downstream.
154- let rows = g.db.queryInt("SELECT count(*) FROM recs.user_similarity").get(0)
155- report("each pair is stored once", rows == 3, &"{rows} rows")
156-
157- block:
158- # Recomputing must replace rather than accumulate.
159- g.computeUserSimilarity()
160- let rows = g.db.queryInt("SELECT count(*) FROM recs.user_similarity").get(0)
161- report("recomputing replaces the previous generation", rows == 3,
162- &"{rows} rows")
163-
164- block:
165- g.computeFeedSimilarity()
166- let f12 = g.feedSimilarity("https://f1.test/feed", "https://f2.test/feed")
167- let f13 = g.feedSimilarity("https://f1.test/feed", "https://f3.test/feed")
168- # f1 and f2 share alice and bob; f1 and f3 share only carol.
169- report("feeds with more shared subscribers score higher", f12 > f13,
170- &"f1/f2={f12:.4f} f1/f3={f13:.4f}")
171- let reversed = g.feedSimilarity("https://f2.test/feed", "https://f1.test/feed")
172- report("feed similarity is symmetric on lookup", abs(reversed - f12) < 1e-12)
173- report("an unshared pair is absent",
174- g.feedSimilarity("https://f2.test/feed", "https://f3.test/feed") == 0.0)
175-
176- # --- decay in the query, not just in Nim --------------------------------
177- block:
178- # Two pairs of people who liked the same articles, one recently and one
179- # long ago. The old overlap must weigh less -- that is the whole reason
180- # EXP is registered as a SQL function.
181- g.upsertFeed(Feed(feedUrl: "https://n.test/feed", title: "News"))
182- let fresh = now().utc
183- let stale = now().utc - initDuration(days = 400)
184- for (who, at, n) in [(alice, fresh, "r"), (bob, fresh, "r"),
185- (carol, stale, "s")]:
186- g.createLike(Like(uri: &"at://{who}/like/{n}", authorDid: who,
187- feedUrl: "https://n.test/feed",
188- articleUrl: &"https://n.test/{n}",
189- createdAt: some(at)))
190- # carol needs a partner on the stale article for a pair to exist.
191- g.createLike(Like(uri: "at://dave/like/s", authorDid: "did:plc:dave",
192- feedUrl: "https://n.test/feed",
193- articleUrl: "https://n.test/s",
194- createdAt: some(stale)))
195-
196- g.computeUserSimilarity()
197- proc overlapOf(a, b: string): float =
198- let (x, y) = if a < b: (a, b) else: (b, a)
199- g.db.queryFirst(
200- "SELECT common_likes FROM recs.user_similarity WHERE user_a = ? AND user_b = ?",
201- [p(x), p(y)], proc(r: Row): float = r.f(0)).get(-1.0)
202-
203- let recent = overlapOf(alice, bob)
204- let old = overlapOf(carol, "did:plc:dave")
205- # The check that was vacuous before: both sides were 0 because the Go
206- # query truncates a decayed sum that is always just under 1.
207- report("a recent shared like counts as one", recent == 1.0, &"{recent}")
208- report("a 400-day-old shared like decays to nothing", old == 0.0, &"{old}")
209- report("decay makes old overlap weigh less than new", recent > old,
210- &"recent={recent} old={old}")
211-
212- echo ""
213- if failures == 0:
214- echo "RESULT: similarity and scoring behave."
215- else:
216- echo &"RESULT: {failures} check(s) failed."
217- quit 1
218-
219-when isMainModule:
220- main()
deleted spike/nim/dpop_probe.nim +0 -23
deleted file mode 100644
@@ -1,23 +0,0 @@
1-## Cross-checks atproto/dpop.nim against an independent Go verifier.
2-##
3-## nim c -r dpop_probe.nim && ./dpop_probe | (cd verify && go run .)
4-
5-import std/json
6-import bearssl/rand
7-import atproto/dpop
8-
9-when isMainModule:
10- let rngRef = HmacDrbgContext.new()
11- doAssert rngRef != nil, "no system randomness"
12- var rng = rngRef[]
13-
14- let key = generateKey(rng)
15- let jwt = key.dpopProof(rng, "POST", "https://bsky.social/oauth/token",
16- nonce = "spike-nonce")
17-
18- # Emitted for verify/verify.go to check independently.
19- echo $(%*{
20- "jwt": jwt,
21- "jwk": key.publicJwk,
22- "thumbprint": key.thumbprint,
23- })
deleted file mode 100644
@@ -1,23 +0,0 @@
1-## Cross-checks atproto/dpop.nim against an independent Go verifier.
2-##
3-## nim c -r dpop_probe.nim && ./dpop_probe | (cd verify && go run .)
4-
5-import std/json
6-import bearssl/rand
7-import atproto/dpop
8-
9-when isMainModule:
10- let rngRef = HmacDrbgContext.new()
11- doAssert rngRef != nil, "no system randomness"
12- var rng = rngRef[]
13-
14- let key = generateKey(rng)
15- let jwt = key.dpopProof(rng, "POST", "https://bsky.social/oauth/token",
16- nonce = "spike-nonce")
17-
18- # Emitted for verify/verify.go to check independently.
19- echo $(%*{
20- "jwt": jwt,
21- "jwk": key.publicJwk,
22- "thumbprint": key.thumbprint,
23- })
modified spike/nim/feedfetcher_probe.nim +1 -1
@@ -9,7 +9,7 @@
99 ## against a real host on demand.
1010
1111 import std/[strformat, strutils, tables]
12-import atproto/[feedfetcher, feedparser]
12+import reader/[feedfetcher, feedparser]
1313
1414 var failures = 0
1515
@@ -9,7 +9,7 @@
9 ## against a real host on demand.9 ## against a real host on demand.
10 10
11 import std/[strformat, strutils, tables]11 import std/[strformat, strutils, tables]
12-import atproto/[feedfetcher, feedparser]12+import reader/[feedfetcher, feedparser]
13 13
14 var failures = 014 var failures = 0
15 15
modified spike/nim/feedparser_probe.nim +1 -1
@@ -8,7 +8,7 @@
88 ## the wild are stranger than any fixture.
99
1010 import std/[httpclient, options, sequtils, strformat, strutils]
11-import atproto/feedparser
11+import reader/feedparser
1212
1313 var failures = 0
1414
@@ -8,7 +8,7 @@
8 ## the wild are stranger than any fixture.8 ## the wild are stranger than any fixture.
9 9
10 import std/[httpclient, options, sequtils, strformat, strutils]10 import std/[httpclient, options, sequtils, strformat, strutils]
11-import atproto/feedparser11+import reader/feedparser
12 12
13 var failures = 013 var failures = 0
14 14
modified spike/nim/glean.nim +1 -1
@@ -12,7 +12,7 @@
1212
1313 import std/[asyncdispatch, asynchttpserver, httpclient, json, options, os,
1414 strformat, strutils, uri]
15-import atproto/[articlestore, feedfetcher, feedparser, feedstore, gleandb,
15+import reader/[articlestore, feedfetcher, feedparser, feedstore, gleandb,
1616 scraper, sqlite]
1717
1818 const
@@ -12,7 +12,7 @@
12 12
13 import std/[asyncdispatch, asynchttpserver, httpclient, json, options, os,13 import std/[asyncdispatch, asynchttpserver, httpclient, json, options, os,
14 strformat, strutils, uri]14 strformat, strutils, uri]
15-import atproto/[articlestore, feedfetcher, feedparser, feedstore, gleandb,15+import reader/[articlestore, feedfetcher, feedparser, feedstore, gleandb,
16 scraper, sqlite]16 scraper, sqlite]
17 17
18 const18 const
renamed spike/nim/glean.nimble +2 -15
similarity index 77%
rename from spike/nim/glean_spike.nimble
rename to spike/nim/glean.nimble
@@ -1,34 +1,21 @@
11 version = "0.1.0"
22 author = "glean"
3-description = "Nim port spike for glean: atproto, storage and the database layer"
3+description = "glean: a small RSS reader"
44 license = "MIT"
55 srcDir = "."
66
77 requires "nim >= 2.2.0"
8-requires "bearssl"
9-requires "nimcrypto"
10-requires "chronos"
11-requires "websock"
128
139 # Probes that touch only this machine. Safe to run anywhere, including CI.
1410 const offlineProbes = @[
15- "sqlite_probe",
16- "dpop_probe",
17- "store_probe",
1811 "gleandb_probe",
1912 "stores_probe",
20- "social_probe",
21- "cluster_probe",
22- "reconnect_probe", # uses a local websocket server, not the network
2313 ]
2414
2515 # Probes that talk to real servers. Kept separate because a failure here can
2616 # mean "bsky.social is having a bad day" rather than "the port is broken", and
2717 # a test suite that cannot tell those apart stops being believed.
2818 const networkProbes = @[
29- "oauth_probe",
30- "xrpc_probe",
31- "jetstream_probe",
3219 "feedparser_probe", # fixtures are offline, but it also fetches real feeds
3320 "feedfetcher_probe", # retry policy is scripted; one check hits the network
3421 "scraper_probe", # fixtures are offline; one check fetches a real article
@@ -58,5 +45,5 @@ task build, "Compile every module without running anything":
5845 exec "nim c --hints:off -d:ssl -c " & probe & ".nim"
5946 echo "\nEverything compiles."
6047
61-task schema, "Regenerate atproto/schema.nim from the Go source":
48+task schema, "Regenerate reader/schema.nim from the Go source":
6249 exec "python3 tools/extract_schema.py"
similarity index 77%
rename from spike/nim/glean_spike.nimble
rename to spike/nim/glean.nimble
@@ -1,34 +1,21 @@
1 version = "0.1.0"1 version = "0.1.0"
2 author = "glean"2 author = "glean"
3-description = "Nim port spike for glean: atproto, storage and the database layer"3+description = "glean: a small RSS reader"
4 license = "MIT"4 license = "MIT"
5 srcDir = "."5 srcDir = "."
6 6
7 requires "nim >= 2.2.0"7 requires "nim >= 2.2.0"
8-requires "bearssl"
9-requires "nimcrypto"
10-requires "chronos"
11-requires "websock"
12 8
13 # Probes that touch only this machine. Safe to run anywhere, including CI.9 # Probes that touch only this machine. Safe to run anywhere, including CI.
14 const offlineProbes = @[10 const offlineProbes = @[
15- "sqlite_probe",
16- "dpop_probe",
17- "store_probe",
18 "gleandb_probe",11 "gleandb_probe",
19 "stores_probe",12 "stores_probe",
20- "social_probe",
21- "cluster_probe",
22- "reconnect_probe", # uses a local websocket server, not the network
23 ]13 ]
24 14
25 # Probes that talk to real servers. Kept separate because a failure here can15 # Probes that talk to real servers. Kept separate because a failure here can
26 # mean "bsky.social is having a bad day" rather than "the port is broken", and16 # mean "bsky.social is having a bad day" rather than "the port is broken", and
27 # a test suite that cannot tell those apart stops being believed.17 # a test suite that cannot tell those apart stops being believed.
28 const networkProbes = @[18 const networkProbes = @[
29- "oauth_probe",
30- "xrpc_probe",
31- "jetstream_probe",
32 "feedparser_probe", # fixtures are offline, but it also fetches real feeds19 "feedparser_probe", # fixtures are offline, but it also fetches real feeds
33 "feedfetcher_probe", # retry policy is scripted; one check hits the network20 "feedfetcher_probe", # retry policy is scripted; one check hits the network
34 "scraper_probe", # fixtures are offline; one check fetches a real article21 "scraper_probe", # fixtures are offline; one check fetches a real article
@@ -58,5 +45,5 @@ task build, "Compile every module without running anything":
58 exec "nim c --hints:off -d:ssl -c " & probe & ".nim"45 exec "nim c --hints:off -d:ssl -c " & probe & ".nim"
59 echo "\nEverything compiles."46 echo "\nEverything compiles."
60 47
61-task schema, "Regenerate atproto/schema.nim from the Go source":48+task schema, "Regenerate reader/schema.nim from the Go source":
62 exec "python3 tools/extract_schema.py"49 exec "python3 tools/extract_schema.py"
modified spike/nim/gleandb_probe.nim +14 -40
@@ -9,7 +9,7 @@
99 ## degrades or diverges later.
1010
1111 import std/[options, os, strformat, strutils]
12-import atproto/[gleandb, schema, sqlite]
12+import reader/[gleandb, schema, sqlite]
1313
1414 var failures = 0
1515
@@ -20,7 +20,7 @@ proc report(name: string, ok: bool, detail = "") =
2020 else: echo &" {label} {name}"
2121
2222 proc cleanup(base: string) =
23- for suffix in ["_users", "_articles", "_recs"]:
23+ for suffix in ["_users", "_articles"]:
2424 for ext in ["", "-wal", "-shm"]:
2525 removeFile(base & suffix & ext)
2626
@@ -35,43 +35,26 @@ proc main() =
3535 var g = gleandb.open(base)
3636 defer: g.close()
3737
38- report("schema statements extracted from the Go source",
39- usersSchema.len + articlesSchema.len + recsSchema.len == 44,
40- &"{usersSchema.len} + {articlesSchema.len} + {recsSchema.len}")
38+ report("the schema covers the reader's tables", readerSchema.len == 15,
39+ &"{readerSchema.len} statements")
4140
4241 # --- the three databases ------------------------------------------------
4342 block:
4443 let schemas = g.attachedSchemas()
45- report("all three databases are attached",
46- "main" in schemas and "articles" in schemas and "recs" in schemas,
47- schemas.join(", "))
44+ report("both databases are attached",
45+ "main" in schemas and "articles" in schemas, schemas.join(", "))
4846 report("each is a separate file on disk",
49- fileExists(base & "_users") and fileExists(base & "_articles") and
50- fileExists(base & "_recs"))
47+ fileExists(base & "_users") and fileExists(base & "_articles"))
5148
5249 # --- pragmas ------------------------------------------------------------
5350 # WAL has to be set per schema: PRAGMA applies to one at a time, so setting
5451 # it on main leaves the others journalling the slow way.
5552 block:
56- for name in ["", "articles", "recs"]:
53+ for name in ["", "articles"]:
5754 let mode = g.journalMode(name).toLowerAscii
5855 report(&"WAL is on for {(if name.len == 0: \"main\" else: name)}",
5956 mode == "wal", mode)
6057
61- # --- custom SQL functions ----------------------------------------------
62- block:
63- let e = g.db.queryText("SELECT exp(1.0)")
64- report("exp() is callable from SQL",
65- e.isSome and e.get.startsWith("2.718"), e.get(""))
66- let l = g.db.queryText("SELECT log(2.718281828459045)")
67- report("log() is callable from SQL",
68- l.isSome and l.get.startsWith("1.0") or l.get("") == "1.0", l.get(""))
69- # log(0) is -inf; returning NULL keeps an ORDER BY sane.
70- report("log() of zero is NULL",
71- g.db.queryText("SELECT log(0)").isNone)
72- report("log() of a negative is NULL",
73- g.db.queryText("SELECT log(-5)").isNone)
74-
7558 # --- schema -------------------------------------------------------------
7659 g.migrate()
7760 report("migrate() is idempotent", (g.migrate(); true))
@@ -81,10 +64,12 @@ proc main() =
8164 let n = g.db.queryInt(
8265 &"SELECT count(*) FROM {schemaName}.sqlite_master WHERE type='table'")
8366 n.get(0).int
84- report("users tables created", tableCount("main") >= 7, $tableCount("main"))
85- report("articles tables created", tableCount("articles") >= 7,
86- $tableCount("articles"))
87- report("recs tables created", tableCount("recs") >= 7, $tableCount("recs"))
67+ # main holds only the local user now; everything that grows lives in
68+ # the articles database.
69+ report("the users database holds just the user table",
70+ tableCount("main") == 1, $tableCount("main"))
71+ report("articles holds feeds, subscriptions, articles and read state",
72+ tableCount("articles") >= 5, $tableCount("articles"))
8873
8974 # --- a real cross-database query ---------------------------------------
9075 # The whole point of attaching rather than opening three connections.
@@ -139,17 +124,6 @@ proc main() =
139124 if afterDelete.isNone: "row gone from the index"
140125 else: "STILL INDEXED: " & afterDelete.get)
141126
142- # --- vec0 ---------------------------------------------------------------
143- block:
144- g.createEmbeddingTables(dimension = 4)
145- report("embedding tables take their dimension from config",
146- g.db.queryInt(
147- "SELECT count(*) FROM recs.sqlite_master WHERE name LIKE '%embeddings%'"
148- ).get(0) > 0)
149- # Proves sqlite-vec is registered on this connection, not just linked.
150- let ver = g.db.queryText("SELECT vec_version()")
151- report("sqlite-vec is live on the glean connection", ver.isSome, ver.get(""))
152-
153127 echo ""
154128 if failures == 0:
155129 echo "RESULT: the database layer matches Glean's layout and behaviour."
@@ -9,7 +9,7 @@
9 ## degrades or diverges later.9 ## degrades or diverges later.
10 10
11 import std/[options, os, strformat, strutils]11 import std/[options, os, strformat, strutils]
12-import atproto/[gleandb, schema, sqlite]12+import reader/[gleandb, schema, sqlite]
13 13
14 var failures = 014 var failures = 0
15 15
@@ -20,7 +20,7 @@ proc report(name: string, ok: bool, detail = "") =
20 else: echo &" {label} {name}"20 else: echo &" {label} {name}"
21 21
22 proc cleanup(base: string) =22 proc cleanup(base: string) =
23- for suffix in ["_users", "_articles", "_recs"]:23+ for suffix in ["_users", "_articles"]:
24 for ext in ["", "-wal", "-shm"]:24 for ext in ["", "-wal", "-shm"]:
25 removeFile(base & suffix & ext)25 removeFile(base & suffix & ext)
26 26
@@ -35,43 +35,26 @@ proc main() =
35 var g = gleandb.open(base)35 var g = gleandb.open(base)
36 defer: g.close()36 defer: g.close()
37 37
38- report("schema statements extracted from the Go source",38+ report("the schema covers the reader's tables", readerSchema.len == 15,
39- usersSchema.len + articlesSchema.len + recsSchema.len == 44,39+ &"{readerSchema.len} statements")
40- &"{usersSchema.len} + {articlesSchema.len} + {recsSchema.len}")
41 40
42 # --- the three databases ------------------------------------------------41 # --- the three databases ------------------------------------------------
43 block:42 block:
44 let schemas = g.attachedSchemas()43 let schemas = g.attachedSchemas()
45- report("all three databases are attached",44+ report("both databases are attached",
46- "main" in schemas and "articles" in schemas and "recs" in schemas,45+ "main" in schemas and "articles" in schemas, schemas.join(", "))
47- schemas.join(", "))
48 report("each is a separate file on disk",46 report("each is a separate file on disk",
49- fileExists(base & "_users") and fileExists(base & "_articles") and47+ fileExists(base & "_users") and fileExists(base & "_articles"))
50- fileExists(base & "_recs"))
51 48
52 # --- pragmas ------------------------------------------------------------49 # --- pragmas ------------------------------------------------------------
53 # WAL has to be set per schema: PRAGMA applies to one at a time, so setting50 # WAL has to be set per schema: PRAGMA applies to one at a time, so setting
54 # it on main leaves the others journalling the slow way.51 # it on main leaves the others journalling the slow way.
55 block:52 block:
56- for name in ["", "articles", "recs"]:53+ for name in ["", "articles"]:
57 let mode = g.journalMode(name).toLowerAscii54 let mode = g.journalMode(name).toLowerAscii
58 report(&"WAL is on for {(if name.len == 0: \"main\" else: name)}",55 report(&"WAL is on for {(if name.len == 0: \"main\" else: name)}",
59 mode == "wal", mode)56 mode == "wal", mode)
60 57
61- # --- custom SQL functions ----------------------------------------------
62- block:
63- let e = g.db.queryText("SELECT exp(1.0)")
64- report("exp() is callable from SQL",
65- e.isSome and e.get.startsWith("2.718"), e.get(""))
66- let l = g.db.queryText("SELECT log(2.718281828459045)")
67- report("log() is callable from SQL",
68- l.isSome and l.get.startsWith("1.0") or l.get("") == "1.0", l.get(""))
69- # log(0) is -inf; returning NULL keeps an ORDER BY sane.
70- report("log() of zero is NULL",
71- g.db.queryText("SELECT log(0)").isNone)
72- report("log() of a negative is NULL",
73- g.db.queryText("SELECT log(-5)").isNone)
74-
75 # --- schema -------------------------------------------------------------58 # --- schema -------------------------------------------------------------
76 g.migrate()59 g.migrate()
77 report("migrate() is idempotent", (g.migrate(); true))60 report("migrate() is idempotent", (g.migrate(); true))
@@ -81,10 +64,12 @@ proc main() =
81 let n = g.db.queryInt(64 let n = g.db.queryInt(
82 &"SELECT count(*) FROM {schemaName}.sqlite_master WHERE type='table'")65 &"SELECT count(*) FROM {schemaName}.sqlite_master WHERE type='table'")
83 n.get(0).int66 n.get(0).int
84- report("users tables created", tableCount("main") >= 7, $tableCount("main"))67+ # main holds only the local user now; everything that grows lives in
85- report("articles tables created", tableCount("articles") >= 7,68+ # the articles database.
86- $tableCount("articles"))69+ report("the users database holds just the user table",
87- report("recs tables created", tableCount("recs") >= 7, $tableCount("recs"))70+ tableCount("main") == 1, $tableCount("main"))
71+ report("articles holds feeds, subscriptions, articles and read state",
72+ tableCount("articles") >= 5, $tableCount("articles"))
88 73
89 # --- a real cross-database query ---------------------------------------74 # --- a real cross-database query ---------------------------------------
90 # The whole point of attaching rather than opening three connections.75 # The whole point of attaching rather than opening three connections.
@@ -139,17 +124,6 @@ proc main() =
139 if afterDelete.isNone: "row gone from the index"124 if afterDelete.isNone: "row gone from the index"
140 else: "STILL INDEXED: " & afterDelete.get)125 else: "STILL INDEXED: " & afterDelete.get)
141 126
142- # --- vec0 ---------------------------------------------------------------
143- block:
144- g.createEmbeddingTables(dimension = 4)
145- report("embedding tables take their dimension from config",
146- g.db.queryInt(
147- "SELECT count(*) FROM recs.sqlite_master WHERE name LIKE '%embeddings%'"
148- ).get(0) > 0)
149- # Proves sqlite-vec is registered on this connection, not just linked.
150- let ver = g.db.queryText("SELECT vec_version()")
151- report("sqlite-vec is live on the glean connection", ver.isSome, ver.get(""))
152-
153 echo ""127 echo ""
154 if failures == 0:128 if failures == 0:
155 echo "RESULT: the database layer matches Glean's layout and behaviour."129 echo "RESULT: the database layer matches Glean's layout and behaviour."
deleted spike/nim/jetstream_probe.nim +0 -150
deleted file mode 100644
@@ -1,150 +0,0 @@
1-## Reads the live ATProto firehose.
2-##
3-## nim c -r -d:ssl jetstream_probe.nim
4-##
5-## Jetstream is public, so unlike the OAuth and XRPC probes this one can check
6-## the real thing end to end. The point of interest is server-side filtering:
7-## an unfiltered subscription is the whole network, so a consumer that loses
8-## its query parameters still receives a healthy-looking torrent of events
9-## that are simply the wrong ones.
10-
11-import std/[json, sequtils, sets, strformat, strutils, tables]
12-import pkg/chronos
13-import atproto/jetstream
14-
15-var failures = 0
16-
17-proc report(name: string, ok: bool, detail = "") =
18- if not ok: inc failures
19- let label = if ok: "PASS" else: "FAIL"
20- if detail.len > 0: echo &" {label} {name} -- {detail}"
21- else: echo &" {label} {name}"
22-
23-proc main() {.async.} =
24- echo "Jetstream probe"
25- echo ""
26-
27- # --- URL construction ---------------------------------------------------
28- block:
29- let cfg = JetstreamConfig(
30- url: DefaultJetstream,
31- wantedCollections: @["app.bsky.feed.post", "app.bsky.feed.like"],
32- cursor: 123,
33- )
34- let u = subscribeUrl(cfg)
35- report("adds /subscribe", "/subscribe" in u)
36- # Repeated keys, not one comma-joined value: Jetstream reads these as a
37- # multi-valued parameter and a joined string matches no collection.
38- report("repeats wantedCollections per value",
39- u.count("wantedCollections=") == 2, u.split('?')[^1])
40- report("carries the cursor", "cursor=123" in u)
41-
42- block:
43- let cfg = JetstreamConfig(url: "wss://example.test/subscribe")
44- report("does not double up /subscribe",
45- subscribeUrl(cfg).count("/subscribe") == 1)
46-
47- # --- parsing ------------------------------------------------------------
48- block:
49- let raw = """{"did":"did:plc:abc","time_us":1720000000000000,
50- "kind":"commit","commit":{"rev":"3k","operation":"create",
51- "collection":"app.bsky.feed.post","rkey":"xyz","cid":"bafy",
52- "record":{"text":"hi","$type":"app.bsky.feed.post"}}}"""
53- let ev = parseEvent(raw)
54- report("parses a commit event", ev.kind == ekCommit and ev.did == "did:plc:abc")
55- report("parses the operation", ev.commit.operation == coCreate)
56- report("parses the record", ev.commit.record{"text"}.getStr == "hi")
57- report("time_us survives as int64", ev.timeUs == 1720000000000000'i64,
58- $ev.timeUs)
59-
60- block:
61- # A delete carries no record; reading one must not crash a consumer.
62- let ev = parseEvent("""{"did":"did:plc:a","time_us":1,"kind":"commit",
63- "commit":{"operation":"delete","collection":"c","rkey":"r"}}""")
64- report("a delete has no record",
65- ev.commit.operation == coDelete and
66- (ev.commit.record == nil or ev.commit.record.kind == JNull))
67-
68- # --- the live firehose --------------------------------------------------
69- const wanted = "app.bsky.feed.post"
70- var events: seq[Event]
71- try:
72- let cfg = JetstreamConfig(
73- url: DefaultJetstream,
74- wantedCollections: @[wanted],
75- )
76- events = await collectEvents(cfg, count = 25, timeout = 30.seconds)
77- except CatchableError as e:
78- report("connected to the live firehose", false, e.msg)
79- echo ""
80- echo "RESULT: could not reach Jetstream."
81- quit 1
82-
83- report("received live events", events.len > 0, &"{events.len} events")
84- if events.len == 0:
85- echo ""
86- echo "RESULT: no events received."
87- quit 1
88-
89- report("every event has a DID", events.allIt(it.did.startsWith("did:")))
90- report("cursors advance",
91- events[^1].timeUs >= events[0].timeUs,
92- &"{events[0].timeUs} -> {events[^1].timeUs}")
93-
94- # The real check: server-side filtering actually applied. Without it this
95- # would be a mix of likes, follows, reposts and everything else.
96- var collections: CountTable[string]
97- for ev in events:
98- if ev.kind == ekCommit:
99- collections.inc ev.commit.collection
100- let others = toSeq(collections.keys).filterIt(it != wanted and it.len > 0)
101- report("filter applied server-side (only " & wanted & ")",
102- others.len == 0,
103- if others.len == 0: $collections.len & " collection"
104- else: "also saw " & others.join(", "))
105-
106- let posts = events.filterIt(
107- it.kind == ekCommit and it.commit.collection == wanted and
108- it.commit.operation == coCreate and it.commit.record != nil)
109- report("create events carry a typed record",
110- posts.len == 0 or posts[0].commit.record{"$type"}.getStr == wanted,
111- if posts.len > 0: posts[0].commit.record{"$type"}.getStr else: "none seen")
112-
113- # --- cursor replay ------------------------------------------------------
114- # This is how a consumer catches up after downtime, so "the cursor is
115- # accepted" is not enough -- it has to actually rewind. Replaying from an
116- # earlier cursor must return events at or after it, including ones already
117- # seen live.
118- block:
119- let rewindTo = events[0].timeUs
120- let seenLive = events.mapIt(it.timeUs).toHashSet
121- try:
122- let cfg = JetstreamConfig(
123- url: DefaultJetstream,
124- wantedCollections: @[wanted],
125- cursor: rewindTo,
126- )
127- let replayed = await collectEvents(cfg, count = 25, timeout = 30.seconds)
128- report("replay from a cursor returns events", replayed.len > 0,
129- &"{replayed.len} events")
130- if replayed.len > 0:
131- report("replayed events start at or after the cursor",
132- replayed[0].timeUs >= rewindTo,
133- &"cursor {rewindTo} -> first {replayed[0].timeUs}")
134- # Proof it rewound rather than resuming live: the replay should
135- # re-deliver timestamps from the first batch.
136- let overlap = replayed.countIt(it.timeUs in seenLive)
137- report("replay re-delivers already-seen events (a real rewind)",
138- overlap > 0, &"{overlap} of {replayed.len} overlap")
139- except CatchableError as e:
140- report("replay from a cursor returns events", false, e.msg)
141-
142- echo ""
143- if failures == 0:
144- echo "RESULT: Jetstream works, with server-side filtering and replay confirmed."
145- else:
146- echo &"RESULT: {failures} check(s) failed."
147- quit 1
148-
149-when isMainModule:
150- waitFor main()
deleted file mode 100644
@@ -1,150 +0,0 @@
1-## Reads the live ATProto firehose.
2-##
3-## nim c -r -d:ssl jetstream_probe.nim
4-##
5-## Jetstream is public, so unlike the OAuth and XRPC probes this one can check
6-## the real thing end to end. The point of interest is server-side filtering:
7-## an unfiltered subscription is the whole network, so a consumer that loses
8-## its query parameters still receives a healthy-looking torrent of events
9-## that are simply the wrong ones.
10-
11-import std/[json, sequtils, sets, strformat, strutils, tables]
12-import pkg/chronos
13-import atproto/jetstream
14-
15-var failures = 0
16-
17-proc report(name: string, ok: bool, detail = "") =
18- if not ok: inc failures
19- let label = if ok: "PASS" else: "FAIL"
20- if detail.len > 0: echo &" {label} {name} -- {detail}"
21- else: echo &" {label} {name}"
22-
23-proc main() {.async.} =
24- echo "Jetstream probe"
25- echo ""
26-
27- # --- URL construction ---------------------------------------------------
28- block:
29- let cfg = JetstreamConfig(
30- url: DefaultJetstream,
31- wantedCollections: @["app.bsky.feed.post", "app.bsky.feed.like"],
32- cursor: 123,
33- )
34- let u = subscribeUrl(cfg)
35- report("adds /subscribe", "/subscribe" in u)
36- # Repeated keys, not one comma-joined value: Jetstream reads these as a
37- # multi-valued parameter and a joined string matches no collection.
38- report("repeats wantedCollections per value",
39- u.count("wantedCollections=") == 2, u.split('?')[^1])
40- report("carries the cursor", "cursor=123" in u)
41-
42- block:
43- let cfg = JetstreamConfig(url: "wss://example.test/subscribe")
44- report("does not double up /subscribe",
45- subscribeUrl(cfg).count("/subscribe") == 1)
46-
47- # --- parsing ------------------------------------------------------------
48- block:
49- let raw = """{"did":"did:plc:abc","time_us":1720000000000000,
50- "kind":"commit","commit":{"rev":"3k","operation":"create",
51- "collection":"app.bsky.feed.post","rkey":"xyz","cid":"bafy",
52- "record":{"text":"hi","$type":"app.bsky.feed.post"}}}"""
53- let ev = parseEvent(raw)
54- report("parses a commit event", ev.kind == ekCommit and ev.did == "did:plc:abc")
55- report("parses the operation", ev.commit.operation == coCreate)
56- report("parses the record", ev.commit.record{"text"}.getStr == "hi")
57- report("time_us survives as int64", ev.timeUs == 1720000000000000'i64,
58- $ev.timeUs)
59-
60- block:
61- # A delete carries no record; reading one must not crash a consumer.
62- let ev = parseEvent("""{"did":"did:plc:a","time_us":1,"kind":"commit",
63- "commit":{"operation":"delete","collection":"c","rkey":"r"}}""")
64- report("a delete has no record",
65- ev.commit.operation == coDelete and
66- (ev.commit.record == nil or ev.commit.record.kind == JNull))
67-
68- # --- the live firehose --------------------------------------------------
69- const wanted = "app.bsky.feed.post"
70- var events: seq[Event]
71- try:
72- let cfg = JetstreamConfig(
73- url: DefaultJetstream,
74- wantedCollections: @[wanted],
75- )
76- events = await collectEvents(cfg, count = 25, timeout = 30.seconds)
77- except CatchableError as e:
78- report("connected to the live firehose", false, e.msg)
79- echo ""
80- echo "RESULT: could not reach Jetstream."
81- quit 1
82-
83- report("received live events", events.len > 0, &"{events.len} events")
84- if events.len == 0:
85- echo ""
86- echo "RESULT: no events received."
87- quit 1
88-
89- report("every event has a DID", events.allIt(it.did.startsWith("did:")))
90- report("cursors advance",
91- events[^1].timeUs >= events[0].timeUs,
92- &"{events[0].timeUs} -> {events[^1].timeUs}")
93-
94- # The real check: server-side filtering actually applied. Without it this
95- # would be a mix of likes, follows, reposts and everything else.
96- var collections: CountTable[string]
97- for ev in events:
98- if ev.kind == ekCommit:
99- collections.inc ev.commit.collection
100- let others = toSeq(collections.keys).filterIt(it != wanted and it.len > 0)
101- report("filter applied server-side (only " & wanted & ")",
102- others.len == 0,
103- if others.len == 0: $collections.len & " collection"
104- else: "also saw " & others.join(", "))
105-
106- let posts = events.filterIt(
107- it.kind == ekCommit and it.commit.collection == wanted and
108- it.commit.operation == coCreate and it.commit.record != nil)
109- report("create events carry a typed record",
110- posts.len == 0 or posts[0].commit.record{"$type"}.getStr == wanted,
111- if posts.len > 0: posts[0].commit.record{"$type"}.getStr else: "none seen")
112-
113- # --- cursor replay ------------------------------------------------------
114- # This is how a consumer catches up after downtime, so "the cursor is
115- # accepted" is not enough -- it has to actually rewind. Replaying from an
116- # earlier cursor must return events at or after it, including ones already
117- # seen live.
118- block:
119- let rewindTo = events[0].timeUs
120- let seenLive = events.mapIt(it.timeUs).toHashSet
121- try:
122- let cfg = JetstreamConfig(
123- url: DefaultJetstream,
124- wantedCollections: @[wanted],
125- cursor: rewindTo,
126- )
127- let replayed = await collectEvents(cfg, count = 25, timeout = 30.seconds)
128- report("replay from a cursor returns events", replayed.len > 0,
129- &"{replayed.len} events")
130- if replayed.len > 0:
131- report("replayed events start at or after the cursor",
132- replayed[0].timeUs >= rewindTo,
133- &"cursor {rewindTo} -> first {replayed[0].timeUs}")
134- # Proof it rewound rather than resuming live: the replay should
135- # re-deliver timestamps from the first batch.
136- let overlap = replayed.countIt(it.timeUs in seenLive)
137- report("replay re-delivers already-seen events (a real rewind)",
138- overlap > 0, &"{overlap} of {replayed.len} overlap")
139- except CatchableError as e:
140- report("replay from a cursor returns events", false, e.msg)
141-
142- echo ""
143- if failures == 0:
144- echo "RESULT: Jetstream works, with server-side filtering and replay confirmed."
145- else:
146- echo &"RESULT: {failures} check(s) failed."
147- quit 1
148-
149-when isMainModule:
150- waitFor main()
modified spike/nim/nim.cfg +0 -10
@@ -1,12 +1,2 @@
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"
71 --passL:"-lsqlite3 -lm"
82
9-# websock pulls in chronicles, which will not compile unless a logging stream
10-# is configured. This spike has no use for its output, and leaving it on makes
11-# every consumer of websock inherit the requirement.
12---define:"chronicles_enabled=off"
@@ -1,12 +1,2 @@
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"1 --passL:"-lsqlite3 -lm"
8 2
9-# websock pulls in chronicles, which will not compile unless a logging stream
10-# is configured. This spike has no use for its output, and leaving it on makes
11-# every consumer of websock inherit the requirement.
12---define:"chronicles_enabled=off"
deleted spike/nim/oauth_probe.nim +0 -140
deleted file mode 100644
@@ -1,140 +0,0 @@
1-## Drives the OAuth flow against a real PDS, as far as it can go without a
2-## human at a browser.
3-##
4-## nim c -r -d:ssl oauth_probe.nim [handle]
5-##
6-## Everything here is unauthenticated and read-only except the PAR, which
7-## creates a short-lived pending authorization the server discards on its own.
8-## Getting a request_uri back is the meaningful result: it means the DPoP
9-## proof was accepted, the nonce retry worked, and PKCE was well-formed.
10-
11-import std/[httpclient, json, os, strformat, strutils, uri]
12-import bearssl/rand
13-import atproto/[dpop, identity, oauth]
14-
15-var failures = 0
16-
17-proc report(name: string, ok: bool, detail = "") =
18- if not ok: inc failures
19- let label = if ok: "PASS" else: "FAIL"
20- if detail.len > 0:
21- echo &" {label} {name} -- {detail}"
22- else:
23- echo &" {label} {name}"
24-
25-when isMainModule:
26- let handle = if paramCount() >= 1: paramStr(1) else: "bsky.app"
27- echo "ATProto OAuth flow probe"
28- echo &" identifier: {handle}"
29- echo ""
30-
31- let rngRef = HmacDrbgContext.new()
32- doAssert rngRef != nil, "no system randomness"
33- var rng = rngRef[]
34-
35- # --- resolution ---------------------------------------------------------
36- var ident: Identity
37- try:
38- ident = lookup(handle)
39- report("handle resolves to a DID", ident.did.isDid, ident.did)
40- report("identity links to a PDS", ident.pdsEndpoint.len > 0, ident.pdsEndpoint)
41- except CatchableError as e:
42- report("handle resolves", false, e.msg)
43- quit 1
44-
45- # --- discovery ----------------------------------------------------------
46- var authServerUrl: string
47- try:
48- authServerUrl = resolveAuthServerUrl(ident.pdsEndpoint)
49- report("PDS advertises an auth server", authServerUrl.len > 0, authServerUrl)
50- except CatchableError as e:
51- report("PDS advertises an auth server", false, e.msg)
52- quit 1
53-
54- var meta: AuthServerMetadata
55- try:
56- meta = resolveAuthServerMetadata(authServerUrl)
57- report("auth server metadata fetched", true, meta.issuer)
58- report("has a PAR endpoint", meta.parEndpoint.len > 0, meta.parEndpoint)
59- report("has a token endpoint", meta.tokenEndpoint.len > 0, meta.tokenEndpoint)
60- except CatchableError as e:
61- report("auth server metadata fetched", false, e.msg)
62- quit 1
63-
64- # --- PKCE ---------------------------------------------------------------
65- block:
66- # RFC 7636 test vector: this verifier must produce this challenge.
67- const
68- verifier = "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
69- want = "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"
70- report("PKCE S256 matches RFC 7636 vector",
71- s256Challenge(verifier) == want, s256Challenge(verifier))
72-
73- # --- the nonce dance ----------------------------------------------------
74- # sendAuthRequest retries transparently, which would hide a server that
75- # never demanded a nonce. Prove the first unnonced attempt really is
76- # rejected, so the retry path is load-bearing rather than decorative.
77- let cfg0 = newLocalhostConfig("http://127.0.0.1:8080/callback")
78- block:
79- let
80- key = generateKey(rng)
81- proof = dpopProof(key, rng, "POST", meta.parEndpoint)
82- client = newHttpClient(userAgent = UserAgent, timeout = 15_000)
83- defer: client.close()
84- client.headers = newHttpHeaders({
85- "Content-Type": "application/x-www-form-urlencoded",
86- "DPoP": proof,
87- })
88- # The body must be fully valid: bsky.social validates parameters before it
89- # checks the nonce, so a stub body returns invalid_request and tells us
90- # nothing about the nonce requirement.
91- let verifier = secureRandomB64(48)
92- let body = @{
93- "client_id": cfg0.clientId,
94- "state": secureRandomB64(16),
95- "redirect_uri": cfg0.callbackUrl,
96- "scope": scopeStr(cfg0.scopes),
97- "response_type": "code",
98- "code_challenge": s256Challenge(verifier),
99- "code_challenge_method": "S256",
100- }.encodeQuery
101- let res = client.post(meta.parEndpoint, body = body)
102- let code = try: parseJson(res.body){"error"}.getStr except CatchableError: ""
103- let nonce = res.headers.getOrDefault("DPoP-Nonce").string
104- report("first PAR without a nonce is rejected",
105- res.code.int == 400 and code == "use_dpop_nonce", &"HTTP {res.code.int} {code}")
106- report("rejection supplies a DPoP-Nonce to retry with", nonce.len > 0)
107-
108- # --- PAR ----------------------------------------------------------------
109- # A localhost client needs no hosted metadata document, so the flow can be
110- # exercised without deploying anything.
111- let cfg = cfg0
112- report("localhost client_id is well-formed",
113- cfg.clientId.startsWith("http://localhost?"), cfg.clientId)
114-
115- try:
116- let info = sendAuthRequest(cfg, meta, rng, loginHint = handle)
117- report("PAR accepted (DPoP proof + nonce retry + PKCE)",
118- info.requestUri.startsWith("urn:ietf:params:oauth:request_uri:"),
119- info.requestUri)
120- report("auth server issued a DPoP nonce", info.dpopNonce.len > 0,
121- info.dpopNonce[0 ..< min(12, info.dpopNonce.len)] & "...")
122- report("PKCE verifier retained for callback", info.pkceVerifier.len >= 43)
123- report("state is unguessable", info.state.len >= 21, &"{info.state.len} chars")
124-
125- let url = authorizeUrl(cfg, meta, info)
126- let q = parseUri(url).query
127- report("authorize URL carries request_uri", "request_uri=" in q)
128- report("authorize URL carries client_id", "client_id=" in q)
129- echo ""
130- echo " authorize URL:"
131- echo " ", url
132- except CatchableError as e:
133- report("PAR accepted", false, e.msg)
134-
135- echo ""
136- if failures == 0:
137- echo "RESULT: the OAuth flow works against a live PDS up to user consent."
138- else:
139- echo &"RESULT: {failures} check(s) failed."
140- quit 1
deleted file mode 100644
@@ -1,140 +0,0 @@
1-## Drives the OAuth flow against a real PDS, as far as it can go without a
2-## human at a browser.
3-##
4-## nim c -r -d:ssl oauth_probe.nim [handle]
5-##
6-## Everything here is unauthenticated and read-only except the PAR, which
7-## creates a short-lived pending authorization the server discards on its own.
8-## Getting a request_uri back is the meaningful result: it means the DPoP
9-## proof was accepted, the nonce retry worked, and PKCE was well-formed.
10-
11-import std/[httpclient, json, os, strformat, strutils, uri]
12-import bearssl/rand
13-import atproto/[dpop, identity, oauth]
14-
15-var failures = 0
16-
17-proc report(name: string, ok: bool, detail = "") =
18- if not ok: inc failures
19- let label = if ok: "PASS" else: "FAIL"
20- if detail.len > 0:
21- echo &" {label} {name} -- {detail}"
22- else:
23- echo &" {label} {name}"
24-
25-when isMainModule:
26- let handle = if paramCount() >= 1: paramStr(1) else: "bsky.app"
27- echo "ATProto OAuth flow probe"
28- echo &" identifier: {handle}"
29- echo ""
30-
31- let rngRef = HmacDrbgContext.new()
32- doAssert rngRef != nil, "no system randomness"
33- var rng = rngRef[]
34-
35- # --- resolution ---------------------------------------------------------
36- var ident: Identity
37- try:
38- ident = lookup(handle)
39- report("handle resolves to a DID", ident.did.isDid, ident.did)
40- report("identity links to a PDS", ident.pdsEndpoint.len > 0, ident.pdsEndpoint)
41- except CatchableError as e:
42- report("handle resolves", false, e.msg)
43- quit 1
44-
45- # --- discovery ----------------------------------------------------------
46- var authServerUrl: string
47- try:
48- authServerUrl = resolveAuthServerUrl(ident.pdsEndpoint)
49- report("PDS advertises an auth server", authServerUrl.len > 0, authServerUrl)
50- except CatchableError as e:
51- report("PDS advertises an auth server", false, e.msg)
52- quit 1
53-
54- var meta: AuthServerMetadata
55- try:
56- meta = resolveAuthServerMetadata(authServerUrl)
57- report("auth server metadata fetched", true, meta.issuer)
58- report("has a PAR endpoint", meta.parEndpoint.len > 0, meta.parEndpoint)
59- report("has a token endpoint", meta.tokenEndpoint.len > 0, meta.tokenEndpoint)
60- except CatchableError as e:
61- report("auth server metadata fetched", false, e.msg)
62- quit 1
63-
64- # --- PKCE ---------------------------------------------------------------
65- block:
66- # RFC 7636 test vector: this verifier must produce this challenge.
67- const
68- verifier = "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
69- want = "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"
70- report("PKCE S256 matches RFC 7636 vector",
71- s256Challenge(verifier) == want, s256Challenge(verifier))
72-
73- # --- the nonce dance ----------------------------------------------------
74- # sendAuthRequest retries transparently, which would hide a server that
75- # never demanded a nonce. Prove the first unnonced attempt really is
76- # rejected, so the retry path is load-bearing rather than decorative.
77- let cfg0 = newLocalhostConfig("http://127.0.0.1:8080/callback")
78- block:
79- let
80- key = generateKey(rng)
81- proof = dpopProof(key, rng, "POST", meta.parEndpoint)
82- client = newHttpClient(userAgent = UserAgent, timeout = 15_000)
83- defer: client.close()
84- client.headers = newHttpHeaders({
85- "Content-Type": "application/x-www-form-urlencoded",
86- "DPoP": proof,
87- })
88- # The body must be fully valid: bsky.social validates parameters before it
89- # checks the nonce, so a stub body returns invalid_request and tells us
90- # nothing about the nonce requirement.
91- let verifier = secureRandomB64(48)
92- let body = @{
93- "client_id": cfg0.clientId,
94- "state": secureRandomB64(16),
95- "redirect_uri": cfg0.callbackUrl,
96- "scope": scopeStr(cfg0.scopes),
97- "response_type": "code",
98- "code_challenge": s256Challenge(verifier),
99- "code_challenge_method": "S256",
100- }.encodeQuery
101- let res = client.post(meta.parEndpoint, body = body)
102- let code = try: parseJson(res.body){"error"}.getStr except CatchableError: ""
103- let nonce = res.headers.getOrDefault("DPoP-Nonce").string
104- report("first PAR without a nonce is rejected",
105- res.code.int == 400 and code == "use_dpop_nonce", &"HTTP {res.code.int} {code}")
106- report("rejection supplies a DPoP-Nonce to retry with", nonce.len > 0)
107-
108- # --- PAR ----------------------------------------------------------------
109- # A localhost client needs no hosted metadata document, so the flow can be
110- # exercised without deploying anything.
111- let cfg = cfg0
112- report("localhost client_id is well-formed",
113- cfg.clientId.startsWith("http://localhost?"), cfg.clientId)
114-
115- try:
116- let info = sendAuthRequest(cfg, meta, rng, loginHint = handle)
117- report("PAR accepted (DPoP proof + nonce retry + PKCE)",
118- info.requestUri.startsWith("urn:ietf:params:oauth:request_uri:"),
119- info.requestUri)
120- report("auth server issued a DPoP nonce", info.dpopNonce.len > 0,
121- info.dpopNonce[0 ..< min(12, info.dpopNonce.len)] & "...")
122- report("PKCE verifier retained for callback", info.pkceVerifier.len >= 43)
123- report("state is unguessable", info.state.len >= 21, &"{info.state.len} chars")
124-
125- let url = authorizeUrl(cfg, meta, info)
126- let q = parseUri(url).query
127- report("authorize URL carries request_uri", "request_uri=" in q)
128- report("authorize URL carries client_id", "client_id=" in q)
129- echo ""
130- echo " authorize URL:"
131- echo " ", url
132- except CatchableError as e:
133- report("PAR accepted", false, e.msg)
134-
135- echo ""
136- if failures == 0:
137- echo "RESULT: the OAuth flow works against a live PDS up to user consent."
138- else:
139- echo &"RESULT: {failures} check(s) failed."
140- quit 1
renamed spike/nim/reader/articlestore.nim +5 -12
similarity index 95%
rename from spike/nim/atproto/articlestore.nim
rename to spike/nim/reader/articlestore.nim
@@ -26,8 +26,6 @@ type
2626 published*: Option[DateTime]
2727 updated*: Option[DateTime]
2828 isRead*: bool
29- likeCount*: int
30- hasLiked*: bool
3129
3230 ArticleFilter* = enum
3331 afAll = "all"
@@ -64,8 +62,6 @@ proc readArticleRow*(r: Row): Article =
6462 published: parseSqliteTime(r.str(11)),
6563 updated: parseSqliteTime(r.str(12)),
6664 isRead: r.b(13),
67- likeCount: r.i(14),
68- hasLiked: r.b(15),
6965 )
7066
7167 ## The column list is shared so every listing returns the same shape and
@@ -75,10 +71,7 @@ const articleColumns = """
7571 a.id, a.feed_url, COALESCE(f.title, a.feed_url), f.favicon_url,
7672 a.guid, a.title, a.url, a.author, a.summary, a.content, a.full_content,
7773 a.published, a.updated,
78- COALESCE(r.is_read, 0),
79- (SELECT count(*) FROM articles.likes l WHERE l.article_url = a.url),
80- EXISTS(SELECT 1 FROM articles.likes l
81- WHERE l.article_url = a.url AND l.author_did = ?)"""
74+ COALESCE(r.is_read, 0)"""
8275
8376 const articleFrom = """
8477 FROM articles.articles a
@@ -87,11 +80,11 @@ const articleFrom = """
8780
8881 proc getArticle*(g: GleanDb, userDid: string, id: int64): Option[Article] =
8982 g.db.queryFirst("SELECT " & articleColumns & articleFrom & " WHERE a.id = ?",
90- [p(userDid), p(userDid), p(id)], readArticleRow)
83+ [p(userDid), p(id)], readArticleRow)
9184
9285 proc getArticleByUrl*(g: GleanDb, userDid, url: string): Option[Article] =
9386 g.db.queryFirst("SELECT " & articleColumns & articleFrom & " WHERE a.url = ?",
94- [p(userDid), p(userDid), p(url)], readArticleRow)
87+ [p(userDid), p(url)], readArticleRow)
9588
9689 proc listArticles*(g: GleanDb, userDid: string, filter = afAll,
9790 feedUrl = "", category = "", sortOldest = false,
@@ -101,7 +94,7 @@ proc listArticles*(g: GleanDb, userDid: string, filter = afAll,
10194 JOIN articles.subscriptions s
10295 ON s.feed_url = a.feed_url AND s.user_did = ?
10396 WHERE 1=1"""
104- var params = @[p(userDid), p(userDid), p(userDid)]
97+ var params = @[p(userDid), p(userDid)]
10598
10699 case filter
107100 of afUnread: sql &= " AND COALESCE(r.is_read, 0) = 0"
@@ -147,7 +140,7 @@ proc searchArticles*(g: GleanDb, userDid, query: string,
147140 WHERE articles_fts MATCH ?) m ON m.rowid = a.id
148141 ORDER BY m.rank
149142 LIMIT ? OFFSET ?""",
150- [p(userDid), p(userDid), p(userDid), p(query),
143+ [p(userDid), p(userDid), p(query),
151144 p(limit.int64), p(offset.int64)], readArticleRow)
152145
153146 proc batchUpsertArticles*(g: GleanDb, articles: seq[NewArticle],
similarity index 95%
rename from spike/nim/atproto/articlestore.nim
rename to spike/nim/reader/articlestore.nim
@@ -26,8 +26,6 @@ type
26 published*: Option[DateTime]26 published*: Option[DateTime]
27 updated*: Option[DateTime]27 updated*: Option[DateTime]
28 isRead*: bool28 isRead*: bool
29- likeCount*: int
30- hasLiked*: bool
31 29
32 ArticleFilter* = enum30 ArticleFilter* = enum
33 afAll = "all"31 afAll = "all"
@@ -64,8 +62,6 @@ proc readArticleRow*(r: Row): Article =
64 published: parseSqliteTime(r.str(11)),62 published: parseSqliteTime(r.str(11)),
65 updated: parseSqliteTime(r.str(12)),63 updated: parseSqliteTime(r.str(12)),
66 isRead: r.b(13),64 isRead: r.b(13),
67- likeCount: r.i(14),
68- hasLiked: r.b(15),
69 )65 )
70 66
71 ## The column list is shared so every listing returns the same shape and67 ## The column list is shared so every listing returns the same shape and
@@ -75,10 +71,7 @@ const articleColumns = """
75 a.id, a.feed_url, COALESCE(f.title, a.feed_url), f.favicon_url,71 a.id, a.feed_url, COALESCE(f.title, a.feed_url), f.favicon_url,
76 a.guid, a.title, a.url, a.author, a.summary, a.content, a.full_content,72 a.guid, a.title, a.url, a.author, a.summary, a.content, a.full_content,
77 a.published, a.updated,73 a.published, a.updated,
78- COALESCE(r.is_read, 0),74+ COALESCE(r.is_read, 0)"""
79- (SELECT count(*) FROM articles.likes l WHERE l.article_url = a.url),
80- EXISTS(SELECT 1 FROM articles.likes l
81- WHERE l.article_url = a.url AND l.author_did = ?)"""
82 75
83 const articleFrom = """76 const articleFrom = """
84 FROM articles.articles a77 FROM articles.articles a
@@ -87,11 +80,11 @@ const articleFrom = """
87 80
88 proc getArticle*(g: GleanDb, userDid: string, id: int64): Option[Article] =81 proc getArticle*(g: GleanDb, userDid: string, id: int64): Option[Article] =
89 g.db.queryFirst("SELECT " & articleColumns & articleFrom & " WHERE a.id = ?",82 g.db.queryFirst("SELECT " & articleColumns & articleFrom & " WHERE a.id = ?",
90- [p(userDid), p(userDid), p(id)], readArticleRow)83+ [p(userDid), p(id)], readArticleRow)
91 84
92 proc getArticleByUrl*(g: GleanDb, userDid, url: string): Option[Article] =85 proc getArticleByUrl*(g: GleanDb, userDid, url: string): Option[Article] =
93 g.db.queryFirst("SELECT " & articleColumns & articleFrom & " WHERE a.url = ?",86 g.db.queryFirst("SELECT " & articleColumns & articleFrom & " WHERE a.url = ?",
94- [p(userDid), p(userDid), p(url)], readArticleRow)87+ [p(userDid), p(url)], readArticleRow)
95 88
96 proc listArticles*(g: GleanDb, userDid: string, filter = afAll,89 proc listArticles*(g: GleanDb, userDid: string, filter = afAll,
97 feedUrl = "", category = "", sortOldest = false,90 feedUrl = "", category = "", sortOldest = false,
@@ -101,7 +94,7 @@ proc listArticles*(g: GleanDb, userDid: string, filter = afAll,
101 JOIN articles.subscriptions s94 JOIN articles.subscriptions s
102 ON s.feed_url = a.feed_url AND s.user_did = ?95 ON s.feed_url = a.feed_url AND s.user_did = ?
103 WHERE 1=1"""96 WHERE 1=1"""
104- var params = @[p(userDid), p(userDid), p(userDid)]97+ var params = @[p(userDid), p(userDid)]
105 98
106 case filter99 case filter
107 of afUnread: sql &= " AND COALESCE(r.is_read, 0) = 0"100 of afUnread: sql &= " AND COALESCE(r.is_read, 0) = 0"
@@ -147,7 +140,7 @@ proc searchArticles*(g: GleanDb, userDid, query: string,
147 WHERE articles_fts MATCH ?) m ON m.rowid = a.id140 WHERE articles_fts MATCH ?) m ON m.rowid = a.id
148 ORDER BY m.rank141 ORDER BY m.rank
149 LIMIT ? OFFSET ?""",142 LIMIT ? OFFSET ?""",
150- [p(userDid), p(userDid), p(userDid), p(query),143+ [p(userDid), p(userDid), p(query),
151 p(limit.int64), p(offset.int64)], readArticleRow)144 p(limit.int64), p(offset.int64)], readArticleRow)
152 145
153 proc batchUpsertArticles*(g: GleanDb, articles: seq[NewArticle],146 proc batchUpsertArticles*(g: GleanDb, articles: seq[NewArticle],
renamed spike/nim/reader/feedfetcher.nim +0 -0
similarity index 100%
rename from spike/nim/atproto/feedfetcher.nim
rename to spike/nim/reader/feedfetcher.nim
similarity index 100%
rename from spike/nim/atproto/feedfetcher.nim
rename to spike/nim/reader/feedfetcher.nim
renamed spike/nim/reader/feedparser.nim +0 -0
similarity index 100%
rename from spike/nim/atproto/feedparser.nim
rename to spike/nim/reader/feedparser.nim
similarity index 100%
rename from spike/nim/atproto/feedparser.nim
rename to spike/nim/reader/feedparser.nim
renamed spike/nim/reader/feedstore.nim +0 -0
similarity index 100%
rename from spike/nim/atproto/feedstore.nim
rename to spike/nim/reader/feedstore.nim
similarity index 100%
rename from spike/nim/atproto/feedstore.nim
rename to spike/nim/reader/feedstore.nim
added spike/nim/reader/gleandb.nim +86 -0
new file mode 100644
@@ -0,0 +1,86 @@
1+## The database: two SQLite files behind one connection.
2+##
3+## `main` holds the local user, `articles` the feeds and their contents. They
4+## are attached together so a query can join across them, while each still
5+## checkpoints and vacuums on its own -- which matters because the article
6+## store is the only one that grows without bound.
7+##
8+## The pragmas are copied from internal/db/db.go. They are not decoration:
9+## WAL is what lets readers run during a write, and the per-database repeats
10+## matter because PRAGMA applies to one schema at a time -- setting WAL on
11+## `main` leaves `articles` journalling the slow way.
12+
13+import std/[options, strformat, strutils]
14+import ./schema
15+import ./sqlite
16+
17+type
18+ GleanDb* = object
19+ db*: Db
20+ basePath*: string
21+
22+const
23+ MainPragmas = [
24+ "journal_mode=WAL",
25+ "wal_autocheckpoint=1000",
26+ "busy_timeout=30000",
27+ "synchronous=NORMAL",
28+ "temp_store=FILE",
29+ "mmap_size=268435456",
30+ "auto_vacuum=INCREMENTAL",
31+ ]
32+ AttachedPragmas = [
33+ "journal_mode=WAL",
34+ "wal_autocheckpoint=1000",
35+ "busy_timeout=30000",
36+ "synchronous=NORMAL",
37+ "temp_store=FILE",
38+ "mmap_size=268435456",
39+ ]
40+
41+proc applyPragmas(db: Db, schemaName: string, pragmas: openArray[string]) =
42+ for pragma in pragmas:
43+ # A pragma that will not take is not worth failing a boot over -- some are
44+ # advisory, and mmap_size in particular is refused on filesystems that
45+ # cannot map. The ones that matter are asserted after opening instead.
46+ try:
47+ if schemaName.len == 0:
48+ db.exec &"PRAGMA {pragma}"
49+ else:
50+ db.exec &"PRAGMA {schemaName}.{pragma}"
51+ except SqliteError:
52+ discard
53+
54+proc open*(basePath: string): GleanDb =
55+ ## Open `<basePath>_users` and attach `_articles`.
56+ let
57+ usersPath = basePath & "_users"
58+ articlesPath = basePath & "_articles"
59+
60+ var db = sqlite.open(usersPath)
61+ applyPragmas(db, "", MainPragmas)
62+
63+ # Quoted and single-quote-escaped: a path is not necessarily tame, and this
64+ # one can come from configuration.
65+ for (name, path) in [("articles", articlesPath)]:
66+ db.exec &"ATTACH DATABASE '{path.replace(\"'\", \"''\")}' AS {name}"
67+ applyPragmas(db, name, AttachedPragmas)
68+
69+ GleanDb(db: db, basePath: basePath)
70+
71+proc migrate*(g: GleanDb) =
72+ ## Apply the schema. Every statement is IF NOT EXISTS, so this runs on every
73+ ## boot rather than being gated on a version.
74+ for stmt in readerSchema: g.db.exec stmt
75+
76+proc close*(g: var GleanDb) =
77+ g.db.close()
78+
79+proc journalMode*(g: GleanDb, schemaName = ""): string =
80+ let sql = if schemaName.len == 0: "PRAGMA journal_mode"
81+ else: &"PRAGMA {schemaName}.journal_mode"
82+ g.db.queryText(sql).get("")
83+
84+proc attachedSchemas*(g: GleanDb): seq[string] =
85+ for row in g.db.rows("PRAGMA database_list"):
86+ result.add row[1]
new file mode 100644
@@ -0,0 +1,86 @@
1+## The database: two SQLite files behind one connection.
2+##
3+## `main` holds the local user, `articles` the feeds and their contents. They
4+## are attached together so a query can join across them, while each still
5+## checkpoints and vacuums on its own -- which matters because the article
6+## store is the only one that grows without bound.
7+##
8+## The pragmas are copied from internal/db/db.go. They are not decoration:
9+## WAL is what lets readers run during a write, and the per-database repeats
10+## matter because PRAGMA applies to one schema at a time -- setting WAL on
11+## `main` leaves `articles` journalling the slow way.
12+
13+import std/[options, strformat, strutils]
14+import ./schema
15+import ./sqlite
16+
17+type
18+ GleanDb* = object
19+ db*: Db
20+ basePath*: string
21+
22+const
23+ MainPragmas = [
24+ "journal_mode=WAL",
25+ "wal_autocheckpoint=1000",
26+ "busy_timeout=30000",
27+ "synchronous=NORMAL",
28+ "temp_store=FILE",
29+ "mmap_size=268435456",
30+ "auto_vacuum=INCREMENTAL",
31+ ]
32+ AttachedPragmas = [
33+ "journal_mode=WAL",
34+ "wal_autocheckpoint=1000",
35+ "busy_timeout=30000",
36+ "synchronous=NORMAL",
37+ "temp_store=FILE",
38+ "mmap_size=268435456",
39+ ]
40+
41+proc applyPragmas(db: Db, schemaName: string, pragmas: openArray[string]) =
42+ for pragma in pragmas:
43+ # A pragma that will not take is not worth failing a boot over -- some are
44+ # advisory, and mmap_size in particular is refused on filesystems that
45+ # cannot map. The ones that matter are asserted after opening instead.
46+ try:
47+ if schemaName.len == 0:
48+ db.exec &"PRAGMA {pragma}"
49+ else:
50+ db.exec &"PRAGMA {schemaName}.{pragma}"
51+ except SqliteError:
52+ discard
53+
54+proc open*(basePath: string): GleanDb =
55+ ## Open `<basePath>_users` and attach `_articles`.
56+ let
57+ usersPath = basePath & "_users"
58+ articlesPath = basePath & "_articles"
59+
60+ var db = sqlite.open(usersPath)
61+ applyPragmas(db, "", MainPragmas)
62+
63+ # Quoted and single-quote-escaped: a path is not necessarily tame, and this
64+ # one can come from configuration.
65+ for (name, path) in [("articles", articlesPath)]:
66+ db.exec &"ATTACH DATABASE '{path.replace(\"'\", \"''\")}' AS {name}"
67+ applyPragmas(db, name, AttachedPragmas)
68+
69+ GleanDb(db: db, basePath: basePath)
70+
71+proc migrate*(g: GleanDb) =
72+ ## Apply the schema. Every statement is IF NOT EXISTS, so this runs on every
73+ ## boot rather than being gated on a version.
74+ for stmt in readerSchema: g.db.exec stmt
75+
76+proc close*(g: var GleanDb) =
77+ g.db.close()
78+
79+proc journalMode*(g: GleanDb, schemaName = ""): string =
80+ let sql = if schemaName.len == 0: "PRAGMA journal_mode"
81+ else: &"PRAGMA {schemaName}.journal_mode"
82+ g.db.queryText(sql).get("")
83+
84+proc attachedSchemas*(g: GleanDb): seq[string] =
85+ for row in g.db.rows("PRAGMA database_list"):
86+ result.add row[1]
added spike/nim/reader/schema.nim +97 -0
new file mode 100644
@@ -0,0 +1,97 @@
1+## The reader's schema.
2+##
3+## Two databases: `main` holds the single local user, `articles` holds feeds
4+## and their contents. They are separate files so the article store -- the
5+## only one that grows without bound -- can be vacuumed or discarded on its
6+## own.
7+##
8+## Originally generated from a Go server's schema and then trimmed to what a
9+## single-user reader needs: no accounts, no OAuth, no social records, no
10+## derived recommendation tables. What is left is feeds, subscriptions,
11+## articles, read state and the full-text index. It is maintained by hand
12+## from here.
13+##
14+## `read_state_history` is not redundant with `read_state`. Articles are
15+## purged on a retention window and can return from the feed later under a
16+## new surrogate id; the history is keyed on (feed_url, guid), which survives
17+## that, so a re-ingested article does not come back unread.
18+
19+const readerSchema* = [
20+ """CREATE TABLE IF NOT EXISTS users (
21+ did TEXT PRIMARY KEY,
22+ indexed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
23+ updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
24+ follows_dirty BOOLEAN NOT NULL DEFAULT 1
25+ )""",
26+ """CREATE TABLE IF NOT EXISTS articles.feeds (
27+ feed_url TEXT PRIMARY KEY,
28+ title TEXT,
29+ site_url TEXT,
30+ description TEXT,
31+ feed_type TEXT CHECK(feed_type IN ('rss', 'atom', 'json', 'atproto')),
32+ last_fetched_at DATETIME,
33+ last_error TEXT,
34+ subscriber_count INTEGER NOT NULL DEFAULT 0,
35+ consecutive_empty_fetches INTEGER NOT NULL DEFAULT 0,
36+ error_count INTEGER NOT NULL DEFAULT 0,
37+ favicon_url TEXT
38+ )""",
39+ """CREATE TABLE IF NOT EXISTS articles.subscriptions (
40+ id INTEGER PRIMARY KEY AUTOINCREMENT,
41+ user_did TEXT NOT NULL,
42+ feed_url TEXT NOT NULL,
43+ title TEXT,
44+ category TEXT,
45+ added_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
46+ uri TEXT,
47+ cid TEXT,
48+ UNIQUE(user_did, feed_url)
49+ )""",
50+ """CREATE TABLE IF NOT EXISTS articles.articles (
51+ id INTEGER PRIMARY KEY AUTOINCREMENT,
52+ feed_url TEXT NOT NULL,
53+ guid TEXT NOT NULL,
54+ title TEXT NOT NULL DEFAULT '',
55+ url TEXT,
56+ author TEXT,
57+ summary TEXT,
58+ content TEXT,
59+ full_content TEXT,
60+ published DATETIME,
61+ updated DATETIME,
62+ fetched_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
63+ language TEXT NOT NULL DEFAULT '',
64+ UNIQUE(feed_url, guid)
65+ )""",
66+ """CREATE TABLE IF NOT EXISTS articles.read_state (
67+ user_did TEXT NOT NULL,
68+ article_id INTEGER NOT NULL,
69+ is_read BOOLEAN NOT NULL DEFAULT 0,
70+ read_at DATETIME,
71+ PRIMARY KEY (user_did, article_id)
72+ )""",
73+ """CREATE TABLE IF NOT EXISTS articles.read_state_history (
74+ user_did TEXT NOT NULL,
75+ feed_url TEXT NOT NULL,
76+ guid TEXT NOT NULL,
77+ is_read BOOLEAN NOT NULL DEFAULT 0,
78+ read_at DATETIME,
79+ PRIMARY KEY (user_did, feed_url, guid)
80+ )""",
81+ """CREATE INDEX IF NOT EXISTS articles.idx_subscriptions_feed_user ON subscriptions(feed_url, user_did)""",
82+ """CREATE INDEX IF NOT EXISTS articles.idx_subscriptions_uri ON subscriptions(uri)""",
83+ """CREATE INDEX IF NOT EXISTS articles.idx_articles_url ON articles(url)""",
84+ """CREATE INDEX IF NOT EXISTS articles.idx_read_state_unread ON read_state(user_did, is_read) WHERE is_read = 0""",
85+ """CREATE INDEX IF NOT EXISTS articles.idx_articles_language ON articles(language)""",
86+ """CREATE VIRTUAL TABLE IF NOT EXISTS articles.articles_fts USING fts5(title, summary, content, author, content=articles, content_rowid=id)""",
87+ """CREATE TRIGGER IF NOT EXISTS articles.articles_ai AFTER INSERT ON articles BEGIN
88+ INSERT INTO articles_fts(rowid, title, summary, content, author) VALUES (new.id, new.title, new.summary, new.content, new.author);
89+ END""",
90+ """CREATE TRIGGER IF NOT EXISTS articles.articles_ad AFTER DELETE ON articles BEGIN
91+ INSERT INTO articles_fts(articles_fts, rowid, title, summary, content, author) VALUES('delete', old.id, old.title, old.summary, old.content, old.author);
92+ END""",
93+ """CREATE TRIGGER IF NOT EXISTS articles.articles_au AFTER UPDATE ON articles BEGIN
94+ INSERT INTO articles_fts(articles_fts, rowid, title, summary, content, author) VALUES('delete', old.id, old.title, old.summary, old.content, old.author);
95+ INSERT INTO articles_fts(rowid, title, summary, content, author) VALUES (new.id, new.title, new.summary, new.content, new.author);
96+ END""",
97+]
new file mode 100644
@@ -0,0 +1,97 @@
1+## The reader's schema.
2+##
3+## Two databases: `main` holds the single local user, `articles` holds feeds
4+## and their contents. They are separate files so the article store -- the
5+## only one that grows without bound -- can be vacuumed or discarded on its
6+## own.
7+##
8+## Originally generated from a Go server's schema and then trimmed to what a
9+## single-user reader needs: no accounts, no OAuth, no social records, no
10+## derived recommendation tables. What is left is feeds, subscriptions,
11+## articles, read state and the full-text index. It is maintained by hand
12+## from here.
13+##
14+## `read_state_history` is not redundant with `read_state`. Articles are
15+## purged on a retention window and can return from the feed later under a
16+## new surrogate id; the history is keyed on (feed_url, guid), which survives
17+## that, so a re-ingested article does not come back unread.
18+
19+const readerSchema* = [
20+ """CREATE TABLE IF NOT EXISTS users (
21+ did TEXT PRIMARY KEY,
22+ indexed_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
23+ updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
24+ follows_dirty BOOLEAN NOT NULL DEFAULT 1
25+ )""",
26+ """CREATE TABLE IF NOT EXISTS articles.feeds (
27+ feed_url TEXT PRIMARY KEY,
28+ title TEXT,
29+ site_url TEXT,
30+ description TEXT,
31+ feed_type TEXT CHECK(feed_type IN ('rss', 'atom', 'json', 'atproto')),
32+ last_fetched_at DATETIME,
33+ last_error TEXT,
34+ subscriber_count INTEGER NOT NULL DEFAULT 0,
35+ consecutive_empty_fetches INTEGER NOT NULL DEFAULT 0,
36+ error_count INTEGER NOT NULL DEFAULT 0,
37+ favicon_url TEXT
38+ )""",
39+ """CREATE TABLE IF NOT EXISTS articles.subscriptions (
40+ id INTEGER PRIMARY KEY AUTOINCREMENT,
41+ user_did TEXT NOT NULL,
42+ feed_url TEXT NOT NULL,
43+ title TEXT,
44+ category TEXT,
45+ added_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
46+ uri TEXT,
47+ cid TEXT,
48+ UNIQUE(user_did, feed_url)
49+ )""",
50+ """CREATE TABLE IF NOT EXISTS articles.articles (
51+ id INTEGER PRIMARY KEY AUTOINCREMENT,
52+ feed_url TEXT NOT NULL,
53+ guid TEXT NOT NULL,
54+ title TEXT NOT NULL DEFAULT '',
55+ url TEXT,
56+ author TEXT,
57+ summary TEXT,
58+ content TEXT,
59+ full_content TEXT,
60+ published DATETIME,
61+ updated DATETIME,
62+ fetched_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
63+ language TEXT NOT NULL DEFAULT '',
64+ UNIQUE(feed_url, guid)
65+ )""",
66+ """CREATE TABLE IF NOT EXISTS articles.read_state (
67+ user_did TEXT NOT NULL,
68+ article_id INTEGER NOT NULL,
69+ is_read BOOLEAN NOT NULL DEFAULT 0,
70+ read_at DATETIME,
71+ PRIMARY KEY (user_did, article_id)
72+ )""",
73+ """CREATE TABLE IF NOT EXISTS articles.read_state_history (
74+ user_did TEXT NOT NULL,
75+ feed_url TEXT NOT NULL,
76+ guid TEXT NOT NULL,
77+ is_read BOOLEAN NOT NULL DEFAULT 0,
78+ read_at DATETIME,
79+ PRIMARY KEY (user_did, feed_url, guid)
80+ )""",
81+ """CREATE INDEX IF NOT EXISTS articles.idx_subscriptions_feed_user ON subscriptions(feed_url, user_did)""",
82+ """CREATE INDEX IF NOT EXISTS articles.idx_subscriptions_uri ON subscriptions(uri)""",
83+ """CREATE INDEX IF NOT EXISTS articles.idx_articles_url ON articles(url)""",
84+ """CREATE INDEX IF NOT EXISTS articles.idx_read_state_unread ON read_state(user_did, is_read) WHERE is_read = 0""",
85+ """CREATE INDEX IF NOT EXISTS articles.idx_articles_language ON articles(language)""",
86+ """CREATE VIRTUAL TABLE IF NOT EXISTS articles.articles_fts USING fts5(title, summary, content, author, content=articles, content_rowid=id)""",
87+ """CREATE TRIGGER IF NOT EXISTS articles.articles_ai AFTER INSERT ON articles BEGIN
88+ INSERT INTO articles_fts(rowid, title, summary, content, author) VALUES (new.id, new.title, new.summary, new.content, new.author);
89+ END""",
90+ """CREATE TRIGGER IF NOT EXISTS articles.articles_ad AFTER DELETE ON articles BEGIN
91+ INSERT INTO articles_fts(articles_fts, rowid, title, summary, content, author) VALUES('delete', old.id, old.title, old.summary, old.content, old.author);
92+ END""",
93+ """CREATE TRIGGER IF NOT EXISTS articles.articles_au AFTER UPDATE ON articles BEGIN
94+ INSERT INTO articles_fts(articles_fts, rowid, title, summary, content, author) VALUES('delete', old.id, old.title, old.summary, old.content, old.author);
95+ INSERT INTO articles_fts(rowid, title, summary, content, author) VALUES (new.id, new.title, new.summary, new.content, new.author);
96+ END""",
97+]
renamed spike/nim/reader/scraper.nim +0 -0
similarity index 100%
rename from spike/nim/atproto/scraper.nim
rename to spike/nim/reader/scraper.nim
similarity index 100%
rename from spike/nim/atproto/scraper.nim
rename to spike/nim/reader/scraper.nim
renamed spike/nim/reader/sqlite.nim +5 -14
similarity index 94%
rename from spike/nim/atproto/sqlite.nim
rename to spike/nim/reader/sqlite.nim
@@ -1,11 +1,8 @@
1-## A small SQLite wrapper, with sqlite-vec compiled in.
1+## A small SQLite wrapper.
22 ##
3-## Deliberately thin: enough to prepare, bind and step, not an ORM. The
4-## sqlite-vec extension is linked rather than loaded at runtime -- see
5-## nim.cfg, where the SQLITE_CORE define lives, and note that getting that
6-## wrong segfaults rather than erroring.
7-
8-{.compile: "../vendor/sqlite-vec.c".}
3+## Deliberately thin: enough to prepare, bind and step, not an ORM. FTS5 is
4+## used for article search and ships with the system library; nothing here
5+## needs an extension.
96
107 import std/options
118
@@ -51,8 +48,6 @@ proc sqlite3_changes(db: pointer): cint
5148 proc sqlite3_clear_bindings(s: pointer): cint
5249 {.pop.}
5350
54-proc sqlite3_vec_init(db: pointer, errMsg: ptr cstring,
55- api: pointer): cint {.importc, cdecl.}
5651
5752 const SQLITE_TRANSIENT = cast[pointer](-1)
5853
@@ -62,7 +57,7 @@ proc check(db: pointer, rc: cint, what: string) =
6257 if rc notin [SQLITE_OK.cint, SQLITE_ROW.cint, SQLITE_DONE.cint]:
6358 raise newException(SqliteError, what & ": " & $sqlite3_errmsg(db))
6459
65-proc open*(path: string, withVec = true): Db =
60+proc open*(path: string): Db =
6661 ## Open (or create) a database.
6762 ##
6863 ## WAL is set here rather than left to callers: it is what makes a reader
@@ -75,10 +70,6 @@ proc open*(path: string, withVec = true): Db =
7570 raise newException(SqliteError, "opening " & path & ": " & msg)
7671 result = Db(handle: handle)
7772
78- if withVec:
79- let rc = sqlite3_vec_init(handle, nil, nil)
80- if rc != SQLITE_OK:
81- raise newException(SqliteError, "registering sqlite-vec failed")
8273
8374 discard sqlite3_exec(handle, "PRAGMA journal_mode=WAL", nil, nil, nil)
8475 discard sqlite3_exec(handle, "PRAGMA synchronous=NORMAL", nil, nil, nil)
similarity index 94%
rename from spike/nim/atproto/sqlite.nim
rename to spike/nim/reader/sqlite.nim
@@ -1,11 +1,8 @@
1-## A small SQLite wrapper, with sqlite-vec compiled in.1+## A small SQLite wrapper.
2 ##2 ##
3-## Deliberately thin: enough to prepare, bind and step, not an ORM. The3+## Deliberately thin: enough to prepare, bind and step, not an ORM. FTS5 is
4-## sqlite-vec extension is linked rather than loaded at runtime -- see4+## used for article search and ships with the system library; nothing here
5-## nim.cfg, where the SQLITE_CORE define lives, and note that getting that5+## needs an extension.
6-## wrong segfaults rather than erroring.
7-
8-{.compile: "../vendor/sqlite-vec.c".}
9 6
10 import std/options7 import std/options
11 8
@@ -51,8 +48,6 @@ proc sqlite3_changes(db: pointer): cint
51 proc sqlite3_clear_bindings(s: pointer): cint48 proc sqlite3_clear_bindings(s: pointer): cint
52 {.pop.}49 {.pop.}
53 50
54-proc sqlite3_vec_init(db: pointer, errMsg: ptr cstring,
55- api: pointer): cint {.importc, cdecl.}
56 51
57 const SQLITE_TRANSIENT = cast[pointer](-1)52 const SQLITE_TRANSIENT = cast[pointer](-1)
58 53
@@ -62,7 +57,7 @@ proc check(db: pointer, rc: cint, what: string) =
62 if rc notin [SQLITE_OK.cint, SQLITE_ROW.cint, SQLITE_DONE.cint]:57 if rc notin [SQLITE_OK.cint, SQLITE_ROW.cint, SQLITE_DONE.cint]:
63 raise newException(SqliteError, what & ": " & $sqlite3_errmsg(db))58 raise newException(SqliteError, what & ": " & $sqlite3_errmsg(db))
64 59
65-proc open*(path: string, withVec = true): Db =60+proc open*(path: string): Db =
66 ## Open (or create) a database.61 ## Open (or create) a database.
67 ##62 ##
68 ## WAL is set here rather than left to callers: it is what makes a reader63 ## WAL is set here rather than left to callers: it is what makes a reader
@@ -75,10 +70,6 @@ proc open*(path: string, withVec = true): Db =
75 raise newException(SqliteError, "opening " & path & ": " & msg)70 raise newException(SqliteError, "opening " & path & ": " & msg)
76 result = Db(handle: handle)71 result = Db(handle: handle)
77 72
78- if withVec:
79- let rc = sqlite3_vec_init(handle, nil, nil)
80- if rc != SQLITE_OK:
81- raise newException(SqliteError, "registering sqlite-vec failed")
82 73
83 discard sqlite3_exec(handle, "PRAGMA journal_mode=WAL", nil, nil, nil)74 discard sqlite3_exec(handle, "PRAGMA journal_mode=WAL", nil, nil, nil)
84 discard sqlite3_exec(handle, "PRAGMA synchronous=NORMAL", nil, nil, nil)75 discard sqlite3_exec(handle, "PRAGMA synchronous=NORMAL", nil, nil, nil)
deleted spike/nim/readme-reader.md +0 -56
deleted file mode 100644
@@ -1,56 +0,0 @@
1-# glean — a small RSS reader
2-
3-One binary, one SQLite database, no accounts.
4-
5-```
6-nim c -d:ssl glean.nim
7-./glean
8-```
9-
10-Listens on `127.0.0.1:8080` and stores its database under
11-`~/.local/share/glean/`. Override with `GLEAN_PORT` and `GLEAN_DB`.
12-
13-## API
14-
15-```
16-GET /feeds subscriptions, with unread counts
17-POST /feeds url=… subscribe and fetch immediately
18-DELETE /feeds?url=… unsubscribe
19-
20-GET /articles unread, newest first
21-GET /articles?status=all|read …or everything, or just the read ones
22-GET /articles?feed=… …scoped to one feed
23-GET /articles?q=… full-text search (FTS5)
24-GET /articles/:id one article
25-
26-POST /read id=… mark read (undo=1 to reverse)
27-POST /read-all [feed=…] mark a feed, or everything, read
28-POST /refresh fetch every subscribed feed
29-POST /fetch-content id=… scrape the full article text
30-
31-GET /unread unread count
32-```
33-
34-Parameters are read from the query string or a form body, whichever the
35-caller used.
36-
37-## What it does not do
38-
39-No accounts, no sessions, no CSRF, no ATProto. It runs on your machine and
40-reads your feeds, and that assumption is what keeps it a few hundred lines
41-rather than a few thousand.
42-
43-The one visible consequence: **do not expose it to a network.** It binds
44-loopback deliberately. Anything that can reach the port can read and change
45-everything, because there is nobody to distinguish from anybody else.
46-
47-## What it is built on
48-
49-The modules under `atproto/` — fetching with a sane retry policy, parsing
50-RSS/RDF/Atom/JSON Feed, SQLite with FTS5, and a readability-style extractor
51-that sanitises what it returns. Those came out of a larger port and carry
52-their own tests (`nimble test`); this file is a few hundred lines of routing
53-on top.
54-
55-Storage keeps a single hard-coded user, because the schema came from a
56-multi-user server and it was cheaper to satisfy that than to rewrite it.
deleted file mode 100644
@@ -1,56 +0,0 @@
1-# glean — a small RSS reader
2-
3-One binary, one SQLite database, no accounts.
4-
5-```
6-nim c -d:ssl glean.nim
7-./glean
8-```
9-
10-Listens on `127.0.0.1:8080` and stores its database under
11-`~/.local/share/glean/`. Override with `GLEAN_PORT` and `GLEAN_DB`.
12-
13-## API
14-
15-```
16-GET /feeds subscriptions, with unread counts
17-POST /feeds url=… subscribe and fetch immediately
18-DELETE /feeds?url=… unsubscribe
19-
20-GET /articles unread, newest first
21-GET /articles?status=all|read …or everything, or just the read ones
22-GET /articles?feed=… …scoped to one feed
23-GET /articles?q=… full-text search (FTS5)
24-GET /articles/:id one article
25-
26-POST /read id=… mark read (undo=1 to reverse)
27-POST /read-all [feed=…] mark a feed, or everything, read
28-POST /refresh fetch every subscribed feed
29-POST /fetch-content id=… scrape the full article text
30-
31-GET /unread unread count
32-```
33-
34-Parameters are read from the query string or a form body, whichever the
35-caller used.
36-
37-## What it does not do
38-
39-No accounts, no sessions, no CSRF, no ATProto. It runs on your machine and
40-reads your feeds, and that assumption is what keeps it a few hundred lines
41-rather than a few thousand.
42-
43-The one visible consequence: **do not expose it to a network.** It binds
44-loopback deliberately. Anything that can reach the port can read and change
45-everything, because there is nobody to distinguish from anybody else.
46-
47-## What it is built on
48-
49-The modules under `atproto/` — fetching with a sane retry policy, parsing
50-RSS/RDF/Atom/JSON Feed, SQLite with FTS5, and a readability-style extractor
51-that sanitises what it returns. Those came out of a larger port and carry
52-their own tests (`nimble test`); this file is a few hundred lines of routing
53-on top.
54-
55-Storage keeps a single hard-coded user, because the schema came from a
56-multi-user server and it was cheaper to satisfy that than to rewrite it.
modified spike/nim/readme.md +54 -556
@@ -1,577 +1,75 @@
1-# Nim port spike
1+# glean — a small RSS reader
22
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:
3+One binary, one SQLite database, no accounts.
94
105 ```
11-cd spike/nim
12-nimble test # everything that needs no network
13-nimble testnet # the probes that talk to real servers
14-nimble testall # both
15-nimble schema # regenerate atproto/schema.nim from the Go source
6+nim c -d:ssl glean.nim
7+./glean
168 ```
179
18-Reusable code lives in `atproto/`; the `*_probe.nim` files are the checks.
19-They are split by whether they touch the network: a failure against
20-`bsky.social` can mean that server is having a bad day rather than that the
21-port is broken, and a suite that cannot tell those apart stops being
22-believed. Individual probes still run on their own:
10+Listens on `127.0.0.1:8080` and stores its database under
11+`~/.local/share/glean/`. Override with `GLEAN_PORT` and `GLEAN_DB`.
12+
13+## API
2314
2415 ```
25-nim c -r -d:ssl <name>_probe.nim
16+GET /feeds subscriptions, with unread counts
17+POST /feeds url=… subscribe and fetch immediately
18+DELETE /feeds?url=… unsubscribe
19+
20+GET /articles unread, newest first
21+GET /articles?status=all|read …or everything, or just the read ones
22+GET /articles?feed=… …scoped to one feed
23+GET /articles?q=… full-text search (FTS5)
24+GET /articles/:id one article
25+
26+POST /read id=… mark read (undo=1 to reverse)
27+POST /read-all [feed=…] mark a feed, or everything, read
28+POST /refresh fetch every subscribed feed
29+POST /fetch-content id=… scrape the full article text
30+
31+GET /unread unread count
2632 ```
2733
28-## 1. sqlite-vec + FTS5 — works
29-
30-Glean needs FTS5 for article search and sqlite-vec for embedding similarity.
31-Both are C-level concerns, so the question was never "does Nim have a SQLite
32-library" but whether the same extension links and registers the same way.
33-
34-`sqlite_probe.nim` compiles the same `sqlite-vec.c` the Go bindings vendor,
35-calls `sqlite3_vec_init` directly, and exercises `vec_version()`, an FTS5
36-`MATCH`, a `vec0` virtual table, float32 blob binding, and a KNN query.
37-All pass; KNN returns the expected nearest neighbour.
38-
39-**One trap, and it is a nasty one.** sqlite-vec must be compiled with
40-`-DSQLITE_CORE`. Without it the file builds as a *loadable extension* whose
41-entry point takes its SQLite API pointer from the caller — pass `nil` there and
42-it segfaults inside `sqlite3_vec_init` rather than returning an error. Nim's
43-`{.compile: (file, flags).}` tuple form did not reliably apply the define, so
44-it lives in `nim.cfg` where it cannot be forgotten.
45-
46-## 2. ES256 / DPoP — works
47-
48-This was the real risk. Glean gets OAuth from `bluesky-social/indigo`; Nim has
49-no ATProto library, so a port has to build the DPoP layer itself, and
50-everything downstream of sign-in depends on it.
51-
52-`dpop_probe.nim` generates a P-256 key, builds the public JWK, computes an
53-RFC 7638 thumbprint, and signs a DPoP proof JWT (RFC 9449).
54-
55-It uses **bearssl** rather than OpenSSL for one specific reason: JOSE requires
56-the signature as raw `R||S`, 64 bytes, and most crypto libraries hand you DER.
57-`br_ecdsa_sign_raw` emits `R||S` directly, so there is no DER-unwrapping step
58-to get subtly wrong.
59-
60-Self-verification would only prove internal consistency, so the token goes to
61-`verify/verify.go`, which re-implements the verifier from Go's standard
62-library — the same stack the current server runs on. It checks the header
63-claims, that the signature is 64 raw bytes and not DER-wrapped, that the
64-signature verifies against the JWK, that a tampered payload is *rejected*, and
65-that the thumbprint matches when recomputed independently.
66-
67-19/19 checks pass, and the whole cycle passes across 30 independently generated
68-keypairs — worth confirming, because ECDSA values with leading zeros are a
69-classic source of intermittent length bugs in `R||S` encoding.
70-
71-## 3. The OAuth flow — works against a live PDS
72-
73-`atproto/identity.nim` and `atproto/oauth.nim` implement resolution and the
74-client flow, ported from the shape of indigo's `atproto/auth/oauth`.
75-`oauth_probe.nim` drives it end to end against real servers:
76-
77- handle -> DID -> PDS -> auth server -> metadata -> PAR -> authorize URL
78-
79-Against `bsky.social` all 15 checks pass, including a **real PAR that the
80-production auth server accepts** and returns a `request_uri` for. That single
81-result exercises the DPoP signature, the nonce retry and PKCE at once — a
82-server-side validation no amount of local testing substitutes for.
83-
84-Two things worth recording:
85-
86-**The nonce retry is load-bearing, and the probe proves it.** A PDS rejects the
87-first DPoP-signed request of a session with `400 use_dpop_nonce` and supplies a
88-`DPoP-Nonce` header to retry with. Since `sendAuthRequest` retries
89-transparently, the probe separately fires an unnonced request and asserts the
90-rejection — otherwise a server that never demanded a nonce would look identical
91-to a working retry.
92-
93-**Order of validation caught me out.** That check first sent a stub body and got
94-`invalid_request`, not `use_dpop_nonce`: bsky.social validates request
95-parameters *before* the nonce. The body has to be fully valid to observe the
96-nonce requirement at all.
97-
98-The flow stops at the authorization URL, which is correct — the next step is a
99-human granting consent in a browser. `exchangeCode` and `refresh` are
100-implemented against that callback but are consequently untested.
101-
102-## 4. XRPC signing — proof accepted by a real PDS
103-
104-`atproto/xrpc.nim` makes authenticated calls with a DPoP-bound access token:
105-`Authorization: DPoP <token>` (not Bearer), a per-request proof, and the retry
106-logic for both failure modes.
107-
108-Resource-server proofs are **not** the same as auth-server proofs, even when
109-one host fills both roles. A PDS request carries `ath` (base64url SHA-256 of
110-the access token, which is what binds the proof to that one credential) and
111-`iss`; the PAR and token requests that mint the token carry neither. The nonce
112-handshake also differs: a PDS answers **401** with `WWW-Authenticate: DPoP
113-error="use_dpop_nonce"`, where the auth server used **400** and a JSON body.
114-A session therefore tracks two independent nonces.
115-
116-`xrpc_probe.nim` cannot get a consented token, so it signs a request with a
117-deliberately invalid one and reads what a real PDS says. Against
118-`bsky.social`'s PDS:
119-
120- WWW-Authenticate: DPoP algs="...", error="invalid_token",
121- error_description="Malformed token"
122-
123-That is the informative result. `invalid_token` rather than
124-`invalid_dpop_proof` means the PDS parsed and accepted the proof — signature,
125-claims, `ath`, `htu`, `htm` — and rejected only the credential. The challenge
126-scheme coming back as `DPoP` confirms the Authorization scheme too, and the
127-response carried a host nonce.
128-
129-`ath` is pinned against a value computed with `openssl dgst` rather than
130-checked for self-consistency, since a self-consistent hash would pass any
131-internal assertion and still be refused by every server.
132-
133-Two bugs the probe caught, both of which would have failed every authenticated
134-call in ways that are hard to read from the outside:
135-
136-- `htm` was built as `($httpMethod)[4..^1]`, assuming the enum stringified to
137- `HttpGet`. It stringifies to `GET`, so this was a range error, not a
138- mis-signed claim — loud, but only once something ran.
139-- `hostOf` returned an empty host for a scheme-less input like
140- `public.api.bsky.app`, because `parseUri` files that under `path`. URLs came
141- out as `https:///xrpc/...`.
142-
143-A third mistake was in the probe rather than the code: the first version
144-pointed at `com.atproto.repo.listRecords`, which is **public**. The PDS served
145-it with a 200 while ignoring the bogus credential entirely — a passing test
146-that tested nothing. `com.atproto.server.getSession` has to evaluate the
147-credential.
148-
149-## 5. Jetstream — works, verified against the live firehose
150-
151-`atproto/jetstream.nim` consumes the ATProto firehose as JSON over a
152-websocket, using `websock` (already vendored via chronos). Jetstream is public
153-and unauthenticated, so unlike the OAuth and XRPC probes this one checks the
154-real thing end to end: 17/17 against the live network.
155-
156-Two capabilities that matter to Glean are confirmed rather than assumed:
157-
158-**Server-side filtering.** The probe subscribes to one collection and asserts
159-that nothing else arrives. This is worth checking explicitly because the
160-failure mode is invisible: an unfiltered subscription is the entire network,
161-so a consumer that loses its query parameters still receives a healthy-looking
162-torrent of events that happen to be the wrong ones.
163-
164-**Cursor replay.** Rewinding to an earlier `time_us` must actually rewind, not
165-quietly resume live — that is how a consumer catches up after downtime. The
166-probe replays from a cursor it saw and asserts the events *overlap* what it
167-already received (25 of 25). Accepting the parameter proves nothing on its
168-own.
169-
170-Two traps:
171-
172-- **websock's `Uri` overload silently drops the query string.** It forwards
173- only `uri.path`, so every Jetstream filter would vanish and the consumer
174- would subscribe to the whole firehose while looking fine. `connect` uses the
175- host/path form and assembles path-plus-query itself.
176-- Jetstream wants `wantedCollections` **repeated per value**, not one
177- comma-joined string. A joined value matches no collection at all.
178-
179-`websock` pulls in chronicles, which does not compile unless a logging stream
180-is configured and resolves it in the *importing* module's scope — hence the
181-`import chronicles` in `jetstream.nim` that nothing appears to use, and the
182-`chronicles_enabled=off` in `nim.cfg`.
183-
184-## 6. Reconnect and backoff — works
185-
186-`atproto/jetstream_run.nim` keeps a subscription alive. Three concerns, all
187-about not losing events rather than about uptime:
188-
189-- **Resume from a rewound cursor.** Cursors are persisted in batches, so the
190- last one written always trails the last one seen. Reconnecting rewinds past
191- that gap and accepts re-delivery: duplicates are cheap when handling is
192- idempotent, a hole in the record is not.
193-- **Rotate connections on a timer.** Server-side filters are fixed for the
194- life of a connection, so a subscription that never reconnects never learns
195- about newly-known DIDs. Glean's Go consumer rotates every 15 minutes for
196- this reason, not for robustness.
197-- **Back off with jitter.** The Go consumer waits a flat 5s, which is fine
198- against one healthy server and unkind to one that is down.
199-
200-A healthy server proves nothing here, so `reconnect_probe.nim` runs a local
201-websocket server that hangs up after three frames. Over two seconds the
202-supervisor made 12 connections and received 36 events across them.
203-
204-The probe checks that the cursor **reaches the server**, not merely that it
205-advances locally: a consumer that tracks a cursor it never sends looks
206-identical from the inside and silently loses the gap on every reconnect. 11
207-of 12 connections carried one, and the first correctly did not — it starts
208-live.
209-
210-One bug worth recording. `delayFor` short-circuited `attempt <= 0` straight
211-to the initial delay, skipping jitter on the first retry. That is precisely
212-the moment jitter exists for: one server restart drops every consumer at
213-once, and un-jittered first retries bring them all back simultaneously.
214-
215-## 7. Session persistence — works
216-
217-`atproto/sqlite.nim` is a thin wrapper over the library the storage spike
218-already proved (sqlite-vec linked in, WAL on). `atproto/store.nim` keeps the
219-three things a restart must not lose, in the same tables the Go backend uses:
220-`oauth_auth_requests`, `oauth_sessions`, `jetstream_cursor`. Rows hold JSON,
221-as on the Go side, so adding a session field does not mean a migration.
222-
223-`atproto/cursor_store.nim` bridges the store and the supervisor, batching
224-writes: the cursor advances on every event, often hundreds a second, and a
225-synchronous write per event would make the database the bottleneck for a
226-firehose. The stored cursor therefore trails the live one — which is exactly
227-the gap the supervisor's rewind exists to cover.
228-
229-The check that earns its keep is the DPoP key. A missing row is loud; a key
230-that round-trips *almost* correctly is not — every request after a restart
231-would fail authentication with nothing pointing at storage. So the reloaded
232-key is checked by **signing with it and verifying against the original public
233-point**, with a second key as a negative control. Comparing fields, or
234-re-signing and comparing bytes, would prove neither thing: ECDSA is
235-randomised.
236-
237-Two other properties worth pinning:
238-
239-- Taking a pending auth request **consumes** it. An authorization code is
240- single-use, and leaving the row behind invites a replayed callback.
241-- Cursors never move backwards. A late event that rewound the stored cursor
242- would rewind the subscription on the next reconnect.
243-
244-The public half of the key is stored rather than re-derived on load. Deriving
245-it needs a curve operation, and a stored key that cannot be read back without
246-one is a restart that silently drops every session.
247-
248-**These rows are key material.** A session holds the DPoP private key and the
249-refresh token; together they are the account. The database file wants the
250-protection of a password store, not of a cache.
251-
252-## 8. The database layer — foundation done
253-
254-`atproto/gleandb.nim` opens Glean's storage the way the Go backend does:
255-`<base>_users` as `main`, with `_articles` and `_recs` **attached**, so a
256-query can join across them while each file still checkpoints and vacuums
257-alone — and so the derived recommendation tables can be dropped and rebuilt
258-without touching anything a user wrote. The suffixes match, so a database
259-written by either implementation is readable by the other.
260-
261-Pragmas are repeated per schema rather than set once, which is the detail
262-worth knowing: `PRAGMA` applies to one database at a time, so setting WAL on
263-`main` leaves `articles` journalling the slow way. The probe asserts WAL on
264-all three rather than trusting the call.
265-
266-`exp()` and `log()` are registered as SQL functions, as the Go connection
267-does, and they are load-bearing: the clustering SQL calls `EXP` and `LOG` in
268-seven places for time decay and popularity normalisation. SQLite only ships
269-maths functions when built with `SQLITE_ENABLE_MATH_FUNCTIONS`, so without
270-these registrations every recommendation query fails at runtime rather than
271-at startup.
272-
273-`log()` returns NULL for zero and negatives where Go's `math.Log` returns
274-`-inf` and NaN. The call sites all pass `1 + count`, so the argument is never
275-below 1 and the two agree in practice -- but it is a divergence, noted here
276-rather than left to be found.
277-
278-**The schema is extracted, not retyped.** `tools/extract_schema.py` pulls all
279-44 DDL statements out of `internal/db/db.go` into `atproto/schema.nim`.
280-Hand-transcribing them would drift, and a column differing by a default or a
281-CHECK constraint surfaces months later as a violation nobody can place.
282-
283-The probe leans on things that fail quietly: an ATTACH that did not happen, a
284-pragma that did not take, FTS5 triggers that never fire. That last one is the
285-sharpest — the index is maintained by triggers, so if they were missed,
286-search returns nothing while every other query looks healthy. Insert, update
287-and delete are each checked against the index.
288-
289-### Scope
34+Parameters are read from the query string or a form body, whichever the
35+caller used.
29036
291-This is the foundation, not the whole layer. Glean's `internal/db` is ~2,800
292-lines of non-test code; what is ported here is the connection, the schema and
293-the primitives, plus a cross-database join, FTS5 and vec0 proven end to end.
37+## What it does not do
29438
295-The remaining query surface — `article.go`, `feed.go`, `social.go`,
296-`retention.go` and friends — is mechanical against this base, and large. It
297-is the first part of the port where the work is volume rather than risk.
39+No accounts, no sessions, no CSRF, no ATProto. It runs on your machine and
40+reads your feeds, and that assumption is what keeps it a few hundred lines
41+rather than a few thousand.
29842
299-## 9. Article and feed queries
43+The one visible consequence: **do not expose it to a network.** It binds
44+loopback deliberately. Anything that can reach the port can read and change
45+everything, because there is nobody to distinguish from anybody else.
30046
301-`atproto/feedstore.nim` and `atproto/articlestore.nim` port the core of
302-`internal/db/feed.go` and `internal/db/article.go`, with `sqlite.nim` grown to
303-carry them: typed nullable row access, transactions, and prepared statements
304-that can be stepped many times.
47+## What it is built on
30548
306-Read state is the shape everything else follows. It is per-user and lives in
307-its own table, so a listing is always a LEFT JOIN and "unread" means *no row*
308-— hence `COALESCE(r.is_read, 0)` throughout rather than a column on the
309-article.
49+Six modules under `reader/`:
31050
311-Two behaviours carried over deliberately, both about not resurfacing things a
312-reader has dealt with:
51+| | |
52+|---|---|
53+| `feedfetcher` | HTTP with a retry policy that retries only 429 and 5xx |
54+| `feedparser` | RSS 2.0, RDF, Atom and JSON Feed |
55+| `scraper` | readability-style extraction, output sanitised |
56+| `sqlite` | a thin wrapper: prepare, bind, step, transactions |
57+| `gleandb` | two attached databases and the schema |
58+| `feedstore` / `articlestore` | the queries |
31359
314-- Articles already older than the retention window are **not ingested**.
315- Writing them only to purge them churns the database.
316-- After ingest, read state is **restored from history**, because a purged
317- article that reappears gets a new surrogate id and would otherwise come
318- back unread.
60+Each carries its own checks:
31961
320-The probe targets the ways these can be quietly wrong: a mark-all-read whose
321-scope clears the account instead of one feed, a subscriber count that drifts,
322-a listing that leaks another user's feeds, a NULL date sorting to the top.
323-
324-### Four bugs the probe caught
325-
326-All four were wrong assumptions about the schema or about SQLite, and none
327-would have raised at compile time:
328-
329-- `likes` keys on **`author_did`**, not `user_did` — a like is an ATProto
330- record with an author, not a user-scoped row.
331-- FTS5 requires the **bare table name** on the left of `MATCH`. Neither an
332- alias nor a schema-qualified name is accepted, so the match moved into a
333- subquery, which also carries `rank` out for ordering.
334-- The shared feed column list used bare names, which work until a query joins
335- `subscriptions` — which also has `feed_url` — and then fail as ambiguous.
336- Now qualified with the table alias.
337-- The probe's own first draft assumed `users` had a `handle` column. It does
338- not; handles are resolved from atproto rather than stored.
339-
340-### Scope
341-
342-Roughly the core of the two files by usage: feeds, subscriptions, listing,
343-filtering, read state, counts, search, retention, dead feeds. Not ported:
344-`BatchReconcileSubscriptions`, OPML-adjacent helpers,
345-`GetNextArticleID`'s navigation logic, and the several list variants the
346-recommendation engine uses. They are mechanical against these patterns.
347-
348-## 10. Likes, annotations, follows and trending
349-
350-`atproto/socialstore.nim` ports `internal/db/social.go`, completing the
351-database layer.
352-
353-Likes and annotations are ATProto records first and rows second. They are
354-keyed by their `at://` URI because that is the identity the PDS owns, and
355-each can arrive twice -- once from the firehose and once from a PDS sync --
356-so every write is idempotent by construction rather than by checking first.
357-An edit keeps the URI and changes the content, which is why annotations
358-upsert rather than insert.
359-
360-Two details worth keeping:
361-
362-- The language filter always keeps `language = ''`. That means *not yet
363- classified*, not *no language*; excluding it would make new articles vanish
364- for anyone with a filter set and reappear minutes later once classified.
365-- Future-dated articles sort last in trending. A feed publishing with a
366- scheduled timestamp would otherwise pin itself to the top indefinitely.
367-
368-The probe weights toward the destructive path. `deleteOrphanedLikes` and its
369-annotation twin delete local rows the PDS no longer has, so the checks assert
370-both halves: the dropped rows go, and *other* users' rows do not. The empty
371-active set is called out in the code, because "you unliked everything" and
372-"the fetch failed" are indistinguishable from inside the store — the caller
373-has to not get there.
374-
375-## 11. Feed parsing
376-
377-`atproto/feedparser.nim` ports `internal/feed/parser.go`: RSS 2.0, RDF
378-(RSS 1.0), Atom and JSON Feed.
379-
380-The interesting part of a feed parser is not the happy path, it is the
381-fallbacks, because feeds in the wild disagree with the spec constantly. A
382-missing `guid` falls back to the link — it is what dedupes an article across
383-fetches, so an item without one would be re-inserted forever. `rel="self"`
384-loses to `rel="alternate"`, or a feed whose first link is itself would point
385-every article back at the feed. An unparseable date leaves the article
386-undated rather than dropping it. Atom's `updated` stands in for a missing
387-`published`, since Atom requires only the former.
388-
389-Namespaces are matched by local name rather than resolved. Go's decoder maps
390-`content:encoded` to its namespace URI; Nim's keeps the literal prefix, and
391-feeds are inconsistent enough about prefixes that suffix matching is both
392-simpler and more forgiving.
393-
394-### Two bugs, one of them quiet and severe
395-
396-**`xmltree.innerText` silently skips CDATA.** RSS ships article bodies inside
397-CDATA *precisely because* they contain HTML, so relying on `innerText` loses
398-the content of most real feeds while every surrounding field parses
399-perfectly — a parser that looks like it works and returns empty articles.
400-The parser now walks the tree itself, including `xnCData`.
401-
402-**Nim's `zzz` requires `-05:00`; `ZZZ` accepts `-0500`.** Feeds use both, so
403-both formats are in the list. Without it, every RSS feed using a numeric
404-offset — which is most of them — parsed as undated.
405-
406-Neither would have failed a compile, and neither is visible without a
407-fixture that contains the case.
408-
409-### Beyond fixtures
410-
411-Fixtures only contain what their author thought of, so the probe also fetches
412-four real feeds across two formats (Rust Blog, nim-lang, LWN, Hacker News)
413-and asserts every item has a guid and a title. All four parse, 65 items
414-between them, all dated and all carrying text.
415-
416-## 12. Fetching feeds
417-
418-`atproto/feedfetcher.nim` ports `internal/feed/fetcher.go`. Retries exist for
419-servers that are briefly unwell -- 429 and 5xx -- and for nothing else. A 404
420-will still be a 404 in a second, and bytes that do not parse will not parse
421-differently on a second read.
422-
423-The retry policy is the substance, so the probe drives it with a scripted
424-transport: 503 twice then success, a 404 that must not be retried, a 429
425-naming its own delay. None of those are reproducible against a real host on
426-demand.
427-
428-### A bug in the Go fetcher
429-
430-The Go retry policy does not do what it reads as doing. `executeRequest`
431-returns its `*http.Response` **only on success**, so in `Fetch` the guard
432-
433-```go
434-if resp != nil && !httpclient.IsRetryable(resp.StatusCode) {
62+```
63+nimble test # no network needed
64+nimble testnet # parses and scrapes real sites
43565 ```
43666
437-can never be true when `err != nil` -- the response is always nil on the
438-error path. Three things follow:
439-
440-- Every failure retries the full four attempts regardless of status, so a
441- permanently-404 feed costs four requests and ~7s of backoff on every
442- refresh cycle.
443-- `retryBackoff`'s `Retry-After` handling is unreachable for the same reason:
444- `lastResp` is only non-nil once the function has already returned, so a
445- rate-limited server's own delay is never honoured.
446-- The HTTP status is never checked at all, so an error page is handed to the
447- parser and fails there instead of failing as an HTTP error.
448-
449-This port checks the status, honours `Retry-After` (capped, so a server
450-asking for an hour cannot hold a worker for one), and retries only what is
451-worth retrying. That is a behaviour change rather than a translation, which
452-is why it is written down both here and at the top of the module.
453-
454-Worth fixing on the Go side too; it is a handful of lines.
455-
456-## 13. The scraper
457-
458-`atproto/scraper.nim` ports `internal/scraper/scraper.go`. Feeds routinely
459-ship a first paragraph and a "read more" link, so the reader view needs the
460-page itself: find the node most likely to be the article, strip the furniture
461-around it, and re-render a whitelisted subset of HTML. Output stays HTML
462-rather than text, because paragraphs, headings, lists and code blocks are
463-most of what makes a long article readable.
464-
465-### Sanitisation is the real test
466-
467-This renders HTML fetched from an arbitrary site into a reader's client, so
468-the probe weights toward what must *not* survive: `script` and `style`
469-content, `iframe`, event handlers, unknown data attributes, `javascript:`
470-URLs. Attributes are whitelisted rather than blacklisted, so an attribute
471-nobody anticipated is dropped by default instead of needing to be added to a
472-list. A page that renders correctly proves nothing about what was stripped,
473-so each of those is asserted separately.
474-
475-Links carry `rel="noopener noreferrer"`, and a link whose target is not
476-http(s) is flattened to its text -- the words survive without the anchor.
477-
478-Two smaller details worth keeping:
479-
480-- A lazy-loaded image resolves through `src` → `data-src` → `data-lazy-src`
481- to the first *reachable* URL. Rendering the placeholder gives a broken
482- image where the picture should be.
483-- Removal walks children back to front. Deleting by index shifts every
484- sibling after it, so a forward walk skips one after each removal -- the
485- probe has four consecutive `<nav>` elements specifically to catch that.
486-
487-The extractor refuses rather than returning something thin. A caller handed a
488-nav menu cannot tell it from a short article; an error lets it fall back to
489-the feed's own summary.
490-
491-The live check takes its URL from the Rust blog's feed rather than
492-hard-coding one, after a hard-coded guess 404'd -- a probe that rots when
493-someone reorganises their permalinks is worse than no probe.
494-
495-## 14. Similarity and scoring
496-
497-`atproto/cluster.nim` ports the similarity computations from
498-`internal/cluster/jaccard.go` and the pure helpers from `scoring.go`.
499-
500-Both similarities are precomputed on a schedule and staged into a temp table
501-before being swapped in, so the expensive pass holds no lock on the live
502-table and readers keep seeing the previous generation rather than an empty
503-one.
504-
505-Feed similarity is time-decayed, `exp(-0.023 * days)` -- a half-life of about
506-a month. Without it a long-lived account's earliest subscriptions dominate
507-its recommendations forever. User similarity over *subscriptions* is
508-deliberately undecayed, because a subscription is a standing choice rather
509-than an event; shared *likes* are events, so those decay on both sides.
510-
511-Recommendations have no obviously correct answer, so the probe checks
512-properties rather than values: that decay decays, that each pair is stored
513-once, that recomputing replaces rather than accumulates, that normalisation
514-of an all-equal set does not produce NaN, and that `samplePeople` reserves
515-half its slots for people the reader does *not* follow -- ranking purely by
516-similarity returns the network they already have, which is not discovery.
517-
518-### A bug inherited from the Go implementation
519-
520-`common_likes` is computed as `CAST(SUM(decay * decay) AS INTEGER)`. Each
521-decay factor is at most 1, so their product is always just under it: a pair
522-sharing one recent like sums to 0.9999 and **truncates to 0**. Two shared
523-likes store 1. The number is shown to readers as the reason for a
524-recommendation -- "N shared likes" -- so a real overlap reads as none.
525-
526-This port uses `ROUND` instead. Worth fixing on the Go side too.
527-
528-The probe caught it only because the assertion looked *vacuous*: both the
529-recent and the stale pair returned 0, so "recent >= old" passed while
530-proving nothing. A test that passes for the wrong reason is worse than one
531-that fails.
532-
533-## What this does and does not prove
534-
535-Proven: the cryptography, the storage layer, and the OAuth flow up to user
536-consent all work from Nim, in the formats ATProto and real servers accept.
537-
538-Not proven, and still ahead:
539-
540-- **Token exchange and refresh** — written, but needs a browser consent to
541- exercise. This is the one place a real end-to-end test still has to happen,
542- and it also covers the `invalid_token` refresh-and-retry branch of
543- `xrpc.request`, which is currently unexercised.
544-
545-### Not needed: DAG-CBOR, CAR and the MST
546-
547-Earlier notes listed these as remaining work. That was wrong, and worth
548-correcting because they are by far the largest format risk left in ATProto.
549-
550-Glean never parses them. `internal/atproto/sync.go` opens by saying so: it
551-reconciles with `com.atproto.repo.listRecords` rather than
552-`com.atproto.sync.getRepo`, because it only syncs known users, Jetstream
553-covers real-time events, and every reconcile is idempotent. Reads go through
554-`listRecords`/`getRecord` and writes through
555-`createRecord`/`putRecord`/`deleteRecord` — all JSON.
556-
557-The imports confirm it. Glean uses four indigo packages — `atclient`,
558-`auth/oauth`, `identity`, `syntax` — and none of `repo`, `mst`, `cbor` or
559-`car`.
560-
561-A Nim port inherits that choice for free. Implementing DAG-CBOR, CIDs, CAR
562-framing and Merkle Search Tree traversal would be real work in service of a
563-code path this application does not have.
564-
565-The honest read: none of the three things that could have sunk the port did.
566-What is left is protocol plumbing against well-specified formats — large, but
567-ordinary.
67+Storage keeps a single hard-coded user. The schema began life in a
68+multi-user server, and satisfying that one column was cheaper than
69+rewriting every query around it.
56870
569-## Caveats in this code
71+## History
57072
571-- `identity.nim` shells out to `dig` for the DNS leg of handle resolution,
572- because Nim's stdlib has no resolver. A missing `dig` degrades to the
573- HTTP `.well-known` fallback rather than failing, but a real port should bind
574- a resolver instead of depending on a binary.
575-- Only public and localhost clients are supported. Confidential clients
576- (`private_key_jwt` client assertions) are not implemented; Glean uses a
577- public client today.
73+This started as a port of a much larger reader — one with accounts, ATProto
74+sync, social recommendations and a clustering engine. That half has been
75+deleted. What remains is the part that reads feeds.
@@ -1,577 +1,75 @@
1-# Nim port spike1+# glean — a small RSS reader
2 2
3-Two questions had to be answered before committing to a Nim rewrite of the Go3+One binary, one SQLite database, no accounts.
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 4
10 ```5 ```
11-cd spike/nim6+nim c -d:ssl glean.nim
12-nimble test # everything that needs no network7+./glean
13-nimble testnet # the probes that talk to real servers
14-nimble testall # both
15-nimble schema # regenerate atproto/schema.nim from the Go source
16 ```8 ```
17 9
18-Reusable code lives in `atproto/`; the `*_probe.nim` files are the checks.10+Listens on `127.0.0.1:8080` and stores its database under
19-They are split by whether they touch the network: a failure against11+`~/.local/share/glean/`. Override with `GLEAN_PORT` and `GLEAN_DB`.
20-`bsky.social` can mean that server is having a bad day rather than that the12+
21-port is broken, and a suite that cannot tell those apart stops being13+## API
22-believed. Individual probes still run on their own:
23 14
24 ```15 ```
25-nim c -r -d:ssl <name>_probe.nim16+GET /feeds subscriptions, with unread counts
17+POST /feeds url=… subscribe and fetch immediately
18+DELETE /feeds?url=… unsubscribe
19+
20+GET /articles unread, newest first
21+GET /articles?status=all|read …or everything, or just the read ones
22+GET /articles?feed=… …scoped to one feed
23+GET /articles?q=… full-text search (FTS5)
24+GET /articles/:id one article
25+
26+POST /read id=… mark read (undo=1 to reverse)
27+POST /read-all [feed=…] mark a feed, or everything, read
28+POST /refresh fetch every subscribed feed
29+POST /fetch-content id=… scrape the full article text
30+
31+GET /unread unread count
26 ```32 ```
27 33
28-## 1. sqlite-vec + FTS5 — works34+Parameters are read from the query string or a form body, whichever the
29-35+caller used.
30-Glean needs FTS5 for article search and sqlite-vec for embedding similarity.
31-Both are C-level concerns, so the question was never "does Nim have a SQLite
32-library" but whether the same extension links and registers the same way.
33-
34-`sqlite_probe.nim` compiles the same `sqlite-vec.c` the Go bindings vendor,
35-calls `sqlite3_vec_init` directly, and exercises `vec_version()`, an FTS5
36-`MATCH`, a `vec0` virtual table, float32 blob binding, and a KNN query.
37-All pass; KNN returns the expected nearest neighbour.
38-
39-**One trap, and it is a nasty one.** sqlite-vec must be compiled with
40-`-DSQLITE_CORE`. Without it the file builds as a *loadable extension* whose
41-entry point takes its SQLite API pointer from the caller — pass `nil` there and
42-it segfaults inside `sqlite3_vec_init` rather than returning an error. Nim's
43-`{.compile: (file, flags).}` tuple form did not reliably apply the define, so
44-it lives in `nim.cfg` where it cannot be forgotten.
45-
46-## 2. ES256 / DPoP — works
47-
48-This was the real risk. Glean gets OAuth from `bluesky-social/indigo`; Nim has
49-no ATProto library, so a port has to build the DPoP layer itself, and
50-everything downstream of sign-in depends on it.
51-
52-`dpop_probe.nim` generates a P-256 key, builds the public JWK, computes an
53-RFC 7638 thumbprint, and signs a DPoP proof JWT (RFC 9449).
54-
55-It uses **bearssl** rather than OpenSSL for one specific reason: JOSE requires
56-the signature as raw `R||S`, 64 bytes, and most crypto libraries hand you DER.
57-`br_ecdsa_sign_raw` emits `R||S` directly, so there is no DER-unwrapping step
58-to get subtly wrong.
59-
60-Self-verification would only prove internal consistency, so the token goes to
61-`verify/verify.go`, which re-implements the verifier from Go's standard
62-library — the same stack the current server runs on. It checks the header
63-claims, that the signature is 64 raw bytes and not DER-wrapped, that the
64-signature verifies against the JWK, that a tampered payload is *rejected*, and
65-that the thumbprint matches when recomputed independently.
66-
67-19/19 checks pass, and the whole cycle passes across 30 independently generated
68-keypairs — worth confirming, because ECDSA values with leading zeros are a
69-classic source of intermittent length bugs in `R||S` encoding.
70-
71-## 3. The OAuth flow — works against a live PDS
72-
73-`atproto/identity.nim` and `atproto/oauth.nim` implement resolution and the
74-client flow, ported from the shape of indigo's `atproto/auth/oauth`.
75-`oauth_probe.nim` drives it end to end against real servers:
76-
77- handle -> DID -> PDS -> auth server -> metadata -> PAR -> authorize URL
78-
79-Against `bsky.social` all 15 checks pass, including a **real PAR that the
80-production auth server accepts** and returns a `request_uri` for. That single
81-result exercises the DPoP signature, the nonce retry and PKCE at once — a
82-server-side validation no amount of local testing substitutes for.
83-
84-Two things worth recording:
85-
86-**The nonce retry is load-bearing, and the probe proves it.** A PDS rejects the
87-first DPoP-signed request of a session with `400 use_dpop_nonce` and supplies a
88-`DPoP-Nonce` header to retry with. Since `sendAuthRequest` retries
89-transparently, the probe separately fires an unnonced request and asserts the
90-rejection — otherwise a server that never demanded a nonce would look identical
91-to a working retry.
92-
93-**Order of validation caught me out.** That check first sent a stub body and got
94-`invalid_request`, not `use_dpop_nonce`: bsky.social validates request
95-parameters *before* the nonce. The body has to be fully valid to observe the
96-nonce requirement at all.
97-
98-The flow stops at the authorization URL, which is correct — the next step is a
99-human granting consent in a browser. `exchangeCode` and `refresh` are
100-implemented against that callback but are consequently untested.
101-
102-## 4. XRPC signing — proof accepted by a real PDS
103-
104-`atproto/xrpc.nim` makes authenticated calls with a DPoP-bound access token:
105-`Authorization: DPoP <token>` (not Bearer), a per-request proof, and the retry
106-logic for both failure modes.
107-
108-Resource-server proofs are **not** the same as auth-server proofs, even when
109-one host fills both roles. A PDS request carries `ath` (base64url SHA-256 of
110-the access token, which is what binds the proof to that one credential) and
111-`iss`; the PAR and token requests that mint the token carry neither. The nonce
112-handshake also differs: a PDS answers **401** with `WWW-Authenticate: DPoP
113-error="use_dpop_nonce"`, where the auth server used **400** and a JSON body.
114-A session therefore tracks two independent nonces.
115-
116-`xrpc_probe.nim` cannot get a consented token, so it signs a request with a
117-deliberately invalid one and reads what a real PDS says. Against
118-`bsky.social`'s PDS:
119-
120- WWW-Authenticate: DPoP algs="...", error="invalid_token",
121- error_description="Malformed token"
122-
123-That is the informative result. `invalid_token` rather than
124-`invalid_dpop_proof` means the PDS parsed and accepted the proof — signature,
125-claims, `ath`, `htu`, `htm` — and rejected only the credential. The challenge
126-scheme coming back as `DPoP` confirms the Authorization scheme too, and the
127-response carried a host nonce.
128-
129-`ath` is pinned against a value computed with `openssl dgst` rather than
130-checked for self-consistency, since a self-consistent hash would pass any
131-internal assertion and still be refused by every server.
132-
133-Two bugs the probe caught, both of which would have failed every authenticated
134-call in ways that are hard to read from the outside:
135-
136-- `htm` was built as `($httpMethod)[4..^1]`, assuming the enum stringified to
137- `HttpGet`. It stringifies to `GET`, so this was a range error, not a
138- mis-signed claim — loud, but only once something ran.
139-- `hostOf` returned an empty host for a scheme-less input like
140- `public.api.bsky.app`, because `parseUri` files that under `path`. URLs came
141- out as `https:///xrpc/...`.
142-
143-A third mistake was in the probe rather than the code: the first version
144-pointed at `com.atproto.repo.listRecords`, which is **public**. The PDS served
145-it with a 200 while ignoring the bogus credential entirely — a passing test
146-that tested nothing. `com.atproto.server.getSession` has to evaluate the
147-credential.
148-
149-## 5. Jetstream — works, verified against the live firehose
150-
151-`atproto/jetstream.nim` consumes the ATProto firehose as JSON over a
152-websocket, using `websock` (already vendored via chronos). Jetstream is public
153-and unauthenticated, so unlike the OAuth and XRPC probes this one checks the
154-real thing end to end: 17/17 against the live network.
155-
156-Two capabilities that matter to Glean are confirmed rather than assumed:
157-
158-**Server-side filtering.** The probe subscribes to one collection and asserts
159-that nothing else arrives. This is worth checking explicitly because the
160-failure mode is invisible: an unfiltered subscription is the entire network,
161-so a consumer that loses its query parameters still receives a healthy-looking
162-torrent of events that happen to be the wrong ones.
163-
164-**Cursor replay.** Rewinding to an earlier `time_us` must actually rewind, not
165-quietly resume live — that is how a consumer catches up after downtime. The
166-probe replays from a cursor it saw and asserts the events *overlap* what it
167-already received (25 of 25). Accepting the parameter proves nothing on its
168-own.
169-
170-Two traps:
171-
172-- **websock's `Uri` overload silently drops the query string.** It forwards
173- only `uri.path`, so every Jetstream filter would vanish and the consumer
174- would subscribe to the whole firehose while looking fine. `connect` uses the
175- host/path form and assembles path-plus-query itself.
176-- Jetstream wants `wantedCollections` **repeated per value**, not one
177- comma-joined string. A joined value matches no collection at all.
178-
179-`websock` pulls in chronicles, which does not compile unless a logging stream
180-is configured and resolves it in the *importing* module's scope — hence the
181-`import chronicles` in `jetstream.nim` that nothing appears to use, and the
182-`chronicles_enabled=off` in `nim.cfg`.
183-
184-## 6. Reconnect and backoff — works
185-
186-`atproto/jetstream_run.nim` keeps a subscription alive. Three concerns, all
187-about not losing events rather than about uptime:
188-
189-- **Resume from a rewound cursor.** Cursors are persisted in batches, so the
190- last one written always trails the last one seen. Reconnecting rewinds past
191- that gap and accepts re-delivery: duplicates are cheap when handling is
192- idempotent, a hole in the record is not.
193-- **Rotate connections on a timer.** Server-side filters are fixed for the
194- life of a connection, so a subscription that never reconnects never learns
195- about newly-known DIDs. Glean's Go consumer rotates every 15 minutes for
196- this reason, not for robustness.
197-- **Back off with jitter.** The Go consumer waits a flat 5s, which is fine
198- against one healthy server and unkind to one that is down.
199-
200-A healthy server proves nothing here, so `reconnect_probe.nim` runs a local
201-websocket server that hangs up after three frames. Over two seconds the
202-supervisor made 12 connections and received 36 events across them.
203-
204-The probe checks that the cursor **reaches the server**, not merely that it
205-advances locally: a consumer that tracks a cursor it never sends looks
206-identical from the inside and silently loses the gap on every reconnect. 11
207-of 12 connections carried one, and the first correctly did not — it starts
208-live.
209-
210-One bug worth recording. `delayFor` short-circuited `attempt <= 0` straight
211-to the initial delay, skipping jitter on the first retry. That is precisely
212-the moment jitter exists for: one server restart drops every consumer at
213-once, and un-jittered first retries bring them all back simultaneously.
214-
215-## 7. Session persistence — works
216-
217-`atproto/sqlite.nim` is a thin wrapper over the library the storage spike
218-already proved (sqlite-vec linked in, WAL on). `atproto/store.nim` keeps the
219-three things a restart must not lose, in the same tables the Go backend uses:
220-`oauth_auth_requests`, `oauth_sessions`, `jetstream_cursor`. Rows hold JSON,
221-as on the Go side, so adding a session field does not mean a migration.
222-
223-`atproto/cursor_store.nim` bridges the store and the supervisor, batching
224-writes: the cursor advances on every event, often hundreds a second, and a
225-synchronous write per event would make the database the bottleneck for a
226-firehose. The stored cursor therefore trails the live one — which is exactly
227-the gap the supervisor's rewind exists to cover.
228-
229-The check that earns its keep is the DPoP key. A missing row is loud; a key
230-that round-trips *almost* correctly is not — every request after a restart
231-would fail authentication with nothing pointing at storage. So the reloaded
232-key is checked by **signing with it and verifying against the original public
233-point**, with a second key as a negative control. Comparing fields, or
234-re-signing and comparing bytes, would prove neither thing: ECDSA is
235-randomised.
236-
237-Two other properties worth pinning:
238-
239-- Taking a pending auth request **consumes** it. An authorization code is
240- single-use, and leaving the row behind invites a replayed callback.
241-- Cursors never move backwards. A late event that rewound the stored cursor
242- would rewind the subscription on the next reconnect.
243-
244-The public half of the key is stored rather than re-derived on load. Deriving
245-it needs a curve operation, and a stored key that cannot be read back without
246-one is a restart that silently drops every session.
247-
248-**These rows are key material.** A session holds the DPoP private key and the
249-refresh token; together they are the account. The database file wants the
250-protection of a password store, not of a cache.
251-
252-## 8. The database layer — foundation done
253-
254-`atproto/gleandb.nim` opens Glean's storage the way the Go backend does:
255-`<base>_users` as `main`, with `_articles` and `_recs` **attached**, so a
256-query can join across them while each file still checkpoints and vacuums
257-alone — and so the derived recommendation tables can be dropped and rebuilt
258-without touching anything a user wrote. The suffixes match, so a database
259-written by either implementation is readable by the other.
260-
261-Pragmas are repeated per schema rather than set once, which is the detail
262-worth knowing: `PRAGMA` applies to one database at a time, so setting WAL on
263-`main` leaves `articles` journalling the slow way. The probe asserts WAL on
264-all three rather than trusting the call.
265-
266-`exp()` and `log()` are registered as SQL functions, as the Go connection
267-does, and they are load-bearing: the clustering SQL calls `EXP` and `LOG` in
268-seven places for time decay and popularity normalisation. SQLite only ships
269-maths functions when built with `SQLITE_ENABLE_MATH_FUNCTIONS`, so without
270-these registrations every recommendation query fails at runtime rather than
271-at startup.
272-
273-`log()` returns NULL for zero and negatives where Go's `math.Log` returns
274-`-inf` and NaN. The call sites all pass `1 + count`, so the argument is never
275-below 1 and the two agree in practice -- but it is a divergence, noted here
276-rather than left to be found.
277-
278-**The schema is extracted, not retyped.** `tools/extract_schema.py` pulls all
279-44 DDL statements out of `internal/db/db.go` into `atproto/schema.nim`.
280-Hand-transcribing them would drift, and a column differing by a default or a
281-CHECK constraint surfaces months later as a violation nobody can place.
282-
283-The probe leans on things that fail quietly: an ATTACH that did not happen, a
284-pragma that did not take, FTS5 triggers that never fire. That last one is the
285-sharpest — the index is maintained by triggers, so if they were missed,
286-search returns nothing while every other query looks healthy. Insert, update
287-and delete are each checked against the index.
288-
289-### Scope
290 36
291-This is the foundation, not the whole layer. Glean's `internal/db` is ~2,80037+## What it does not do
292-lines of non-test code; what is ported here is the connection, the schema and
293-the primitives, plus a cross-database join, FTS5 and vec0 proven end to end.
294 38
295-The remaining query surface — `article.go`, `feed.go`, `social.go`,39+No accounts, no sessions, no CSRF, no ATProto. It runs on your machine and
296-`retention.go` and friends — is mechanical against this base, and large. It40+reads your feeds, and that assumption is what keeps it a few hundred lines
297-is the first part of the port where the work is volume rather than risk.41+rather than a few thousand.
298 42
299-## 9. Article and feed queries43+The one visible consequence: **do not expose it to a network.** It binds
44+loopback deliberately. Anything that can reach the port can read and change
45+everything, because there is nobody to distinguish from anybody else.
300 46
301-`atproto/feedstore.nim` and `atproto/articlestore.nim` port the core of47+## What it is built on
302-`internal/db/feed.go` and `internal/db/article.go`, with `sqlite.nim` grown to
303-carry them: typed nullable row access, transactions, and prepared statements
304-that can be stepped many times.
305 48
306-Read state is the shape everything else follows. It is per-user and lives in49+Six modules under `reader/`:
307-its own table, so a listing is always a LEFT JOIN and "unread" means *no row*
308-— hence `COALESCE(r.is_read, 0)` throughout rather than a column on the
309-article.
310 50
311-Two behaviours carried over deliberately, both about not resurfacing things a51+| | |
312-reader has dealt with:52+|---|---|
53+| `feedfetcher` | HTTP with a retry policy that retries only 429 and 5xx |
54+| `feedparser` | RSS 2.0, RDF, Atom and JSON Feed |
55+| `scraper` | readability-style extraction, output sanitised |
56+| `sqlite` | a thin wrapper: prepare, bind, step, transactions |
57+| `gleandb` | two attached databases and the schema |
58+| `feedstore` / `articlestore` | the queries |
313 59
314-- Articles already older than the retention window are **not ingested**.60+Each carries its own checks:
315- Writing them only to purge them churns the database.
316-- After ingest, read state is **restored from history**, because a purged
317- article that reappears gets a new surrogate id and would otherwise come
318- back unread.
319 61
320-The probe targets the ways these can be quietly wrong: a mark-all-read whose62+```
321-scope clears the account instead of one feed, a subscriber count that drifts,63+nimble test # no network needed
322-a listing that leaks another user's feeds, a NULL date sorting to the top.64+nimble testnet # parses and scrapes real sites
323-
324-### Four bugs the probe caught
325-
326-All four were wrong assumptions about the schema or about SQLite, and none
327-would have raised at compile time:
328-
329-- `likes` keys on **`author_did`**, not `user_did` — a like is an ATProto
330- record with an author, not a user-scoped row.
331-- FTS5 requires the **bare table name** on the left of `MATCH`. Neither an
332- alias nor a schema-qualified name is accepted, so the match moved into a
333- subquery, which also carries `rank` out for ordering.
334-- The shared feed column list used bare names, which work until a query joins
335- `subscriptions` — which also has `feed_url` — and then fail as ambiguous.
336- Now qualified with the table alias.
337-- The probe's own first draft assumed `users` had a `handle` column. It does
338- not; handles are resolved from atproto rather than stored.
339-
340-### Scope
341-
342-Roughly the core of the two files by usage: feeds, subscriptions, listing,
343-filtering, read state, counts, search, retention, dead feeds. Not ported:
344-`BatchReconcileSubscriptions`, OPML-adjacent helpers,
345-`GetNextArticleID`'s navigation logic, and the several list variants the
346-recommendation engine uses. They are mechanical against these patterns.
347-
348-## 10. Likes, annotations, follows and trending
349-
350-`atproto/socialstore.nim` ports `internal/db/social.go`, completing the
351-database layer.
352-
353-Likes and annotations are ATProto records first and rows second. They are
354-keyed by their `at://` URI because that is the identity the PDS owns, and
355-each can arrive twice -- once from the firehose and once from a PDS sync --
356-so every write is idempotent by construction rather than by checking first.
357-An edit keeps the URI and changes the content, which is why annotations
358-upsert rather than insert.
359-
360-Two details worth keeping:
361-
362-- The language filter always keeps `language = ''`. That means *not yet
363- classified*, not *no language*; excluding it would make new articles vanish
364- for anyone with a filter set and reappear minutes later once classified.
365-- Future-dated articles sort last in trending. A feed publishing with a
366- scheduled timestamp would otherwise pin itself to the top indefinitely.
367-
368-The probe weights toward the destructive path. `deleteOrphanedLikes` and its
369-annotation twin delete local rows the PDS no longer has, so the checks assert
370-both halves: the dropped rows go, and *other* users' rows do not. The empty
371-active set is called out in the code, because "you unliked everything" and
372-"the fetch failed" are indistinguishable from inside the store — the caller
373-has to not get there.
374-
375-## 11. Feed parsing
376-
377-`atproto/feedparser.nim` ports `internal/feed/parser.go`: RSS 2.0, RDF
378-(RSS 1.0), Atom and JSON Feed.
379-
380-The interesting part of a feed parser is not the happy path, it is the
381-fallbacks, because feeds in the wild disagree with the spec constantly. A
382-missing `guid` falls back to the link — it is what dedupes an article across
383-fetches, so an item without one would be re-inserted forever. `rel="self"`
384-loses to `rel="alternate"`, or a feed whose first link is itself would point
385-every article back at the feed. An unparseable date leaves the article
386-undated rather than dropping it. Atom's `updated` stands in for a missing
387-`published`, since Atom requires only the former.
388-
389-Namespaces are matched by local name rather than resolved. Go's decoder maps
390-`content:encoded` to its namespace URI; Nim's keeps the literal prefix, and
391-feeds are inconsistent enough about prefixes that suffix matching is both
392-simpler and more forgiving.
393-
394-### Two bugs, one of them quiet and severe
395-
396-**`xmltree.innerText` silently skips CDATA.** RSS ships article bodies inside
397-CDATA *precisely because* they contain HTML, so relying on `innerText` loses
398-the content of most real feeds while every surrounding field parses
399-perfectly — a parser that looks like it works and returns empty articles.
400-The parser now walks the tree itself, including `xnCData`.
401-
402-**Nim's `zzz` requires `-05:00`; `ZZZ` accepts `-0500`.** Feeds use both, so
403-both formats are in the list. Without it, every RSS feed using a numeric
404-offset — which is most of them — parsed as undated.
405-
406-Neither would have failed a compile, and neither is visible without a
407-fixture that contains the case.
408-
409-### Beyond fixtures
410-
411-Fixtures only contain what their author thought of, so the probe also fetches
412-four real feeds across two formats (Rust Blog, nim-lang, LWN, Hacker News)
413-and asserts every item has a guid and a title. All four parse, 65 items
414-between them, all dated and all carrying text.
415-
416-## 12. Fetching feeds
417-
418-`atproto/feedfetcher.nim` ports `internal/feed/fetcher.go`. Retries exist for
419-servers that are briefly unwell -- 429 and 5xx -- and for nothing else. A 404
420-will still be a 404 in a second, and bytes that do not parse will not parse
421-differently on a second read.
422-
423-The retry policy is the substance, so the probe drives it with a scripted
424-transport: 503 twice then success, a 404 that must not be retried, a 429
425-naming its own delay. None of those are reproducible against a real host on
426-demand.
427-
428-### A bug in the Go fetcher
429-
430-The Go retry policy does not do what it reads as doing. `executeRequest`
431-returns its `*http.Response` **only on success**, so in `Fetch` the guard
432-
433-```go
434-if resp != nil && !httpclient.IsRetryable(resp.StatusCode) {
435 ```65 ```
436 66
437-can never be true when `err != nil` -- the response is always nil on the67+Storage keeps a single hard-coded user. The schema began life in a
438-error path. Three things follow:68+multi-user server, and satisfying that one column was cheaper than
439-69+rewriting every query around it.
440-- Every failure retries the full four attempts regardless of status, so a
441- permanently-404 feed costs four requests and ~7s of backoff on every
442- refresh cycle.
443-- `retryBackoff`'s `Retry-After` handling is unreachable for the same reason:
444- `lastResp` is only non-nil once the function has already returned, so a
445- rate-limited server's own delay is never honoured.
446-- The HTTP status is never checked at all, so an error page is handed to the
447- parser and fails there instead of failing as an HTTP error.
448-
449-This port checks the status, honours `Retry-After` (capped, so a server
450-asking for an hour cannot hold a worker for one), and retries only what is
451-worth retrying. That is a behaviour change rather than a translation, which
452-is why it is written down both here and at the top of the module.
453-
454-Worth fixing on the Go side too; it is a handful of lines.
455-
456-## 13. The scraper
457-
458-`atproto/scraper.nim` ports `internal/scraper/scraper.go`. Feeds routinely
459-ship a first paragraph and a "read more" link, so the reader view needs the
460-page itself: find the node most likely to be the article, strip the furniture
461-around it, and re-render a whitelisted subset of HTML. Output stays HTML
462-rather than text, because paragraphs, headings, lists and code blocks are
463-most of what makes a long article readable.
464-
465-### Sanitisation is the real test
466-
467-This renders HTML fetched from an arbitrary site into a reader's client, so
468-the probe weights toward what must *not* survive: `script` and `style`
469-content, `iframe`, event handlers, unknown data attributes, `javascript:`
470-URLs. Attributes are whitelisted rather than blacklisted, so an attribute
471-nobody anticipated is dropped by default instead of needing to be added to a
472-list. A page that renders correctly proves nothing about what was stripped,
473-so each of those is asserted separately.
474-
475-Links carry `rel="noopener noreferrer"`, and a link whose target is not
476-http(s) is flattened to its text -- the words survive without the anchor.
477-
478-Two smaller details worth keeping:
479-
480-- A lazy-loaded image resolves through `src` → `data-src` → `data-lazy-src`
481- to the first *reachable* URL. Rendering the placeholder gives a broken
482- image where the picture should be.
483-- Removal walks children back to front. Deleting by index shifts every
484- sibling after it, so a forward walk skips one after each removal -- the
485- probe has four consecutive `<nav>` elements specifically to catch that.
486-
487-The extractor refuses rather than returning something thin. A caller handed a
488-nav menu cannot tell it from a short article; an error lets it fall back to
489-the feed's own summary.
490-
491-The live check takes its URL from the Rust blog's feed rather than
492-hard-coding one, after a hard-coded guess 404'd -- a probe that rots when
493-someone reorganises their permalinks is worse than no probe.
494-
495-## 14. Similarity and scoring
496-
497-`atproto/cluster.nim` ports the similarity computations from
498-`internal/cluster/jaccard.go` and the pure helpers from `scoring.go`.
499-
500-Both similarities are precomputed on a schedule and staged into a temp table
501-before being swapped in, so the expensive pass holds no lock on the live
502-table and readers keep seeing the previous generation rather than an empty
503-one.
504-
505-Feed similarity is time-decayed, `exp(-0.023 * days)` -- a half-life of about
506-a month. Without it a long-lived account's earliest subscriptions dominate
507-its recommendations forever. User similarity over *subscriptions* is
508-deliberately undecayed, because a subscription is a standing choice rather
509-than an event; shared *likes* are events, so those decay on both sides.
510-
511-Recommendations have no obviously correct answer, so the probe checks
512-properties rather than values: that decay decays, that each pair is stored
513-once, that recomputing replaces rather than accumulates, that normalisation
514-of an all-equal set does not produce NaN, and that `samplePeople` reserves
515-half its slots for people the reader does *not* follow -- ranking purely by
516-similarity returns the network they already have, which is not discovery.
517-
518-### A bug inherited from the Go implementation
519-
520-`common_likes` is computed as `CAST(SUM(decay * decay) AS INTEGER)`. Each
521-decay factor is at most 1, so their product is always just under it: a pair
522-sharing one recent like sums to 0.9999 and **truncates to 0**. Two shared
523-likes store 1. The number is shown to readers as the reason for a
524-recommendation -- "N shared likes" -- so a real overlap reads as none.
525-
526-This port uses `ROUND` instead. Worth fixing on the Go side too.
527-
528-The probe caught it only because the assertion looked *vacuous*: both the
529-recent and the stale pair returned 0, so "recent >= old" passed while
530-proving nothing. A test that passes for the wrong reason is worse than one
531-that fails.
532-
533-## What this does and does not prove
534-
535-Proven: the cryptography, the storage layer, and the OAuth flow up to user
536-consent all work from Nim, in the formats ATProto and real servers accept.
537-
538-Not proven, and still ahead:
539-
540-- **Token exchange and refresh** — written, but needs a browser consent to
541- exercise. This is the one place a real end-to-end test still has to happen,
542- and it also covers the `invalid_token` refresh-and-retry branch of
543- `xrpc.request`, which is currently unexercised.
544-
545-### Not needed: DAG-CBOR, CAR and the MST
546-
547-Earlier notes listed these as remaining work. That was wrong, and worth
548-correcting because they are by far the largest format risk left in ATProto.
549-
550-Glean never parses them. `internal/atproto/sync.go` opens by saying so: it
551-reconciles with `com.atproto.repo.listRecords` rather than
552-`com.atproto.sync.getRepo`, because it only syncs known users, Jetstream
553-covers real-time events, and every reconcile is idempotent. Reads go through
554-`listRecords`/`getRecord` and writes through
555-`createRecord`/`putRecord`/`deleteRecord` — all JSON.
556-
557-The imports confirm it. Glean uses four indigo packages — `atclient`,
558-`auth/oauth`, `identity`, `syntax` — and none of `repo`, `mst`, `cbor` or
559-`car`.
560-
561-A Nim port inherits that choice for free. Implementing DAG-CBOR, CIDs, CAR
562-framing and Merkle Search Tree traversal would be real work in service of a
563-code path this application does not have.
564-
565-The honest read: none of the three things that could have sunk the port did.
566-What is left is protocol plumbing against well-specified formats — large, but
567-ordinary.
568 70
569-## Caveats in this code71+## History
570 72
571-- `identity.nim` shells out to `dig` for the DNS leg of handle resolution,73+This started as a port of a much larger reader — one with accounts, ATProto
572- because Nim's stdlib has no resolver. A missing `dig` degrades to the74+sync, social recommendations and a clustering engine. That half has been
573- HTTP `.well-known` fallback rather than failing, but a real port should bind75+deleted. What remains is the part that reads feeds.
574- a resolver instead of depending on a binary.
575-- Only public and localhost clients are supported. Confidential clients
576- (`private_key_jwt` client assertions) are not implemented; Glean uses a
577- public client today.
deleted spike/nim/reconnect_probe.nim +0 -182
deleted file mode 100644
@@ -1,182 +0,0 @@
1-## Checks the reconnect/backoff supervisor.
2-##
3-## nim c -r -d:ssl reconnect_probe.nim
4-##
5-## The backoff curve and cursor arithmetic are pure and tested directly. The
6-## reconnect behaviour itself is tested against a local websocket server that
7-## is deliberately hostile -- it drops every connection after a few frames --
8-## because a consumer that "works" against a healthy server proves nothing
9-## about the case this module exists for.
10-
11-import std/[json, sequtils, strformat, strutils]
12-import pkg/chronicles
13-import pkg/chronos
14-import pkg/websock/websock
15-import atproto/[jetstream, jetstream_run]
16-
17-var failures = 0
18-
19-proc report(name: string, ok: bool, detail = "") =
20- if not ok: inc failures
21- let label = if ok: "PASS" else: "FAIL"
22- if detail.len > 0: echo &" {label} {name} -- {detail}"
23- else: echo &" {label} {name}"
24-
25-# --- a server that hangs up ------------------------------------------------
26-
27-type FlakyServer = ref object
28- http: HttpServer
29- connections*: int
30- framesPerConn*: int
31- cursorsSeen*: seq[string] ## the cursor query param of each connection
32-
33-proc start(framesPerConn: int): FlakyServer =
34- let srv = FlakyServer(framesPerConn: framesPerConn)
35-
36- proc handle(request: HttpRequest) {.async, gcsafe.} =
37- srv.connections.inc
38- srv.cursorsSeen.add request.uri.query
39- let server = WSServer.new(protos = ["", "jetstream"])
40- let ws = await server.handleRequest(request)
41- # Send a few well-formed events, then drop the connection mid-stream.
42- for i in 0 ..< srv.framesPerConn:
43- let ev = %*{
44- "did": "did:plc:probe",
45- "time_us": 1_700_000_000_000_000'i64 + srv.connections * 1000 + i,
46- "kind": "commit",
47- "commit": {"rev": "r", "operation": "create",
48- "collection": "app.test.item", "rkey": $i,
49- "record": {"n": i}},
50- }
51- await ws.send($ev)
52- await ws.stream.closeWait()
53-
54- srv.http = HttpServer.create(initTAddress("127.0.0.1:0"), handle)
55- srv.http.start()
56- srv
57-
58-proc port(s: FlakyServer): int = s.http.localAddress.port.int
59-
60-# --- checks ----------------------------------------------------------------
61-
62-proc main() {.async.} =
63- echo "Jetstream reconnect probe"
64- echo ""
65-
66- # --- backoff curve ------------------------------------------------------
67- block:
68- let b = BackoffConfig(initial: 1.seconds, maxDelay: 30.seconds,
69- factor: 2.0, jitter: 0.0)
70- report("first retry uses the initial delay",
71- delayFor(b, 0) == 1.seconds, $delayFor(b, 0).milliseconds & "ms")
72- report("the first retry is jittered too (herd control)",
73- delayFor(BackoffConfig(initial: 1.seconds, maxDelay: 10.seconds,
74- factor: 2.0, jitter: 0.5),
75- 0, rand01 = 1.0) == 500.milliseconds)
76- report("delay grows geometrically",
77- delayFor(b, 1) == 2.seconds and delayFor(b, 2) == 4.seconds,
78- &"{delayFor(b, 1).milliseconds}ms, {delayFor(b, 2).milliseconds}ms")
79- report("delay is capped",
80- delayFor(b, 20) == 30.seconds, $delayFor(b, 20).milliseconds & "ms")
81-
82- block:
83- # Jitter must reduce, never exceed, so the cap stays a real ceiling.
84- let b = BackoffConfig(initial: 1.seconds, maxDelay: 10.seconds,
85- factor: 2.0, jitter: 0.5)
86- let atMax = delayFor(b, 0, rand01 = 0.0)
87- let atMin = delayFor(b, 0, rand01 = 1.0)
88- report("jitter only shortens the delay",
89- atMax == 1.seconds and atMin == 500.milliseconds,
90- &"{atMin.milliseconds}..{atMax.milliseconds}ms")
91- report("a jittered capped delay never exceeds the cap",
92- delayFor(b, 20, rand01 = 0.0) <= 10.seconds)
93-
94- # --- cursor rewind ------------------------------------------------------
95- block:
96- let store = CursorStore(cursor: 1_700_000_000_000_000'i64)
97- let c = rewoundCursor(store, 5.seconds)
98- report("resume rewinds past the unflushed window",
99- c == 1_700_000_000_000_000'i64 - 5_000_000,
100- &"{store.cursor} -> {c}")
101-
102- # Starting live must stay live; rewinding zero would ask for 1970.
103- report("no cursor stays live", rewoundCursor(CursorStore(), 5.seconds) == 0)
104-
105- let young = CursorStore(cursor: 1000)
106- report("rewind clamps at zero", rewoundCursor(young, 5.seconds) == 0)
107-
108- # --- reconnect against a server that hangs up ---------------------------
109- let server = start(framesPerConn = 3)
110- var received: seq[Event]
111-
112- proc onEvent(ev: Event) {.gcsafe, raises: [].} =
113- received.add ev
114-
115- var cfg = newSupervisorConfig(JetstreamConfig(
116- url: &"ws://127.0.0.1:{server.port}",
117- wantedCollections: @["app.test.item"],
118- ))
119- # Tight timings so the probe finishes; rotation off, so every reconnect
120- # here is a genuine recovery from the server hanging up.
121- cfg.backoff = BackoffConfig(initial: 50.milliseconds,
122- maxDelay: 200.milliseconds,
123- factor: 2.0, jitter: 0.1)
124- cfg.rotateEvery = ZeroDuration
125- cfg.rewind = 0.seconds
126-
127- let sup = newSupervisor(cfg, onEvent)
128- let runner = sup.run()
129- await sleepAsync(2.seconds)
130- sup.stop()
131- await runner.cancelAndWait()
132-
133- report("survived the server hanging up", server.connections > 1,
134- &"{server.connections} connections")
135- report("kept receiving across reconnects", received.len > 3,
136- &"{received.len} events over {server.connections} connections")
137- report("events arrived from more than one connection",
138- received.len >= server.connections,
139- &"{received.len} events")
140-
141- # The cursor must advance across connections, or a restart would replay
142- # from the beginning forever.
143- report("cursor advanced to the newest event seen",
144- sup.store.load() == received[^1].timeUs,
145- &"{sup.store.load()}")
146-
147- # And it must actually reach the server. Tracking a cursor that is never
148- # sent looks identical from the client side and loses every event in the
149- # gap on each reconnect.
150- block:
151- let withCursor = server.cursorsSeen.filterIt("cursor=" in it)
152- report("reconnects send a cursor",
153- withCursor.len > 0,
154- &"{withCursor.len} of {server.cursorsSeen.len} connections")
155- report("the first connection has no cursor (starts live)",
156- server.cursorsSeen.len > 0 and "cursor=" notin server.cursorsSeen[0])
157- if withCursor.len >= 2:
158- proc cursorOf(q: string): int64 =
159- for part in q.split('&'):
160- if part.startsWith("cursor="):
161- return try: parseBiggestInt(part[7 .. ^1]).int64
162- except ValueError: 0
163- 0
164- report("the cursor advances between reconnects",
165- cursorOf(withCursor[^1]) > cursorOf(withCursor[0]),
166- &"{cursorOf(withCursor[0])} -> {cursorOf(withCursor[^1])}")
167-
168- # Stop accepting before closing: a connection still being torn down by the
169- # supervisor would otherwise keep closeWait pending indefinitely.
170- server.http.stop()
171- if not await server.http.closeWait().withTimeout(5.seconds):
172- report("server shut down cleanly", false, "closeWait timed out")
173-
174- echo ""
175- if failures == 0:
176- echo "RESULT: reconnect, backoff and cursor resume all work."
177- else:
178- echo &"RESULT: {failures} check(s) failed."
179- quit 1
180-
181-when isMainModule:
182- waitFor main()
deleted file mode 100644
@@ -1,182 +0,0 @@
1-## Checks the reconnect/backoff supervisor.
2-##
3-## nim c -r -d:ssl reconnect_probe.nim
4-##
5-## The backoff curve and cursor arithmetic are pure and tested directly. The
6-## reconnect behaviour itself is tested against a local websocket server that
7-## is deliberately hostile -- it drops every connection after a few frames --
8-## because a consumer that "works" against a healthy server proves nothing
9-## about the case this module exists for.
10-
11-import std/[json, sequtils, strformat, strutils]
12-import pkg/chronicles
13-import pkg/chronos
14-import pkg/websock/websock
15-import atproto/[jetstream, jetstream_run]
16-
17-var failures = 0
18-
19-proc report(name: string, ok: bool, detail = "") =
20- if not ok: inc failures
21- let label = if ok: "PASS" else: "FAIL"
22- if detail.len > 0: echo &" {label} {name} -- {detail}"
23- else: echo &" {label} {name}"
24-
25-# --- a server that hangs up ------------------------------------------------
26-
27-type FlakyServer = ref object
28- http: HttpServer
29- connections*: int
30- framesPerConn*: int
31- cursorsSeen*: seq[string] ## the cursor query param of each connection
32-
33-proc start(framesPerConn: int): FlakyServer =
34- let srv = FlakyServer(framesPerConn: framesPerConn)
35-
36- proc handle(request: HttpRequest) {.async, gcsafe.} =
37- srv.connections.inc
38- srv.cursorsSeen.add request.uri.query
39- let server = WSServer.new(protos = ["", "jetstream"])
40- let ws = await server.handleRequest(request)
41- # Send a few well-formed events, then drop the connection mid-stream.
42- for i in 0 ..< srv.framesPerConn:
43- let ev = %*{
44- "did": "did:plc:probe",
45- "time_us": 1_700_000_000_000_000'i64 + srv.connections * 1000 + i,
46- "kind": "commit",
47- "commit": {"rev": "r", "operation": "create",
48- "collection": "app.test.item", "rkey": $i,
49- "record": {"n": i}},
50- }
51- await ws.send($ev)
52- await ws.stream.closeWait()
53-
54- srv.http = HttpServer.create(initTAddress("127.0.0.1:0"), handle)
55- srv.http.start()
56- srv
57-
58-proc port(s: FlakyServer): int = s.http.localAddress.port.int
59-
60-# --- checks ----------------------------------------------------------------
61-
62-proc main() {.async.} =
63- echo "Jetstream reconnect probe"
64- echo ""
65-
66- # --- backoff curve ------------------------------------------------------
67- block:
68- let b = BackoffConfig(initial: 1.seconds, maxDelay: 30.seconds,
69- factor: 2.0, jitter: 0.0)
70- report("first retry uses the initial delay",
71- delayFor(b, 0) == 1.seconds, $delayFor(b, 0).milliseconds & "ms")
72- report("the first retry is jittered too (herd control)",
73- delayFor(BackoffConfig(initial: 1.seconds, maxDelay: 10.seconds,
74- factor: 2.0, jitter: 0.5),
75- 0, rand01 = 1.0) == 500.milliseconds)
76- report("delay grows geometrically",
77- delayFor(b, 1) == 2.seconds and delayFor(b, 2) == 4.seconds,
78- &"{delayFor(b, 1).milliseconds}ms, {delayFor(b, 2).milliseconds}ms")
79- report("delay is capped",
80- delayFor(b, 20) == 30.seconds, $delayFor(b, 20).milliseconds & "ms")
81-
82- block:
83- # Jitter must reduce, never exceed, so the cap stays a real ceiling.
84- let b = BackoffConfig(initial: 1.seconds, maxDelay: 10.seconds,
85- factor: 2.0, jitter: 0.5)
86- let atMax = delayFor(b, 0, rand01 = 0.0)
87- let atMin = delayFor(b, 0, rand01 = 1.0)
88- report("jitter only shortens the delay",
89- atMax == 1.seconds and atMin == 500.milliseconds,
90- &"{atMin.milliseconds}..{atMax.milliseconds}ms")
91- report("a jittered capped delay never exceeds the cap",
92- delayFor(b, 20, rand01 = 0.0) <= 10.seconds)
93-
94- # --- cursor rewind ------------------------------------------------------
95- block:
96- let store = CursorStore(cursor: 1_700_000_000_000_000'i64)
97- let c = rewoundCursor(store, 5.seconds)
98- report("resume rewinds past the unflushed window",
99- c == 1_700_000_000_000_000'i64 - 5_000_000,
100- &"{store.cursor} -> {c}")
101-
102- # Starting live must stay live; rewinding zero would ask for 1970.
103- report("no cursor stays live", rewoundCursor(CursorStore(), 5.seconds) == 0)
104-
105- let young = CursorStore(cursor: 1000)
106- report("rewind clamps at zero", rewoundCursor(young, 5.seconds) == 0)
107-
108- # --- reconnect against a server that hangs up ---------------------------
109- let server = start(framesPerConn = 3)
110- var received: seq[Event]
111-
112- proc onEvent(ev: Event) {.gcsafe, raises: [].} =
113- received.add ev
114-
115- var cfg = newSupervisorConfig(JetstreamConfig(
116- url: &"ws://127.0.0.1:{server.port}",
117- wantedCollections: @["app.test.item"],
118- ))
119- # Tight timings so the probe finishes; rotation off, so every reconnect
120- # here is a genuine recovery from the server hanging up.
121- cfg.backoff = BackoffConfig(initial: 50.milliseconds,
122- maxDelay: 200.milliseconds,
123- factor: 2.0, jitter: 0.1)
124- cfg.rotateEvery = ZeroDuration
125- cfg.rewind = 0.seconds
126-
127- let sup = newSupervisor(cfg, onEvent)
128- let runner = sup.run()
129- await sleepAsync(2.seconds)
130- sup.stop()
131- await runner.cancelAndWait()
132-
133- report("survived the server hanging up", server.connections > 1,
134- &"{server.connections} connections")
135- report("kept receiving across reconnects", received.len > 3,
136- &"{received.len} events over {server.connections} connections")
137- report("events arrived from more than one connection",
138- received.len >= server.connections,
139- &"{received.len} events")
140-
141- # The cursor must advance across connections, or a restart would replay
142- # from the beginning forever.
143- report("cursor advanced to the newest event seen",
144- sup.store.load() == received[^1].timeUs,
145- &"{sup.store.load()}")
146-
147- # And it must actually reach the server. Tracking a cursor that is never
148- # sent looks identical from the client side and loses every event in the
149- # gap on each reconnect.
150- block:
151- let withCursor = server.cursorsSeen.filterIt("cursor=" in it)
152- report("reconnects send a cursor",
153- withCursor.len > 0,
154- &"{withCursor.len} of {server.cursorsSeen.len} connections")
155- report("the first connection has no cursor (starts live)",
156- server.cursorsSeen.len > 0 and "cursor=" notin server.cursorsSeen[0])
157- if withCursor.len >= 2:
158- proc cursorOf(q: string): int64 =
159- for part in q.split('&'):
160- if part.startsWith("cursor="):
161- return try: parseBiggestInt(part[7 .. ^1]).int64
162- except ValueError: 0
163- 0
164- report("the cursor advances between reconnects",
165- cursorOf(withCursor[^1]) > cursorOf(withCursor[0]),
166- &"{cursorOf(withCursor[0])} -> {cursorOf(withCursor[^1])}")
167-
168- # Stop accepting before closing: a connection still being torn down by the
169- # supervisor would otherwise keep closeWait pending indefinitely.
170- server.http.stop()
171- if not await server.http.closeWait().withTimeout(5.seconds):
172- report("server shut down cleanly", false, "closeWait timed out")
173-
174- echo ""
175- if failures == 0:
176- echo "RESULT: reconnect, backoff and cursor resume all work."
177- else:
178- echo &"RESULT: {failures} check(s) failed."
179- quit 1
180-
181-when isMainModule:
182- waitFor main()
modified spike/nim/scraper_probe.nim +1 -1
@@ -9,7 +9,7 @@
99 ## page that renders correctly proves nothing about what was stripped.
1010
1111 import std/[htmlparser, httpclient, options, strformat, strutils, xmltree]
12-import atproto/[feedparser, scraper]
12+import reader/[feedparser, scraper]
1313
1414 var failures = 0
1515
@@ -9,7 +9,7 @@
9 ## page that renders correctly proves nothing about what was stripped.9 ## page that renders correctly proves nothing about what was stripped.
10 10
11 import std/[htmlparser, httpclient, options, strformat, strutils, xmltree]11 import std/[htmlparser, httpclient, options, strformat, strutils, xmltree]
12-import atproto/[feedparser, scraper]12+import reader/[feedparser, scraper]
13 13
14 var failures = 014 var failures = 0
15 15
deleted spike/nim/social_probe.nim +0 -216
deleted file mode 100644
@@ -1,216 +0,0 @@
1-## Exercises likes, annotations, follows and trending.
2-##
3-## nim c -r -d:ssl social_probe.nim
4-##
5-## Weighted towards the destructive and the social-graph paths: orphan
6-## reconciliation deletes rows, and trending decides what a reader sees, so
7-## both are worth more than "the insert worked".
8-
9-import std/[options, os, strformat, strutils, times]
10-import atproto/[articlestore, feedstore, gleandb, socialstore, sqlite]
11-
12-var failures = 0
13-
14-proc report(name: string, ok: bool, detail = "") =
15- if not ok: inc failures
16- let label = if ok: "PASS" else: "FAIL"
17- if detail.len > 0: echo &" {label} {name} -- {detail}"
18- else: echo &" {label} {name}"
19-
20-proc cleanup(base: string) =
21- for suffix in ["_users", "_articles", "_recs"]:
22- for ext in ["", "-wal", "-shm"]:
23- removeFile(base & suffix & ext)
24-
25-const
26- alice = "did:plc:alice"
27- bob = "did:plc:bob"
28- carol = "did:plc:carol"
29- feedA = "https://a.test/feed"
30- urlA1 = "https://a.test/1"
31- urlA2 = "https://a.test/2"
32-
33-proc main() =
34- echo "Social store probe"
35- echo ""
36-
37- let base = getTempDir() / "glean_social_probe"
38- cleanup(base)
39- defer: cleanup(base)
40-
41- var g = gleandb.open(base)
42- defer: g.close()
43- g.migrate()
44-
45- for did in [alice, bob, carol]:
46- g.db.run("INSERT INTO users (did) VALUES (?)", p(did))
47- g.upsertFeed(Feed(feedUrl: feedA, title: "Feed A", feedType: "rss"))
48- for did in [alice, bob, carol]:
49- g.createSubscription(did, feedA, "", "", "", "")
50-
51- let t = now().utc
52- g.batchUpsertArticles(@[
53- NewArticle(feedUrl: feedA, guid: "a1", title: "First", url: urlA1,
54- published: some(t - initDuration(hours = 2))),
55- NewArticle(feedUrl: feedA, guid: "a2", title: "Second", url: urlA2,
56- published: some(t - initDuration(hours = 1))),
57- ])
58-
59- # --- likes --------------------------------------------------------------
60- block:
61- g.createLike(Like(uri: "at://alice/like/1", authorDid: alice,
62- feedUrl: feedA, articleUrl: urlA1, createdAt: some(t)))
63- report("a like round-trips", g.hasLiked(alice, feedA, urlA1))
64- report("a like is per author", not g.hasLiked(bob, feedA, urlA1))
65-
66- # The same record arrives from Jetstream and from a PDS sync.
67- g.createLike(Like(uri: "at://alice/like/1", authorDid: alice,
68- feedUrl: feedA, articleUrl: urlA1, createdAt: some(t)))
69- report("re-delivering a like does not duplicate it",
70- g.getLikeCount(feedA, urlA1) == 1, $g.getLikeCount(feedA, urlA1))
71-
72- g.batchCreateLikes(@[
73- Like(uri: "at://bob/like/1", authorDid: bob, feedUrl: feedA,
74- articleUrl: urlA1, createdAt: some(t)),
75- Like(uri: "at://carol/like/1", authorDid: carol, feedUrl: feedA,
76- articleUrl: urlA1, createdAt: some(t)),
77- Like(uri: "at://bob/like/2", authorDid: bob, feedUrl: feedA,
78- articleUrl: urlA2, createdAt: some(t)),
79- ])
80- report("batch insert works", g.getLikeCount(feedA, urlA1) == 3,
81- $g.getLikeCount(feedA, urlA1))
82- report("likes are listed per author",
83- g.listLikes(bob).len == 2, $g.listLikes(bob).len)
84- report("likes can be scoped to a feed",
85- g.listLikes(bob, feedUrl = feedA).len == 2)
86-
87- g.deleteLikeByUserArticle(alice, feedA, urlA1)
88- report("unliking removes only that author's like",
89- g.getLikeCount(feedA, urlA1) == 2 and not g.hasLiked(alice, feedA, urlA1),
90- $g.getLikeCount(feedA, urlA1))
91-
92- # --- annotations --------------------------------------------------------
93- block:
94- g.createAnnotation(Annotation(
95- uri: "at://alice/ann/1", authorDid: alice, feedUrl: feedA,
96- articleUrl: urlA1, quote: "a quote", note: "a note",
97- tags: @["nim", "systems"], rating: some(4), createdAt: some(t)))
98- let listed = g.listAnnotations(articleUrl = urlA1)
99- report("an annotation round-trips", listed.len == 1, $listed.len)
100- if listed.len > 0:
101- report("tags round-trip as a list", listed[0].tags == @["nim", "systems"],
102- listed[0].tags.join("|"))
103- report("rating round-trips", listed[0].rating == some(4))
104-
105- # An edit keeps the URI and changes the content.
106- g.createAnnotation(Annotation(
107- uri: "at://alice/ann/1", authorDid: alice, feedUrl: feedA,
108- articleUrl: urlA1, note: "edited", tags: @["nim"], createdAt: some(t)))
109- let edited = g.listAnnotations(articleUrl = urlA1)
110- report("re-delivering an annotation updates rather than duplicates",
111- edited.len == 1 and edited[0].note == "edited",
112- &"{edited.len} rows, note={edited[0].note}")
113- report("a cleared rating round-trips as none", edited[0].rating.isNone)
114-
115- report("annotationExists finds it", g.annotationExists("at://alice/ann/1"))
116- report("annotationExists rejects an unknown uri",
117- not g.annotationExists("at://nobody/ann/9"))
118-
119- g.batchCreateAnnotations(@[
120- Annotation(uri: "at://bob/ann/1", authorDid: bob, feedUrl: feedA,
121- articleUrl: urlA1, note: "bob's", createdAt: some(t))])
122- report("annotations are listed per author",
123- g.listAnnotations(authorDid = bob).len == 1)
124-
125- # --- follows ------------------------------------------------------------
126- block:
127- g.upsertFollow(alice, bob, "at://alice/follow/1", "cid1")
128- g.upsertFollow(alice, bob, "at://alice/follow/1", "cid1")
129- report("following is idempotent", g.listFollows(alice) == @[bob],
130- g.listFollows(alice).join(","))
131- g.upsertFollow(alice, carol, "", "")
132- report("follows accumulate", g.listFollows(alice).len == 2)
133- g.deleteFollow(alice, carol)
134- report("unfollowing removes one edge", g.listFollows(alice) == @[bob])
135-
136- # --- trending -----------------------------------------------------------
137- block:
138- let since = formatSqliteTime(t - initDuration(days = 1))
139- let global = g.listTrendingArticles(alice, since)
140- report("global trending ranks by likes",
141- global.len == 2 and global[0].url == urlA1,
142- if global.len > 0: &"top={global[0].url} likes={global[0].likeCount}"
143- else: "empty")
144- report("trending counts annotations too",
145- global.len > 0 and global[0].annotationCount == 2,
146- if global.len > 0: $global[0].annotationCount else: "-")
147- report("hasLiked is per viewer",
148- not global[0].hasLiked, "alice unliked hers earlier")
149-
150- let viewedByBob = g.listTrendingArticles(bob, since)
151- report("hasLiked is true for a viewer who liked it",
152- viewedByBob[0].hasLiked)
153-
154- # Alice follows only bob, so carol's like should not pull an article in
155- # that bob did not also like.
156- let scoped = g.listTrendingArticlesForUser(alice, since)
157- report("graph-scoped trending only counts followed authors",
158- scoped.len > 0, $scoped.len)
159- let scopedTop = scoped[0]
160- report("a followed author's likes are counted",
161- scopedTop.likeCount >= 1, $scopedTop.likeCount)
162-
163- # An author nobody follows contributes nothing to a graph-scoped view.
164- let carolOnly = g.listTrendingArticlesForUser(carol, since)
165- report("a reader with no follows still sees their own likes",
166- carolOnly.len == 1, $carolOnly.len)
167-
168- # --- language filter ----------------------------------------------------
169- block:
170- let (sql, params) = langFilter(@["en", "fr"], "ar.")
171- report("language filter builds placeholders",
172- sql.count("?") == 2 and params.len == 2, sql.strip())
173- # Unclassified articles must survive the filter, or new items vanish for
174- # a few minutes and then reappear.
175- report("language filter keeps unclassified articles",
176- "language = ''" in sql)
177- report("an empty language list adds no filter",
178- langFilter(@[], "ar.")[0].len == 0)
179-
180- # --- liked articles -----------------------------------------------------
181- block:
182- let liked = g.listLikedArticles(bob)
183- report("liked articles are listed newest first", liked.len == 2, $liked.len)
184- report("liked articles carry article fields",
185- liked[0].title.len > 0 and liked[0].feedTitle == "Feed A",
186- liked[0].title)
187-
188- # --- reconciliation -----------------------------------------------------
189- # The destructive path: local rows the PDS no longer has must go, and
190- # everyone else's must not.
191- block:
192- report("bob starts with 2 likes", g.listLikes(bob).len == 2)
193- g.deleteOrphanedLikes(bob, @["at://bob/like/1"])
194- report("reconciling deletes likes the PDS dropped",
195- g.listLikes(bob).len == 1, $g.listLikes(bob).len)
196- report("reconciling leaves other users alone",
197- g.listLikes(carol).len == 1, $g.listLikes(carol).len)
198-
199- g.deleteOrphanedAnnotations(alice, @["at://alice/ann/1"])
200- report("reconciling keeps annotations the PDS still has",
201- g.listAnnotations(authorDid = alice).len == 1)
202- g.deleteOrphanedAnnotations(alice, @[])
203- report("an empty active set clears that author's annotations",
204- g.listAnnotations(authorDid = alice).len == 0)
205- report("and leaves another author's annotations",
206- g.listAnnotations(authorDid = bob).len == 1)
207-
208- echo ""
209- if failures == 0:
210- echo "RESULT: likes, annotations, follows and trending behave."
211- else:
212- echo &"RESULT: {failures} check(s) failed."
213- quit 1
214-
215-when isMainModule:
216- main()
deleted file mode 100644
@@ -1,216 +0,0 @@
1-## Exercises likes, annotations, follows and trending.
2-##
3-## nim c -r -d:ssl social_probe.nim
4-##
5-## Weighted towards the destructive and the social-graph paths: orphan
6-## reconciliation deletes rows, and trending decides what a reader sees, so
7-## both are worth more than "the insert worked".
8-
9-import std/[options, os, strformat, strutils, times]
10-import atproto/[articlestore, feedstore, gleandb, socialstore, sqlite]
11-
12-var failures = 0
13-
14-proc report(name: string, ok: bool, detail = "") =
15- if not ok: inc failures
16- let label = if ok: "PASS" else: "FAIL"
17- if detail.len > 0: echo &" {label} {name} -- {detail}"
18- else: echo &" {label} {name}"
19-
20-proc cleanup(base: string) =
21- for suffix in ["_users", "_articles", "_recs"]:
22- for ext in ["", "-wal", "-shm"]:
23- removeFile(base & suffix & ext)
24-
25-const
26- alice = "did:plc:alice"
27- bob = "did:plc:bob"
28- carol = "did:plc:carol"
29- feedA = "https://a.test/feed"
30- urlA1 = "https://a.test/1"
31- urlA2 = "https://a.test/2"
32-
33-proc main() =
34- echo "Social store probe"
35- echo ""
36-
37- let base = getTempDir() / "glean_social_probe"
38- cleanup(base)
39- defer: cleanup(base)
40-
41- var g = gleandb.open(base)
42- defer: g.close()
43- g.migrate()
44-
45- for did in [alice, bob, carol]:
46- g.db.run("INSERT INTO users (did) VALUES (?)", p(did))
47- g.upsertFeed(Feed(feedUrl: feedA, title: "Feed A", feedType: "rss"))
48- for did in [alice, bob, carol]:
49- g.createSubscription(did, feedA, "", "", "", "")
50-
51- let t = now().utc
52- g.batchUpsertArticles(@[
53- NewArticle(feedUrl: feedA, guid: "a1", title: "First", url: urlA1,
54- published: some(t - initDuration(hours = 2))),
55- NewArticle(feedUrl: feedA, guid: "a2", title: "Second", url: urlA2,
56- published: some(t - initDuration(hours = 1))),
57- ])
58-
59- # --- likes --------------------------------------------------------------
60- block:
61- g.createLike(Like(uri: "at://alice/like/1", authorDid: alice,
62- feedUrl: feedA, articleUrl: urlA1, createdAt: some(t)))
63- report("a like round-trips", g.hasLiked(alice, feedA, urlA1))
64- report("a like is per author", not g.hasLiked(bob, feedA, urlA1))
65-
66- # The same record arrives from Jetstream and from a PDS sync.
67- g.createLike(Like(uri: "at://alice/like/1", authorDid: alice,
68- feedUrl: feedA, articleUrl: urlA1, createdAt: some(t)))
69- report("re-delivering a like does not duplicate it",
70- g.getLikeCount(feedA, urlA1) == 1, $g.getLikeCount(feedA, urlA1))
71-
72- g.batchCreateLikes(@[
73- Like(uri: "at://bob/like/1", authorDid: bob, feedUrl: feedA,
74- articleUrl: urlA1, createdAt: some(t)),
75- Like(uri: "at://carol/like/1", authorDid: carol, feedUrl: feedA,
76- articleUrl: urlA1, createdAt: some(t)),
77- Like(uri: "at://bob/like/2", authorDid: bob, feedUrl: feedA,
78- articleUrl: urlA2, createdAt: some(t)),
79- ])
80- report("batch insert works", g.getLikeCount(feedA, urlA1) == 3,
81- $g.getLikeCount(feedA, urlA1))
82- report("likes are listed per author",
83- g.listLikes(bob).len == 2, $g.listLikes(bob).len)
84- report("likes can be scoped to a feed",
85- g.listLikes(bob, feedUrl = feedA).len == 2)
86-
87- g.deleteLikeByUserArticle(alice, feedA, urlA1)
88- report("unliking removes only that author's like",
89- g.getLikeCount(feedA, urlA1) == 2 and not g.hasLiked(alice, feedA, urlA1),
90- $g.getLikeCount(feedA, urlA1))
91-
92- # --- annotations --------------------------------------------------------
93- block:
94- g.createAnnotation(Annotation(
95- uri: "at://alice/ann/1", authorDid: alice, feedUrl: feedA,
96- articleUrl: urlA1, quote: "a quote", note: "a note",
97- tags: @["nim", "systems"], rating: some(4), createdAt: some(t)))
98- let listed = g.listAnnotations(articleUrl = urlA1)
99- report("an annotation round-trips", listed.len == 1, $listed.len)
100- if listed.len > 0:
101- report("tags round-trip as a list", listed[0].tags == @["nim", "systems"],
102- listed[0].tags.join("|"))
103- report("rating round-trips", listed[0].rating == some(4))
104-
105- # An edit keeps the URI and changes the content.
106- g.createAnnotation(Annotation(
107- uri: "at://alice/ann/1", authorDid: alice, feedUrl: feedA,
108- articleUrl: urlA1, note: "edited", tags: @["nim"], createdAt: some(t)))
109- let edited = g.listAnnotations(articleUrl = urlA1)
110- report("re-delivering an annotation updates rather than duplicates",
111- edited.len == 1 and edited[0].note == "edited",
112- &"{edited.len} rows, note={edited[0].note}")
113- report("a cleared rating round-trips as none", edited[0].rating.isNone)
114-
115- report("annotationExists finds it", g.annotationExists("at://alice/ann/1"))
116- report("annotationExists rejects an unknown uri",
117- not g.annotationExists("at://nobody/ann/9"))
118-
119- g.batchCreateAnnotations(@[
120- Annotation(uri: "at://bob/ann/1", authorDid: bob, feedUrl: feedA,
121- articleUrl: urlA1, note: "bob's", createdAt: some(t))])
122- report("annotations are listed per author",
123- g.listAnnotations(authorDid = bob).len == 1)
124-
125- # --- follows ------------------------------------------------------------
126- block:
127- g.upsertFollow(alice, bob, "at://alice/follow/1", "cid1")
128- g.upsertFollow(alice, bob, "at://alice/follow/1", "cid1")
129- report("following is idempotent", g.listFollows(alice) == @[bob],
130- g.listFollows(alice).join(","))
131- g.upsertFollow(alice, carol, "", "")
132- report("follows accumulate", g.listFollows(alice).len == 2)
133- g.deleteFollow(alice, carol)
134- report("unfollowing removes one edge", g.listFollows(alice) == @[bob])
135-
136- # --- trending -----------------------------------------------------------
137- block:
138- let since = formatSqliteTime(t - initDuration(days = 1))
139- let global = g.listTrendingArticles(alice, since)
140- report("global trending ranks by likes",
141- global.len == 2 and global[0].url == urlA1,
142- if global.len > 0: &"top={global[0].url} likes={global[0].likeCount}"
143- else: "empty")
144- report("trending counts annotations too",
145- global.len > 0 and global[0].annotationCount == 2,
146- if global.len > 0: $global[0].annotationCount else: "-")
147- report("hasLiked is per viewer",
148- not global[0].hasLiked, "alice unliked hers earlier")
149-
150- let viewedByBob = g.listTrendingArticles(bob, since)
151- report("hasLiked is true for a viewer who liked it",
152- viewedByBob[0].hasLiked)
153-
154- # Alice follows only bob, so carol's like should not pull an article in
155- # that bob did not also like.
156- let scoped = g.listTrendingArticlesForUser(alice, since)
157- report("graph-scoped trending only counts followed authors",
158- scoped.len > 0, $scoped.len)
159- let scopedTop = scoped[0]
160- report("a followed author's likes are counted",
161- scopedTop.likeCount >= 1, $scopedTop.likeCount)
162-
163- # An author nobody follows contributes nothing to a graph-scoped view.
164- let carolOnly = g.listTrendingArticlesForUser(carol, since)
165- report("a reader with no follows still sees their own likes",
166- carolOnly.len == 1, $carolOnly.len)
167-
168- # --- language filter ----------------------------------------------------
169- block:
170- let (sql, params) = langFilter(@["en", "fr"], "ar.")
171- report("language filter builds placeholders",
172- sql.count("?") == 2 and params.len == 2, sql.strip())
173- # Unclassified articles must survive the filter, or new items vanish for
174- # a few minutes and then reappear.
175- report("language filter keeps unclassified articles",
176- "language = ''" in sql)
177- report("an empty language list adds no filter",
178- langFilter(@[], "ar.")[0].len == 0)
179-
180- # --- liked articles -----------------------------------------------------
181- block:
182- let liked = g.listLikedArticles(bob)
183- report("liked articles are listed newest first", liked.len == 2, $liked.len)
184- report("liked articles carry article fields",
185- liked[0].title.len > 0 and liked[0].feedTitle == "Feed A",
186- liked[0].title)
187-
188- # --- reconciliation -----------------------------------------------------
189- # The destructive path: local rows the PDS no longer has must go, and
190- # everyone else's must not.
191- block:
192- report("bob starts with 2 likes", g.listLikes(bob).len == 2)
193- g.deleteOrphanedLikes(bob, @["at://bob/like/1"])
194- report("reconciling deletes likes the PDS dropped",
195- g.listLikes(bob).len == 1, $g.listLikes(bob).len)
196- report("reconciling leaves other users alone",
197- g.listLikes(carol).len == 1, $g.listLikes(carol).len)
198-
199- g.deleteOrphanedAnnotations(alice, @["at://alice/ann/1"])
200- report("reconciling keeps annotations the PDS still has",
201- g.listAnnotations(authorDid = alice).len == 1)
202- g.deleteOrphanedAnnotations(alice, @[])
203- report("an empty active set clears that author's annotations",
204- g.listAnnotations(authorDid = alice).len == 0)
205- report("and leaves another author's annotations",
206- g.listAnnotations(authorDid = bob).len == 1)
207-
208- echo ""
209- if failures == 0:
210- echo "RESULT: likes, annotations, follows and trending behave."
211- else:
212- echo &"RESULT: {failures} check(s) failed."
213- quit 1
214-
215-when isMainModule:
216- main()
deleted spike/nim/sqlite_probe.nim +0 -151
deleted file mode 100644
@@ -1,151 +0,0 @@
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()
deleted file mode 100644
@@ -1,151 +0,0 @@
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()
deleted spike/nim/store_probe.nim +0 -185
deleted file mode 100644
@@ -1,185 +0,0 @@
1-## Checks that a session survives a restart intact.
2-##
3-## nim c -r -d:ssl store_probe.nim
4-##
5-## The interesting failure is not "the row is missing" -- that is loud. It is
6-## a DPoP key that round-trips *almost* correctly: every request after a
7-## restart then fails authentication with no clue that storage is at fault.
8-## So the key is checked by signing with the reloaded copy and verifying
9-## against the original public point, not by comparing fields.
10-
11-import std/[json, options, os, strformat, strutils]
12-import pkg/chronos
13-import bearssl/rand
14-import atproto/[cursor_store, dpop, jetstream_run, oauth, sqlite, store]
15-
16-var failures = 0
17-
18-proc report(name: string, ok: bool, detail = "") =
19- if not ok: inc failures
20- let label = if ok: "PASS" else: "FAIL"
21- if detail.len > 0: echo &" {label} {name} -- {detail}"
22- else: echo &" {label} {name}"
23-
24-proc main() =
25- echo "Session persistence probe"
26- echo ""
27-
28- let rngRef = HmacDrbgContext.new()
29- doAssert rngRef != nil
30- var rng = rngRef[]
31-
32- let path = getTempDir() / "glean_store_probe.db"
33- removeFile(path)
34- removeFile(path & "-wal")
35- removeFile(path & "-shm")
36- defer:
37- removeFile(path)
38- removeFile(path & "-wal")
39- removeFile(path & "-shm")
40-
41- let key = generateKey(rng)
42- let originalThumb = key.thumbprint
43-
44- # --- write, then close the database entirely ---------------------------
45- block:
46- var db = sqlite.open(path)
47- migrate(db)
48- report("schema applies", true, "sqlite " & libVersion())
49-
50- saveSession(db, SessionRecord(
51- accountDid: "did:plc:probe",
52- sessionId: "sess-1",
53- pdsUrl: "https://pds.example",
54- authServerUrl: "https://bsky.social",
55- tokenEndpoint: "https://bsky.social/oauth/token",
56- accessToken: "access-abc",
57- refreshToken: "refresh-xyz",
58- dpopKey: key,
59- hostNonce: "host-n",
60- authNonce: "auth-n",
61- scopes: @["atproto", "transition:generic"],
62- createdAt: 1_700_000_000,
63- ))
64- saveCursor(db, 1_700_000_000_000_000'i64)
65- saveAuthRequest(db, AuthRequestData(
66- state: "state-1",
67- authServerUrl: "https://bsky.social",
68- scopes: @["atproto"],
69- pkceVerifier: "verifier-1",
70- requestUri: "urn:ietf:params:oauth:request_uri:req-1",
71- tokenEndpoint: "https://bsky.social/oauth/token",
72- dpopNonce: "n-1",
73- dpopKey: key,
74- accountDid: "did:plc:probe",
75- ))
76- db.close()
77-
78- # --- reopen, as a restart would ----------------------------------------
79- var db = sqlite.open(path)
80- defer: db.close()
81- migrate(db) # idempotent, as it is on every boot
82-
83- block:
84- let got = getSession(db, "did:plc:probe", "sess-1")
85- report("session survives a restart", got.isSome)
86- if got.isSome:
87- let s = got.get
88- report("tokens round-trip",
89- s.accessToken == "access-abc" and s.refreshToken == "refresh-xyz")
90- report("nonces round-trip",
91- s.hostNonce == "host-n" and s.authNonce == "auth-n")
92- report("scopes round-trip", s.scopes == @["atproto", "transition:generic"])
93-
94- # The check that matters.
95- report("reloaded key has the same JWK thumbprint",
96- s.dpopKey.thumbprint == originalThumb, s.dpopKey.thumbprint)
97-
98- const msg = "eyJhbGciOiJFUzI1NiJ9.eyJodG0iOiJHRVQifQ"
99- let sig = s.dpopKey.signEs256(rng, msg)
100- report("reloaded key still signs verifiably",
101- key.verifyEs256(msg, sig),
102- "signed with the reloaded key, verified against the original")
103-
104- # Negative control: a verifier that accepts anything proves nothing.
105- let other = generateKey(rng)
106- report("a different key does not verify",
107- not other.verifyEs256(msg, sig))
108-
109- block:
110- report("unknown session reads as none",
111- getSession(db, "did:plc:probe", "nope").isNone)
112- report("sessions are listed per account",
113- listSessions(db, "did:plc:probe") == @["sess-1"])
114- report("accounts are listed", listAccounts(db) == @["did:plc:probe"])
115-
116- block:
117- report("cursor survives a restart",
118- loadCursor(db) == 1_700_000_000_000_000'i64, $loadCursor(db))
119-
120- # --- auth requests are single-use --------------------------------------
121- block:
122- let first = takeAuthRequest(db, "state-1")
123- report("pending auth request survives a restart", first.isSome)
124- if first.isSome:
125- report("PKCE verifier round-trips", first.get.pkceVerifier == "verifier-1")
126- report("auth request key round-trips",
127- first.get.dpopKey.thumbprint == originalThumb)
128- # An authorization code is single-use; leaving the row behind invites a
129- # replayed callback.
130- report("taking an auth request consumes it",
131- takeAuthRequest(db, "state-1").isNone)
132-
133- # --- batched cursor store ----------------------------------------------
134- block:
135- let cs = newDbCursorStore(db, flushEvery = 3600) # never auto-flush here
136- report("cursor store loads the persisted value",
137- cs.load() == 1_700_000_000_000_000'i64)
138-
139- cs.save(1_700_000_000_500_000'i64)
140- # Batching means the write is deferred, not lost.
141- report("a fresh cursor is not written immediately",
142- loadCursor(db) == 1_700_000_000_000_000'i64)
143- cs.close()
144- report("closing flushes the pending cursor",
145- loadCursor(db) == 1_700_000_000_500_000'i64, $loadCursor(db))
146-
147- # Cursors must never move backwards, or a late event would rewind the
148- # subscription on the next reconnect.
149- let cs2 = newDbCursorStore(db, flushEvery = 0)
150- cs2.save(1'i64)
151- cs2.close()
152- report("an older cursor does not overwrite a newer one",
153- loadCursor(db) == 1_700_000_000_500_000'i64, $loadCursor(db))
154-
155- # --- the seam with the supervisor --------------------------------------
156- # jetstream_run takes a CursorStore base reference, so this only works if
157- # the overrides dispatch dynamically. A silent fall-through to the base
158- # method would read the in-memory field instead of the database and look
159- # entirely healthy while resuming from the wrong place.
160- block:
161- let cs = newDbCursorStore(db, flushEvery = 3600)
162- let asBase: CursorStore = cs
163- report("DbCursorStore dispatches through the CursorStore base",
164- asBase.load() == loadCursor(db), $asBase.load())
165-
166- let rewound = rewoundCursor(asBase, 5.seconds)
167- report("the supervisor rewinds the persisted cursor",
168- rewound == loadCursor(db) - 5_000_000, $rewound)
169-
170- # --- deletion ----------------------------------------------------------
171- block:
172- deleteSession(db, "did:plc:probe", "sess-1")
173- report("sessions can be deleted",
174- getSession(db, "did:plc:probe", "sess-1").isNone and
175- listAccounts(db).len == 0)
176-
177- echo ""
178- if failures == 0:
179- echo "RESULT: sessions, keys and cursors survive a restart."
180- else:
181- echo &"RESULT: {failures} check(s) failed."
182- quit 1
183-
184-when isMainModule:
185- main()
deleted file mode 100644
@@ -1,185 +0,0 @@
1-## Checks that a session survives a restart intact.
2-##
3-## nim c -r -d:ssl store_probe.nim
4-##
5-## The interesting failure is not "the row is missing" -- that is loud. It is
6-## a DPoP key that round-trips *almost* correctly: every request after a
7-## restart then fails authentication with no clue that storage is at fault.
8-## So the key is checked by signing with the reloaded copy and verifying
9-## against the original public point, not by comparing fields.
10-
11-import std/[json, options, os, strformat, strutils]
12-import pkg/chronos
13-import bearssl/rand
14-import atproto/[cursor_store, dpop, jetstream_run, oauth, sqlite, store]
15-
16-var failures = 0
17-
18-proc report(name: string, ok: bool, detail = "") =
19- if not ok: inc failures
20- let label = if ok: "PASS" else: "FAIL"
21- if detail.len > 0: echo &" {label} {name} -- {detail}"
22- else: echo &" {label} {name}"
23-
24-proc main() =
25- echo "Session persistence probe"
26- echo ""
27-
28- let rngRef = HmacDrbgContext.new()
29- doAssert rngRef != nil
30- var rng = rngRef[]
31-
32- let path = getTempDir() / "glean_store_probe.db"
33- removeFile(path)
34- removeFile(path & "-wal")
35- removeFile(path & "-shm")
36- defer:
37- removeFile(path)
38- removeFile(path & "-wal")
39- removeFile(path & "-shm")
40-
41- let key = generateKey(rng)
42- let originalThumb = key.thumbprint
43-
44- # --- write, then close the database entirely ---------------------------
45- block:
46- var db = sqlite.open(path)
47- migrate(db)
48- report("schema applies", true, "sqlite " & libVersion())
49-
50- saveSession(db, SessionRecord(
51- accountDid: "did:plc:probe",
52- sessionId: "sess-1",
53- pdsUrl: "https://pds.example",
54- authServerUrl: "https://bsky.social",
55- tokenEndpoint: "https://bsky.social/oauth/token",
56- accessToken: "access-abc",
57- refreshToken: "refresh-xyz",
58- dpopKey: key,
59- hostNonce: "host-n",
60- authNonce: "auth-n",
61- scopes: @["atproto", "transition:generic"],
62- createdAt: 1_700_000_000,
63- ))
64- saveCursor(db, 1_700_000_000_000_000'i64)
65- saveAuthRequest(db, AuthRequestData(
66- state: "state-1",
67- authServerUrl: "https://bsky.social",
68- scopes: @["atproto"],
69- pkceVerifier: "verifier-1",
70- requestUri: "urn:ietf:params:oauth:request_uri:req-1",
71- tokenEndpoint: "https://bsky.social/oauth/token",
72- dpopNonce: "n-1",
73- dpopKey: key,
74- accountDid: "did:plc:probe",
75- ))
76- db.close()
77-
78- # --- reopen, as a restart would ----------------------------------------
79- var db = sqlite.open(path)
80- defer: db.close()
81- migrate(db) # idempotent, as it is on every boot
82-
83- block:
84- let got = getSession(db, "did:plc:probe", "sess-1")
85- report("session survives a restart", got.isSome)
86- if got.isSome:
87- let s = got.get
88- report("tokens round-trip",
89- s.accessToken == "access-abc" and s.refreshToken == "refresh-xyz")
90- report("nonces round-trip",
91- s.hostNonce == "host-n" and s.authNonce == "auth-n")
92- report("scopes round-trip", s.scopes == @["atproto", "transition:generic"])
93-
94- # The check that matters.
95- report("reloaded key has the same JWK thumbprint",
96- s.dpopKey.thumbprint == originalThumb, s.dpopKey.thumbprint)
97-
98- const msg = "eyJhbGciOiJFUzI1NiJ9.eyJodG0iOiJHRVQifQ"
99- let sig = s.dpopKey.signEs256(rng, msg)
100- report("reloaded key still signs verifiably",
101- key.verifyEs256(msg, sig),
102- "signed with the reloaded key, verified against the original")
103-
104- # Negative control: a verifier that accepts anything proves nothing.
105- let other = generateKey(rng)
106- report("a different key does not verify",
107- not other.verifyEs256(msg, sig))
108-
109- block:
110- report("unknown session reads as none",
111- getSession(db, "did:plc:probe", "nope").isNone)
112- report("sessions are listed per account",
113- listSessions(db, "did:plc:probe") == @["sess-1"])
114- report("accounts are listed", listAccounts(db) == @["did:plc:probe"])
115-
116- block:
117- report("cursor survives a restart",
118- loadCursor(db) == 1_700_000_000_000_000'i64, $loadCursor(db))
119-
120- # --- auth requests are single-use --------------------------------------
121- block:
122- let first = takeAuthRequest(db, "state-1")
123- report("pending auth request survives a restart", first.isSome)
124- if first.isSome:
125- report("PKCE verifier round-trips", first.get.pkceVerifier == "verifier-1")
126- report("auth request key round-trips",
127- first.get.dpopKey.thumbprint == originalThumb)
128- # An authorization code is single-use; leaving the row behind invites a
129- # replayed callback.
130- report("taking an auth request consumes it",
131- takeAuthRequest(db, "state-1").isNone)
132-
133- # --- batched cursor store ----------------------------------------------
134- block:
135- let cs = newDbCursorStore(db, flushEvery = 3600) # never auto-flush here
136- report("cursor store loads the persisted value",
137- cs.load() == 1_700_000_000_000_000'i64)
138-
139- cs.save(1_700_000_000_500_000'i64)
140- # Batching means the write is deferred, not lost.
141- report("a fresh cursor is not written immediately",
142- loadCursor(db) == 1_700_000_000_000_000'i64)
143- cs.close()
144- report("closing flushes the pending cursor",
145- loadCursor(db) == 1_700_000_000_500_000'i64, $loadCursor(db))
146-
147- # Cursors must never move backwards, or a late event would rewind the
148- # subscription on the next reconnect.
149- let cs2 = newDbCursorStore(db, flushEvery = 0)
150- cs2.save(1'i64)
151- cs2.close()
152- report("an older cursor does not overwrite a newer one",
153- loadCursor(db) == 1_700_000_000_500_000'i64, $loadCursor(db))
154-
155- # --- the seam with the supervisor --------------------------------------
156- # jetstream_run takes a CursorStore base reference, so this only works if
157- # the overrides dispatch dynamically. A silent fall-through to the base
158- # method would read the in-memory field instead of the database and look
159- # entirely healthy while resuming from the wrong place.
160- block:
161- let cs = newDbCursorStore(db, flushEvery = 3600)
162- let asBase: CursorStore = cs
163- report("DbCursorStore dispatches through the CursorStore base",
164- asBase.load() == loadCursor(db), $asBase.load())
165-
166- let rewound = rewoundCursor(asBase, 5.seconds)
167- report("the supervisor rewinds the persisted cursor",
168- rewound == loadCursor(db) - 5_000_000, $rewound)
169-
170- # --- deletion ----------------------------------------------------------
171- block:
172- deleteSession(db, "did:plc:probe", "sess-1")
173- report("sessions can be deleted",
174- getSession(db, "did:plc:probe", "sess-1").isNone and
175- listAccounts(db).len == 0)
176-
177- echo ""
178- if failures == 0:
179- echo "RESULT: sessions, keys and cursors survive a restart."
180- else:
181- echo &"RESULT: {failures} check(s) failed."
182- quit 1
183-
184-when isMainModule:
185- main()
modified spike/nim/stores_probe.nim +1 -1
@@ -7,7 +7,7 @@
77 ## returns other people's feeds, a NULL that sorts to the top.
88
99 import std/[options, os, strformat, strutils, times]
10-import atproto/[articlestore, feedstore, gleandb, sqlite]
10+import reader/[articlestore, feedstore, gleandb, sqlite]
1111
1212 var failures = 0
1313
@@ -7,7 +7,7 @@
7 ## returns other people's feeds, a NULL that sorts to the top.7 ## returns other people's feeds, a NULL that sorts to the top.
8 8
9 import std/[options, os, strformat, strutils, times]9 import std/[options, os, strformat, strutils, times]
10-import atproto/[articlestore, feedstore, gleandb, sqlite]10+import reader/[articlestore, feedstore, gleandb, sqlite]
11 11
12 var failures = 012 var failures = 0
13 13
deleted spike/nim/tools/extract_schema.py +0 -72
deleted file mode 100755
@@ -1,72 +0,0 @@
1-#!/usr/bin/env python3
2-"""Regenerate atproto/schema.nim from the Go schema.
3-
4-The DDL lives in internal/db/db.go as three []string literals. Retyping 44
5-statements into Nim would drift -- a differing default or CHECK constraint
6-surfaces months later as a violation nobody can place -- so they are copied
7-mechanically instead.
8-
9- python3 tools/extract_schema.py
10-"""
11-import pathlib
12-import re
13-import sys
14-
15-GO = pathlib.Path(__file__).resolve().parents[3] / "internal" / "db" / "db.go"
16-OUT = pathlib.Path(__file__).resolve().parents[1] / "atproto" / "schema.nim"
17-
18-HEADER = '''## Glean's schema, in the three databases it uses.
19-##
20-## Extracted mechanically from internal/db/db.go rather than retyped: 44 DDL
21-## statements transcribed by hand would drift from the Go definition, and a
22-## column that differs by a default or a CHECK is the kind of thing that only
23-## shows up as a constraint violation months later.
24-##
25-## Regenerate with tools/extract_schema.py after changing the Go schema.
26-
27-'''
28-
29-
30-def groups(src):
31- """Locate each `var <name>Schema = []string{ ... }` block by line."""
32- lines = src.splitlines()
33- found = []
34- for i, line in enumerate(lines):
35- m = re.match(r"var (\w+Schema) = \[\]string\{", line)
36- if not m:
37- continue
38- for j in range(i + 1, len(lines)):
39- if lines[j] == "}":
40- found.append((m.group(1), i + 1, j))
41- break
42- return lines, found
43-
44-
45-def main():
46- if not GO.exists():
47- sys.exit(f"cannot find {GO}")
48- src = GO.read_text()
49- lines, found = groups(src)
50- if not found:
51- sys.exit("no schema blocks found; did db.go change shape?")
52-
53- out = [HEADER]
54- total = 0
55- for name, start, end in found:
56- body = "\n".join(lines[start:end])
57- stmts = re.findall(r"`([^`]*)`", body)
58- total += len(stmts)
59- out.append(f"const {name}* = [\n")
60- for st in stmts:
61- if '"""' in st:
62- sys.exit(f"statement contains a triple quote: {st[:60]}")
63- out.append(f' """{st}""",\n')
64- out.append("]\n\n")
65- print(f"{name}: {len(stmts)} statements")
66-
67- OUT.write_text("".join(out))
68- print(f"wrote {OUT} ({total} statements)")
69-
70-
71-if __name__ == "__main__":
72- main()
deleted file mode 100755
@@ -1,72 +0,0 @@
1-#!/usr/bin/env python3
2-"""Regenerate atproto/schema.nim from the Go schema.
3-
4-The DDL lives in internal/db/db.go as three []string literals. Retyping 44
5-statements into Nim would drift -- a differing default or CHECK constraint
6-surfaces months later as a violation nobody can place -- so they are copied
7-mechanically instead.
8-
9- python3 tools/extract_schema.py
10-"""
11-import pathlib
12-import re
13-import sys
14-
15-GO = pathlib.Path(__file__).resolve().parents[3] / "internal" / "db" / "db.go"
16-OUT = pathlib.Path(__file__).resolve().parents[1] / "atproto" / "schema.nim"
17-
18-HEADER = '''## Glean's schema, in the three databases it uses.
19-##
20-## Extracted mechanically from internal/db/db.go rather than retyped: 44 DDL
21-## statements transcribed by hand would drift from the Go definition, and a
22-## column that differs by a default or a CHECK is the kind of thing that only
23-## shows up as a constraint violation months later.
24-##
25-## Regenerate with tools/extract_schema.py after changing the Go schema.
26-
27-'''
28-
29-
30-def groups(src):
31- """Locate each `var <name>Schema = []string{ ... }` block by line."""
32- lines = src.splitlines()
33- found = []
34- for i, line in enumerate(lines):
35- m = re.match(r"var (\w+Schema) = \[\]string\{", line)
36- if not m:
37- continue
38- for j in range(i + 1, len(lines)):
39- if lines[j] == "}":
40- found.append((m.group(1), i + 1, j))
41- break
42- return lines, found
43-
44-
45-def main():
46- if not GO.exists():
47- sys.exit(f"cannot find {GO}")
48- src = GO.read_text()
49- lines, found = groups(src)
50- if not found:
51- sys.exit("no schema blocks found; did db.go change shape?")
52-
53- out = [HEADER]
54- total = 0
55- for name, start, end in found:
56- body = "\n".join(lines[start:end])
57- stmts = re.findall(r"`([^`]*)`", body)
58- total += len(stmts)
59- out.append(f"const {name}* = [\n")
60- for st in stmts:
61- if '"""' in st:
62- sys.exit(f"statement contains a triple quote: {st[:60]}")
63- out.append(f' """{st}""",\n')
64- out.append("]\n\n")
65- print(f"{name}: {len(stmts)} statements")
66-
67- OUT.write_text("".join(out))
68- print(f"wrote {OUT} ({total} statements)")
69-
70-
71-if __name__ == "__main__":
72- main()
deleted spike/nim/vendor/sqlite-vec.c +0 -9749
deleted file mode 100644
@@ -1,9749 +0,0 @@
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-}
deleted file mode 100644
@@ -1,9749 +0,0 @@
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-}
deleted spike/nim/vendor/sqlite-vec.h +0 -41
deleted file mode 100644
@@ -1,41 +0,0 @@
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 */
deleted file mode 100644
@@ -1,41 +0,0 @@
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 */
deleted spike/nim/verify/go.mod +0 -3
deleted file mode 100644
@@ -1,3 +0,0 @@
1-module glean.spike/verify
2-
3-go 1.26.2
deleted file mode 100644
@@ -1,3 +0,0 @@
1-module glean.spike/verify
2-
3-go 1.26.2
deleted spike/nim/verify/verify.go +0 -178
deleted file mode 100644
@@ -1,178 +0,0 @@
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-}
deleted file mode 100644
@@ -1,178 +0,0 @@
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-}
deleted spike/nim/xrpc_probe.nim +0 -123
deleted file mode 100644
@@ -1,123 +0,0 @@
1-## Exercises the XRPC layer as far as it goes without a consented session.
2-##
3-## nim c -r -d:ssl xrpc_probe.nim
4-##
5-## A valid access token needs a human in a browser, so the authenticated path
6-## is probed with a deliberately invalid one. That still proves the parts most
7-## likely to be wrong: that a real PDS accepts the *shape* of our request and
8-## rejects it for the right reason, and that the 401/WWW-Authenticate nonce
9-## handshake behaves the way the retry logic assumes.
10-
11-import std/[httpclient, json, strformat, strutils]
12-import bearssl/rand
13-import atproto/[dpop, identity, xrpc]
14-
15-var failures = 0
16-
17-proc report(name: string, ok: bool, detail = "") =
18- if not ok: inc failures
19- let label = if ok: "PASS" else: "FAIL"
20- if detail.len > 0: echo &" {label} {name} -- {detail}"
21- else: echo &" {label} {name}"
22-
23-when isMainModule:
24- echo "XRPC signing probe"
25- echo ""
26-
27- let rngRef = HmacDrbgContext.new()
28- doAssert rngRef != nil
29- var rng = rngRef[]
30-
31- # --- claim construction -------------------------------------------------
32- block:
33- report("htu strips the query string",
34- dpopUrl("https://pds.example/xrpc/a.b.c?limit=5#frag") ==
35- "https://pds.example/xrpc/a.b.c",
36- dpopUrl("https://pds.example/xrpc/a.b.c?limit=5#frag"))
37-
38- block:
39- # Pinned against an independently computed value:
40- # printf 'test-access-token' | openssl dgst -sha256 -binary \
41- # | basenc --base64url | tr -d '='
42- # A self-consistent ath would pass any internal check and still be
43- # rejected by every PDS.
44- const
45- token = "test-access-token"
46- want = "WXSA1LYsphIZPxnnP-TMOtF_C_nPwWp8v0tQZBMcSAU"
47- let ath = accessTokenHash(token)
48- report("ath matches an independently computed SHA-256", ath == want, ath)
49- report("ath is unpadded base64url", ath.len == 43 and '=' notin ath)
50-
51- block:
52- let key = generateKey(rng)
53- let proof = dpopProof(key, rng, "GET", "https://pds.example/xrpc/a.b.c",
54- accessToken = "tok", issuer = "https://auth.example")
55- let payload = parseJson(b64uDecode(proof.split('.')[1]))
56- report("resource proof carries ath", payload.hasKey("ath"))
57- report("resource proof carries iss", payload{"iss"}.getStr == "https://auth.example")
58- report("resource proof carries exp", payload.hasKey("exp"))
59-
60- let authProof = dpopProof(key, rng, "POST", "https://auth.example/oauth/par")
61- let authPayload = parseJson(b64uDecode(authProof.split('.')[1]))
62- # An auth-server proof must NOT claim to be bound to a token.
63- report("auth-server proof omits ath", not authPayload.hasKey("ath"))
64-
65- # --- unauthenticated XRPC -----------------------------------------------
66- # Separates "can we speak XRPC" from "is our signing right".
67- block:
68- try:
69- let j = publicQuery("public.api.bsky.app", "app.bsky.actor.getProfile",
70- {"actor": "bsky.app"})
71- report("unauthenticated XRPC query works", j{"did"}.getStr.isDid,
72- j{"handle"}.getStr)
73- except CatchableError as e:
74- report("unauthenticated XRPC query works", false, e.msg)
75-
76- # --- the authenticated path, against a real PDS -------------------------
77- var ident: Identity
78- try:
79- ident = lookup("bsky.app")
80- except CatchableError as e:
81- report("resolve a PDS to talk to", false, e.msg)
82- quit 1
83- report("resolved a PDS", true, ident.pdsEndpoint)
84-
85- block:
86- # A well-formed DPoP request carrying a token the PDS has never issued.
87- # The interesting question is *which* error comes back: a shape problem
88- # would read as invalid_dpop_proof, whereas invalid_token means the PDS
89- # parsed and accepted everything except the credential itself.
90- var sess = Session(
91- did: ident.did,
92- pdsUrl: ident.pdsEndpoint,
93- authServerUrl: "https://bsky.social",
94- accessToken: "not-a-real-token",
95- dpopKey: generateKey(rng),
96- clientId: "http://localhost",
97- )
98- # getSession, not listRecords: repo reads are public, so a PDS serves them
99- # happily while ignoring the Authorization header entirely -- which looks
100- # like a pass and tests nothing. This endpoint exists to answer "who am
101- # I", so it has to evaluate the credential.
102- let url = &"{sess.pdsUrl}/xrpc/com.atproto.server.getSession"
103- let res = request(sess, rng, HttpGet, url)
104-
105- report("PDS rejects an unknown token with 401", res.code.int == 401,
106- &"HTTP {res.code.int}")
107-
108- let wwwAuth = res.headers.getOrDefault("WWW-Authenticate").string
109- report("401 carries WWW-Authenticate", wwwAuth.len > 0, wwwAuth)
110- report("challenge is DPoP, not Bearer",
111- wwwAuth.startsWith("DPoP"), wwwAuth.split(' ')[0])
112- report("the proof itself was accepted (error is about the token, "&
113- "not the proof)",
114- wwwAuth.len > 0 and "invalid_dpop_proof" notin wwwAuth, wwwAuth)
115- report("PDS issued a host nonce for the retry", sess.hostNonce.len > 0,
116- sess.hostNonce[0 ..< min(12, sess.hostNonce.len)] & "...")
117-
118- echo ""
119- if failures == 0:
120- echo "RESULT: XRPC signing is correct as far as an unconsented session can show."
121- else:
122- echo &"RESULT: {failures} check(s) failed."
123- quit 1
deleted file mode 100644
@@ -1,123 +0,0 @@
1-## Exercises the XRPC layer as far as it goes without a consented session.
2-##
3-## nim c -r -d:ssl xrpc_probe.nim
4-##
5-## A valid access token needs a human in a browser, so the authenticated path
6-## is probed with a deliberately invalid one. That still proves the parts most
7-## likely to be wrong: that a real PDS accepts the *shape* of our request and
8-## rejects it for the right reason, and that the 401/WWW-Authenticate nonce
9-## handshake behaves the way the retry logic assumes.
10-
11-import std/[httpclient, json, strformat, strutils]
12-import bearssl/rand
13-import atproto/[dpop, identity, xrpc]
14-
15-var failures = 0
16-
17-proc report(name: string, ok: bool, detail = "") =
18- if not ok: inc failures
19- let label = if ok: "PASS" else: "FAIL"
20- if detail.len > 0: echo &" {label} {name} -- {detail}"
21- else: echo &" {label} {name}"
22-
23-when isMainModule:
24- echo "XRPC signing probe"
25- echo ""
26-
27- let rngRef = HmacDrbgContext.new()
28- doAssert rngRef != nil
29- var rng = rngRef[]
30-
31- # --- claim construction -------------------------------------------------
32- block:
33- report("htu strips the query string",
34- dpopUrl("https://pds.example/xrpc/a.b.c?limit=5#frag") ==
35- "https://pds.example/xrpc/a.b.c",
36- dpopUrl("https://pds.example/xrpc/a.b.c?limit=5#frag"))
37-
38- block:
39- # Pinned against an independently computed value:
40- # printf 'test-access-token' | openssl dgst -sha256 -binary \
41- # | basenc --base64url | tr -d '='
42- # A self-consistent ath would pass any internal check and still be
43- # rejected by every PDS.
44- const
45- token = "test-access-token"
46- want = "WXSA1LYsphIZPxnnP-TMOtF_C_nPwWp8v0tQZBMcSAU"
47- let ath = accessTokenHash(token)
48- report("ath matches an independently computed SHA-256", ath == want, ath)
49- report("ath is unpadded base64url", ath.len == 43 and '=' notin ath)
50-
51- block:
52- let key = generateKey(rng)
53- let proof = dpopProof(key, rng, "GET", "https://pds.example/xrpc/a.b.c",
54- accessToken = "tok", issuer = "https://auth.example")
55- let payload = parseJson(b64uDecode(proof.split('.')[1]))
56- report("resource proof carries ath", payload.hasKey("ath"))
57- report("resource proof carries iss", payload{"iss"}.getStr == "https://auth.example")
58- report("resource proof carries exp", payload.hasKey("exp"))
59-
60- let authProof = dpopProof(key, rng, "POST", "https://auth.example/oauth/par")
61- let authPayload = parseJson(b64uDecode(authProof.split('.')[1]))
62- # An auth-server proof must NOT claim to be bound to a token.
63- report("auth-server proof omits ath", not authPayload.hasKey("ath"))
64-
65- # --- unauthenticated XRPC -----------------------------------------------
66- # Separates "can we speak XRPC" from "is our signing right".
67- block:
68- try:
69- let j = publicQuery("public.api.bsky.app", "app.bsky.actor.getProfile",
70- {"actor": "bsky.app"})
71- report("unauthenticated XRPC query works", j{"did"}.getStr.isDid,
72- j{"handle"}.getStr)
73- except CatchableError as e:
74- report("unauthenticated XRPC query works", false, e.msg)
75-
76- # --- the authenticated path, against a real PDS -------------------------
77- var ident: Identity
78- try:
79- ident = lookup("bsky.app")
80- except CatchableError as e:
81- report("resolve a PDS to talk to", false, e.msg)
82- quit 1
83- report("resolved a PDS", true, ident.pdsEndpoint)
84-
85- block:
86- # A well-formed DPoP request carrying a token the PDS has never issued.
87- # The interesting question is *which* error comes back: a shape problem
88- # would read as invalid_dpop_proof, whereas invalid_token means the PDS
89- # parsed and accepted everything except the credential itself.
90- var sess = Session(
91- did: ident.did,
92- pdsUrl: ident.pdsEndpoint,
93- authServerUrl: "https://bsky.social",
94- accessToken: "not-a-real-token",
95- dpopKey: generateKey(rng),
96- clientId: "http://localhost",
97- )
98- # getSession, not listRecords: repo reads are public, so a PDS serves them
99- # happily while ignoring the Authorization header entirely -- which looks
100- # like a pass and tests nothing. This endpoint exists to answer "who am
101- # I", so it has to evaluate the credential.
102- let url = &"{sess.pdsUrl}/xrpc/com.atproto.server.getSession"
103- let res = request(sess, rng, HttpGet, url)
104-
105- report("PDS rejects an unknown token with 401", res.code.int == 401,
106- &"HTTP {res.code.int}")
107-
108- let wwwAuth = res.headers.getOrDefault("WWW-Authenticate").string
109- report("401 carries WWW-Authenticate", wwwAuth.len > 0, wwwAuth)
110- report("challenge is DPoP, not Bearer",
111- wwwAuth.startsWith("DPoP"), wwwAuth.split(' ')[0])
112- report("the proof itself was accepted (error is about the token, "&
113- "not the proof)",
114- wwwAuth.len > 0 and "invalid_dpop_proof" notin wwwAuth, wwwAuth)
115- report("PDS issued a host nonce for the retry", sess.hostNonce.len > 0,
116- sess.hostNonce[0 ..< min(12, sess.hostNonce.len)] & "...")
117-
118- echo ""
119- if failures == 0:
120- echo "RESULT: XRPC signing is correct as far as an unconsented session can show."
121- else:
122- echo &"RESULT: {failures} check(s) failed."
123- quit 1