Repository navigation
Expand file tree
/
Copy pathtext_index.cpp
More file actions
476 lines (440 loc) · 18.7 KB
/
Copy pathtext_index.cpp
File metadata and controls
476 lines (440 loc) · 18.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
/**
* sqlite_text_index.cpp — SQLite storage hooks for the custom terms index.
*/
#include "sqlite/text_index.h"
#include <format>
#include <iostream>
#include <stdexcept>
#include <unordered_set>
#include "util/sqlite.h"
namespace ragger::sqlite {
// Doc embed text formula (title appended when present) — mirrors the constant
// Backend uses for embeddings so the index sees the same document text.
namespace {
constexpr const char* kDocEmbedTextSQL =
"text || CASE WHEN (SELECT ifnull(title,'') FROM document_sources ds "
"WHERE ds.document_source_id = documents.document_source_id) = '' "
"THEN '' ELSE char(10) || (SELECT title FROM document_sources ds "
"WHERE ds.document_source_id = documents.document_source_id) END";
} // namespace
void TextIndex::validate_table(const std::string& table) {
static const std::unordered_set<std::string> valid = {
"turns", "turn_summaries", "summaries", "documents", "decisions"};
if (valid.find(table) == valid.end())
throw std::runtime_error("Invalid table name: " + table);
}
std::string TextIndex::id_col_of(const std::string& table) {
if (table == "turns") return "turn_id";
if (table == "turn_summaries") return "turn_summary_id";
if (table == "summaries") return "summary_id";
if (table == "documents") return "document_id";
if (table == "decisions") return "decision_id";
throw std::runtime_error("Invalid table name: " + table);
}
std::string TextIndex::text_select_of(const std::string& table) {
const std::string id = id_col_of(table);
if (table == "documents")
return std::format("SELECT {}, {} AS text FROM documents ORDER BY {}",
id, kDocEmbedTextSQL, id);
if (table == "turns")
return "SELECT turn_id, user_text, assistant_text FROM turns ORDER BY turn_id";
return std::format("SELECT {}, text FROM {} ORDER BY {}", id, table, id);
}
void TextIndex::create_schema() {
auto exec = [&](const std::string& sql) {
char* err = nullptr;
if (sqlite3_exec(db_, sql.c_str(), nullptr, nullptr, &err) != SQLITE_OK) {
std::string msg = err ? err : "unknown error";
sqlite3_free(err);
throw std::runtime_error("TextIndex schema: " + msg);
}
};
exec(R"(
CREATE TABLE IF NOT EXISTS terms (
term_id INTEGER PRIMARY KEY AUTOINCREMENT,
term TEXT NOT NULL UNIQUE
)
)");
// One junction table per text table. FK ON DELETE CASCADE means a deleted
// record's index rows go away automatically — nothing to maintain in the
// write path beyond delete-then-insert for updates.
struct T { const char* junc; const char* base; const char* pk; };
const T tables[] = {
{"turns_terms", "turns", "turn_id"},
{"turn_summaries_terms", "turn_summaries", "turn_summary_id"},
{"summaries_terms", "summaries", "summary_id"},
{"documents_terms", "documents", "document_id"},
{"decisions_terms", "decisions", "decision_id"},
};
for (const auto& t : tables) {
exec(std::format(
"CREATE TABLE IF NOT EXISTS {} (\n"
" {} INTEGER NOT NULL REFERENCES {}({}) ON DELETE CASCADE,\n"
" term_id INTEGER NOT NULL REFERENCES terms(term_id) ON DELETE CASCADE,\n"
" count INTEGER NOT NULL DEFAULT 1,\n"
" PRIMARY KEY ({}, term_id)\n)",
t.junc, t.pk, t.base, t.pk, t.pk));
exec(std::format(
"CREATE INDEX IF NOT EXISTS idx_{}_term ON {}(term_id)",
t.junc, t.junc));
}
}
// ---- cached-statement plumbing --------------------------------------------
Stmt& TextIndex::cached(std::optional<Stmt>& slot, const char* sql) {
if (!slot) {
slot.emplace(db_, sql);
} else {
sqlite3_reset(slot->raw());
sqlite3_clear_bindings(slot->raw());
}
return *slot;
}
Stmt& TextIndex::select_term_id() {
return cached(select_term_id_, "SELECT term_id FROM terms WHERE term = ?");
}
Stmt& TextIndex::insert_term() {
return cached(insert_term_, "INSERT OR IGNORE INTO terms(term) VALUES (?)");
}
Stmt& TextIndex::doc_freq_turns() {
return cached(doc_freq_turns_, "SELECT COUNT(*) FROM turns_terms WHERE term_id = ?");
}
Stmt& TextIndex::doc_freq_turn_summaries() {
return cached(doc_freq_turn_summaries_,
"SELECT COUNT(*) FROM turn_summaries_terms WHERE term_id = ?");
}
Stmt& TextIndex::doc_freq_summaries() {
return cached(doc_freq_summaries_,
"SELECT COUNT(*) FROM summaries_terms WHERE term_id = ?");
}
Stmt& TextIndex::doc_freq_documents() {
return cached(doc_freq_documents_,
"SELECT COUNT(*) FROM documents_terms WHERE term_id = ?");
}
Stmt& TextIndex::doc_freq_decisions() {
return cached(doc_freq_decisions_,
"SELECT COUNT(*) FROM decisions_terms WHERE term_id = ?");
}
Stmt& TextIndex::count_turns() {
return cached(count_turns_, "SELECT COUNT(*) FROM turns");
}
Stmt& TextIndex::count_turn_summaries() {
return cached(count_turn_summaries_, "SELECT COUNT(*) FROM turn_summaries");
}
Stmt& TextIndex::count_summaries() {
return cached(count_summaries_, "SELECT COUNT(*) FROM summaries");
}
Stmt& TextIndex::count_documents() {
return cached(count_documents_, "SELECT COUNT(*) FROM documents");
}
Stmt& TextIndex::count_decisions() {
return cached(count_decisions_, "SELECT COUNT(*) FROM decisions");
}
Stmt& TextIndex::postings_turns() {
return cached(postings_turns_,
"SELECT j.turn_id, j.count, b.unigram_count, b.bigram_count "
"FROM turns_terms j JOIN turns b ON b.turn_id = j.turn_id "
"WHERE j.term_id = ?");
}
Stmt& TextIndex::postings_turn_summaries() {
return cached(postings_turn_summaries_,
"SELECT j.turn_summary_id, j.count, b.unigram_count, b.bigram_count "
"FROM turn_summaries_terms j JOIN turn_summaries b "
"ON b.turn_summary_id = j.turn_summary_id "
"WHERE j.term_id = ?");
}
Stmt& TextIndex::postings_summaries() {
return cached(postings_summaries_,
"SELECT j.summary_id, j.count, b.unigram_count, b.bigram_count "
"FROM summaries_terms j JOIN summaries b ON b.summary_id = j.summary_id "
"WHERE j.term_id = ?");
}
Stmt& TextIndex::postings_documents() {
return cached(postings_documents_,
"SELECT j.document_id, j.count, b.unigram_count, b.bigram_count "
"FROM documents_terms j JOIN documents b ON b.document_id = j.document_id "
"WHERE j.term_id = ?");
}
Stmt& TextIndex::postings_decisions() {
return cached(postings_decisions_,
"SELECT j.decision_id, j.count, b.unigram_count, b.bigram_count "
"FROM decisions_terms j JOIN decisions b ON b.decision_id = j.decision_id "
"WHERE j.term_id = ?");
}
Stmt& TextIndex::delete_terms_turns() {
return cached(delete_terms_turns_, "DELETE FROM turns_terms WHERE turn_id = ?");
}
Stmt& TextIndex::delete_terms_turn_summaries() {
return cached(delete_terms_turn_summaries_,
"DELETE FROM turn_summaries_terms WHERE turn_summary_id = ?");
}
Stmt& TextIndex::delete_terms_summaries() {
return cached(delete_terms_summaries_, "DELETE FROM summaries_terms WHERE summary_id = ?");
}
Stmt& TextIndex::delete_terms_documents() {
return cached(delete_terms_documents_, "DELETE FROM documents_terms WHERE document_id = ?");
}
Stmt& TextIndex::delete_terms_decisions() {
return cached(delete_terms_decisions_, "DELETE FROM decisions_terms WHERE decision_id = ?");
}
Stmt& TextIndex::insert_terms_turns() {
return cached(insert_terms_turns_,
"INSERT OR REPLACE INTO turns_terms (turn_id, term_id, count) VALUES (?, ?, ?)");
}
Stmt& TextIndex::insert_terms_turn_summaries() {
return cached(insert_terms_turn_summaries_,
"INSERT OR REPLACE INTO turn_summaries_terms (turn_summary_id, term_id, count) "
"VALUES (?, ?, ?)");
}
Stmt& TextIndex::insert_terms_summaries() {
return cached(insert_terms_summaries_,
"INSERT OR REPLACE INTO summaries_terms (summary_id, term_id, count) VALUES (?, ?, ?)");
}
Stmt& TextIndex::insert_terms_documents() {
return cached(insert_terms_documents_,
"INSERT OR REPLACE INTO documents_terms (document_id, term_id, count) "
"VALUES (?, ?, ?)");
}
Stmt& TextIndex::insert_terms_decisions() {
return cached(insert_terms_decisions_,
"INSERT OR REPLACE INTO decisions_terms (decision_id, term_id, count) "
"VALUES (?, ?, ?)");
}
Stmt& TextIndex::update_counts_turns() {
return cached(update_counts_turns_,
"UPDATE turns SET unigram_count = ?, bigram_count = ? WHERE turn_id = ?");
}
Stmt& TextIndex::update_counts_turn_summaries() {
return cached(update_counts_turn_summaries_,
"UPDATE turn_summaries SET unigram_count = ?, bigram_count = ? "
"WHERE turn_summary_id = ?");
}
Stmt& TextIndex::update_counts_summaries() {
return cached(update_counts_summaries_,
"UPDATE summaries SET unigram_count = ?, bigram_count = ? WHERE summary_id = ?");
}
Stmt& TextIndex::update_counts_documents() {
return cached(update_counts_documents_,
"UPDATE documents SET unigram_count = ?, bigram_count = ? WHERE document_id = ?");
}
Stmt& TextIndex::update_counts_decisions() {
return cached(update_counts_decisions_,
"UPDATE decisions SET unigram_count = ?, bigram_count = ? WHERE decision_id = ?");
}
// ---- runtime dispatch over the fixed 5-table set --------------------------
Stmt& TextIndex::doc_freq_stmt(const std::string& table) {
if (table == "turns") return doc_freq_turns();
if (table == "turn_summaries") return doc_freq_turn_summaries();
if (table == "summaries") return doc_freq_summaries();
if (table == "documents") return doc_freq_documents();
return doc_freq_decisions();
}
Stmt& TextIndex::count_stmt(const std::string& table) {
if (table == "turns") return count_turns();
if (table == "turn_summaries") return count_turn_summaries();
if (table == "summaries") return count_summaries();
if (table == "documents") return count_documents();
return count_decisions();
}
Stmt& TextIndex::postings_stmt(const std::string& table) {
if (table == "turns") return postings_turns();
if (table == "turn_summaries") return postings_turn_summaries();
if (table == "summaries") return postings_summaries();
if (table == "documents") return postings_documents();
return postings_decisions();
}
Stmt& TextIndex::delete_terms_stmt(const std::string& table) {
if (table == "turns") return delete_terms_turns();
if (table == "turn_summaries") return delete_terms_turn_summaries();
if (table == "summaries") return delete_terms_summaries();
if (table == "documents") return delete_terms_documents();
return delete_terms_decisions();
}
Stmt& TextIndex::insert_terms_stmt(const std::string& table) {
if (table == "turns") return insert_terms_turns();
if (table == "turn_summaries") return insert_terms_turn_summaries();
if (table == "summaries") return insert_terms_summaries();
if (table == "documents") return insert_terms_documents();
return insert_terms_decisions();
}
Stmt& TextIndex::update_counts_stmt(const std::string& table) {
if (table == "turns") return update_counts_turns();
if (table == "turn_summaries") return update_counts_turn_summaries();
if (table == "summaries") return update_counts_summaries();
if (table == "documents") return update_counts_documents();
return update_counts_decisions();
}
// ---- TextIndex overrides ---------------------------------------------------
void TextIndex::reset_terms_table() {
// Drop-and-recreate instead of DELETE: on ~900K+ rows, DELETE FROM terms
// has to walk every row, update the UNIQUE index B-tree, and write a WAL
// frame per row -- O(n) work. DROP TABLE is a catalog-only operation
// (SQLite just forgets the table/index exist), so it's near-instant
// regardless of row count; CREATE TABLE then reuses create_schema()'s
// exact DDL, so the two paths can never drift apart.
//
// Wrapped in one transaction: a crash mid-sequence rolls back to the
// pre-reset state instead of leaving tables half-gone.
//
// No manual sqlite_sequence cleanup needed -- a freshly CREATE'd
// AUTOINCREMENT table has no sqlite_sequence row until its first insert.
//
// Any cached Stmt bound to `terms` or a `<table>_terms` junction (e.g.
// select_term_id_, insert_term_, the per-table junction insert/delete
// members) is prepared via sqlite3_prepare_v2, which SQLite transparently
// reprepares against the new schema on its next step() -- no code
// elsewhere needs to change for this.
char* err = nullptr;
auto exec = [&](const char* sql) {
if (sqlite3_exec(db_, sql, nullptr, nullptr, &err) != SQLITE_OK) {
std::string msg = err ? err : "unknown error";
sqlite3_free(err);
throw std::runtime_error("TextIndex reset_terms_table: " + msg);
}
};
exec("BEGIN IMMEDIATE");
try {
// Junction tables reference terms(term_id); drop them first so no
// statement can observe a junction row pointing at a gone table
// (though within one transaction this is moot for correctness, it
// keeps drop order mirroring the FK dependency direction).
exec("DROP TABLE IF EXISTS turns_terms");
exec("DROP TABLE IF EXISTS turn_summaries_terms");
exec("DROP TABLE IF EXISTS summaries_terms");
exec("DROP TABLE IF EXISTS documents_terms");
exec("DROP TABLE IF EXISTS decisions_terms");
exec("DROP TABLE IF EXISTS terms");
create_schema(); // CREATE TABLE/INDEX IF NOT EXISTS -- recreates all 6, empty
exec("COMMIT");
} catch (...) {
exec("ROLLBACK");
throw;
}
}
int TextIndex::upsert_term(const std::string& term) {
// SELECT-first, not INSERT-OR-IGNORE-first: AUTOINCREMENT allocates the
// next term_id BEFORE the uniqueness check runs, so an "ignored" insert
// (term already exists) still permanently burns an id -- with upsert_term
// called on every tokenize of every store/reindex/backfill, that leaked
// ~13M ids against ~900K real terms in the live DB. Checking first avoids
// ever attempting the redundant insert for an existing term.
Stmt& sel = select_term_id();
sel.bind(1, term);
int existing = sel.step() ? sel.column_int(0) : 0;
sqlite3_reset(sel.raw()); // release the read snapshot immediately, see below
if (existing != 0) return existing;
Stmt& ins = insert_term();
ins.bind(1, term);
ins.exec();
return static_cast<int>(sqlite3_last_insert_rowid(db_));
}
void TextIndex::replace_record_terms(
const std::string& table, int id,
const std::unordered_map<int, int>& term_counts) {
validate_table(table);
{
Stmt& del = delete_terms_stmt(table);
del.bind(1, id);
del.exec();
}
for (const auto& [term_id, count] : term_counts) {
Stmt& ins = insert_terms_stmt(table);
ins.bind(1, id).bind(2, term_id).bind(3, count);
ins.exec();
}
}
void TextIndex::set_counts(const std::string& table, int id,
int unigram_count, int bigram_count) {
validate_table(table);
Stmt& upd = update_counts_stmt(table);
upd.bind(1, unigram_count).bind(2, bigram_count).bind(3, id);
upd.exec();
}
int TextIndex::corpus_size(const std::string& table) {
validate_table(table);
Stmt& s = count_stmt(table);
int result = s.step() ? s.column_int(0) : 0;
sqlite3_reset(s.raw()); // release the read snapshot immediately, see upsert_term
return result;
}
int TextIndex::doc_freq(const std::string& table, int term_id) {
validate_table(table);
Stmt& s = doc_freq_stmt(table);
s.bind(1, term_id);
int result = s.step() ? s.column_int(0) : 0;
sqlite3_reset(s.raw()); // release the read snapshot immediately, see upsert_term
return result;
}
int TextIndex::lookup_term(const std::string& term) {
Stmt& s = select_term_id();
s.bind(1, term);
int result = s.step() ? s.column_int(0) : 0;
sqlite3_reset(s.raw()); // release the read snapshot immediately, see upsert_term
return result;
}
std::vector<TextIndex::Posting> TextIndex::postings(
const std::string& table, int term_id) {
validate_table(table);
Stmt& s = postings_stmt(table);
s.bind(1, term_id);
std::vector<Posting> out;
while (s.step()) {
out.push_back({s.column_int(0), s.column_int(1),
s.column_int(2), s.column_int(3)});
}
return out;
}
int TextIndex::reindex_table(const std::string& table, bool progress) {
validate_table(table);
const std::string id_col = id_col_of(table);
const std::string select_sql = text_select_of(table);
int total = 0;
if (progress) {
Stmt& c = count_stmt(table);
if (c.step()) total = c.column_int(0);
}
// Two-pass batches: read a bounded batch (finalize the SELECT), then write.
// Walking a SELECT while UPDATE-ing the same table on one connection can
// invalidate the cursor mid-scan (see rebuild_phon's note).
//
// This SELECT's shape is loop-invariant (select_sql doesn't change across
// outer iterations) but reindex_table() itself is a rare admin/CLI
// operation, not a per-row hot path -- so it's hoisted above the loop for
// clarity, not promoted to a cached instance member.
constexpr int kBatch = 200;
struct Row { int id; std::string text; };
int done = 0;
int last_id = 0;
for (;;) {
std::vector<Row> batch;
batch.reserve(kBatch);
{
// select_sql is ORDER BY id; skip ids already processed to page.
Stmt s(db_, select_sql);
while (s.step()) {
int id = s.column_int(0);
if (id <= last_id) continue;
std::string text = s.column_text(1);
if (table == "turns") {
std::string a = s.column_text(2);
if (!a.empty()) text += "\n" + a;
}
batch.push_back({id, std::move(text)});
if (static_cast<int>(batch.size()) >= kBatch) break;
}
}
if (batch.empty()) break;
for (const auto& row : batch) {
index_record(table, row.id, row.text);
last_id = row.id;
++done;
if (progress) {
std::cout << std::format("\rReindexing {}: {}/{}",
table, done, total);
std::cout.flush();
}
}
}
if (progress) std::cout << "\n";
return done;
}
} // namespace ragger::sqlite