A vector store in SQLite
Step 5 of 8. Step 4 ended at tag llm-ch04 with the index-side half of
RAG working end to end — a document cut into chunks, every chunk
embedded, a question scored against each one — but the vectors lived in
memory and died with the process. This step gives them a home: a SQLite
database where each vector is a BLOB, a brute-force search that reads
every row, and then an index on top of the same table — k-means clusters
the stored vectors, and a query only visits the clusters nearest to it.
The store ends up at 405 rows, and the index returns the same top-3 as
the full scan after reading 119 of them. The step ends at tag
llm-ch05.
One store, four tables
Section titled “One store, four tables”SQLite is the storage engine the rest of this track will keep using:
one file, transactional, no server, and rusqlite speaks to it from
ordinary Rust. The bundled feature compiles SQLite itself into the
binary, so there is nothing to install.
[package]name = "llm-apps"version = "0.1.0"edition = "2024"
[dependencies]anyhow = "1.0.104"# The runtime the whole track runs on. `default-features = false` keeps the# CPU path only - no Metal, no CUDA - which is everything this track needs# and the shortest dependency tree candle has.candle-core = { version = "0.11.0", default-features = false }candle-nn = "0.11.0"candle-transformers = "0.11.0"hf-hub = "0.3.2"rusqlite = { version = "0.41.1", features = ["bundled"] }serde_json = "1.0.154"tokenizers = "0.23.2"The schema is two tables of data and two tables of index. chunks and
vectors are the store: a chunk’s text, and its embedding as a BLOB.
centroids and assignments are the index this step builds halfway
through — empty at first, filled by the same program that fills the
store.
A vector enters SQLite as bytes. all-MiniLM-L6-v2 produces 384 f32s,
each to_le_bytes is 4 of them, so one row carries a 1,536-byte BLOB —
and the same trick runs backwards to decode:
/// A vector becomes a BLOB: 384 little-endian f32s, 1536 bytes.fn to_blob(v: &[f32]) -> Vec<u8> { v.iter().flat_map(|f| f.to_le_bytes()).collect()}
/// And back. `chunks_exact` panics on a corrupted row rather than/// silently mis-decoding it.fn from_blob(bytes: &[u8]) -> Vec<f32> { bytes .chunks_exact(4) .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) .collect()}Note what is not in the schema: any vector arithmetic. SQLite stores the vectors and filters rows; every dot product in this step happens in Rust, after the bytes are decoded. That split holds all the way up — even sqlite-vec, the extension built for exactly this job, is at heart a faster version of the scan written below.
Brute force first
Section titled “Brute force first”The first search reads every row, decodes every BLOB, and cosines against the question vector in Rust:
SELECT v.chunk_id, v.embedding, c.textFROM vectors v JOIN chunks c ON c.id = v.chunk_idNothing smarter is available, and nothing smarter is needed yet: the
function collects (id, score, text) for every row, sorts descending,
and truncates. It also counts what it read, because the scanned-row
count is the number this whole step is about. On the five real chunks
from step 4, it is instant and correct — chunk 2, the one holding “Its
peak memory footprint is about 340 megabytes”, wins at 0.5836, the same
score that step’s overlap run printed.
Five rows is not a store worth indexing. A brute-force scan is O(N) per query, and the index built later in this step costs something on every query too — at five rows the scan wins every time. So the program scales the store before building anything on top of it: 80 near-copies of each real chunk, 400 extra rows, each copy being its original plus a deterministic wobble of at most 0.01 per dimension:
fn jitter(seed: usize, dim: usize) -> f32 { ((seed * 31 + dim * 7) as f32 * 0.173).sin() * 0.01}No rand crate, no seed file — the same command prints the same table
on every machine. The copies stand in for the near-neighbours a real
corpus produces: on a real store the top hits would be different
documents about memory, and here they are copies of one chunk, but at
this layer of the system both are just rows of bytes.
Then an index
Section titled “Then an index”The index is built out of the store itself, not the vectors still in
memory: every row is read back as a BLOB, decoded, clustered with
k-means (K = 8), and the result is written back into centroids and
assignments. One row of assignments per vector says which cluster it
belongs to.
The k-means is Lloyd’s algorithm in forty lines, seeded deterministically:
cluster 0 starts at vector 0, cluster 1 at the vector farthest from it,
and each further seed is the vector farthest from everything chosen so
far. No rand, so the cluster sizes print identically every run.
Searching the index inverts the work. Instead of reading 405 rows, the
query reads 8 — one per centroid — cosines against each, and visits only
the PROBE = 2 nearest clusters. The SQL filters by those clusters
before any vector bytes are read:
SELECT v.chunk_id, v.embedding, c.textFROM vectors vJOIN chunks c ON c.id = v.chunk_idJOIN assignments a ON a.chunk_id = v.chunk_idWHERE a.centroid_id IN (6, 7)PROBE is the dial every vector index exposes under some name. One
cluster is the cheapest and riskiest query; more clusters cost more rows
and buy recall when a question sits between two of them.
The whole program
Section titled “The whole program”//! Step 5: a vector store in SQLite.//!//! Same model, chunker and embed pipeline as steps 3 and 4. The chunk//! vectors become BLOBs in a SQLite database; a question is answered by//! scanning every row - brute force - and then by an index: k-means//! clusters the stored vectors, and the query only visits the clusters//! nearest to it. On this step's 405-row store the index reads under a//! third of the rows and returns the same top hits.
use anyhow::Result;use candle_core::{Device, Tensor};use candle_nn::VarBuilder;use candle_transformers::models::bert::{BertModel, Config, DTYPE};use hf_hub::api::sync::Api;use rusqlite::{params, params_from_iter, Connection};use tokenizers::Tokenizer;
const MODEL_ID: &str = "sentence-transformers/all-MiniLM-L6-v2";const DOC: &str = "Running a sentence encoder on your own machine is mostly an exercise in memory budgeting. The runtime in this project is candle, which loads the safetensors weights through memory mapping, so nothing is copied until a layer actually runs. Startup cost is dominated entirely by paging the model in on the first forward pass. Its peak memory footprint is about 340 megabytes. That is with three sentences in the batch; a long document pushes the number up. Tokenization itself is cheap, microseconds per sentence, and it happens before anything reaches the model. Embedding a whole document at once is the mistake this step fixes: the vector averages every topic together and retrieval degrades quietly.";const QUESTION: &str = "How much memory does inference use at peak?";
/// Index size (clusters) and how many of them a query visits.const K: usize = 8;const PROBE: usize = 2;/// Synthetic near-copies per real chunk, to give the store a realistic size.const COPIES: usize = 80;
/// Cut `doc` into word windows of `window` words, advancing `stride`/// words each time. Same chunker as step 4.fn chunk(doc: &str, window: usize, stride: usize) -> Vec<String> { let words: Vec<&str> = doc.split_whitespace().collect(); let mut out = Vec::new(); let mut start = 0; while start < words.len() { let end = (start + window).min(words.len()); out.push(words[start..end].join(" ")); if end == words.len() { break; } start += stride; } out}
fn cosine(a: &[f32], b: &[f32]) -> f32 { let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum(); let len_a = a.iter().map(|x| x * x).sum::<f32>().sqrt(); let len_b = b.iter().map(|x| x * x).sum::<f32>().sqrt(); dot / (len_a * len_b)}
/// Tokenize, forward, mean-pool: one Vec<f32> of 384 numbers per input.fn embed( model: &BertModel, device: &Device, tokenizer: &Tokenizer, texts: &[String],) -> Result<Vec<Vec<f32>>> { let encodings = tokenizer .encode_batch(texts.to_vec(), true) .map_err(|e| anyhow::anyhow!("{e}"))?; let len = encodings.iter().map(|e| e.get_ids().len()).max().unwrap(); let token_ids: Vec<Vec<u32>> = encodings .iter() .map(|e| { let mut ids = e.get_ids().to_vec(); ids.resize(len, 0); ids }) .collect(); let attention: Vec<Vec<u32>> = encodings .iter() .map(|e| { let mut mask = e.get_attention_mask().to_vec(); mask.resize(len, 0); mask }) .collect(); let types = vec![vec![0u32; len]; texts.len()];
let ids = Tensor::new(token_ids, device)?; let mask = Tensor::new(attention.clone(), device)?; let types = Tensor::new(types, device)?;
let hidden = model.forward(&ids, &types, Some(&mask))?.to_vec3::<f32>()?; Ok(hidden .iter() .zip(attention.iter()) .map(|(tokens, mask)| { let mut v = vec![0f32; tokens[0].len()]; let mut n = 0f32; for (t, &keep) in tokens.iter().zip(mask) { if keep == 1 { n += 1.0; for (v, &x) in v.iter_mut().zip(t) { *v += x; } } } v.iter_mut().for_each(|x| *x /= n); v }) .collect())}
/// A vector becomes a BLOB: 384 little-endian f32s, 1536 bytes.fn to_blob(v: &[f32]) -> Vec<u8> { v.iter().flat_map(|f| f.to_le_bytes()).collect()}
/// And back. `chunks_exact` panics on a corrupted row rather than/// silently mis-decoding it.fn from_blob(bytes: &[u8]) -> Vec<f32> { bytes .chunks_exact(4) .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) .collect()}
/// The first 10 words of a chunk, for the printed table.fn preview(s: &str) -> String { let words: Vec<&str> = s.split_whitespace().take(10).collect(); let mut out = words.join(" "); if s.split_whitespace().count() > 10 { out.push_str(" ..."); } out}
/// Deterministic pseudo-noise, so every run of this program prints the/// same numbers. Amplitude 0.01 keeps each copy within a hair of its/// original - a stand-in for the near-neighbours a real corpus produces.fn jitter(seed: usize, dim: usize) -> f32 { ((seed * 31 + dim * 7) as f32 * 0.173).sin() * 0.01}
/// Lloyd's k-means with a deterministic start: seed 0, then the vector/// farthest from it, then repeatedly the vector farthest from everything/// chosen so far. Runs to convergence or `iters` passes, whichever first.fn kmeans(vectors: &[Vec<f32>], k: usize, iters: usize) -> (Vec<Vec<f32>>, Vec<usize>) { let dist = |a: &Vec<f32>, b: &Vec<f32>| 1.0 - cosine(a, b); let mut seeds: Vec<usize> = vec![0]; while seeds.len() < k { let (mut best, mut bd) = (0usize, -1.0f32); for (i, v) in vectors.iter().enumerate() { if seeds.contains(&i) { continue; } let d: f32 = seeds.iter().map(|&s| dist(v, &vectors[s])).sum(); if d > bd { bd = d; best = i; } } seeds.push(best); } let mut centroids: Vec<Vec<f32>> = seeds.iter().map(|&s| vectors[s].clone()).collect(); let mut assign = vec![0usize; vectors.len()]; for _ in 0..iters { let mut changed = false; for (i, v) in vectors.iter().enumerate() { let nearest = centroids .iter() .enumerate() .max_by(|a, b| cosine(v, a.1).total_cmp(&cosine(v, b.1))) .map(|(ci, _)| ci) .unwrap(); if assign[i] != nearest { changed = true; assign[i] = nearest; } } for (ci, c) in centroids.iter_mut().enumerate() { let members: Vec<&Vec<f32>> = vectors .iter() .enumerate() .filter(|(i, _)| assign[*i] == ci) .map(|(_, v)| v) .collect(); if members.is_empty() { continue; } for d in 0..c.len() { c[d] = members.iter().map(|v| v[d]).sum::<f32>() / members.len() as f32; } } if !changed { break; } } (centroids, assign)}
type Hit = (i64, f32, String);
/// Brute force: read every row, decode the blob, cosine in Rust.fn search_brute(conn: &Connection, q: &[f32], top: usize) -> Result<(usize, Vec<Hit>)> { let mut stmt = conn.prepare( "SELECT v.chunk_id, v.embedding, c.text FROM vectors v JOIN chunks c ON c.id = v.chunk_id", )?; let mut rows = stmt.query([])?; let mut scored: Vec<Hit> = Vec::new(); while let Some(row) = rows.next()? { let id: i64 = row.get(0)?; let blob: Vec<u8> = row.get(1)?; let text: String = row.get(2)?; scored.push((id, cosine(&from_blob(&blob), q), text)); } let scanned = scored.len(); scored.sort_by(|a, b| b.1.total_cmp(&a.1)); scored.truncate(top); Ok((scanned, scored))}
/// Indexed search: rank the centroids, then read only the rows assigned/// to the `probe` nearest clusters. Returns the clusters visited, the/// rows read, and the hits.fn search_index( conn: &Connection, q: &[f32], probe: usize, top: usize,) -> Result<(Vec<(i64, f32)>, usize, Vec<Hit>)> { let mut stmt = conn.prepare("SELECT id, vector FROM centroids")?; let mut rows = stmt.query([])?; let mut clusters: Vec<(i64, f32)> = Vec::new(); while let Some(row) = rows.next()? { let id: i64 = row.get(0)?; let c = from_blob(&row.get::<_, Vec<u8>>(1)?); clusters.push((id, cosine(&c, q))); } clusters.sort_by(|a, b| b.1.total_cmp(&a.1)); let chosen: Vec<i64> = clusters[..probe].iter().map(|c| c.0).collect();
let list = chosen .iter() .enumerate() .map(|(i, _)| format!("?{}", i + 1)) .collect::<Vec<_>>() .join(", "); let mut stmt = conn.prepare(&format!( "SELECT v.chunk_id, v.embedding, c.text FROM vectors v JOIN chunks c ON c.id = v.chunk_id JOIN assignments a ON a.chunk_id = v.chunk_id WHERE a.centroid_id IN ({list})" ))?; let mut rows = stmt.query(params_from_iter(chosen.iter()))?; let mut scored: Vec<Hit> = Vec::new(); while let Some(row) = rows.next()? { let id: i64 = row.get(0)?; let blob: Vec<u8> = row.get(1)?; let text: String = row.get(2)?; scored.push((id, cosine(&from_blob(&blob), q), text)); } let scanned = scored.len(); scored.sort_by(|a, b| b.1.total_cmp(&a.1)); scored.truncate(top); let visited: Vec<(i64, f32)> = clusters[..probe].to_vec(); Ok((visited, scanned, scored))}
fn print_hits(hits: &[Hit]) { for (rank, (_, score, text)) in hits.iter().enumerate() { println!(" {}. [{score:.4}] {text}", rank + 1); }}
fn main() -> Result<()> { let device = Device::Cpu; let started = std::time::Instant::now();
let api = Api::new()?; let repo = api.model(MODEL_ID.to_string()); let config: Config = serde_json::from_str(&std::fs::read_to_string(repo.get("config.json")?)?)?; let tokenizer = Tokenizer::from_file(repo.get("tokenizer.json")?) .map_err(|e| anyhow::anyhow!("{e}"))?; let weights = repo.get("model.safetensors")?; let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights], DTYPE, &device)? }; let model = BertModel::load(vb, &config)?; println!("model loaded in {:.1?} on {device:?}", started.elapsed());
// Same chunker, document and question as step 4: 30-word windows, // 6-word overlap. let chunks = chunk(DOC, 30, 24);
// The store lives in the temp directory; delete any copy from an // earlier run so the numbers below start from an empty database. let db_path = std::env::temp_dir().join("llm-step5-vectors.db"); let _ = std::fs::remove_file(&db_path); let conn = Connection::open(&db_path)?; conn.execute_batch( "CREATE TABLE chunks (id INTEGER PRIMARY KEY, text TEXT NOT NULL); CREATE TABLE vectors ( chunk_id INTEGER PRIMARY KEY REFERENCES chunks(id), embedding BLOB NOT NULL); CREATE TABLE centroids (id INTEGER PRIMARY KEY, vector BLOB NOT NULL); CREATE TABLE assignments ( chunk_id INTEGER PRIMARY KEY REFERENCES chunks(id), centroid_id INTEGER NOT NULL REFERENCES centroids(id));", )?;
let mut texts = chunks.clone(); texts.push(QUESTION.to_string()); let vectors = embed(&model, &device, &tokenizer, &texts)?; let question = vectors.last().unwrap().clone();
for (i, v) in vectors[..vectors.len() - 1].iter().enumerate() { conn.execute( "INSERT INTO chunks (id, text) VALUES (?1, ?2)", params![i as i64, format!("chunk {i}: {}", preview(&chunks[i]))], )?; conn.execute( "INSERT INTO vectors (chunk_id, embedding) VALUES (?1, ?2)", params![i as i64, to_blob(v)], )?; } println!("stored {} chunk vectors in the database", vectors.len() - 1);
let (scanned, hits) = search_brute(&conn, &question, 3)?; println!("\nbrute force over {scanned} rows:"); print_hits(&hits);
// Scale the store up: 80 near-copies of each real chunk, 400 rows, // standing in for the neighbours a real corpus produces. for i in 0..5usize { for j in 0..COPIES { let id = (5 + i * COPIES + j) as i64; let noisy: Vec<f32> = vectors[i] .iter() .enumerate() .map(|(d, x)| x + jitter(id as usize, d)) .collect(); conn.execute( "INSERT INTO chunks (id, text) VALUES (?1, ?2)", params![id, format!("chunk {i} copy #{j}")], )?; conn.execute( "INSERT INTO vectors (chunk_id, embedding) VALUES (?1, ?2)", params![id, to_blob(&noisy)], )?; } } let total: i64 = conn.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))?; let (scanned, brute_hits) = search_brute(&conn, &question, 3)?; println!("\nscaled to {total} rows; brute force over {scanned} rows:"); print_hits(&brute_hits);
// The index is built from what the store holds: every vector is read // back as a blob, clustered, and the cluster map is written back. let mut stmt = conn.prepare("SELECT chunk_id, embedding FROM vectors")?; let mut rows = stmt.query([])?; let mut ids: Vec<i64> = Vec::new(); let mut all: Vec<Vec<f32>> = Vec::new(); while let Some(row) = rows.next()? { ids.push(row.get(0)?); all.push(from_blob(&row.get::<_, Vec<u8>>(1)?)); } let (centroids, assign) = kmeans(&all, K, 20); for (ci, c) in centroids.iter().enumerate() { conn.execute( "INSERT INTO centroids (id, vector) VALUES (?1, ?2)", params![ci as i64, to_blob(c)], )?; } for (&chunk_id, &ci) in ids.iter().zip(assign.iter()) { conn.execute( "INSERT INTO assignments (chunk_id, centroid_id) VALUES (?1, ?2)", params![chunk_id, ci as i64], )?; } let mut sizes = vec![0usize; K]; for &ci in &assign { sizes[ci] += 1; } sizes.sort(); println!("\nindex built: {K} clusters, sizes {sizes:?}");
let (visited, scanned, hits) = search_index(&conn, &question, PROBE, 3)?; let pct = 100.0 * scanned as f64 / total as f64; let visited = visited .iter() .map(|(id, s)| format!("{id} ({s:.4})")) .collect::<Vec<_>>() .join(", "); println!( "\nindexed search: cluster {visited} nearest to the question, scanned {scanned} of {total} rows ({pct:.0}%)" ); print_hits(&hits); let overlap = hits .iter() .filter(|(id, _, _)| brute_hits.iter().any(|(bid, _, _)| bid == id)) .count(); println!("top-3 overlap with brute force: {overlap} of 3");
Ok(())}Running it (cargo run -q, the model weights come from the local
Hugging Face cache, and the database goes to the system temp directory)
prints:
model loaded in 76.1ms on Cpustored 5 chunk vectors in the database
brute force over 5 rows: 1. [0.5836] chunk 2: in on the first forward pass. Its peak memory footprint ... 2. [0.2620] chunk 0: Running a sentence encoder on your own machine is mostly ... 3. [0.2545] chunk 1: the safetensors weights through memory mapping, so nothing is copied ...
scaled to 405 rows; brute force over 405 rows: 1. [0.5847] chunk 2 copy #19 2. [0.5847] chunk 2 copy #60 3. [0.5847] chunk 2 copy #53
index built: 8 clusters, sizes [38, 38, 40, 41, 43, 43, 81, 81]
indexed search: cluster 6 (0.5836), 7 (0.2629) nearest to the question, scanned 119 of 405 rows (29%) 1. [0.5847] chunk 2 copy #19 2. [0.5847] chunk 2 copy #60 3. [0.5847] chunk 2 copy #53top-3 overlap with brute force: 3 of 3What the code does, move by move
Section titled “What the code does, move by move”The store and the index are the same database. There is no second
file, no in-process structure to rebuild on startup. centroids and
assignments are ordinary tables, which means the index survives a
restart, can be inspected with any SQLite client, and — because
assignments has a foreign key — cannot drift out of sync with the
vectors it points at.
The index is built from the store, not from memory. Halfway through, the program re-reads every vector as a BLOB and clusters what comes back. That ordering is deliberate: if the clustering only worked on vectors it happened to still have in memory, the code would work here and fail the moment the index was rebuilt by a second process.
Sorting, not argmax, on purpose — the same choice step 4 made. The top-3 lists are what make the overlap check possible: the last line compares the index’s hits against the full scan’s by row id, which is the only way to say the index lost nothing.
The scoring happens outside SQL. search_brute and search_index
both decode bytes and run cosine in Rust. rusqlite offers no vector
functions, and writing the similarity as a SQL expression over 384
columns would be its own kind of unreadable.
Reading the numbers
Section titled “Reading the numbers”The cluster sizes are the first thing worth reading. The 405 rows form
five natural clumps of 81 — one real chunk plus its 80 copies — and
k-means was asked for 8 clusters, so it kept two clumps whole and split
the other three: 81 + 81 for the intact ones, and 38 + 43, 38 + 43,
40 + 41 for the split pairs. Splitting a clump is not an error; it just
means a query landing near that clump may need both halves, which is
exactly what PROBE is for.
The query’s nearest cluster (6, similarity 0.5836) is one of the intact 81-row clumps — the chunk-2 copies — and the second nearest (7, 0.2629) is the chunk-0 clump, matching the brute-force ranking from the five-row store. Visiting both reads 119 of 405 rows and returns the same top-3, ids and scores. Probing only cluster 6 would have read 81 rows and also returned 3 of 3 this time — but with no margin: had the question sat between two clumps, or had the best clump been one of the split ones, one cluster would not have been enough. Two probes is the cheap insurance this store size can afford.
The scores deserve a word too. The copies score 0.5847, a hair above the original chunk 2 at 0.5836 — the wobble happened to nudge them marginally closer to the question, so the original ranks fourth. Nothing is wrong: on a real corpus the neighbours really would be different documents, and the store cannot tell a synthetic copy from a real row.
And the honest limit: at 405 rows both searches finish in well under a millisecond, so the stopwatch is silent and the scanned-row count does the talking. That count is what scales — brute force reads all N rows on every query, the index reads 8 centroids plus the visited clusters. At 405 rows the difference is 29%; at a million rows it is the difference between reading half a gigabyte of BLOBs per query and reading a fraction of one, which is why dedicated vector indexes exist. This one is small, honest about its recall, and built entirely out of tables the reader can open and inspect.
The tag
Section titled “The tag”The step’s end state is tagged in the examples repository:
git checkout llm-ch05At that tag the project embeds chunks, writes them to SQLite as BLOBs, answers a question by brute force over the whole store, then builds a k-means index on top and answers it again by visiting two clusters. Step 6 turns the store into the R of RAG: retrieving chunks and building the prompt a model will answer from. The track overview keeps the full list current as articles land.
Related
Section titled “Related”- Chunking documents for retrieval — the previous step, whose chunker, document and question this step stores
- Embeddings and vector similarity — the
cosinearithmetic every search in this step runs