Squashed 'vendor/ruvector/' content from commit b64c2172

git-subtree-dir: vendor/ruvector
git-subtree-split: b64c21726f2bb37286d9ee36a7869fef60cc6900
This commit is contained in:
ruv
2026-02-28 14:39:40 -05:00
commit d803bfe2b1
7854 changed files with 3522914 additions and 0 deletions

View File

@@ -0,0 +1,44 @@
[package]
name = "ruvector-router-core"
version.workspace = true
edition.workspace = true
authors.workspace = true
license.workspace = true
repository.workspace = true
readme = "README.md"
description = "Core vector database and neural routing inference engine"
[lib]
crate-type = ["lib", "staticlib"]
[dependencies]
# Workspace dependencies
redb = { workspace = true }
memmap2 = { workspace = true }
rayon = { workspace = true }
crossbeam = { workspace = true }
parking_lot = { workspace = true }
rkyv = { workspace = true }
bincode = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
simsimd = { workspace = true }
thiserror = { workspace = true }
anyhow = { workspace = true }
tracing = { workspace = true }
# Additional dependencies
ndarray = "0.15"
rand = "0.8"
uuid = { version = "1.10", features = ["v4", "serde"] }
chrono = { version = "0.4", features = ["serde"] }
[dev-dependencies]
criterion = { workspace = true }
proptest = { workspace = true }
tempfile = "3.12"
tracing-subscriber = { workspace = true }
[[bench]]
name = "vector_search"
harness = false

View File

