diff --git a/docs/fetch-and-storage.md b/docs/fetch-and-storage.md index 18437fb6..fdf1a280 100644 --- a/docs/fetch-and-storage.md +++ b/docs/fetch-and-storage.md @@ -288,6 +288,16 @@ so (`Rows (est.)`, a per-row `(est.)` marker, a footer pointing at `--exact`, an `db status` printed `Entities: 0 / Filings: 0` under a column labelled "Rows". Zero now means "no statistics yet"; a genuinely empty table pays one cheap `COUNT(*)`. +### A newly declared index needs one `db setup`, and nothing else + +`setupDatabase()` emits every index a registry entry declares as +`CREATE INDEX IF NOT EXISTS`, on every run and not only on a fresh database — so an index +added to `indexes` in `storageRegistry.ts` appears on an existing deployment the next time +`sec db setup` (or `sec setup`) runs, and no rows move. That is the opposite of a new +*column*, which `CREATE TABLE IF NOT EXISTS` cannot add and which needs the catch-up pass +below. On a large table the `CREATE INDEX` itself is the cost: it reads the table once and +holds a write lock while it builds. + ### Three schema catch-up passes `db setup` finishes with these, in order, all after the extension loop. diff --git a/src/cli/groups/ask.ts b/src/cli/groups/ask.ts index e0ba7756..a754cfad 100644 --- a/src/cli/groups/ask.ts +++ b/src/cli/groups/ask.ts @@ -66,12 +66,13 @@ export function addAskCommands(program: Command): void { defaults: { ...scope, limit: options.limit, force: options.force === true }, }), ]); + const skipped = out.skipped ?? 0; console.log( `indexed ${out.indexed} filing(s) · ${out.sections} sections` + - (out.skipped > 0 ? ` · ${out.skipped} already indexed` : "") + + (skipped > 0 ? ` · ${skipped} already indexed` : "") + (out.truncated ? " · stopped at --limit, run again for more" : "") ); - if (out.indexed > 0 || out.skipped > 0) { + if (out.indexed > 0 || skipped > 0) { suggest({ command: 'sec ask "..."', why: "ask a question about what is indexed" }); } else { suggest({ @@ -111,8 +112,12 @@ export function addAskCommands(program: Command): void { // question is so much as read. if (options.index !== false) { const limit = options.indexLimit ?? DEFAULT_ASK_INDEX_LIMIT; + // This pre-index reports what it indexed and nothing about what it + // skipped, so it asks for no count of the latter: that count is a + // join over every converted filing, and a question pays for it on + // the way to an answer that never shows it. const indexed = await runWorkflowCli([ - new IndexFilingSectionsTask({ defaults: { ...scope, limit } }), + new IndexFilingSectionsTask({ defaults: { ...scope, limit, countSkipped: false } }), ]); if (indexed.indexed > 0) { console.log( diff --git a/src/config/storageRegistry.ts b/src/config/storageRegistry.ts index ccfe37d2..ce8b895d 100644 --- a/src/config/storageRegistry.ts +++ b/src/config/storageRegistry.ts @@ -236,7 +236,17 @@ export const SEC_STORAGE_REGISTRY: readonly StorageDefinition[] = [ // at the current converter version" — the primary is written last, so its // presence is what means the whole submission landed. `converted_at` serves // the recency listing. - indexes: [["form", "converter_version", "is_primary"], ["converted_at"]], + // + // `(filing_date, accession_number)` is the order the index selection reads + // this table in — newest filing first, ties broken by accession — and the + // column order matches that ORDER BY exactly, so SQLite walks the index + // backwards instead of sorting the whole table into a temp B-tree to take + // the first page of it. + indexes: [ + ["form", "converter_version", "is_primary"], + ["converted_at"], + ["filing_date", "accession_number"], + ], }), defineStorage({ token: FILING_SECTION_REPOSITORY_TOKEN, diff --git a/src/kb/PagedChunkVectorStorage.ts b/src/kb/PagedChunkVectorStorage.ts index d1587547..9bc4eabe 100644 --- a/src/kb/PagedChunkVectorStorage.ts +++ b/src/kb/PagedChunkVectorStorage.ts @@ -7,6 +7,7 @@ import type { ChunkVectorPrimaryKey, ChunkVectorStorageSchema, + PageCursor, TypedArray, VectorSearchOptions, } from "workglow"; @@ -24,8 +25,11 @@ import { * A chunk carries its text and a JSON-encoded vector, so a page is the working * set: large enough that the scan is a few hundred statements over a corpus of * hundreds of thousands of chunks, small enough to stay a fixed cost. + * + * Exported so a test can seed exactly one page boundary rather than a number + * that happens to cross one today. */ -const SCAN_PAGE = 512; +export const SCAN_PAGE = 512; /** A scored row, as {@link SqliteVectorStorage.similaritySearch} returns them. */ interface Scored { @@ -65,6 +69,14 @@ function keepBest(kept: T[], row: T, topK: number): void { * index here, so every question still scores every chunk and latency grows with * the corpus. What this removes is the heap ceiling that made a large index * unqueryable rather than slow. + * + * The pages are keyset (seek) pages, not `OFFSET` pages. `OFFSET n` makes the + * database walk and discard the first `n` rows on every page, so paging a table + * that way costs O(rows²) and a large index is slower to read a page at a time + * than to read whole — which is the opposite of what paging it is for. Resuming + * from the last key seen reads each row once: the ordering is the primary key, + * which is the index SQLite already keeps, so a page is a seek into it and the + * scan is linear again. */ export class PagedChunkVectorStorage extends SqliteVectorStorage< ChunkVectorStorageSchema, @@ -77,27 +89,33 @@ export class PagedChunkVectorStorage extends SqliteVectorStorage< assertVectorShape(query, this.getVectorDimensions(), "query"); const { topK = 10, filter, scoreThreshold = 0 } = options; - type Row = NonNullable>>[number]; + type Row = Awaited>["items"][number]; const kept: (Row & Scored)[] = []; if (topK <= 0) return emitSimilaritySearch(this.events, query, kept); - // Ordered by the primary key so the pages partition the table: LIMIT with - // OFFSET and no ORDER BY is free to hand back a row twice and skip another. - for (let offset = 0; ; offset += SCAN_PAGE) { - const page = - (await this.getAll({ - orderBy: [{ column: "chunk_id", direction: "ASC" }], - limit: SCAN_PAGE, - offset, - })) ?? []; - for (const entity of page) { + // Ordered by the primary key so the pages partition the table, and resumed + // from the last key of the previous page rather than from a row count. The + // cursor is the store's own, which encodes that key and is refused if the + // ordering it was built under stops matching. + let cursor: PageCursor | undefined; + for (;;) { + const page = await this.getPage({ + orderBy: [{ column: "chunk_id", direction: "ASC" }], + limit: SCAN_PAGE, + cursor, + }); + for (const entity of page.items) { const metadata = (entity.metadata ?? {}) as Record; if (filter && !matchesFilter(metadata, filter)) continue; const score = cosineSimilarity(query, toVector(entity.vector)); if (score < scoreThreshold) continue; keepBest(kept, { ...entity, score }, topK); } - if (page.length < SCAN_PAGE) break; + // Both conditions. A table whose size is an exact multiple of the page + // hands back a cursor for its last full page, and the page after it is + // empty rather than absent; looping on the cursor alone would not end. + if (page.nextCursor === undefined || page.items.length === 0) break; + cursor = page.nextCursor; } return emitSimilaritySearch(this.events, query, kept); @@ -109,6 +127,15 @@ export class PagedChunkVectorStorage extends SqliteVectorStorage< * * SQLite holds it as a JSON string and the tabular read usually decodes it * already; a row written by an older release, or handed back raw, does not. + * + * `Float32Array` is hardcoded, and so are the `vector` and `metadata` property + * names the scan reads, where the base class resolves all three from the schema + * it was constructed with. That is a duplication, not a choice: the base holds + * them in private fields with no accessor, and every instance of this class is + * the one `ChunkVectorStorageSchema` store `getSecKnowledgeBase` builds, which + * declares exactly those. Constructed against a different schema this would be + * silently wrong, so widening it means reading those three off the base rather + * than restating them here. */ function toVector(stored: unknown): TypedArray { if (typeof stored === "string") return new Float32Array(JSON.parse(stored) as number[]); diff --git a/src/kb/secKnowledgeBaseSearch.test.ts b/src/kb/secKnowledgeBaseSearch.test.ts index 13b4019e..deae5627 100644 --- a/src/kb/secKnowledgeBaseSearch.test.ts +++ b/src/kb/secKnowledgeBaseSearch.test.ts @@ -8,24 +8,61 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { ChunkVectorPrimaryKey, ChunkVectorStorageSchema, SqliteVectorStorage } from "workglow"; import { withSqliteDb } from "../config/testing/withSqliteDb"; import { getDb } from "../util/db"; +import { SCAN_PAGE } from "./PagedChunkVectorStorage"; import { KB_CHUNK_TABLE } from "./secKbTables"; import { getSecKnowledgeBase, resetSecKnowledgeBaseForTesting } from "./secKnowledgeBase"; /** Narrow enough to seed a few thousand rows without a real embedding model. */ const DIMENSIONS = 4; -const CHUNKS = 1200; + +function chunkId(index: number): string { + return `c${String(index).padStart(5, "0")}`; +} /** * Similarity to `[1, 0, 0, 0]` is `1 / sqrt(1 + spread²)`, so the best chunks * are the ones with the smallest spread — and the seeding puts those LAST in * `chunk_id` order, where a scan that stops early would never reach them. */ -function seedVector(index: number): Float32Array { - return new Float32Array([1, CHUNKS - 1 - index, 0, 0]); +function seedVector(index: number, total: number): Float32Array { + return new Float32Array([1, total - 1 - index, 0, 0]); } -function chunkId(index: number): string { - return `c${String(index).padStart(5, "0")}`; +/** Writes `count` chunks onto the table the knowledge base just created. */ +async function seedChunks(count: number): Promise { + const seeder = new SqliteVectorStorage( + getDb(), + KB_CHUNK_TABLE, + ChunkVectorStorageSchema, + ChunkVectorPrimaryKey, + [], + DIMENSIONS + ); + await seeder.putBulk( + Array.from({ length: count }, (_unused, index) => ({ + chunk_id: chunkId(index), + doc_id: "0000320193-26-000001:primary.htm", + vector: seedVector(index, count), + metadata: { text: `chunk ${index}` }, + })) as never + ); +} + +/** Records every statement the shared connection prepares while `run` executes. */ +async function recordSql(run: () => Promise): Promise<{ result: T; prepared: string[] }> { + const db = getDb(); + const prepared: string[] = []; + const realPrepare = db.prepare.bind(db); + vi.spyOn(db, "prepare").mockImplementation(((sql: string) => { + prepared.push(sql); + return realPrepare(sql); + }) as never); + const result = await run(); + return { result, prepared }; +} + +function chunkReads(prepared: readonly string[]): string[] { + return prepared.filter((sql) => new RegExp(`FROM \`?${KB_CHUNK_TABLE}\``).test(sql)); } /** @@ -55,47 +92,78 @@ describe("the SEC knowledge base's chunk search", () => { // Opening the base creates the three tables; the seeding then writes // through a plain vector storage onto the very same table. const kb = await getSecKnowledgeBase(); - const seeder = new SqliteVectorStorage( - getDb(), - KB_CHUNK_TABLE, - ChunkVectorStorageSchema, - ChunkVectorPrimaryKey, - [], - DIMENSIONS - ); - await seeder.putBulk( - Array.from({ length: CHUNKS }, (_unused, index) => ({ - chunk_id: chunkId(index), - doc_id: "0000320193-26-000001:primary.htm", - vector: seedVector(index), - metadata: { text: `chunk ${index}` }, - })) as never - ); + const total = SCAN_PAGE * 2 + 37; + await seedChunks(total); - const db = getDb(); - const prepared: string[] = []; - const realPrepare = db.prepare.bind(db); - vi.spyOn(db, "prepare").mockImplementation(((sql: string) => { - prepared.push(sql); - return realPrepare(sql); - }) as never); - - const hits = await kb.similaritySearch(new Float32Array([1, 0, 0, 0]), { topK: 3 }); + const { result: hits, prepared } = await recordSql(() => + kb.similaritySearch(new Float32Array([1, 0, 0, 0]), { topK: 3 }) + ); // The whole index is still ranked — the three best chunks are the last // three rows, which only a scan that reaches the end can find. expect(hits.map((hit) => (hit as { chunk_id: string }).chunk_id)).toEqual([ - chunkId(CHUNKS - 1), - chunkId(CHUNKS - 2), - chunkId(CHUNKS - 3), + chunkId(total - 1), + chunkId(total - 2), + chunkId(total - 3), ]); - const chunkReads = prepared.filter((sql) => - new RegExp(`FROM \`?${KB_CHUNK_TABLE}\``).test(sql) - ); - expect(chunkReads.length).toBeGreaterThan(1); + const reads = chunkReads(prepared); + expect(reads.length).toBeGreaterThan(1); // Not one of them may be the unbounded read: a corpus of any size is then // in the heap at once, every row hydrated before a single score is taken. - for (const sql of chunkReads) expect(sql).toMatch(/LIMIT/); + for (const sql of reads) expect(sql).toMatch(/LIMIT/); + }); + + it("seeks to the last key seen rather than counting rows past", async () => { + // The distinction this pins: `OFFSET n` makes SQLite walk and discard the + // first n rows of every page, so scanning the table a page at a time costs + // O(rows²) — at a few hundred thousand chunks, slower than the unbounded + // read the paging replaced. A keyset page is a seek into the primary-key + // index and reads each row once. + const kb = await getSecKnowledgeBase(); + await seedChunks(SCAN_PAGE * 2 + 5); + + const { prepared } = await recordSql(() => + kb.similaritySearch(new Float32Array([1, 0, 0, 0]), { topK: 3 }) + ); + + const reads = chunkReads(prepared); + expect(reads.length).toBe(3); + for (const sql of reads) expect(sql).not.toMatch(/OFFSET/i); + // The first page starts at the beginning; every page after it resumes from + // the last `chunk_id` of the one before. + expect(reads[0]).not.toMatch(/chunk_id`? >/); + for (const sql of reads.slice(1)) expect(sql).toMatch(/`chunk_id` >/); + }); + + it("terminates on a corpus that is an exact multiple of the page", async () => { + // The empty-final-page case: the last full page still hands back a cursor, + // because nothing about it says it was the last. A loop that trusts the + // cursor alone asks for one more page, gets none, and asks again forever. + const kb = await getSecKnowledgeBase(); + const total = SCAN_PAGE * 2; + await seedChunks(total); + + const { result: hits, prepared } = await recordSql(() => + kb.similaritySearch(new Float32Array([1, 0, 0, 0]), { topK: 2 }) + ); + + expect(hits.map((hit) => (hit as { chunk_id: string }).chunk_id)).toEqual([ + chunkId(total - 1), + chunkId(total - 2), + ]); + // Two full pages plus the empty one that ends the scan. + expect(chunkReads(prepared).length).toBe(3); + }); + + it("ranks an empty index without reading a page twice", async () => { + const kb = await getSecKnowledgeBase(); + + const { result: hits, prepared } = await recordSql(() => + kb.similaritySearch(new Float32Array([1, 0, 0, 0]), { topK: 3 }) + ); + + expect(hits).toEqual([]); + expect(chunkReads(prepared).length).toBe(1); }); }); diff --git a/src/task/kb/IndexFilingSectionsTask.test.ts b/src/task/kb/IndexFilingSectionsTask.test.ts index 823996f1..e46fe74b 100644 --- a/src/task/kb/IndexFilingSectionsTask.test.ts +++ b/src/task/kb/IndexFilingSectionsTask.test.ts @@ -15,6 +15,7 @@ import { type FilingDocument, } from "../../storage/document/FilingDocumentSchema"; import { FILING_SECTION_REPOSITORY_TOKEN } from "../../storage/document/FilingSectionSchema"; +import { getDb } from "../../util/db"; import { IndexFilingSectionsTask } from "./IndexFilingSectionsTask"; import { kbDocIdFor } from "./selectDocumentsToIndex"; @@ -124,6 +125,30 @@ describe("IndexFilingSectionsTask selection", () => { expect(out).toMatchObject({ indexed: 0, sections: 0, skipped: 0, truncated: false }); }); + + it("does not count what is already indexed for a caller that will not show it", async () => { + // The count is a join over every converted filing. `sec ask` pre-indexes + // before every question and reports only what it indexed, so it pays for a + // scan of the largest table in the database to produce a number nothing + // renders. + await seed(3); + for (const index of [0, 1, 2]) await markIndexed(index); + const counted: string[] = []; + const db = getDb(); + const realPrepare = db.prepare.bind(db); + vi.spyOn(db, "prepare").mockImplementation(((sql: string) => { + if (/COUNT\(\*\)/.test(sql)) counted.push(sql); + return realPrepare(sql); + }) as never); + + const out = await new IndexFilingSectionsTask().run({ limit: 5, countSkipped: false }); + + expect(counted).toEqual([]); + // Absent rather than zero: zero says the index holds nothing in scope, + // which is the opposite of what this run knows. + expect(out.skipped).toBeUndefined(); + expect(out).toMatchObject({ indexed: 0, sections: 0, truncated: false }); + }); }); /** diff --git a/src/task/kb/IndexFilingSectionsTask.ts b/src/task/kb/IndexFilingSectionsTask.ts index 519136fb..b37fba2c 100644 --- a/src/task/kb/IndexFilingSectionsTask.ts +++ b/src/task/kb/IndexFilingSectionsTask.ts @@ -38,6 +38,15 @@ export interface IndexFilingSectionsTaskInput { readonly limit?: number | undefined; /** Re-index filings already in the knowledge base. */ readonly force?: boolean | undefined; + /** + * Count the filings already in the index, for the `skipped` output. + * + * Its own COUNT over `filing_document` joined to the knowledge base, so a + * caller that does not show the number pays for a join it never reads — + * `sec ask` pre-indexes on every question and reports only what it indexed. + * Defaults to counting, so a caller has to have decided not to. + */ + readonly countSkipped?: boolean | undefined; } export interface IndexFilingSectionsTaskOutput { @@ -46,8 +55,12 @@ export interface IndexFilingSectionsTaskOutput { readonly indexed: number; /** Sections embedded across them. */ readonly sections: number; - /** Filings already in the index, skipped. */ - readonly skipped: number; + /** + * Filings already in the index, skipped — `undefined` when this run did not + * count them. Absent rather than zero, because zero is an answer: it says the + * index holds nothing in scope, which is not what "nobody asked" means. + */ + readonly skipped: number | undefined; /** True when the run hit its limit with filings still unexamined. */ readonly truncated: boolean; } @@ -138,6 +151,7 @@ export class IndexFilingSectionsTask extends Task< accession: Type.Optional(Type.String()), limit: Type.Optional(Type.Number()), force: Type.Optional(Type.Boolean()), + countSkipped: Type.Optional(Type.Boolean()), }); } @@ -146,7 +160,7 @@ export class IndexFilingSectionsTask extends Task< success: Type.Boolean(), indexed: Type.Integer(), sections: Type.Integer(), - skipped: Type.Integer(), + skipped: Type.Optional(Type.Integer()), truncated: Type.Boolean(), }); } @@ -175,7 +189,13 @@ export class IndexFilingSectionsTask extends Task< }); const truncated = limit !== undefined && candidates.length > limit; const work = truncated ? candidates.slice(0, limit) : candidates; - const skipped = input.force === true ? 0 : await countAlreadyIndexed(scope); + // Left unset for a caller that will not render it: the count is a join over + // every converted filing, which is the slowest read this task takes. Under + // `--force` nothing is skipped, so the answer is zero without asking. + let skipped: number | undefined; + if (input.countSkipped !== false) { + skipped = input.force === true ? 0 : await countAlreadyIndexed(scope); + } const denominator = Math.max(1, work.length); // The knowledge base's storages are built against `getDb()` rather than diff --git a/src/task/kb/selectDocumentsToIndex.sqlite.test.ts b/src/task/kb/selectDocumentsToIndex.sqlite.test.ts index 2558827b..8451fe63 100644 --- a/src/task/kb/selectDocumentsToIndex.sqlite.test.ts +++ b/src/task/kb/selectDocumentsToIndex.sqlite.test.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { DocumentNode } from "workglow"; import { Document, globalServiceRegistry, NodeKind } from "workglow"; import { withSqliteDb } from "../../config/testing/withSqliteDb"; @@ -13,6 +13,7 @@ import { FILING_DOCUMENT_REPOSITORY_TOKEN, type FilingDocument, } from "../../storage/document/FilingDocumentSchema"; +import { getDb } from "../../util/db"; import { countAlreadyIndexed, kbDocIdFor, selectDocumentsToIndex } from "./selectDocumentsToIndex"; const doc = (index: number, over: Partial = {}): FilingDocument => ({ @@ -148,4 +149,51 @@ describe("selectDocumentsToIndex (sqlite)", () => { await seed(2); expect(await selectDocumentsToIndex({ limit: 0 })).toEqual([]); }); + + /** + * The selection reads `filing_document` newest first and takes a page of it. + * Without an index in that order SQLite reads every row of the largest table + * in the database and sorts the lot into a temp B-tree to hand back the first + * few — on every `sec ask`, which pre-indexes before it answers. + * + * Asserted through the plan of the statement the selection actually prepares, + * rather than by looking the index up in `sqlite_master`: an index whose + * column order does not match the `ORDER BY` exists and removes nothing. + */ + it("reads the newest filings from an index rather than sorting the table", async () => { + await seed(3); + // The knowledge-base tables exist here, so the plan is the one with the + // anti-join — the shape `sec index` and `sec ask` both run. + await markIndexed(1); + + const db = getDb(); + const prepared: { sql: string; params: unknown[] }[] = []; + const realPrepare = db.prepare.bind(db); + vi.spyOn(db, "prepare").mockImplementation(((sql: string) => { + const statement = realPrepare(sql); + if (!/FROM `filing_document`/.test(sql)) return statement; + const realAll = statement.all.bind(statement); + return new Proxy(statement, { + get(target, property, receiver) { + if (property !== "all") return Reflect.get(target, property, receiver); + return (...params: unknown[]) => { + prepared.push({ sql, params }); + return realAll(...(params as never[])); + }; + }, + }); + }) as never); + + await selectDocumentsToIndex({ limit: 10 }); + + expect(prepared).toHaveLength(1); + const plan = ( + db + .prepare(`EXPLAIN QUERY PLAN ${prepared[0]!.sql}`) + .all(...(prepared[0]!.params as never[])) as { detail: string }[] + ).map((row) => row.detail); + + expect(plan.join("\n")).toContain("filing_document_filing_date_accession_number"); + expect(plan.some((step) => /TEMP B-TREE FOR ORDER BY/.test(step))).toBe(false); + }); });