Skip to content

perf: float vectors in a plain table, allocation-free exact scan #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Verdenroz merged 5 commits into from
Sep 4, 2026
Merged

perf: float vectors in a plain table, allocation-free exact scan #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9dfe36c
test: tie-aware recall and recall@10 in refiner_curve
Verdenroz Sep 4, 2026
c986c2e
perf: float vectors in a plain table, coarse pass rescored from it
Verdenroz Sep 4, 2026
ef33d53
perf: allocation-free exact scan and schema guards
Verdenroz Sep 4, 2026
8681ffe
docs: allocation claim for the dense search path
Verdenroz Sep 4, 2026
1600c35
perf: one allocation for the rescore id list, exact alloc claim
Verdenroz Sep 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 68 additions & 29 deletions crates/scry-core/examples/refiner_curve.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,13 @@
//! Copy the live DB first (`sqlite3 index.db "VACUUM INTO 'copy.db'"`) so
//! WAL pages are included and the server is never touched.

use std::collections::HashSet;
use std::time::{Duration, Instant};

use scry_core::config::Config;
use scry_core::store::Store;

const K: usize = 50;
const K_SHOWN: usize = 10;
const SAMPLES: usize = 200;
const COARSE_KS: &[usize] = &[100, 200, 400, 800, 1600];