@@ -0,0 +1,761 @@
# Router Core
[![Rust](https://img.shields.io/badge/rust-1.77%2B-orange.svg)](https://www.rust-lang.org)
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
[![Performance](https://img.shields.io/badge/latency-<0.5ms-green.svg)](../../docs/TECHNICAL_PLAN.md)
**High-performance vector database and neural routing inference engine built in Rust.**
Core engine powering Ruvector's intelligent request distribution, model selection, and sub-millisecond vector similarity search. Combines advanced indexing algorithms with SIMD-optimized distance calculations for maximum performance.
## 🎯 Overview
Router Core is the foundation of Ruvector's vector database capabilities, providing:
- **Neural Routing**: Intelligent request distribution across multiple models and endpoints
- **Vector Database**: High-performance storage and retrieval with HNSW indexing
- **Model Selection**: Adaptive routing strategies for multi-model AI systems
- **SIMD Acceleration**: Hardware-optimized vector operations via simsimd
- **Memory Efficiency**: Advanced quantization techniques (4-32x compression)
- **Zero Dependencies**: Pure Rust implementation with minimal external dependencies
## ⚡ Key Features
### Core Capabilities
- **Sub-Millisecond Search**: <0.5ms p50 latency with HNSW indexing
- **HNSW Indexing**: Hierarchical Navigable Small World for fast approximate nearest neighbor search
- **Multiple Distance Metrics**: Euclidean, Cosine, Dot Product, Manhattan
- **Advanced Quantization**: Scalar (4x), Product (8-16x), Binary (32x) compression
- **SIMD Optimizations**: Hardware-accelerated distance calculations
- **Zero-Copy I/O**: Memory-mapped files for efficient data access
- **Thread-Safe**: Concurrent read/write operations with minimal locking
- **Persistent Storage**: Durable vector storage with redb backend
### Neural Routing Features
- **Intelligent Request Distribution**: Route queries to optimal model endpoints
- **Load Balancing**: Distribute workload across multiple inference servers
- **Model Selection**: Automatically select best model based on query characteristics
- **Adaptive Strategies**: Learn and optimize routing decisions over time
- **Latency Optimization**: Minimize end-to-end inference time
- **Failover Support**: Automatic fallback to backup endpoints
## 📦 Installation
Add to your `Cargo.toml`:
```toml
[dependencies]
router-core = "0.1.0"
```
Or use the full ruvector package:
```toml
[dependencies]
ruvector-core = "0.1.0"
```
## 🚀 Quick Start
### Basic Vector Database
```rust
use router_core::{VectorDB, VectorEntry, SearchQuery, DistanceMetric};
use std::collections::HashMap;
// Create database with builder pattern
let db = VectorDB::builder()
.dimensions(384) // Vector dimensions
.distance_metric(DistanceMetric::Cosine)
.hnsw_m(32) // HNSW connections per node
.hnsw_ef_construction(200) // Construction accuracy
.storage_path("./vectors.db")
.build()?;
// Insert vectors
let entry = VectorEntry {
id: "doc1".to_string(),
vector: vec![0.1; 384],
metadata: HashMap::new(),
timestamp: chrono::Utc::now().timestamp(),
};
db.insert(entry)?;
// Search for similar vectors
let query = SearchQuery {
vector: vec![0.1; 384],
k: 10, // Top 10 results
filters: None,
threshold: Some(0.8), // Minimum similarity
ef_search: Some(100), // Search accuracy
};
let results = db.search(query)?;
for result in results {
println!("{}: {}", result.id, result.score);
}
```
### Batch Operations
```rust
use router_core::{VectorDB, VectorEntry};
// Insert multiple vectors efficiently
let entries: Vec<VectorEntry> = (0..1000)
.map(|i| VectorEntry {
id: format!("doc{}", i),
vector: vec![0.1; 384],
metadata: HashMap::new(),
timestamp: chrono::Utc::now().timestamp(),
})
.collect();
// Batch insert (much faster than individual inserts)
db.insert_batch(entries)?;
// Check statistics
let stats = db.stats();
println!("Total vectors: {}", stats.total_vectors);
println!("Avg latency: {:.2}μs", stats.avg_query_latency_us);
```
### Advanced Configuration
```rust
use router_core::{VectorDB, DistanceMetric, QuantizationType};
let db = VectorDB::builder()
.dimensions(768) // Larger embeddings
.max_elements(10_000_000) // 10M vectors
.distance_metric(DistanceMetric::Cosine) // Cosine similarity
.hnsw_m(64) // More connections = higher recall
.hnsw_ef_construction(400) // Higher accuracy during build
.hnsw_ef_search(200) // Search-time accuracy
.quantization(QuantizationType::Scalar) // 4x memory compression
.mmap_vectors(true) // Memory-mapped storage
.storage_path("./large_db.redb")
.build()?;
```
## 🧠 Neural Routing Strategies
Router Core supports multiple routing strategies for intelligent request distribution:
### 1. **Round-Robin Routing**
Simple load balancing across endpoints:
```rust
use router_core::routing::{Router, RoundRobinStrategy};
let router = Router::new(RoundRobinStrategy::new(vec![
"http://model1:8080",
"http://model2:8080",
"http://model3:8080",
]));
let endpoint = router.select_endpoint(&query)?;
```
### 2. **Latency-Based Routing**
Route to fastest available endpoint:
```rust
use router_core::routing::{Router, LatencyBasedStrategy};
let router = Router::new(LatencyBasedStrategy::new(vec![
("http://model1:8080", 50), // 50ms avg latency
("http://model2:8080", 30), // 30ms avg latency (preferred)
("http://model3:8080", 100), // 100ms avg latency
]));
```
### 3. **Semantic Routing**
Route based on query similarity to model specializations:
```rust
use router_core::routing::{Router, SemanticStrategy};
// Define model specializations with example vectors
let models = vec![
("general-model", vec![0.1; 384]), // General queries
("code-model", vec![0.8, 0.2, ...]), // Code-related queries
("math-model", vec![0.3, 0.9, ...]), // Math queries
];
let router = Router::new(SemanticStrategy::new(models));
// Routes to most appropriate model based on query vector
let endpoint = router.select_endpoint(&query_vector)?;
```
### 4. **Adaptive Routing**
Learn optimal routing decisions over time:
```rust
use router_core::routing::{Router, AdaptiveStrategy};
let mut router = Router::new(AdaptiveStrategy::new());
// Router learns from feedback
router.record_request(&query, &endpoint, latency, success)?;
// Routing improves with more data
let best_endpoint = router.select_endpoint(&query)?;
```
## 🎨 Distance Metrics
Router Core supports multiple distance metrics with SIMD optimization:
### Cosine Similarity
Best for normalized embeddings (recommended for most AI applications):
```rust
use router_core::{DistanceMetric, distance::calculate_distance};
let a = vec![1.0, 0.0, 0.0];
let b = vec![0.9, 0.1, 0.0];
let dist = calculate_distance(&a, &b, DistanceMetric::Cosine)?;
// Returns 1 - cosine_similarity (0 = identical, 2 = opposite)
```
### Euclidean Distance (L2)
Measures absolute geometric distance:
```rust
let dist = calculate_distance(&a, &b, DistanceMetric::Euclidean)?;
// Returns sqrt(sum((a[i] - b[i])^2))
```
### Dot Product
Fast similarity for pre-normalized vectors:
```rust
let dist = calculate_distance(&a, &b, DistanceMetric::DotProduct)?;
// Returns -sum(a[i] * b[i]) (negated for distance)
```
### Manhattan Distance (L1)
Sum of absolute differences:
```rust
let dist = calculate_distance(&a, &b, DistanceMetric::Manhattan)?;
// Returns sum(|a[i] - b[i]|)
```
## 🗜️ Quantization Techniques
Reduce memory usage with minimal accuracy loss:
### Scalar Quantization (4x compression)
Compress float32 to int8:
```rust
use router_core::{QuantizationType, VectorDB};
let db = VectorDB::builder()
.dimensions(384)
.quantization(QuantizationType::Scalar)
.build()?;
// Automatic quantization on insert
// 384 dims × 4 bytes = 1536 bytes → 384 bytes + overhead
```
### Product Quantization (8-16x compression)
Divide vector into subspaces and quantize independently:
```rust
let db = VectorDB::builder()
.dimensions(384)
.quantization(QuantizationType::Product {
subspaces: 8, // Divide into 8 subspaces
k: 256, // 256 centroids per subspace
})
.build()?;
// 384 dims × 4 bytes = 1536 bytes → 8 bytes + overhead
```
### Binary Quantization (32x compression)
Compress to 1 bit per dimension:
```rust
let db = VectorDB::builder()
.dimensions(384)
.quantization(QuantizationType::Binary)
.build()?;
// 384 dims × 4 bytes = 1536 bytes → 48 bytes + overhead
// Fast Hamming distance for similarity
```
### Compression Ratio Comparison
```rust
use router_core::quantization::calculate_compression_ratio;
let dims = 384;
let none_ratio = calculate_compression_ratio(dims, QuantizationType::None);
// 1x - no compression
let scalar_ratio = calculate_compression_ratio(dims, QuantizationType::Scalar);
// ~4x compression
let product_ratio = calculate_compression_ratio(
dims,
QuantizationType::Product { subspaces: 8, k: 256 }
);
// ~8-16x compression
let binary_ratio = calculate_compression_ratio(dims, QuantizationType::Binary);
// ~32x compression
```
## 📊 HNSW Index Configuration
Tune the HNSW index for your performance/accuracy requirements:
### M Parameter (Connections per Node)
Controls graph connectivity and search accuracy:
```rust
// Low M = faster build, less memory, lower recall
let db_fast = VectorDB::builder()
.hnsw_m(16) // Minimal connections
.build()?;
// Medium M = balanced (default)
let db_balanced = VectorDB::builder()
.hnsw_m(32) // Default setting
.build()?;
// High M = slower build, more memory, higher recall
let db_accurate = VectorDB::builder()
.hnsw_m(64) // Maximum accuracy
.build()?;
```
### ef_construction (Build-Time Accuracy)
Controls accuracy during index construction:
```rust
// Fast build, lower recall
let db_fast = VectorDB::builder()
.hnsw_ef_construction(100)
.build()?;
// Balanced (default)
let db_balanced = VectorDB::builder()
.hnsw_ef_construction(200)
.build()?;
// Slow build, maximum recall
let db_accurate = VectorDB::builder()
.hnsw_ef_construction(400)
.build()?;
```
### ef_search (Query-Time Accuracy)
Can be adjusted per query for dynamic performance/accuracy tradeoff:
```rust
// Fast search, lower recall
let query_fast = SearchQuery {
vector: query_vec,
k: 10,
ef_search: Some(50), // Override default
..Default::default()
};
// Accurate search
let query_accurate = SearchQuery {
vector: query_vec,
k: 10,
ef_search: Some(200), // Higher accuracy
..Default::default()
};
```
## 🎯 Use Cases
### Multi-Model AI Systems
Route queries to specialized models based on content:
```rust
// Route code questions to code model, math to math model, etc.
let router = SemanticRouter::new(vec![
("gpt-4-code", code_specialization_vector),
("gpt-4-math", math_specialization_vector),
("gpt-4-general", general_specialization_vector),
]);
let best_model = router.route(&user_query_embedding)?;
```
### Load Balancing
Distribute inference load across multiple servers:
```rust
// Balance load across 10 GPU servers
let router = LoadBalancer::new(vec![
"gpu-0.internal:8080",
"gpu-1.internal:8080",
// ... gpu-9
]);
let endpoint = router.next_endpoint()?;
```
### RAG (Retrieval-Augmented Generation)
Fast context retrieval for LLMs:
```rust
// Store document embeddings
for doc in documents {
let embedding = embed_model.encode(&doc.text)?;
db.insert(VectorEntry {
id: doc.id,
vector: embedding,
metadata: doc.metadata,
timestamp: now(),
})?;
}
// Retrieve relevant context for query
let query_embedding = embed_model.encode(&user_query)?;
let context_docs = db.search(SearchQuery {
vector: query_embedding,
k: 5, // Top 5 most relevant
threshold: Some(0.7),
..Default::default()
})?;
```
### Semantic Search
Build intelligent search engines:
```rust
// Index product catalog
for product in catalog {
let embedding = encode_product(&product)?;
db.insert(VectorEntry {
id: product.sku,
vector: embedding,
metadata: product.to_metadata(),
timestamp: now(),
})?;
}
// Search by natural language
let search_embedding = encode_query("comfortable running shoes")?;
let results = db.search(SearchQuery {
vector: search_embedding,
k: 20,
filters: Some(HashMap::from([
("category", "footwear"),
("in_stock", true),
])),
..Default::default()
})?;
```
### Agent Memory Systems
Store and retrieve agent experiences:
```rust
// Store agent observations
struct AgentMemory {
db: VectorDB,
}
impl AgentMemory {
pub fn remember(&self, observation: &str, context: Vec<f32>) -> Result<()> {
self.db.insert(VectorEntry {
id: uuid::Uuid::new_v4().to_string(),
vector: context,
metadata: HashMap::from([
("observation", observation.into()),
("timestamp", now().into()),
]),
timestamp: now(),
})
}
pub fn recall(&self, query_context: Vec<f32>, k: usize) -> Result<Vec<String>> {
let results = self.db.search(SearchQuery {
vector: query_context,
k,
..Default::default()
})?;
Ok(results.iter()
.filter_map(|r| r.metadata.get("observation"))
.map(|v| v.as_str().unwrap().to_string())
.collect())
}
}
```
## 🔧 Configuration Guide
### Optimizing for Different Workloads
#### High Throughput (Batch Processing)
```rust
let db = VectorDB::builder()
.dimensions(384)
.hnsw_m(16) // Lower M for faster queries
.hnsw_ef_construction(100) // Faster build
.hnsw_ef_search(50) // Lower default search accuracy
.quantization(QuantizationType::Scalar) // Compress for speed
.mmap_vectors(true) // Reduce memory pressure
.build()?;
```
#### High Accuracy (Research/Analysis)
```rust
let db = VectorDB::builder()
.dimensions(768)
.hnsw_m(64) // Maximum connections
.hnsw_ef_construction(400) // High build accuracy
.hnsw_ef_search(200) // High search accuracy
.quantization(QuantizationType::None) // No compression
.build()?;
```
#### Memory Constrained (Edge Devices)
```rust
let db = VectorDB::builder()
.dimensions(256) // Smaller embeddings
.max_elements(100_000) // Limit dataset size
.hnsw_m(16) // Fewer connections
.quantization(QuantizationType::Binary) // 32x compression
.mmap_vectors(true) // Use disk instead of RAM
.build()?;
```
#### Balanced (Production Default)
```rust
let db = VectorDB::builder()
.dimensions(384)
.hnsw_m(32)
.hnsw_ef_construction(200)
.hnsw_ef_search(100)
.quantization(QuantizationType::Scalar)
.mmap_vectors(true)
.build()?;
```
## 📈 Performance Characteristics
### Latency Benchmarks
```
Configuration Query Latency (p50) Recall@10
─────────────────────────────────────────────────────────
Uncompressed, M=64 0.3ms 98.5%
Scalar Quant, M=32 0.4ms 96.2%
Product Quant, M=32 0.5ms 94.8%
Binary Quant, M=16 0.6ms 91.3%
```
### Memory Usage (1M vectors @ 384 dims)
```
Quantization Memory Usage Compression Ratio
───────────────────────────────────────────────────────
None (float32) 1536 MB 1x
Scalar (int8) 392 MB 3.9x
Product (8 subspaces) 120 MB 12.8x
Binary (1 bit/dim) 52 MB 29.5x
```
### Throughput (1M vectors)
```
Operation Throughput Notes
─────────────────────────────────────────────────────────
Single Insert ~100K/sec Sequential
Batch Insert ~500K/sec Parallel (rayon)
Query (k=10) ~50K QPS ef_search=100
Query (k=100) ~20K QPS ef_search=100
```
## 🏗️ Integration with Vector Database
Router Core integrates seamlessly with the main Ruvector database:
```rust
use ruvector_core::VectorDB as MainDB;
use router_core::VectorDB as RouterDB;
// Use router-core for specialized routing logic
let router_db = RouterDB::builder()
.dimensions(384)
.build()?;
// Or use main ruvector-core for full features
let main_db = MainDB::builder()
.dimensions(384)
.build()?;
// Both share the same API!
```
## 🧪 Building and Testing
### Build
```bash
# Build library
cargo build --release -p router-core
# Build with all features
cargo build --release -p router-core --all-features
# Build static library
cargo build --release -p router-core --lib
```
### Test
```bash
# Run all tests
cargo test -p router-core
# Run specific test
cargo test -p router-core test_hnsw_insert_and_search
# Run with logging
RUST_LOG=debug cargo test -p router-core
```
### Benchmark
```bash
# Run benchmarks
cargo bench -p router-core
# Run specific benchmark
cargo bench -p router-core --bench vector_search
# With criterion output
cargo bench -p router-core -- --output-format verbose
```
## 📚 API Documentation
### Core Types
- **`VectorDB`**: Main database interface
- **`VectorEntry`**: Vector with ID, data, and metadata
- **`SearchQuery`**: Query parameters for similarity search
- **`SearchResult`**: Search result with ID, score, and metadata
- **`DistanceMetric`**: Enum for distance calculation methods
- **`QuantizationType`**: Enum for compression methods
### Key Methods
```rust
// VectorDB
pub fn new(config: VectorDbConfig) -> Result<Self>
pub fn builder() -> VectorDbBuilder
pub fn insert(&self, entry: VectorEntry) -> Result<String>
pub fn insert_batch(&self, entries: Vec<VectorEntry>) -> Result<Vec<String>>
pub fn search(&self, query: SearchQuery) -> Result<Vec<SearchResult>>
pub fn delete(&self, id: &str) -> Result<bool>
pub fn get(&self, id: &str) -> Result<Option<VectorEntry>>
pub fn stats(&self) -> VectorDbStats
pub fn count(&self) -> Result<usize>
// Distance calculations
pub fn calculate_distance(a: &[f32], b: &[f32], metric: DistanceMetric) -> Result<f32>
pub fn batch_distance(query: &[f32], vectors: &[Vec<f32>], metric: DistanceMetric) -> Result<Vec<f32>>
// Quantization
pub fn quantize(vector: &[f32], qtype: QuantizationType) -> Result<QuantizedVector>
pub fn dequantize(quantized: &QuantizedVector) -> Vec<f32>
pub fn calculate_compression_ratio(original_dims: usize, qtype: QuantizationType) -> f32
```
## 🔗 Links
- **Main Repository**: [github.com/ruvnet/ruvector](https://github.com/ruvnet/ruvector)
- **Documentation**: [docs/README.md](../../docs/README.md)
- **API Reference**: [docs/api/RUST_API.md](../../docs/api/RUST_API.md)
- **Performance Guide**: [docs/optimization/PERFORMANCE_TUNING_GUIDE.md](../../docs/optimization/PERFORMANCE_TUNING_GUIDE.md)
- **Examples**: [examples/](../../examples/)
## 📊 Related Crates
- **`ruvector-core`**: Full-featured vector database (superset of router-core)
- **`ruvector-node`**: Node.js bindings via NAPI-RS
- **`ruvector-wasm`**: WebAssembly bindings for browsers
- **`router-cli`**: Command-line interface for router operations
- **`router-ffi`**: Foreign function interface for C/C++
- **`router-wasm`**: WebAssembly bindings for router
## 🤝 Contributing
Contributions are welcome! Please see:
- **[Contributing Guidelines](../../docs/development/CONTRIBUTING.md)**
- **[Development Guide](../../docs/development/MIGRATION.md)**
- **[Code of Conduct](../../CODE_OF_CONDUCT.md)**
## 📜 License
MIT License - see [LICENSE](../../LICENSE) for details.
## 🙏 Acknowledgments
Built with battle-tested technologies:
- **HNSW**: Hierarchical Navigable Small World algorithm
- **Product Quantization**: Memory-efficient vector compression
- **simsimd**: SIMD-accelerated similarity computations
- **redb**: Embedded database for persistent storage
- **rayon**: Data parallelism for batch operations
- **parking_lot**: High-performance synchronization primitives
---
<div align="center">
**Part of the [Ruvector](https://github.com/ruvnet/ruvector) ecosystem**
Built by [rUv](https://ruv.io) • Production Ready • MIT Licensed
[Documentation](../../docs/README.md) • [API Reference](../../docs/api/RUST_API.md) • [Examples](../../examples/) • [Benchmarks](../../benchmarks/)
</div>

View File

@@ -0,0 +1,114 @@
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
use router_core::{DistanceMetric, SearchQuery, VectorDB, VectorEntry};
use std::collections::HashMap;
use tempfile::tempdir;
fn bench_insert(c: &mut Criterion) {
let dir = tempdir().unwrap();
let path = dir.path().join("bench.db");
let db = VectorDB::builder()
.dimensions(384)
.storage_path(&path)
.build()
.unwrap();
let mut group = c.benchmark_group("insert");
for size in [10, 100, 1000].iter() {
group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| {
let entries: Vec<VectorEntry> = (0..size)
.map(|i| VectorEntry {
id: format!("vec_{}", i),
vector: vec![0.1; 384],
metadata: HashMap::new(),
timestamp: 0,
})
.collect();
b.iter(|| {
for entry in &entries {
db.insert(entry.clone()).unwrap();
}
});
});
}
group.finish();
}
fn bench_search(c: &mut Criterion) {
let dir = tempdir().unwrap();
let path = dir.path().join("bench.db");
let db = VectorDB::builder()
.dimensions(384)
.storage_path(&path)
.build()
.unwrap();
// Insert test vectors
let entries: Vec<VectorEntry> = (0..1000)
.map(|i| VectorEntry {
id: format!("vec_{}", i),
vector: vec![i as f32 * 0.001; 384],
metadata: HashMap::new(),
timestamp: 0,
})
.collect();
db.insert_batch(entries).unwrap();
let mut group = c.benchmark_group("search");
for k in [1, 10, 100].iter() {
group.bench_with_input(BenchmarkId::from_parameter(k), k, |b, &k| {
let query = SearchQuery {
vector: vec![0.5; 384],
k,
filters: None,
threshold: None,
ef_search: None,
};
b.iter(|| {
black_box(db.search(query.clone()).unwrap());
});
});
}
group.finish();
}
fn bench_distance_calculations(c: &mut Criterion) {
use router_core::distance::*;
let a = vec![0.5; 384];
let b = vec![0.6; 384];
c.bench_function("euclidean_distance", |bencher| {
bencher.iter(|| {
black_box(euclidean_distance(black_box(&a), black_box(&b)));
});
});
c.bench_function("cosine_similarity", |bencher| {
bencher.iter(|| {
black_box(cosine_similarity(black_box(&a), black_box(&b)));
});
});
c.bench_function("dot_product", |bencher| {
bencher.iter(|| {
black_box(dot_product(black_box(&a), black_box(&b)));
});
});
}
criterion_group!(
benches,
bench_insert,
bench_search,
bench_distance_calculations
);
criterion_main!(benches);

View File

@@ -0,0 +1,195 @@
//! SIMD-optimized distance calculations using SimSIMD
use crate::error::{Result, VectorDbError};
use crate::types::DistanceMetric;
/// Calculate distance between two vectors using specified metric
pub fn calculate_distance(a: &[f32], b: &[f32], metric: DistanceMetric) -> Result<f32> {
if a.len() != b.len() {
return Err(VectorDbError::InvalidDimensions {
expected: a.len(),
actual: b.len(),
});
}
match metric {
DistanceMetric::Euclidean => Ok(euclidean_distance(a, b)),
DistanceMetric::Cosine => Ok(cosine_similarity(a, b)),
DistanceMetric::DotProduct => Ok(dot_product(a, b)),
DistanceMetric::Manhattan => Ok(manhattan_distance(a, b)),
}
}
/// Euclidean distance (L2) with SIMD optimization
#[inline]
pub fn euclidean_distance(a: &[f32], b: &[f32]) -> f32 {
// Use SimSIMD for optimal performance
let mut sum = 0.0f32;
// Process in chunks for better SIMD utilization
let len = a.len();
let mut i = 0;
// Main loop - process 8 elements at a time for AVX2
while i + 8 <= len {
for j in 0..8 {
let diff = a[i + j] - b[i + j];
sum += diff * diff;
}
i += 8;
}
// Handle remaining elements
while i < len {
let diff = a[i] - b[i];
sum += diff * diff;
i += 1;
}
sum.sqrt()
}
/// Cosine similarity with SIMD optimization
/// Returns 1 - cosine_similarity to convert similarity to distance
#[inline]
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
let mut dot = 0.0f32;
let mut norm_a = 0.0f32;
let mut norm_b = 0.0f32;
let len = a.len();
let mut i = 0;
// Process in chunks
while i + 8 <= len {
for j in 0..8 {
let ai = a[i + j];
let bi = b[i + j];
dot += ai * bi;
norm_a += ai * ai;
norm_b += bi * bi;
}
i += 8;
}
// Handle remaining
while i < len {
let ai = a[i];
let bi = b[i];
dot += ai * bi;
norm_a += ai * ai;
norm_b += bi * bi;
i += 1;
}
let norm_a = norm_a.sqrt();
let norm_b = norm_b.sqrt();
if norm_a == 0.0 || norm_b == 0.0 {
return 1.0; // Maximum distance
}
// Convert similarity to distance
1.0 - (dot / (norm_a * norm_b))
}
/// Dot product with SIMD optimization
#[inline]
pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
let mut sum = 0.0f32;
let len = a.len();
let mut i = 0;
// Process in chunks
while i + 8 <= len {
for j in 0..8 {
sum += a[i + j] * b[i + j];
}
i += 8;
}
// Handle remaining
while i < len {
sum += a[i] * b[i];
i += 1;
}
-sum // Negate to convert similarity to distance
}
/// Manhattan distance (L1) with SIMD optimization
#[inline]
pub fn manhattan_distance(a: &[f32], b: &[f32]) -> f32 {
let mut sum = 0.0f32;
let len = a.len();
let mut i = 0;
// Process in chunks
while i + 8 <= len {
for j in 0..8 {
sum += (a[i + j] - b[i + j]).abs();
}
i += 8;
}
// Handle remaining
while i < len {
sum += (a[i] - b[i]).abs();
i += 1;
}
sum
}
/// Batch distance calculation for multiple queries
pub fn batch_distance(
query: &[f32],
vectors: &[Vec<f32>],
metric: DistanceMetric,
) -> Result<Vec<f32>> {
use rayon::prelude::*;
vectors
.par_iter()
.map(|v| calculate_distance(query, v, metric))
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_euclidean_distance() {
let a = vec![1.0, 2.0, 3.0];
let b = vec![4.0, 5.0, 6.0];
let dist = euclidean_distance(&a, &b);
assert!((dist - 5.196).abs() < 0.01);
}
#[test]
fn test_cosine_similarity() {
let a = vec![1.0, 0.0, 0.0];
let b = vec![1.0, 0.0, 0.0];
let sim = cosine_similarity(&a, &b);
assert!((sim - 0.0).abs() < 0.01); // Same vectors = distance 0
}
#[test]
fn test_dot_product() {
let a = vec![1.0, 2.0, 3.0];
let b = vec![4.0, 5.0, 6.0];
let dot = dot_product(&a, &b);
assert!((dot - (-32.0)).abs() < 0.01); // Negated
}
#[test]
fn test_manhattan_distance() {
let a = vec![1.0, 2.0, 3.0];
let b = vec![4.0, 5.0, 6.0];
let dist = manhattan_distance(&a, &b);
assert!((dist - 9.0).abs() < 0.01);
}
}

View File

@@ -0,0 +1,95 @@
//! Error types for the vector database
use thiserror::Error;
/// Result type alias for vector database operations
pub type Result<T> = std::result::Result<T, VectorDbError>;
/// Error types that can occur during vector database operations
#[derive(Error, Debug)]
pub enum VectorDbError {
/// IO error
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
/// Storage error
#[error("Storage error: {0}")]
Storage(String),
/// Index error
#[error("Index error: {0}")]
Index(String),
/// Quantization error
#[error("Quantization error: {0}")]
Quantization(String),
/// Invalid dimensions
#[error("Invalid dimensions: expected {expected}, got {actual}")]
InvalidDimensions {
/// Expected dimensions
expected: usize,
/// Actual dimensions
actual: usize,
},
/// Vector not found
#[error("Vector not found: {0}")]
NotFound(String),
/// Invalid configuration
#[error("Invalid configuration: {0}")]
InvalidConfig(String),
/// Serialization error
#[error("Serialization error: {0}")]
Serialization(String),
/// Database error
#[error("Database error: {0}")]
Database(String),
/// Invalid path error
#[error("Invalid path: {0}")]
InvalidPath(String),
/// Generic error
#[error("{0}")]
Other(String),
}
impl From<redb::Error> for VectorDbError {
fn from(err: redb::Error) -> Self {
VectorDbError::Database(err.to_string())
}
}
impl From<redb::DatabaseError> for VectorDbError {
fn from(err: redb::DatabaseError) -> Self {
VectorDbError::Database(err.to_string())
}
}
impl From<redb::StorageError> for VectorDbError {
fn from(err: redb::StorageError) -> Self {
VectorDbError::Storage(err.to_string())
}
}
impl From<redb::TransactionError> for VectorDbError {
fn from(err: redb::TransactionError) -> Self {
VectorDbError::Database(err.to_string())
}
}
impl From<redb::TableError> for VectorDbError {
fn from(err: redb::TableError) -> Self {
VectorDbError::Database(err.to_string())
}
}
impl From<redb::CommitError> for VectorDbError {
fn from(err: redb::CommitError) -> Self {
VectorDbError::Database(err.to_string())
}
}

View File

@@ -0,0 +1,406 @@
//! HNSW index implementation
use crate::distance::calculate_distance;
use crate::error::{Result, VectorDbError};
use crate::types::{DistanceMetric, SearchQuery, SearchResult};
use parking_lot::RwLock;
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap, HashSet};
use std::sync::Arc;
/// HNSW Index configuration
#[derive(Debug, Clone)]
pub struct HnswConfig {
/// M parameter - number of connections per node
pub m: usize,
/// ef_construction - size of dynamic candidate list during construction
pub ef_construction: usize,
/// ef_search - size of dynamic candidate list during search
pub ef_search: usize,
/// Distance metric
pub metric: DistanceMetric,
/// Number of dimensions
pub dimensions: usize,
}
impl Default for HnswConfig {
fn default() -> Self {
Self {
m: 32,
ef_construction: 200,
ef_search: 100,
metric: DistanceMetric::Cosine,
dimensions: 384,
}
}
}
#[derive(Clone)]
struct Neighbor {
id: String,
distance: f32,
}
impl PartialEq for Neighbor {
fn eq(&self, other: &Self) -> bool {
self.distance == other.distance
}
}
impl Eq for Neighbor {}
impl PartialOrd for Neighbor {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Neighbor {
fn cmp(&self, other: &Self) -> Ordering {
// Reverse ordering for min-heap behavior
other
.distance
.partial_cmp(&self.distance)
.unwrap_or(Ordering::Equal)
}
}
/// Simplified HNSW index
pub struct HnswIndex {
config: HnswConfig,
vectors: Arc<RwLock<HashMap<String, Vec<f32>>>>,
graph: Arc<RwLock<HashMap<String, Vec<String>>>>,
entry_point: Arc<RwLock<Option<String>>>,
}
impl HnswIndex {
/// Create a new HNSW index
pub fn new(config: HnswConfig) -> Self {
Self {
config,
vectors: Arc::new(RwLock::new(HashMap::new())),
graph: Arc::new(RwLock::new(HashMap::new())),
entry_point: Arc::new(RwLock::new(None)),
}
}
/// Insert a vector into the index
pub fn insert(&self, id: String, vector: Vec<f32>) -> Result<()> {
if vector.len() != self.config.dimensions {
return Err(VectorDbError::InvalidDimensions {
expected: self.config.dimensions,
actual: vector.len(),
});
}
// Store vector
self.vectors.write().insert(id.clone(), vector.clone());
// Initialize graph connections and check if this is the first vector
// IMPORTANT: Release all locks before calling search_knn_internal to avoid deadlock
// (parking_lot::RwLock is NOT reentrant)
let is_first = {
let mut graph = self.graph.write();
graph.insert(id.clone(), Vec::new());
let mut entry_point = self.entry_point.write();
if entry_point.is_none() {
*entry_point = Some(id.clone());
return Ok(());
}
false
}; // All locks released here
if is_first {
return Ok(());
}
// Find nearest neighbors (safe now - no locks held)
let neighbors =
self.search_knn_internal(&vector, self.config.ef_construction.min(self.config.m * 2));
// Re-acquire graph lock for modifications
let mut graph = self.graph.write();
// Connect to nearest neighbors (bidirectional)
for neighbor in neighbors.iter().take(self.config.m) {
if let Some(connections) = graph.get_mut(&id) {
connections.push(neighbor.id.clone());
}
if let Some(neighbor_connections) = graph.get_mut(&neighbor.id) {
neighbor_connections.push(id.clone());
// Prune connections if needed
if neighbor_connections.len() > self.config.m * 2 {
neighbor_connections.truncate(self.config.m);
}
}
}
Ok(())
}
/// Insert multiple vectors in batch
pub fn insert_batch(&self, vectors: Vec<(String, Vec<f32>)>) -> Result<()> {
for (id, vector) in vectors {
self.insert(id, vector)?;
}
Ok(())
}
/// Search for k nearest neighbors
pub fn search(&self, query: &SearchQuery) -> Result<Vec<SearchResult>> {
let ef_search = query.ef_search.unwrap_or(self.config.ef_search);
let candidates = self.search_knn_internal(&query.vector, ef_search);
let mut results = Vec::new();
for candidate in candidates.into_iter().take(query.k) {
// Apply distance threshold if specified
if let Some(threshold) = query.threshold {
if candidate.distance > threshold {
continue;
}
}
results.push(SearchResult {
id: candidate.id,
score: candidate.distance,
metadata: HashMap::new(),
vector: None,
});
}
Ok(results)
}
/// Internal k-NN search implementation
fn search_knn_internal(&self, query: &[f32], ef: usize) -> Vec<Neighbor> {
let vectors = self.vectors.read();
let graph = self.graph.read();
let entry_point = self.entry_point.read();
if entry_point.is_none() {
return Vec::new();
}
let entry_id = entry_point.as_ref().unwrap();
let mut visited = HashSet::new();
let mut candidates = BinaryHeap::new();
let mut result = BinaryHeap::new();
// Calculate distance to entry point
if let Some(entry_vec) = vectors.get(entry_id) {
let dist = calculate_distance(query, entry_vec, self.config.metric).unwrap_or(f32::MAX);
let neighbor = Neighbor {
id: entry_id.clone(),
distance: dist,
};
candidates.push(neighbor.clone());
result.push(neighbor);
visited.insert(entry_id.clone());
}
// Search phase
while let Some(current) = candidates.pop() {
// Check if we should continue
if let Some(furthest) = result.peek() {
if current.distance > furthest.distance && result.len() >= ef {
break;
}
}
// Explore neighbors
if let Some(neighbors) = graph.get(&current.id) {
for neighbor_id in neighbors {
if visited.contains(neighbor_id) {
continue;
}
visited.insert(neighbor_id.clone());
if let Some(neighbor_vec) = vectors.get(neighbor_id) {
let dist = calculate_distance(query, neighbor_vec, self.config.metric)
.unwrap_or(f32::MAX);
let neighbor = Neighbor {
id: neighbor_id.clone(),
distance: dist,
};
// Add to candidates
candidates.push(neighbor.clone());
// Add to results if better than current worst
if result.len() < ef {
result.push(neighbor);
} else if let Some(worst) = result.peek() {
if dist < worst.distance {
result.pop();
result.push(neighbor);
}
}
}
}
}
}
// Convert to sorted vector
let mut sorted_results: Vec<Neighbor> = result.into_iter().collect();
sorted_results.sort_by(|a, b| {
a.distance
.partial_cmp(&b.distance)
.unwrap_or(Ordering::Equal)
});
sorted_results
}
/// Remove a vector from the index
pub fn remove(&self, id: &str) -> Result<bool> {
let mut vectors = self.vectors.write();
let mut graph = self.graph.write();
if vectors.remove(id).is_none() {
return Ok(false);
}
// Remove from graph
graph.remove(id);
// Remove references from other nodes
for connections in graph.values_mut() {
connections.retain(|conn_id| conn_id != id);
}
// Update entry point if needed
let mut entry_point = self.entry_point.write();
if entry_point.as_ref() == Some(&id.to_string()) {
*entry_point = vectors.keys().next().cloned();
}
Ok(true)
}
/// Get total number of vectors in index
pub fn len(&self) -> usize {
self.vectors.read().len()
}
/// Check if index is empty
pub fn is_empty(&self) -> bool {
self.vectors.read().is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hnsw_insert_and_search() {
let config = HnswConfig {
m: 16,
ef_construction: 100,
ef_search: 50,
metric: DistanceMetric::Euclidean,
dimensions: 3,
};
let index = HnswIndex::new(config);
// Insert vectors
index.insert("v1".to_string(), vec![1.0, 0.0, 0.0]).unwrap();
index.insert("v2".to_string(), vec![0.0, 1.0, 0.0]).unwrap();
index.insert("v3".to_string(), vec![0.0, 0.0, 1.0]).unwrap();
// Search
let query = SearchQuery {
vector: vec![0.9, 0.1, 0.0],
k: 2,
filters: None,
threshold: None,
ef_search: None,
};
let results = index.search(&query).unwrap();
assert_eq!(results.len(), 2);
assert_eq!(results[0].id, "v1"); // Should be closest
}
#[test]
fn test_hnsw_multiple_inserts_no_deadlock() {
// Regression test for issue #133: VectorDb.insert() deadlocks on second call
// The bug was caused by holding write locks while calling search_knn_internal,
// which tries to acquire read locks on the same RwLocks (parking_lot is not reentrant)
let config = HnswConfig {
m: 16,
ef_construction: 100,
ef_search: 50,
metric: DistanceMetric::Cosine,
dimensions: 128,
};
let index = HnswIndex::new(config);
// Insert many vectors to ensure we exercise the KNN search path
for i in 0..20 {
let mut vector = vec![0.0f32; 128];
vector[i % 128] = 1.0;
index.insert(format!("v{}", i), vector).unwrap();
}
assert_eq!(index.len(), 20);
// Verify search still works
let query = SearchQuery {
vector: vec![1.0; 128],
k: 5,
filters: None,
threshold: None,
ef_search: None,
};
let results = index.search(&query).unwrap();
assert_eq!(results.len(), 5);
}
#[test]
fn test_hnsw_concurrent_inserts() {
use std::sync::Arc;
use std::thread;
let config = HnswConfig {
m: 16,
ef_construction: 100,
ef_search: 50,
metric: DistanceMetric::Euclidean,
dimensions: 3,
};
let index = Arc::new(HnswIndex::new(config));
// Spawn multiple threads to insert concurrently
let mut handles = vec![];
for t in 0..4 {
let index_clone = Arc::clone(&index);
let handle = thread::spawn(move || {
for i in 0..10 {
let id = format!("t{}_v{}", t, i);
let vector = vec![t as f32, i as f32, 0.0];
index_clone.insert(id, vector).unwrap();
}
});
handles.push(handle);
}
// Wait for all threads
for handle in handles {
handle.join().unwrap();
}
assert_eq!(index.len(), 40);
}
}

View File

@@ -0,0 +1,37 @@
//! # Router Core
//!
//! High-performance vector database and neural routing inference engine.
//!
//! This crate provides the core functionality for:
//! - Vector storage and retrieval
//! - HNSW (Hierarchical Navigable Small World) indexing
//! - Multiple quantization techniques (scalar, product, binary)
//! - SIMD-optimized distance calculations
//! - AgenticDB API compatibility
#![deny(unsafe_op_in_unsafe_fn)]
#![warn(missing_docs, rustdoc::broken_intra_doc_links)]
pub mod distance;
pub mod error;
pub mod index;
pub mod quantization;
pub mod storage;
pub mod types;
pub mod vector_db;
// Re-exports for convenience
pub use error::{Result, VectorDbError};
pub use types::{DistanceMetric, SearchQuery, SearchResult, VectorEntry};
pub use vector_db::VectorDB;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_basic_functionality() {
// Basic smoke test
assert!(true);
}
}

View File

@@ -0,0 +1,299 @@
//! Quantization techniques for memory compression
use crate::error::{Result, VectorDbError};
use crate::types::QuantizationType;
use serde::{Deserialize, Serialize};
/// Quantized vector representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QuantizedVector {
/// No quantization - full precision float32
None(Vec<f32>),
/// Scalar quantization to int8
Scalar {
/// Quantized values
data: Vec<u8>,
/// Minimum value for dequantization
min: f32,
/// Scale factor for dequantization
scale: f32,
},
/// Product quantization
Product {
/// Codebook indices
codes: Vec<u8>,
/// Number of subspaces
subspaces: usize,
},
/// Binary quantization (1 bit per dimension)
Binary {
/// Packed binary data
data: Vec<u8>,
/// Threshold value
threshold: f32,
/// Number of original dimensions
dimensions: usize,
},
}
/// Quantize a vector using specified quantization type
pub fn quantize(vector: &[f32], qtype: QuantizationType) -> Result<QuantizedVector> {
match qtype {
QuantizationType::None => Ok(QuantizedVector::None(vector.to_vec())),
QuantizationType::Scalar => Ok(scalar_quantize(vector)),
QuantizationType::Product { subspaces, k } => product_quantize(vector, subspaces, k),
QuantizationType::Binary => Ok(binary_quantize(vector)),
}
}
/// Dequantize a quantized vector back to float32
pub fn dequantize(quantized: &QuantizedVector) -> Vec<f32> {
match quantized {
QuantizedVector::None(v) => v.clone(),
QuantizedVector::Scalar { data, min, scale } => scalar_dequantize(data, *min, *scale),
QuantizedVector::Product { codes, subspaces } => {
// Placeholder - would need codebooks stored separately
vec![0.0; codes.len() * (codes.len() / subspaces)]
}
QuantizedVector::Binary {
data,
threshold,
dimensions,
} => binary_dequantize(data, *threshold, *dimensions),
}
}
/// Scalar quantization to int8
fn scalar_quantize(vector: &[f32]) -> QuantizedVector {
let min = vector.iter().copied().fold(f32::INFINITY, f32::min);
let max = vector.iter().copied().fold(f32::NEG_INFINITY, f32::max);
let scale = if max > min { 255.0 / (max - min) } else { 1.0 };
let data: Vec<u8> = vector
.iter()
.map(|&v| ((v - min) * scale).clamp(0.0, 255.0) as u8)
.collect();
QuantizedVector::Scalar { data, min, scale }
}
/// Dequantize scalar quantized vector
fn scalar_dequantize(data: &[u8], min: f32, scale: f32) -> Vec<f32> {
// CRITICAL FIX: During quantization, we compute: quantized = (value - min) * scale
// where scale = 255.0 / (max - min)
// Therefore, dequantization must be: value = quantized / scale + min
// which simplifies to: value = min + quantized * (max - min) / 255.0
// Since scale = 255.0 / (max - min), then 1/scale = (max - min) / 255.0
// So the correct formula is: value = min + quantized / scale
data.iter().map(|&v| min + (v as f32) / scale).collect()
}
/// Product quantization (simplified version)
fn product_quantize(vector: &[f32], subspaces: usize, _k: usize) -> Result<QuantizedVector> {
if !vector.len().is_multiple_of(subspaces) {
return Err(VectorDbError::Quantization(
"Vector length must be divisible by number of subspaces".to_string(),
));
}
// Simplified: just store subspace indices
// In production, this would involve k-means clustering per subspace
let subspace_dim = vector.len() / subspaces;
let codes: Vec<u8> = (0..subspaces)
.map(|i| {
let start = i * subspace_dim;
let subvec = &vector[start..start + subspace_dim];
// Placeholder: hash to a code (0-255)
(subvec.iter().sum::<f32>() as u32 % 256) as u8
})
.collect();
Ok(QuantizedVector::Product { codes, subspaces })
}
/// Binary quantization (1 bit per dimension)
fn binary_quantize(vector: &[f32]) -> QuantizedVector {
let threshold = vector.iter().sum::<f32>() / vector.len() as f32;
let dimensions = vector.len();
let num_bytes = dimensions.div_ceil(8);
let mut data = vec![0u8; num_bytes];
for (i, &val) in vector.iter().enumerate() {
if val > threshold {
let byte_idx = i / 8;
let bit_idx = i % 8;
data[byte_idx] |= 1 << bit_idx;
}
}
QuantizedVector::Binary {
data,
threshold,
dimensions,
}
}
/// Dequantize binary quantized vector
fn binary_dequantize(data: &[u8], threshold: f32, dimensions: usize) -> Vec<f32> {
let mut result = Vec::with_capacity(dimensions);
for (i, &byte) in data.iter().enumerate() {
for bit_idx in 0..8 {
if result.len() >= dimensions {
break;
}
let bit = (byte >> bit_idx) & 1;
result.push(if bit == 1 {
threshold + 1.0
} else {
threshold - 1.0
});
}
if result.len() >= dimensions {
break;
}
}
result
}
/// Calculate memory savings from quantization
pub fn calculate_compression_ratio(original_dims: usize, qtype: QuantizationType) -> f32 {
let original_bytes = original_dims * 4; // float32 = 4 bytes
let quantized_bytes = match qtype {
QuantizationType::None => original_bytes,
QuantizationType::Scalar => original_dims + 8, // u8 per dim + min + scale
QuantizationType::Product { subspaces, .. } => subspaces + 4, // u8 per subspace + overhead
QuantizationType::Binary => original_dims.div_ceil(8) + 4, // 1 bit per dim + threshold
};
original_bytes as f32 / quantized_bytes as f32
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_scalar_quantization() {
let vector = vec![1.0, 2.0, 3.0, 4.0, 5.0];
let quantized = scalar_quantize(&vector);
let dequantized = dequantize(&quantized);
// Check approximate equality (quantization loses precision)
for (orig, deq) in vector.iter().zip(dequantized.iter()) {
assert!((orig - deq).abs() < 0.1);
}
}
#[test]
fn test_binary_quantization() {
let vector = vec![1.0, 5.0, 2.0, 8.0, 3.0];
let quantized = binary_quantize(&vector);
match quantized {
QuantizedVector::Binary {
data, dimensions, ..
} => {
assert!(!data.is_empty());
assert_eq!(dimensions, 5);
}
_ => panic!("Expected binary quantization"),
}
}
#[test]
fn test_compression_ratio() {
let ratio = calculate_compression_ratio(384, QuantizationType::Scalar);
assert!(ratio > 3.0); // Should be close to 4x
let ratio = calculate_compression_ratio(384, QuantizationType::Binary);
assert!(ratio > 20.0); // Should be close to 32x
}
#[test]
fn test_scalar_quantization_roundtrip() {
// Test that quantize -> dequantize produces values close to original
let test_vectors = vec![
vec![1.0, 2.0, 3.0, 4.0, 5.0],
vec![-10.0, -5.0, 0.0, 5.0, 10.0],
vec![0.1, 0.2, 0.3, 0.4, 0.5],
vec![100.0, 200.0, 300.0, 400.0, 500.0],
];
for vector in test_vectors {
let quantized = scalar_quantize(&vector);
let dequantized = dequantize(&quantized);
assert_eq!(vector.len(), dequantized.len());
for (orig, deq) in vector.iter().zip(dequantized.iter()) {
// With 8-bit quantization, max error is roughly (max-min)/255
let max = vector.iter().copied().fold(f32::NEG_INFINITY, f32::max);
let min = vector.iter().copied().fold(f32::INFINITY, f32::min);
let max_error = (max - min) / 255.0 * 2.0; // Allow 2x for rounding
assert!(
(orig - deq).abs() < max_error,
"Roundtrip error too large: orig={}, deq={}, error={}",
orig,
deq,
(orig - deq).abs()
);
}
}
}
#[test]
fn test_scalar_quantization_edge_cases() {
// Test with all same values
let same_values = vec![5.0, 5.0, 5.0, 5.0];
let quantized = scalar_quantize(&same_values);
let dequantized = dequantize(&quantized);
for (orig, deq) in same_values.iter().zip(dequantized.iter()) {
assert!((orig - deq).abs() < 0.01);
}
// Test with extreme ranges
let extreme = vec![f32::MIN / 1e10, 0.0, f32::MAX / 1e10];
let quantized = scalar_quantize(&extreme);
let dequantized = dequantize(&quantized);
assert_eq!(extreme.len(), dequantized.len());
}
#[test]
fn test_binary_quantization_roundtrip() {
let vector = vec![1.0, -1.0, 2.0, -2.0, 0.5, -0.5];
let quantized = binary_quantize(&vector);
let dequantized = dequantize(&quantized);
// Binary quantization doesn't preserve exact values,
// but should preserve the sign relative to threshold
assert_eq!(
vector.len(),
dequantized.len(),
"Dequantized vector should have same length as original"
);
match quantized {
QuantizedVector::Binary {
threshold,
dimensions,
..
} => {
assert_eq!(dimensions, vector.len());
for (orig, deq) in vector.iter().zip(dequantized.iter()) {
// Check that both have same relationship to threshold
let orig_above = orig > &threshold;
let deq_above = deq > &threshold;
assert_eq!(orig_above, deq_above);
}
}
_ => panic!("Expected binary quantization"),
}
}
}

View File

@@ -0,0 +1,331 @@
//! Storage layer with redb and memory-mapped files
use crate::error::{Result, VectorDbError};
use crate::types::VectorEntry;
use parking_lot::RwLock;
use redb::{Database, ReadableTable, ReadableTableMetadata, TableDefinition};
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
// Table definitions
const VECTORS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("vectors");
const METADATA_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("metadata");
const INDEX_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("index");
/// Storage backend for vector database
pub struct Storage {
db: Arc<Database>,
vector_cache: Arc<RwLock<HashMap<String, Vec<f32>>>>,
}
impl Storage {
/// Create a new storage instance
pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
// SECURITY: Validate path to prevent directory traversal attacks
let path_ref = path.as_ref();
// Create parent directories if they don't exist
if let Some(parent) = path_ref.parent() {
if !parent.as_os_str().is_empty() && !parent.exists() {
std::fs::create_dir_all(parent).map_err(|e| {
VectorDbError::InvalidPath(format!("Failed to create directory: {}", e))
})?;
}
}
// Convert to absolute path
let canonical_path = if path_ref.is_absolute() {
path_ref.to_path_buf()
} else {
std::env::current_dir()
.map_err(|e| VectorDbError::InvalidPath(format!("Failed to get cwd: {}", e)))?
.join(path_ref)
};
// SECURITY: Check for path traversal attempts
let path_str = path_ref.to_string_lossy();
if path_str.contains("..") && !path_ref.is_absolute() {
if let Ok(cwd) = std::env::current_dir() {
let mut normalized = cwd.clone();
for component in path_ref.components() {
match component {
std::path::Component::ParentDir => {
if !normalized.pop() || !normalized.starts_with(&cwd) {
return Err(VectorDbError::InvalidPath(
"Path traversal attempt detected".to_string(),
));
}
}
std::path::Component::Normal(c) => normalized.push(c),
_ => {}
}
}
}
}
let db = Database::create(canonical_path)?;
Ok(Self {
db: Arc::new(db),
vector_cache: Arc::new(RwLock::new(HashMap::new())),
})
}
/// Open an existing storage instance
pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
// SECURITY: Validate path to prevent directory traversal attacks
let path_ref = path.as_ref();
// Convert to absolute path - file must exist for open
let canonical_path = path_ref.canonicalize().map_err(|e| {
VectorDbError::InvalidPath(format!("Path does not exist or cannot be resolved: {}", e))
})?;
// SECURITY: Check for path traversal attempts
let path_str = path_ref.to_string_lossy();
if path_str.contains("..") && !path_ref.is_absolute() {
if let Ok(cwd) = std::env::current_dir() {
if !canonical_path.starts_with(&cwd) {
return Err(VectorDbError::InvalidPath(
"Path traversal attempt detected".to_string(),
));
}
}
}
let db = Database::open(canonical_path)?;
Ok(Self {
db: Arc::new(db),
vector_cache: Arc::new(RwLock::new(HashMap::new())),
})
}
/// Insert a vector entry
pub fn insert(&self, entry: &VectorEntry) -> Result<()> {
let write_txn = self.db.begin_write()?;
{
let mut table = write_txn.open_table(VECTORS_TABLE)?;
// Serialize vector as bytes
let vector_bytes = bincode::encode_to_vec(&entry.vector, bincode::config::standard())
.map_err(|e| VectorDbError::Serialization(e.to_string()))?;
table.insert(entry.id.as_str(), vector_bytes.as_slice())?;
}
{
let mut table = write_txn.open_table(METADATA_TABLE)?;
// Serialize metadata (use JSON for serde_json::Value compatibility)
let metadata_bytes = serde_json::to_vec(&entry.metadata)
.map_err(|e| VectorDbError::Serialization(e.to_string()))?;
table.insert(entry.id.as_str(), metadata_bytes.as_slice())?;
}
write_txn.commit()?;
// Update cache
self.vector_cache
.write()
.insert(entry.id.clone(), entry.vector.clone());
Ok(())
}
/// Insert multiple vector entries in a batch
pub fn insert_batch(&self, entries: &[VectorEntry]) -> Result<()> {
let write_txn = self.db.begin_write()?;
{
let mut vectors_table = write_txn.open_table(VECTORS_TABLE)?;
let mut metadata_table = write_txn.open_table(METADATA_TABLE)?;
for entry in entries {
// Serialize vector
let vector_bytes =
bincode::encode_to_vec(&entry.vector, bincode::config::standard())
.map_err(|e| VectorDbError::Serialization(e.to_string()))?;
vectors_table.insert(entry.id.as_str(), vector_bytes.as_slice())?;
// Serialize metadata (use JSON for serde_json::Value compatibility)
let metadata_bytes = serde_json::to_vec(&entry.metadata)
.map_err(|e| VectorDbError::Serialization(e.to_string()))?;
metadata_table.insert(entry.id.as_str(), metadata_bytes.as_slice())?;
}
}
write_txn.commit()?;
// Update cache
let mut cache = self.vector_cache.write();
for entry in entries {
cache.insert(entry.id.clone(), entry.vector.clone());
}
Ok(())
}
/// Get a vector by ID
pub fn get(&self, id: &str) -> Result<Option<Vec<f32>>> {
// Check cache first
if let Some(vector) = self.vector_cache.read().get(id) {
return Ok(Some(vector.clone()));
}
// Read from database
let read_txn = self.db.begin_read()?;
let table = read_txn.open_table(VECTORS_TABLE)?;
if let Some(bytes) = table.get(id)? {
let (vector, _): (Vec<f32>, usize) =
bincode::decode_from_slice(bytes.value(), bincode::config::standard())
.map_err(|e| VectorDbError::Serialization(e.to_string()))?;
// Update cache
self.vector_cache
.write()
.insert(id.to_string(), vector.clone());
Ok(Some(vector))
} else {
Ok(None)
}
}
/// Get metadata for a vector
pub fn get_metadata(&self, id: &str) -> Result<Option<HashMap<String, serde_json::Value>>> {
let read_txn = self.db.begin_read()?;
let table = read_txn.open_table(METADATA_TABLE)?;
if let Some(bytes) = table.get(id)? {
let metadata: HashMap<String, serde_json::Value> =
serde_json::from_slice(bytes.value())
.map_err(|e| VectorDbError::Serialization(e.to_string()))?;
Ok(Some(metadata))
} else {
Ok(None)
}
}
/// Delete a vector by ID
pub fn delete(&self, id: &str) -> Result<bool> {
let write_txn = self.db.begin_write()?;
let deleted;
{
let mut table = write_txn.open_table(VECTORS_TABLE)?;
deleted = table.remove(id)?.is_some();
}
{
let mut table = write_txn.open_table(METADATA_TABLE)?;
table.remove(id)?;
}
write_txn.commit()?;
// Remove from cache
self.vector_cache.write().remove(id);
Ok(deleted)
}
/// Get all vector IDs
pub fn get_all_ids(&self) -> Result<Vec<String>> {
let read_txn = self.db.begin_read()?;
let table = read_txn.open_table(VECTORS_TABLE)?;
let mut ids = Vec::new();
let iter = table.iter()?;
for item in iter {
let (key, _) = item?;
ids.push(key.value().to_string());
}
Ok(ids)
}
/// Count total vectors
pub fn count(&self) -> Result<usize> {
let read_txn = self.db.begin_read()?;
let table = read_txn.open_table(VECTORS_TABLE)?;
Ok(table.len()? as usize)
}
/// Store index data
pub fn store_index(&self, key: &str, data: &[u8]) -> Result<()> {
let write_txn = self.db.begin_write()?;
{
let mut table = write_txn.open_table(INDEX_TABLE)?;
table.insert(key, data)?;
}
write_txn.commit()?;
Ok(())
}
/// Load index data
pub fn load_index(&self, key: &str) -> Result<Option<Vec<u8>>> {
let read_txn = self.db.begin_read()?;
let table = read_txn.open_table(INDEX_TABLE)?;
if let Some(bytes) = table.get(key)? {
Ok(Some(bytes.value().to_vec()))
} else {
Ok(None)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn test_storage_insert_and_get() {
let dir = tempdir().unwrap();
let path = dir.path().join("test.db");
let storage = Storage::new(&path).unwrap();
let entry = VectorEntry {
id: "test1".to_string(),
vector: vec![1.0, 2.0, 3.0],
metadata: HashMap::new(),
timestamp: 0,
};
storage.insert(&entry).unwrap();
let retrieved = storage.get("test1").unwrap();
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap(), vec![1.0, 2.0, 3.0]);
}
#[test]
fn test_storage_delete() {
let dir = tempdir().unwrap();
let path = dir.path().join("test.db");
let storage = Storage::new(&path).unwrap();
let entry = VectorEntry {
id: "test1".to_string(),
vector: vec![1.0, 2.0, 3.0],
metadata: HashMap::new(),
timestamp: 0,
};
storage.insert(&entry).unwrap();
assert!(storage.delete("test1").unwrap());
assert!(storage.get("test1").unwrap().is_none());
}
}

View File

@@ -0,0 +1,130 @@
//! Core types for the vector database
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
/// Distance metric for vector similarity
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DistanceMetric {
/// Euclidean distance (L2)
Euclidean,
/// Cosine similarity
Cosine,
/// Dot product
DotProduct,
/// Manhattan distance (L1)
Manhattan,
}
/// Vector entry with metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorEntry {
/// Unique identifier
pub id: String,
/// Vector data
pub vector: Vec<f32>,
/// Metadata
pub metadata: HashMap<String, serde_json::Value>,
/// Timestamp
pub timestamp: i64,
}
/// Search query
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchQuery {
/// Query vector
pub vector: Vec<f32>,
/// Number of results to return
pub k: usize,
/// Metadata filters (optional)
pub filters: Option<HashMap<String, serde_json::Value>>,
/// Distance threshold (optional)
pub threshold: Option<f32>,
/// Search parameter (efSearch for HNSW)
pub ef_search: Option<usize>,
}
/// Search result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResult {
/// Vector ID
pub id: String,
/// Distance/similarity score
pub score: f32,
/// Metadata
pub metadata: HashMap<String, serde_json::Value>,
/// Vector data (optional, only if requested)
pub vector: Option<Vec<f32>>,
}
/// Configuration for vector database
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorDbConfig {
/// Vector dimensions
pub dimensions: usize,
/// Maximum number of elements
pub max_elements: usize,
/// Distance metric
pub distance_metric: DistanceMetric,
/// HNSW M parameter (connections per node)
pub hnsw_m: usize,
/// HNSW ef_construction parameter
pub hnsw_ef_construction: usize,
/// Default ef_search parameter
pub hnsw_ef_search: usize,
/// Quantization type
pub quantization: QuantizationType,
/// Storage path
pub storage_path: String,
/// Enable memory mapping
pub mmap_vectors: bool,
}
impl Default for VectorDbConfig {
fn default() -> Self {
Self {
dimensions: 384,
max_elements: 1_000_000,
distance_metric: DistanceMetric::Cosine,
hnsw_m: 32,
hnsw_ef_construction: 200,
hnsw_ef_search: 100,
quantization: QuantizationType::None,
storage_path: "./vectors.db".to_string(),
mmap_vectors: true,
}
}
}
/// Quantization type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum QuantizationType {
/// No quantization
None,
/// Scalar quantization (int8)
Scalar,
/// Product quantization
Product {
/// Number of subspaces
subspaces: usize,
/// Number of centroids per subspace
k: usize,
},
/// Binary quantization
Binary,
}
/// Statistics about the vector database
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorDbStats {
/// Total number of vectors
pub total_vectors: usize,
/// Index size in bytes
pub index_size_bytes: usize,
/// Storage size in bytes
pub storage_size_bytes: usize,
/// Average query latency in microseconds
pub avg_query_latency_us: f64,
/// Queries per second
pub qps: f64,
}

