use crate::hnsw::{DistanceMetric, SearchResult, VectorIndex};
use crate::metadata::{Metadata, MetadataFilter, MetadataStore, TemporalOptions};
use crate::stats::{IndexHealth, IndexStats, MemoryUsage, PerfTimer, StatsSnapshot};
use ipfrs_core::{Cid, Error, Result};
use lru::LruCache;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::num::NonZeroUsize;
use std::sync::{Arc, RwLock};
#[derive(Debug, Clone)]
pub struct HybridConfig {
pub dimension: usize,
pub metric: DistanceMetric,
pub max_connections: usize,
pub ef_construction: usize,
pub ef_search: usize,
pub cache_size: usize,
pub collect_stats: bool,
pub filter_strategy: FilterStrategy,
}
impl Default for HybridConfig {
fn default() -> Self {
Self {
dimension: 768,
metric: DistanceMetric::Cosine,
max_connections: 16,
ef_construction: 200,
ef_search: 50,
cache_size: 1000,
collect_stats: true,
filter_strategy: FilterStrategy::Auto,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FilterStrategy {
Auto,
PreFilter,
PostFilter,
}
#[derive(Debug, Clone)]
pub struct HybridQuery {
pub vector: Vec<f32>,
pub k: usize,
pub filter: Option<MetadataFilter>,
pub temporal: Option<TemporalOptions>,
pub min_score: Option<f32>,
pub ef_search: Option<usize>,
pub include_metadata: bool,
}
impl HybridQuery {
pub fn knn(vector: Vec<f32>, k: usize) -> Self {
Self {
vector,
k,
filter: None,
temporal: None,
min_score: None,
ef_search: None,
include_metadata: false,
}
}
pub fn with_filter(mut self, filter: MetadataFilter) -> Self {
self.filter = Some(filter);
self
}
pub fn with_temporal(mut self, temporal: TemporalOptions) -> Self {
self.temporal = Some(temporal);
self
}
pub fn with_min_score(mut self, min_score: f32) -> Self {
self.min_score = Some(min_score);
self
}
pub fn with_metadata(mut self) -> Self {
self.include_metadata = true;
self
}
pub fn with_ef_search(mut self, ef_search: usize) -> Self {
self.ef_search = Some(ef_search);
self
}
}
#[derive(Debug, Clone)]
pub struct HybridResult {
pub cid: Cid,
pub score: f32,
pub metadata: Option<Metadata>,
}
impl From<SearchResult> for HybridResult {
fn from(result: SearchResult) -> Self {
Self {
cid: result.cid,
score: result.score,
metadata: None,
}
}
}
#[derive(Debug, Clone)]
pub struct HybridResponse {
pub results: Vec<HybridResult>,
pub total_evaluated: usize,
pub latency_us: u64,
pub strategy_used: FilterStrategy,
}
pub struct HybridIndex {
vector_index: Arc<RwLock<VectorIndex>>,
metadata_store: Arc<MetadataStore>,
config: HybridConfig,
stats: Arc<IndexStats>,
cache: Arc<RwLock<LruCache<u64, Vec<HybridResult>>>>,
}
impl HybridIndex {
pub fn new(config: HybridConfig) -> Result<Self> {
let vector_index = VectorIndex::new(
config.dimension,
config.metric,
config.max_connections,
config.ef_construction,
)?;
let cache_size = NonZeroUsize::new(config.cache_size)
.unwrap_or(NonZeroUsize::new(1000).expect("1000 > 0"));
Ok(Self {
vector_index: Arc::new(RwLock::new(vector_index)),
metadata_store: Arc::new(MetadataStore::new()),
config,
stats: Arc::new(IndexStats::new()),
cache: Arc::new(RwLock::new(LruCache::new(cache_size))),
})
}
pub fn with_defaults() -> Result<Self> {
Self::new(HybridConfig::default())
}
pub fn insert(&self, cid: &Cid, vector: &[f32], metadata: Option<Metadata>) -> Result<()> {
let timer = PerfTimer::start();
self.vector_index
.write()
.unwrap_or_else(|e| e.into_inner())
.insert(cid, vector)?;
if let Some(meta) = metadata {
self.metadata_store.insert(*cid, meta)?;
} else {
self.metadata_store.insert(*cid, Metadata::new())?;
}
if self.config.collect_stats {
self.stats.record_insert(timer.stop());
}
self.cache
.write()
.unwrap_or_else(|e| e.into_inner())
.clear();
Ok(())
}
pub fn insert_batch(&self, items: &[(Cid, Vec<f32>, Option<Metadata>)]) -> Result<()> {
for (cid, vector, metadata) in items {
self.insert(cid, vector, metadata.clone())?;
}
Ok(())
}
pub fn delete(&self, cid: &Cid) -> Result<()> {
self.vector_index
.write()
.unwrap_or_else(|e| e.into_inner())
.delete(cid)?;
self.metadata_store.remove(cid)?;
if self.config.collect_stats {
self.stats.record_delete();
}
self.cache
.write()
.unwrap_or_else(|e| e.into_inner())
.clear();
Ok(())
}
pub async fn search(&self, query: HybridQuery) -> Result<HybridResponse> {
let timer = PerfTimer::start();
let strategy = self.determine_strategy(&query);
let mut total_evaluated = 0;
let results = match strategy {
FilterStrategy::PreFilter => {
self.search_pre_filter(&query, &mut total_evaluated).await?
}
FilterStrategy::PostFilter | FilterStrategy::Auto => {
self.search_post_filter(&query, &mut total_evaluated)
.await?
}
};
let latency = timer.stop();
if self.config.collect_stats {
self.stats.record_search(latency, query.k, results.len());
}
Ok(HybridResponse {
results,
total_evaluated,
latency_us: latency.as_micros() as u64,
strategy_used: strategy,
})
}
async fn search_pre_filter(
&self,
query: &HybridQuery,
total_evaluated: &mut usize,
) -> Result<Vec<HybridResult>> {
let candidates: HashSet<Cid> = if let Some(ref filter) = query.filter {
self.metadata_store.filter(filter).into_iter().collect()
} else {
self.metadata_store.cids().into_iter().collect()
};
let candidates = if let Some(ref temporal) = query.temporal {
let time_filtered = self
.metadata_store
.get_by_time_range(temporal.start, temporal.end);
candidates
.intersection(&time_filtered.into_iter().collect())
.copied()
.collect()
} else {
candidates
};
*total_evaluated = candidates.len();
if candidates.is_empty() {
return Ok(Vec::new());
}
let ef_search = query.ef_search.unwrap_or(self.config.ef_search);
let fetch_k = (query.k * 3).max(100);
let search_results = self
.vector_index
.read()
.unwrap_or_else(|e| e.into_inner())
.search(&query.vector, fetch_k, ef_search)?;
let mut results: Vec<HybridResult> = search_results
.into_iter()
.filter(|r| candidates.contains(&r.cid))
.map(|r| {
let mut hr = HybridResult::from(r);
if let Some(ref temporal) = query.temporal {
if let Some(meta) = self.metadata_store.get(&hr.cid) {
let boost = temporal.recency_multiplier(meta.created_at);
hr.score *= boost;
}
}
hr
})
.collect();
if let Some(min_score) = query.min_score {
results.retain(|r| r.score >= min_score);
}
results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
results.truncate(query.k);
if query.include_metadata {
for result in &mut results {
result.metadata = self.metadata_store.get(&result.cid);
}
}
Ok(results)
}
async fn search_post_filter(
&self,
query: &HybridQuery,
total_evaluated: &mut usize,
) -> Result<Vec<HybridResult>> {
let ef_search = query.ef_search.unwrap_or(self.config.ef_search);
let fetch_k = if query.filter.is_some() || query.temporal.is_some() {
(query.k * 5).max(100)
} else {
query.k
};
let search_results = self
.vector_index
.read()
.unwrap_or_else(|e| e.into_inner())
.search(&query.vector, fetch_k, ef_search)?;
*total_evaluated = search_results.len();
let mut results: Vec<HybridResult> = search_results
.into_iter()
.filter_map(|r| {
if let Some(ref filter) = query.filter {
if let Some(meta) = self.metadata_store.get(&r.cid) {
if !filter.matches(&meta) {
return None;
}
} else {
return None; }
}
if let Some(ref temporal) = query.temporal {
if let Some(meta) = self.metadata_store.get(&r.cid) {
if let (Some(start), Some(end)) = (temporal.start, temporal.end) {
if meta.created_at < start || meta.created_at > end {
return None;
}
}
}
}
let mut hr = HybridResult::from(r);
if let Some(ref temporal) = query.temporal {
if let Some(meta) = self.metadata_store.get(&hr.cid) {
let boost = temporal.recency_multiplier(meta.created_at);
hr.score *= boost;
}
}
Some(hr)
})
.collect();
if let Some(min_score) = query.min_score {
results.retain(|r| r.score >= min_score);
}
if query.temporal.is_some() {
results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
}
results.truncate(query.k);
if query.include_metadata {
for result in &mut results {
result.metadata = self.metadata_store.get(&result.cid);
}
}
Ok(results)
}
fn determine_strategy(&self, query: &HybridQuery) -> FilterStrategy {
if self.config.filter_strategy != FilterStrategy::Auto {
return self.config.filter_strategy;
}
let total_count = self.metadata_store.len();
if total_count == 0 {
return FilterStrategy::PostFilter;
}
if query.filter.is_none() && query.temporal.is_none() {
return FilterStrategy::PostFilter;
}
let filtered_count = if let Some(ref filter) = query.filter {
self.metadata_store.filter(filter).len()
} else {
total_count
};
let selectivity = filtered_count as f64 / total_count as f64;
if selectivity < 0.1 {
FilterStrategy::PreFilter
} else {
FilterStrategy::PostFilter
}
}
pub fn len(&self) -> usize {
self.vector_index
.read()
.unwrap_or_else(|e| e.into_inner())
.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn contains(&self, cid: &Cid) -> bool {
self.vector_index
.read()
.unwrap_or_else(|e| e.into_inner())
.contains(cid)
}
pub fn get_metadata(&self, cid: &Cid) -> Option<Metadata> {
self.metadata_store.get(cid)
}
pub fn update_metadata(&self, cid: &Cid, metadata: Metadata) -> Result<()> {
if !self.contains(cid) {
return Err(Error::NotFound(format!("CID not in index: {}", cid)));
}
self.metadata_store.insert(*cid, metadata)?;
Ok(())
}
pub fn stats(&self) -> StatsSnapshot {
self.stats.snapshot()
}
pub fn health(&self) -> IndexHealth {
let stats = self.stats.snapshot();
IndexHealth::analyze(self.len(), self.config.dimension, Some(&stats))
}
pub fn memory_usage(&self) -> MemoryUsage {
MemoryUsage::estimate(
self.len(),
self.config.dimension,
self.metadata_store.len(),
self.config.cache_size,
)
}
pub fn facet_counts(&self, field: &str) -> std::collections::HashMap<String, usize> {
self.metadata_store.get_facet_counts(field)
}
pub fn clear_cache(&self) {
self.cache
.write()
.unwrap_or_else(|e| e.into_inner())
.clear();
}
pub fn reset_stats(&self) {
self.stats.reset();
}
pub async fn save(&self, path: impl AsRef<std::path::Path>) -> Result<()> {
self.vector_index
.read()
.unwrap_or_else(|e| e.into_inner())
.save(path)
}
pub fn clear(&self) -> Result<()> {
let new_index = VectorIndex::new(
self.config.dimension,
self.config.metric,
self.config.max_connections,
self.config.ef_construction,
)?;
*self.vector_index.write().unwrap_or_else(|e| e.into_inner()) = new_index;
self.metadata_store.clear();
self.cache
.write()
.unwrap_or_else(|e| e.into_inner())
.clear();
self.stats.reset();
Ok(())
}
pub fn prune_by_ttl(&self, ttl_seconds: u64) -> Result<usize> {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let cutoff = now.saturating_sub(ttl_seconds);
self.prune_older_than(cutoff)
}
pub fn prune_older_than(&self, timestamp: u64) -> Result<usize> {
let cids_to_remove: Vec<Cid> = self
.metadata_store
.cids()
.into_iter()
.filter(|cid| {
self.metadata_store
.get(cid)
.map(|m| m.created_at < timestamp)
.unwrap_or(false)
})
.collect();
let count = cids_to_remove.len();
for cid in &cids_to_remove {
let _ = self
.vector_index
.write()
.unwrap_or_else(|e| e.into_inner())
.delete(cid);
let _ = self.metadata_store.remove(cid);
}
self.cache
.write()
.unwrap_or_else(|e| e.into_inner())
.clear();
Ok(count)
}
pub fn prune_to_max_entries(&self, max_entries: usize) -> Result<usize> {
let current_count = self.len();
if current_count <= max_entries {
return Ok(0);
}
let mut entries: Vec<(Cid, u64)> = self
.metadata_store
.cids()
.into_iter()
.filter_map(|cid| self.metadata_store.get(&cid).map(|m| (cid, m.created_at)))
.collect();
entries.sort_by_key(|(_, ts)| *ts);
let to_remove = current_count - max_entries;
for (cid, _) in entries.iter().take(to_remove) {
let _ = self
.vector_index
.write()
.unwrap_or_else(|e| e.into_inner())
.delete(cid);
let _ = self.metadata_store.remove(cid);
}
self.cache
.write()
.unwrap_or_else(|e| e.into_inner())
.clear();
Ok(to_remove)
}
pub fn prune_lru(&self, max_entries: usize) -> Result<usize> {
let current_count = self.len();
if current_count <= max_entries {
return Ok(0);
}
let mut entries: Vec<(Cid, u64)> = self
.metadata_store
.cids()
.into_iter()
.filter_map(|cid| self.metadata_store.get(&cid).map(|m| (cid, m.updated_at)))
.collect();
entries.sort_by_key(|(_, ts)| *ts);
let to_remove = current_count - max_entries;
for (cid, _) in entries.iter().take(to_remove) {
let _ = self
.vector_index
.write()
.unwrap_or_else(|e| e.into_inner())
.delete(cid);
let _ = self.metadata_store.remove(cid);
}
self.cache
.write()
.unwrap_or_else(|e| e.into_inner())
.clear();
Ok(to_remove)
}
pub fn pruning_stats(&self) -> PruningStats {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let entries: Vec<(u64, u64)> = self
.metadata_store
.cids()
.into_iter()
.filter_map(|cid| {
self.metadata_store
.get(&cid)
.map(|m| (m.created_at, m.updated_at))
})
.collect();
if entries.is_empty() {
return PruningStats::default();
}
let oldest_created = entries.iter().map(|(c, _)| *c).min().unwrap_or(now);
let newest_created = entries.iter().map(|(c, _)| *c).max().unwrap_or(now);
let oldest_updated = entries.iter().map(|(_, u)| *u).min().unwrap_or(now);
let age_1day = entries.iter().filter(|(c, _)| now - *c < 86400).count();
let age_7days = entries.iter().filter(|(c, _)| now - *c < 86400 * 7).count();
let age_30days = entries
.iter()
.filter(|(c, _)| now - *c < 86400 * 30)
.count();
PruningStats {
total_entries: entries.len(),
oldest_entry_age: now.saturating_sub(oldest_created),
newest_entry_age: now.saturating_sub(newest_created),
oldest_update_age: now.saturating_sub(oldest_updated),
entries_last_day: age_1day,
entries_last_week: age_7days,
entries_last_month: age_30days,
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct PruningStats {
pub total_entries: usize,
pub oldest_entry_age: u64,
pub newest_entry_age: u64,
pub oldest_update_age: u64,
pub entries_last_day: usize,
pub entries_last_week: usize,
pub entries_last_month: usize,
}
impl PruningStats {
pub fn summary(&self) -> String {
format!(
"Total: {}, Last day: {}, Last week: {}, Last month: {}, Oldest: {}s ago",
self.total_entries,
self.entries_last_day,
self.entries_last_week,
self.entries_last_month,
self.oldest_entry_age
)
}
pub fn would_prune_for_ttl(&self, ttl_seconds: u64) -> usize {
if ttl_seconds < 86400 {
self.total_entries - self.entries_last_day
} else if ttl_seconds < 86400 * 7 {
self.total_entries - self.entries_last_week
} else if ttl_seconds < 86400 * 30 {
self.total_entries - self.entries_last_month
} else {
0
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metadata::MetadataValue;
fn test_cid(n: u8) -> Cid {
let cids = [
"bafybeigdyrzt5sfp7udm7hu76uh7y26nf3efuylqabf3oclgtqy55fbzdi",
"bafybeiczsscdsbs7ffqz55asqdf3smv6klcw3gofszvwlyarci47bgf354",
"bafybeibvfkifsqbapirjrj7zbfwddz5qz5awvbftjgktpcqcxjkzstszlm",
];
cids[n as usize % cids.len()]
.parse()
.expect("test: parse test cid")
}
#[tokio::test]
async fn test_hybrid_index_basic() {
let config = HybridConfig {
dimension: 4,
..Default::default()
};
let index = HybridIndex::new(config).expect("test: create hybrid index basic");
let cid1 = test_cid(0);
let vec1 = vec![1.0, 0.0, 0.0, 0.0];
let meta1 = Metadata::new().with_string("type", "image");
index
.insert(&cid1, &vec1, Some(meta1))
.expect("test: insert cid1 basic");
assert_eq!(index.len(), 1);
assert!(index.contains(&cid1));
}
#[tokio::test]
async fn test_hybrid_search() {
let config = HybridConfig {
dimension: 4,
..Default::default()
};
let index = HybridIndex::new(config).expect("test: create hybrid index for search");
let cid1 = test_cid(0);
let vec1 = vec![1.0, 0.0, 0.0, 0.0];
let meta1 = Metadata::new()
.with_string("type", "image")
.with_integer("size", 1024);
let cid2 = test_cid(1);
let vec2 = vec![0.9, 0.1, 0.0, 0.0];
let meta2 = Metadata::new()
.with_string("type", "document")
.with_integer("size", 2048);
let cid3 = test_cid(2);
let vec3 = vec![0.0, 1.0, 0.0, 0.0];
let meta3 = Metadata::new()
.with_string("type", "audio")
.with_integer("size", 512);
index
.insert(&cid1, &vec1, Some(meta1))
.expect("test: insert cid1 for search");
index
.insert(&cid2, &vec2, Some(meta2))
.expect("test: insert cid2 for search");
index
.insert(&cid3, &vec3, Some(meta3))
.expect("test: insert cid3 for search");
let mut query = HybridQuery::knn(vec![1.0, 0.0, 0.0, 0.0], 2);
query.ef_search = Some(50); let response = index.search(query).await.expect("test: hybrid search");
assert!(
!response.results.is_empty(),
"Expected at least 1 result, got {}",
response.results.len()
);
assert!(
!response.results.is_empty() && response.results.len() <= 2,
"Expected 1-2 results, got {}",
response.results.len()
);
assert_eq!(response.results[0].cid, cid1);
}
#[tokio::test]
async fn test_filtered_search() {
let config = HybridConfig {
dimension: 4,
..Default::default()
};
let index =
HybridIndex::new(config).expect("test: create hybrid index for filtered search");
let cid1 = test_cid(0);
let vec1 = vec![1.0, 0.0, 0.0, 0.0];
let meta1 = Metadata::new().with_string("category", "tech");
let cid2 = test_cid(1);
let vec2 = vec![0.9, 0.1, 0.0, 0.0];
let meta2 = Metadata::new().with_string("category", "science");
index
.insert(&cid1, &vec1, Some(meta1))
.expect("test: insert cid1 filtered");
index
.insert(&cid2, &vec2, Some(meta2))
.expect("test: insert cid2 filtered");
let filter = MetadataFilter::eq("category", MetadataValue::String("tech".to_string()));
let query = HybridQuery::knn(vec![0.9, 0.1, 0.0, 0.0], 10).with_filter(filter);
let response = index.search(query).await.expect("test: filtered search");
assert_eq!(response.results.len(), 1);
assert_eq!(response.results[0].cid, cid1);
}
#[tokio::test]
async fn test_search_with_metadata() {
let config = HybridConfig {
dimension: 4,
..Default::default()
};
let index = HybridIndex::new(config).expect("test: create hybrid index with metadata");
let cid1 = test_cid(0);
let vec1 = vec![1.0, 0.0, 0.0, 0.0];
let meta1 = Metadata::new().with_string("title", "Test Document");
index
.insert(&cid1, &vec1, Some(meta1))
.expect("test: insert cid1 with metadata");
let query = HybridQuery::knn(vec![1.0, 0.0, 0.0, 0.0], 1).with_metadata();
let response = index
.search(query)
.await
.expect("test: search with metadata");
assert_eq!(response.results.len(), 1);
assert!(response.results[0].metadata.is_some());
let meta = response.results[0]
.metadata
.as_ref()
.expect("test: result should have metadata");
assert_eq!(
meta.get("title"),
Some(&MetadataValue::String("Test Document".to_string()))
);
}
#[test]
fn test_health_and_stats() {
let config = HybridConfig {
dimension: 4,
..Default::default()
};
let index = HybridIndex::new(config).expect("test: create hybrid index for health stats");
let health = index.health();
assert_eq!(health.size, 0);
let stats = index.stats();
assert_eq!(stats.search_count, 0);
}
#[test]
fn test_pruning_to_max_entries() {
let config = HybridConfig {
dimension: 4,
..Default::default()
};
let index = HybridIndex::new(config).expect("test: create hybrid index for pruning");
for i in 0..3 {
let cid = test_cid(i);
let vec = vec![i as f32, 0.0, 0.0, 0.0];
let meta = Metadata::new().with_integer("order", i as i64);
index
.insert(&cid, &vec, Some(meta))
.expect("test: insert vector for pruning");
}
assert_eq!(index.len(), 3);
let pruned = index
.prune_to_max_entries(2)
.expect("test: prune to max entries");
assert_eq!(pruned, 1);
assert_eq!(index.len(), 2);
}
#[test]
fn test_pruning_stats() {
let config = HybridConfig {
dimension: 4,
..Default::default()
};
let index = HybridIndex::new(config).expect("test: create hybrid index for pruning stats");
for i in 0..3 {
let cid = test_cid(i);
let vec = vec![i as f32, 0.0, 0.0, 0.0];
index
.insert(&cid, &vec, None)
.expect("test: insert vector for pruning stats");
}
let stats = index.pruning_stats();
assert_eq!(stats.total_entries, 3);
assert_eq!(stats.entries_last_day, 3);
assert_eq!(stats.entries_last_week, 3);
}
}