Expand All @@ -19,57 +19,96 @@ fn main() {
let db = args.next().expect("index.db path");
let repo_key = args.next();
let config = Config::load(None).unwrap();
let opened = Instant::now();
let store = Store::open(
std::path::Path::new(&db),
&config.embedding.model,
config.embedding.dim,
)
.unwrap();
println!(
"open (with any migration) {:.1}ms",
opened.elapsed().as_secs_f64() * 1000.0
);
let repo_id = repo_key.and_then(|key| store.repo_id(&key).unwrap());
let samples = store.sample_chunk_vectors(repo_id, SAMPLES).unwrap();
println!("repo {:?} samples {} k {}", repo_id, samples.len(), K);

let mut exact_time = Duration::ZERO;
let exact: Vec<HashSet<i64>> = samples
let exact: Vec<Truth> = samples
.iter()
.map(|(id, query)| {
let started = Instant::now();
let hits = store.dense_search_exact(repo_id, query, K + 1).unwrap();
exact_time += started.elapsed();
hits.iter()
.map(|h| h.chunk_id)
.filter(|c| c != id)
let others: Vec<f64> = hits
.iter()
.filter(|h| h.chunk_id != *id)
.map(|h| h.distance)
.take(K)
.collect()
.collect();
Truth {
at_shown: others[K_SHOWN - 1],
at_k: others[K - 1],
}
})
.collect();
println!(
"exact recall 1.000 {:>6.1}ms/query",
"exact recall@50 1.000 recall@10 1.000 {:>6.1}ms/query",
exact_time.as_secs_f64() * 1000.0 / samples.len() as f64
);

for &coarse_k in COARSE_KS {
let mut time = Duration::ZERO;
let mut overlap = 0usize;
for ((id, query), truth) in samples.iter().zip(&exact) {
let started = Instant::now();
let hits = store
.dense_search_coarse(repo_id, query, K + 1, coarse_k + 1)
.unwrap();
time += started.elapsed();
overlap += hits
.iter()
.map(|h| h.chunk_id)
.filter(|c| c != id)
.take(K)
.filter(|c| truth.contains(c))
.count();
}
println!(
"coarse {:>5} recall {:.3} {:>6.1}ms/query",
coarse_k,
overlap as f64 / (samples.len() * K) as f64,
time.as_secs_f64() * 1000.0 / samples.len() as f64
report(
&format!("coarse {coarse_k:>5}"),
&samples,
&exact,
|query| {
store
.dense_search_coarse(repo_id, query, K + 1, coarse_k + 1)
.unwrap()
},
);
}
}

/// Recall counts a hit when its distance is within the exact k-th distance,
/// so ties between duplicate vectors do not read as misses.
struct Truth {
at_shown: f64,
at_k: f64,
}

const TIE: f64 = 1e-6;

fn report(
label: &str,
samples: &[(i64, Vec<f32>)],
exact: &[Truth],
mut search: impl FnMut(&[f32]) -> Vec<scry_core::store::DenseHit>,
) {
let mut time = Duration::ZERO;
let (mut overlap, mut overlap_shown) = (0usize, 0usize);
for ((id, query), truth) in samples.iter().zip(exact) {
let started = Instant::now();
let hits = search(query);
time += started.elapsed();
let got: Vec<f64> = hits
.iter()
.filter(|h| h.chunk_id != *id)
.map(|h| h.distance)
.take(K)
.collect();
overlap += got.iter().filter(|d| **d <= truth.at_k + TIE).count();
overlap_shown += got
.iter()
.take(K_SHOWN)
.filter(|d| **d <= truth.at_shown + TIE)
.count();
}
println!(
"{label} recall@50 {:.3} recall@10 {:.3} {:>6.1}ms/query",
overlap as f64 / (samples.len() * K) as f64,
overlap_shown as f64 / (samples.len() * K_SHOWN) as f64,
time.as_secs_f64() * 1000.0 / samples.len() as f64
);
}
154 changes: 91 additions & 63 deletions crates/scry-core/src/store/chunks.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,17 +54,80 @@ fn bytes_to_vector(bytes: &[u8]) -> Vec<f32> {
.collect()
}

/// Cosine distance between the query and a little-endian f32 blob,
/// computed in place so a full scan allocates nothing per row.
fn cosine_distance(query: &[f32], query_norm: f64, blob: &[u8]) -> f64 {
let (mut dot, mut norm) = (0f64, 0f64);
for (x, bytes) in query.iter().zip(blob.as_chunks::<4>().0) {
let y = f64::from(f32::from_le_bytes(*bytes));
dot += f64::from(*x) * y;
norm += y * y;
}
if query_norm == 0.0 || norm == 0.0 {
return 1.0;
}
1.0 - dot / (query_norm * norm.sqrt())
}

fn norm(vector: &[f32]) -> f64 {
vector
.iter()
.map(|x| f64::from(*x) * f64::from(*x))
.sum::<f64>()
.sqrt()
}

/// Keeps `top` sorted by distance and at most `k` long.
fn keep_best(top: &mut Vec<DenseHit>, hit: DenseHit, k: usize) {
if top.len() == k && hit.distance >= top[k - 1].distance {
return;
}
let at = top.partition_point(|h| h.distance <= hit.distance);
top.insert(at, hit);
top.truncate(k);
}

/// Scans `stmt` rows of (chunk_id, embedding blob) and keeps the k nearest.
fn nearest(
stmt: &mut rusqlite::Statement<'_>,
params: impl rusqlite::Params,
query: &[f32],
k: usize,
) -> Result<Vec<DenseHit>> {
let query_norm = norm(query);
let mut top = Vec::with_capacity(k + 1);
let mut rows = stmt.query(params)?;
while let Some(row) = rows.next()? {
let blob = row.get_ref(1)?.as_blob().map_err(rusqlite::Error::from)?;
let hit = DenseHit {
chunk_id: row.get(0)?,
distance: cosine_distance(query, query_norm, blob),
};
keep_best(&mut top, hit, k);
}
Ok(top)
}

fn json_id_list(ids: &[i64]) -> String {
let inner: Vec<String> = ids.iter().map(i64::to_string).collect();
format!("[{}]", inner.join(","))
use std::fmt::Write;
let mut list = String::with_capacity(ids.len() * 8 + 2);
list.push('[');
for (i, id) in ids.iter().enumerate() {
if i > 0 {
list.push(',');
}
write!(list, "{id}").unwrap();
}
list.push(']');
list
}

impl Store {
pub fn vector_for_hash(&self, content_hash: &str) -> Result<Option<Vec<f32>>> {
let bytes: Option<Vec<u8>> = self
.conn
.query_row(
"SELECT v.embedding FROM vec_chunks v
"SELECT v.embedding FROM chunk_vectors v
JOIN chunks c ON c.id = v.chunk_id
WHERE c.content_hash = ?1 LIMIT 1",
[content_hash],
Expand All @@ -85,11 +148,6 @@ impl Store {
chunks: &[NewChunk],
) -> Result<()> {
let tx = self.conn.transaction()?;
tx.execute(
"DELETE FROM vec_chunks WHERE chunk_id IN
(SELECT id FROM chunks WHERE file_id = ?1)",
[file_id],
)?;
tx.execute(
"DELETE FROM vec_chunks_bit WHERE chunk_id IN
(SELECT id FROM chunks WHERE file_id = ?1)",
Expand All @@ -111,15 +169,15 @@ impl Store {
)?;
let chunk_id = tx.last_insert_rowid();
let bytes = vector_bytes(&chunk.embedding);
tx.execute(
"INSERT INTO vec_chunks (chunk_id, repo_id, embedding) VALUES (?1, ?2, ?3)",
params![chunk_id, repo_id, bytes],
)?;
tx.execute(
"INSERT INTO vec_chunks_bit (chunk_id, repo_id, embedding)
VALUES (?1, ?2, vec_quantize_binary(?3))",
params![chunk_id, repo_id, bytes],
)?;
tx.execute(
"INSERT INTO chunk_vectors (chunk_id, embedding) VALUES (?1, ?2)",
params![chunk_id, bytes],
)?;
}
tx.commit()?;
Ok(())
Expand All @@ -132,8 +190,7 @@ impl Store {
k: usize,
) -> Result<Vec<DenseHit>> {
let chunk_count: i64 = self.conn.query_row(
"SELECT count(*) FROM chunks c JOIN files f ON f.id = c.file_id
WHERE ?1 IS NULL OR f.repo_id = ?1",
"SELECT coalesce(sum(chunk_count), 0) FROM repos WHERE ?1 IS NULL OR id = ?1",
[repo_id],
|row| row.get(0),
)?;
Expand All @@ -144,42 +201,6 @@ impl Store {
}
}

pub fn dense_search_exact(
&self,
repo_id: Option<i64>,
query: &[f32],
k: usize,
) -> Result<Vec<DenseHit>> {
// vec0 KNN cannot take an optional partition constraint in one
// statement; the filter must be present or absent in the SQL.
let sql = match repo_id {
Some(_) => {
"SELECT chunk_id, distance FROM vec_chunks
WHERE embedding MATCH ?1 AND k = ?2 AND repo_id = ?3
ORDER BY distance"
}
None => {
"SELECT chunk_id, distance FROM vec_chunks
WHERE embedding MATCH ?1 AND k = ?2
ORDER BY distance"
}
};
let mut stmt = self.conn.prepare(sql)?;
let map = |row: &rusqlite::Row<'_>| {
Ok(DenseHit {
chunk_id: row.get(0)?,
distance: row.get(1)?,
})
};
let rows = match repo_id {
Some(repo_id) => {
stmt.query_map(params![vector_bytes(query), k as i64, repo_id], map)?
}
None => stmt.query_map(params![vector_bytes(query), k as i64], map)?,
};
Ok(rows.collect::<std::result::Result<_, _>>()?)
}

/// Hamming pass over `vec_chunks_bit` keeps `coarse_k` candidates,
/// then the float table rescores them by cosine and keeps `k`.
pub fn dense_search_coarse(
Expand Down Expand Up @@ -214,20 +235,27 @@ impl Store {
.collect::<std::result::Result<_, _>>()?,
};

let ids = json_id_list(&candidates);
let mut rescore = self.conn.prepare(
"SELECT chunk_id, distance FROM vec_chunks
WHERE embedding MATCH ?1 AND k = ?2
AND chunk_id IN (SELECT value FROM json_each(?3))
ORDER BY distance",
"SELECT chunk_id, embedding FROM chunk_vectors
WHERE chunk_id IN (SELECT value FROM json_each(?1))",
)?;
let rows = rescore.query_map(params![vector_bytes(query), k as i64, ids], |row| {
Ok(DenseHit {
chunk_id: row.get(0)?,
distance: row.get(1)?,
})
})?;
Ok(rows.collect::<std::result::Result<_, _>>()?)
nearest(&mut rescore, [json_id_list(&candidates)], query, k)
}

/// Exact cosine KNN as a scan over the plain float table.
pub fn dense_search_exact(
&self,
repo_id: Option<i64>,
query: &[f32],
k: usize,
) -> Result<Vec<DenseHit>> {
let mut stmt = self.conn.prepare(
"SELECT v.chunk_id, v.embedding FROM chunk_vectors v
JOIN chunks c ON c.id = v.chunk_id
JOIN files f ON f.id = c.file_id
WHERE ?1 IS NULL OR f.repo_id = ?1",
)?;
nearest(&mut stmt, [repo_id], query, k)
}

pub fn sample_chunk_vectors(
Expand All @@ -236,7 +264,7 @@ impl Store {
n: usize,
) -> Result<Vec<(i64, Vec<f32>)>> {
let mut stmt = self.conn.prepare(
"SELECT v.chunk_id, v.embedding FROM vec_chunks v
"SELECT v.chunk_id, v.embedding FROM chunk_vectors v
JOIN chunks c ON c.id = v.chunk_id
JOIN files f ON f.id = c.file_id
WHERE ?1 IS NULL OR f.repo_id = ?1
Expand Down
5 changes: 0 additions & 5 deletions crates/scry-core/src/store/files.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,11 +110,6 @@ impl Store {
let Some(file_id) = file_id else {
return Ok(());
};
tx.execute(
"DELETE FROM vec_chunks WHERE chunk_id IN
(SELECT id FROM chunks WHERE file_id = ?1)",
[file_id],
)?;
tx.execute(
"DELETE FROM vec_chunks_bit WHERE chunk_id IN
(SELECT id FROM chunks WHERE file_id = ?1)",
Expand Down
Loading