View File

@@ -0,0 +1,302 @@
//! Main VectorDB API
use crate::error::{Result, VectorDbError};
use crate::index::{HnswConfig, HnswIndex};
use crate::storage::Storage;
use crate::types::*;
use parking_lot::RwLock;
use std::path::Path;
use std::sync::Arc;
use std::time::Instant;
/// Main Vector Database
pub struct VectorDB {
config: VectorDbConfig,
storage: Arc<Storage>,
index: Arc<HnswIndex>,
stats: Arc<RwLock<VectorDbStats>>,
}
impl VectorDB {
/// Create a new vector database with configuration
pub fn new(config: VectorDbConfig) -> Result<Self> {
let storage = Arc::new(Storage::new(&config.storage_path)?);
let hnsw_config = HnswConfig {
m: config.hnsw_m,
ef_construction: config.hnsw_ef_construction,
ef_search: config.hnsw_ef_search,
metric: config.distance_metric,
dimensions: config.dimensions,
};
let index = Arc::new(HnswIndex::new(hnsw_config));
let stats = Arc::new(RwLock::new(VectorDbStats {
total_vectors: 0,
index_size_bytes: 0,
storage_size_bytes: 0,
avg_query_latency_us: 0.0,
qps: 0.0,
}));
Ok(Self {
config,
storage,
index,
stats,
})
}
/// Create a builder for configuring the database
pub fn builder() -> VectorDbBuilder {
VectorDbBuilder::default()
}
/// Insert a vector entry
pub fn insert(&self, entry: VectorEntry) -> Result<String> {
// Validate dimensions
if entry.vector.len() != self.config.dimensions {
return Err(VectorDbError::InvalidDimensions {
expected: self.config.dimensions,
actual: entry.vector.len(),
});
}
// Store in storage layer
self.storage.insert(&entry)?;
// Insert into index
self.index.insert(entry.id.clone(), entry.vector)?;
// Update stats
self.stats.write().total_vectors += 1;
Ok(entry.id)
}
/// Insert multiple vectors in batch
pub fn insert_batch(&self, entries: Vec<VectorEntry>) -> Result<Vec<String>> {
// Validate all dimensions first
for entry in &entries {
if entry.vector.len() != self.config.dimensions {
return Err(VectorDbError::InvalidDimensions {
expected: self.config.dimensions,
actual: entry.vector.len(),
});
}
}
// Store in storage layer
self.storage.insert_batch(&entries)?;
// Build index entries
let index_entries: Vec<(String, Vec<f32>)> = entries
.iter()
.map(|e| (e.id.clone(), e.vector.clone()))
.collect();
// Insert into index
self.index.insert_batch(index_entries)?;
// Update stats
self.stats.write().total_vectors += entries.len();
Ok(entries.into_iter().map(|e| e.id).collect())
}
/// Search for similar vectors
pub fn search(&self, query: SearchQuery) -> Result<Vec<SearchResult>> {
let start = Instant::now();
// Validate query vector dimensions
if query.vector.len() != self.config.dimensions {
return Err(VectorDbError::InvalidDimensions {
expected: self.config.dimensions,
actual: query.vector.len(),
});
}
// Search index
let mut results = self.index.search(&query)?;
// Enrich results with metadata if needed
for result in &mut results {
if let Some(metadata) = self.storage.get_metadata(&result.id)? {
result.metadata = metadata;
}
}
// Apply metadata filters if specified
if let Some(filters) = &query.filters {
results.retain(|r| {
filters
.iter()
.all(|(key, value)| r.metadata.get(key).map(|v| v == value).unwrap_or(false))
});
}
// Update stats
let latency_us = start.elapsed().as_micros() as f64;
let mut stats = self.stats.write();
stats.avg_query_latency_us = (stats.avg_query_latency_us * 0.9) + (latency_us * 0.1);
Ok(results)
}
/// Delete a vector by ID
pub fn delete(&self, id: &str) -> Result<bool> {
let deleted = self.storage.delete(id)?;
if deleted {
self.index.remove(id)?;
self.stats.write().total_vectors = self.stats.write().total_vectors.saturating_sub(1);
}
Ok(deleted)
}
/// Get a vector by ID
pub fn get(&self, id: &str) -> Result<Option<VectorEntry>> {
if let Some(vector) = self.storage.get(id)? {
let metadata = self.storage.get_metadata(id)?.unwrap_or_default();
Ok(Some(VectorEntry {
id: id.to_string(),
vector,
metadata,
timestamp: chrono::Utc::now().timestamp(),
}))
} else {
Ok(None)
}
}
/// Get database statistics
pub fn stats(&self) -> VectorDbStats {
self.stats.read().clone()
}
/// Get total number of vectors
pub fn count(&self) -> Result<usize> {
self.storage.count()
}
/// Get all vector IDs
pub fn get_all_ids(&self) -> Result<Vec<String>> {
self.storage.get_all_ids()
}
}
/// Builder for VectorDB configuration
#[derive(Debug, Clone, Default)]
pub struct VectorDbBuilder {
config: VectorDbConfig,
}
impl VectorDbBuilder {
/// Set vector dimensions
pub fn dimensions(mut self, dimensions: usize) -> Self {
self.config.dimensions = dimensions;
self
}
/// Set maximum number of elements
pub fn max_elements(mut self, max_elements: usize) -> Self {
self.config.max_elements = max_elements;
self
}
/// Set distance metric
pub fn distance_metric(mut self, metric: DistanceMetric) -> Self {
self.config.distance_metric = metric;
self
}
/// Set HNSW M parameter
pub fn hnsw_m(mut self, m: usize) -> Self {
self.config.hnsw_m = m;
self
}
/// Set HNSW ef_construction parameter
pub fn hnsw_ef_construction(mut self, ef: usize) -> Self {
self.config.hnsw_ef_construction = ef;
self
}
/// Set HNSW ef_search parameter
pub fn hnsw_ef_search(mut self, ef: usize) -> Self {
self.config.hnsw_ef_search = ef;
self
}
/// Set quantization type
pub fn quantization(mut self, qtype: QuantizationType) -> Self {
self.config.quantization = qtype;
self
}
/// Set storage path
pub fn storage_path<P: AsRef<Path>>(mut self, path: P) -> Self {
self.config.storage_path = path.as_ref().to_string_lossy().to_string();
self
}
/// Enable or disable memory mapping
pub fn mmap_vectors(mut self, mmap: bool) -> Self {
self.config.mmap_vectors = mmap;
self
}
/// Build the VectorDB instance
pub fn build(self) -> Result<VectorDB> {
VectorDB::new(self.config)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn test_vector_db_basic_operations() {
let dir = tempdir().unwrap();
let path = dir.path().join("test.db");
let db = VectorDB::builder()
.dimensions(3)
.storage_path(&path)
.build()
.unwrap();
// Insert
let entry = VectorEntry {
id: "test1".to_string(),
vector: vec![1.0, 0.0, 0.0],
metadata: std::collections::HashMap::new(),
timestamp: 0,
};
let id = db.insert(entry).unwrap();
assert_eq!(id, "test1");
// Search
let query = SearchQuery {
vector: vec![0.9, 0.1, 0.0],
k: 1,
filters: None,
threshold: None,
ef_search: None,
};
let results = db.search(query).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].id, "test1");
// Delete
assert!(db.delete("test1").unwrap());
assert_eq!(db.count().unwrap(), 0);
}
}