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:
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user