HNSW vector index: 100K search in 27ms (58x faster than brute-force)

- instant-distance HNSW implementation for approximate nearest neighbors
- HnswStore: build from stored embeddings, in-memory index, thread-safe
- POST /vectors/hnsw/build — build index from Parquet (100K in 35s release)
- POST /vectors/hnsw/search — fast ANN search
- GET /vectors/hnsw/list — list loaded indexes

Benchmark (100K × 768d, release build):
  Brute-force: 1,567ms
  HNSW:           31ms (50x)
  HNSW warm:      27ms (58x)

Build cost: 35s one-time for 100K vectors (release mode)
ef_construction=40, ef_search=50 — good recall/speed balance

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
root
2026-03-27 20:00:50 -05:00
co-authored by Claude Opus 4.6
parent 8282842eaf
commit 04770c97eb
7 changed files with 289 additions and 2 deletions
+77 -1
View File
@@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize};
use std::sync::Arc;
use aibridge::client::{AiClient, EmbedRequest};
use crate::{chunker, index_registry, jobs, rag, search, store, supervisor};
use crate::{chunker, hnsw, index_registry, jobs, rag, search, store, supervisor};
#[derive(Clone)]
pub struct VectorState {
@@ -18,6 +18,7 @@ pub struct VectorState {
pub ai_client: AiClient,
pub job_tracker: jobs::JobTracker,
pub index_registry: index_registry::IndexRegistry,
pub hnsw_store: hnsw::HnswStore,
}
pub fn router(state: VectorState) -> Router {
@@ -30,6 +31,9 @@ pub fn router(state: VectorState) -> Router {
.route("/jobs/{id}", get(get_job))
.route("/search", post(search_index))
.route("/rag", post(rag_query))
.route("/hnsw/build", post(build_hnsw))
.route("/hnsw/search", post(search_hnsw))
.route("/hnsw/list", get(list_hnsw))
.with_state(state)
}
@@ -296,3 +300,75 @@ async fn rag_query(
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
// --- HNSW Fast Search ---
#[derive(Deserialize)]
struct BuildHnswRequest {
/// Name of the stored vector index to build HNSW from
index_name: String,
}
/// Build an HNSW index from an existing stored vector index.
/// Loads embeddings from Parquet, builds HNSW in memory.
async fn build_hnsw(
State(state): State<VectorState>,
Json(req): Json<BuildHnswRequest>,
) -> impl IntoResponse {
tracing::info!("building HNSW for '{}'", req.index_name);
// Load embeddings from Parquet
let embeddings = store::load_embeddings(&state.store, &req.index_name)
.await
.map_err(|e| (StatusCode::NOT_FOUND, format!("index not found: {e}")))?;
let n = embeddings.len();
tracing::info!("loaded {} embeddings, building HNSW...", n);
// Build HNSW
match state.hnsw_store.build_index(&req.index_name, embeddings).await {
Ok(stats) => Ok(Json(stats)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Deserialize)]
struct HnswSearchRequest {
index_name: String,
query: String,
top_k: Option<usize>,
}
/// Search using HNSW — approximate nearest neighbors, much faster than brute-force.
async fn search_hnsw(
State(state): State<VectorState>,
Json(req): Json<HnswSearchRequest>,
) -> impl IntoResponse {
let top_k = req.top_k.unwrap_or(5);
// Embed query
let embed_resp = state.ai_client.embed(EmbedRequest {
texts: vec![req.query.clone()],
model: None,
}).await.map_err(|e| (StatusCode::BAD_GATEWAY, format!("embed error: {e}")))?;
if embed_resp.embeddings.is_empty() {
return Err((StatusCode::BAD_GATEWAY, "no embedding returned".to_string()));
}
let query_vec: Vec<f32> = embed_resp.embeddings[0].iter().map(|&x| x as f32).collect();
// Search HNSW
match state.hnsw_store.search(&req.index_name, &query_vec, top_k).await {
Ok(results) => Ok(Json(serde_json::json!({
"results": results,
"query": req.query,
"method": "hnsw",
}))),
Err(e) => Err((StatusCode::NOT_FOUND, e)),
}
}
async fn list_hnsw(State(state): State<VectorState>) -> impl IntoResponse {
Json(state.hnsw_store.list().await)
}