diff --git a/Cargo.lock b/Cargo.lock index 3471623..238f72d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -53,6 +53,8 @@ version = "0.1.0" dependencies = [ "defs", "index", + "serde", + "serde_json", "snafu", "snapshot", "storage", @@ -490,6 +492,7 @@ version = "0.1.0" dependencies = [ "axum", "serde", + "serde_json", "snafu", "uuid", ] @@ -2449,6 +2452,7 @@ dependencies = [ "flate2", "rocksdb", "serde", + "serde_json", "snafu", "tar", "tempfile", diff --git a/crates/api/Cargo.toml b/crates/api/Cargo.toml index 7c198a3..08756c8 100644 --- a/crates/api/Cargo.toml +++ b/crates/api/Cargo.toml @@ -9,6 +9,8 @@ license.workspace = true [dependencies] defs.workspace = true index.workspace = true +serde.workspace = true +serde_json.workspace = true snafu.workspace = true snapshot.workspace = true storage.workspace = true diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index 5377dde..5ff22a5 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -1,7 +1,10 @@ -use defs::{DbError, Dimension, IndexedVector, SearchQueryInput, Similarity, SnapshottableDb}; +use defs::{ + DbError, Dimension, FilterMode, IndexedVector, SearchQueryInput, Similarity, SnapshottableDb, +}; use defs::{DenseVector, Payload, Point, PointId, PointInput}; use index::hnsw::{HnswConfig, HnswIndex}; use index::kd_tree::{KDTree, KDTreeConfig}; +use std::iter::FilterMap; use std::path::{Path, PathBuf}; use tempfile::tempdir; // use std::sync::atomic::{AtomicU64, Ordering}; @@ -142,9 +145,12 @@ impl VectorDb { // Use vector index to find similar vectors let index = self.index.read().map_err(|_| ApiError::LockError)?; - let vectors = + let mut vectors = index.search_with_ef(query.vector, query.similarity, query.limit, query.ef)?; + if query.filter_mode == FilterMode::Post { + vectors = self.filter_ids(vectors, query.filter.as_ref())?; + } Ok(vectors) } @@ -153,8 +159,11 @@ impl VectorDb { let index = self.index.read().unwrap(); for query in queries { - let found = + let mut found = index.search_with_ef(query.vector, query.similarity, query.limit, query.ef)?; + if query.filter_mode == FilterMode::Post { + found = self.filter_ids(found, query.filter.as_ref())?; + } results.push(found); } @@ -190,6 +199,30 @@ impl VectorDb { Ok(inserted) } + + fn payload_matches(&self, id: PointId, filter: &Filter) -> Result { + Ok(self + .storage + .get_payload(id)? + .is_some_and(|p| filter.matches(&p.metadata))) + } + + fn filter_ids( + &self, + ids: impl IntoIterator, + filter: Option<&Filter>, + ) -> Result> { + let Some(f) = filter else { + return Ok(ids.into_iter().collect()); + }; + let mut out = Vec::new(); + for id in ids { + if self.payload_matches(id.0, f)? { + out.push(id); + } + } + Ok(out) + } } impl SnapshottableDb for VectorDb { diff --git a/crates/api/src/tests.rs b/crates/api/src/tests.rs index f99d77b..94b1891 100644 --- a/crates/api/src/tests.rs +++ b/crates/api/src/tests.rs @@ -2,10 +2,10 @@ use super::*; // TODO: Add more exhaustive tests -use std::sync::Mutex; - -use defs::ContentType; +use defs::{ContentType, FilterMode::Post, Metadata}; +use serde_json::{Value, json}; use snapshot::{engine::SnapshotEngine, registry::local::LocalRegistry}; +use std::sync::Mutex; use tempfile::{TempDir, tempdir}; // Helper function to create a test database @@ -27,13 +27,18 @@ fn create_test_db_with_storage(storage_type: StorageType) -> (VectorDb, TempDir) (init_api(config).unwrap(), temp_dir) } -fn test_payload(content: &str) -> Payload { +fn test_payload(content: &str, metadata: Metadata) -> Payload { Payload { content_type: ContentType::Text, content: content.to_string(), + metadata: metadata, } } +fn meta(v: serde_json::Value) -> Metadata { + v.as_object().unwrap().clone() +} + #[test] fn test_insert_and_get() { let (db, _temp_dir) = create_test_db(); @@ -41,6 +46,7 @@ fn test_insert_and_get() { let payload = Payload { content_type: ContentType::Text, content: "Test content".to_string(), + metadata: meta(json!({"age": 23})), }; // Test insert @@ -63,7 +69,7 @@ fn test_insert_and_get() { fn test_insert_and_get_with_in_memory_storage() { let (db, _temp_dir) = create_test_db_with_storage(StorageType::InMemory); let vector = vec![1.0, 2.0, 3.0]; - let payload = test_payload("Test content"); + let payload = test_payload("Test content", meta(json!({"age":12,"height":121}))); let id = db.insert(vector.clone(), payload.clone()).unwrap(); let point = db.get(id).unwrap().unwrap(); @@ -81,6 +87,7 @@ fn test_dimension_mismatch() { let payload = defs::Payload { content_type: ContentType::Text, content: "tester".to_string(), + metadata: meta(json!({"age": 23})), }; let res1 = db.insert(v1, payload.clone()); @@ -105,6 +112,7 @@ fn test_delete() { let payload = Payload { content_type: ContentType::Text, content: "Test content".to_string(), + metadata: meta(json!({"age":23})), }; // Insert a point @@ -137,6 +145,7 @@ fn test_search() { let payload = Payload { content_type: ContentType::Text, content: format!("Test content {vector:?}"), + metadata: meta(json!({"age":23})), }; let id = db.insert(vector, payload).unwrap(); ids.push(id); @@ -150,6 +159,8 @@ fn test_search() { similarity: Similarity::Cosine, limit: 1, ef: None, + filter: None, + filter_mode: Post, }) .unwrap(); @@ -173,6 +184,7 @@ fn test_search_limit() { Payload { content_type: ContentType::Text, content: format!("Test content {i}"), + metadata: meta(json!({"age":23})), }, ) .unwrap(); @@ -187,6 +199,8 @@ fn test_search_limit() { similarity: Similarity::Euclidean, limit: 3, ef: None, + filter: None, + filter_mode: Post, }) .unwrap(); @@ -203,6 +217,8 @@ fn test_search_zero_limit() { similarity: Similarity::Cosine, limit: 0, ef: None, + filter: None, + filter_mode: Post, }); assert!(result.is_err()); @@ -228,6 +244,8 @@ fn test_empty_database() { similarity: Similarity::Cosine, limit: 10, ef: None, + filter: None, + filter_mode: Post, }) .unwrap(); assert_eq!(results.len(), 0); @@ -247,6 +265,7 @@ fn test_list_vectors() { Payload { content_type: ContentType::Text, content: format!("Test content {i}"), + metadata: meta(json!({"age":23})), }, ) .unwrap(); @@ -277,6 +296,7 @@ fn test_build_index() { Payload { content_type: ContentType::Text, content: format!("Test content {i}"), + metadata: meta(json!({"age":23})), }, ) .unwrap(); @@ -301,6 +321,7 @@ fn test_create_and_load_snapshot() { Payload { content_type: ContentType::Text, content: "test".to_string(), + metadata: meta(json!({"age":23})), }, ) .unwrap(); @@ -311,6 +332,7 @@ fn test_create_and_load_snapshot() { Payload { content_type: ContentType::Text, content: "test".to_string(), + metadata: meta(json!({"age":23})), }, ) .unwrap(); @@ -325,6 +347,7 @@ fn test_create_and_load_snapshot() { Payload { content_type: ContentType::Text, content: "test".to_string(), + metadata: meta(json!({"age":23})), }, ) .unwrap(); @@ -354,13 +377,19 @@ fn test_create_and_load_snapshot_with_in_memory_storage() { let v2 = vec![3.0, 4.0, 5.0]; let v3 = vec![6.0, 7.0, 8.0]; - let id1 = old_db.insert(v1.clone(), test_payload("one")).unwrap(); - let id2 = old_db.insert(v2.clone(), test_payload("two")).unwrap(); + let id1 = old_db + .insert(v1.clone(), test_payload("one", meta(json!({"age": 12})))) + .unwrap(); + let id2 = old_db + .insert(v2.clone(), test_payload("two", meta(json!({"age": 32})))) + .unwrap(); let temp_snapshot_dir = tempdir().unwrap(); let snapshot_path = old_db.create_snapshot(temp_snapshot_dir.path()).unwrap(); - let id3 = old_db.insert(v3, test_payload("three")).unwrap(); + let id3 = old_db + .insert(v3, test_payload("three", meta(json!({"age": 22})))) + .unwrap(); let reload_config = DbRestoreConfig { data_path: temp_dir.path().to_path_buf(), @@ -405,6 +434,7 @@ fn test_snapshot_engine() { Payload { content_type: ContentType::Text, content: format!("{}", i), + metadata: meta(json!({"age":23})), }, ) .unwrap(); diff --git a/crates/defs/Cargo.toml b/crates/defs/Cargo.toml index 81d8dd9..c388b57 100644 --- a/crates/defs/Cargo.toml +++ b/crates/defs/Cargo.toml @@ -9,5 +9,6 @@ license.workspace = true [dependencies] axum.workspace = true serde.workspace = true +serde_json.workspace = true snafu.workspace = true uuid.workspace = true diff --git a/crates/defs/src/types.rs b/crates/defs/src/types.rs index 7e67399..947e3a7 100644 --- a/crates/defs/src/types.rs +++ b/crates/defs/src/types.rs @@ -11,6 +11,8 @@ pub type ElementByte = u8; pub type Dimension = usize; +pub type Metadata = serde_json::Map; + // Dense Vector and Vector are considered same // Sparse vector implementation not supported yet. Refer lib/sparse/src/common/sparse_vector.rs pub type DenseVector = Vec; @@ -32,6 +34,8 @@ pub enum ContentType { pub struct Payload { pub content_type: ContentType, pub content: String, + #[serde(default)] + pub metadata: Metadata, } #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] @@ -48,6 +52,43 @@ pub struct Point { pub payload: Option, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum FilterOp { + Eq, + Neq, + Lt, + Leq, + Gt, + Gteq, + In, + Exists, + NotExists, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FieldCondition { + pub key: String, + pub op: FilterOp, + #[serde(default)] + pub value: Option, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum FilterMode { + #[default] + Post, + Inline, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum Filter { + Condition(FieldCondition), + And(Vec), + Or(Vec), + Not(Box), + Xor(Vec), +} + /// Struct which will be stored in the vector index #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] pub struct IndexedVector { @@ -86,7 +127,11 @@ pub struct SearchQueryInput { pub similarity: Similarity, pub limit: usize, #[serde(default)] - pub ef: Option, + pub ef: Option, // adding 2 terms in this Struct, filter and method of filteration + #[serde(default)] + pub filter: Option, + #[serde(default)] + pub filter_mode: FilterMode, } #[derive(Clone, Serialize, Deserialize, Debug)] @@ -139,6 +184,32 @@ impl<'q> Eq for DistanceOrderedVector<'q> {} // Context(ContextQuery), // } +impl Filter { + pub fn matches(&self, metadata: &Metadata) -> bool { + match self { + Filter::Condition(c) => c.matches(metadata), + Filter::And(fs) => fs.iter().all(|f| f.matches(metadata)), + Filter::Or(fs) => fs.iter().any(|f| f.matches(metadata)), + Filter::Not(f) => !f.matches(metadata), + Filter::Xor(fs) => fs.iter().filter(|f| f.matches(metadata)).count() % 2 == 1, + } + } +} + +impl FieldCondition { + fn matches(&self, metadata: &Metadata) -> bool { + let found = resolve_path(metadata, &self.key); + match self.op { + FilterOp::Exists => found.is_some(), + FilterOp::NotExists => found.is_none(), + _ => match (found, self.value.as_ref()) { + (Some(lhs), Some(rhs)) => compare_op(&self.op, lhs, rhs), + _ => false, + }, + } + } +} + #[derive(Clone, Copy, Debug, PartialEq, Default)] pub struct OrdF32(pub f32); @@ -170,3 +241,50 @@ impl From for OrdF32 { Self(x) } } + +fn resolve_path<'a>(metadata: &'a Metadata, key: &str) -> Option<&'a serde_json::Value> { + let mut segments = key.split('.'); + let first = segments.next()?; + let mut cur = metadata.get(first)?; + for seg in segments { + cur = cur.get(seg)?; + } + Some(cur) +} + +fn value_eq(a: &serde_json::Value, b: &serde_json::Value) -> bool { + match (a, b) { + (serde_json::Value::Number(x), serde_json::Value::Number(y)) => { + matches!((x.as_f64(), y.as_f64()), (Some(x), Some(y)) if x == y) + } + _ => a == b, + } +} + +fn ordered_cmp(a: &serde_json::Value, b: &serde_json::Value) -> Option { + match (a, b) { + (serde_json::Value::Number(x), serde_json::Value::Number(y)) => { + x.as_f64()?.partial_cmp(&y.as_f64()?) + } + (serde_json::Value::String(x), serde_json::Value::String(y)) => Some(x.cmp(y)), + (serde_json::Value::Bool(x), serde_json::Value::Bool(y)) => Some(x.cmp(y)), + _ => None, + } +} + +fn compare_op(op: &FilterOp, a: &serde_json::Value, b: &serde_json::Value) -> bool { + match op { + FilterOp::Eq => value_eq(a, b), + FilterOp::Neq => !value_eq(a, b), + FilterOp::Lt | FilterOp::Leq | FilterOp::Gt | FilterOp::Gteq => match ordered_cmp(a, b) { + Some(Ordering::Less) => matches!(op, FilterOp::Lt | FilterOp::Leq), + Some(Ordering::Equal) => matches!(op, FilterOp::Leq | FilterOp::Gteq), + Some(Ordering::Greater) => matches!(op, FilterOp::Gt | FilterOp::Gteq), + None => false, + }, + FilterOp::In => { + matches!(b, serde_json::Value::Array(xs) if xs.iter().any(|i| value_eq(i, a))) + } + FilterOp::Exists | FilterOp::NotExists => true, + } +} diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml index 9fe9322..16d0183 100644 --- a/crates/storage/Cargo.toml +++ b/crates/storage/Cargo.toml @@ -12,6 +12,7 @@ defs.workspace = true flate2 = "1.1.5" rocksdb.workspace = true serde.workspace = true +serde_json.workspace = true snafu.workspace = true tar = "0.4.44" tempfile.workspace = true diff --git a/crates/storage/src/in_memory/tests.rs b/crates/storage/src/in_memory/tests.rs index 9d8510e..4555a28 100644 --- a/crates/storage/src/in_memory/tests.rs +++ b/crates/storage/src/in_memory/tests.rs @@ -1,6 +1,7 @@ use super::*; -use defs::ContentType; +use defs::{ContentType, Metadata}; +use serde_json::{Value, json}; use std::io::{Read, Write}; use tempfile::{TempDir, tempdir}; use uuid::Uuid; @@ -9,19 +10,24 @@ fn create_test_storage() -> MemoryStorage { MemoryStorage::new() } -fn test_payload(content: &str) -> Payload { +fn test_payload(content: &str, metadata: Metadata) -> Payload { Payload { content_type: ContentType::Text, content: content.to_string(), + metadata: meta(json!({"age": 12})), } } +fn meta(v: serde_json::Value) -> Metadata { + v.as_object().unwrap().clone() +} + #[test] fn test_insert_and_get_vector() { let storage = create_test_storage(); let id = Uuid::new_v4(); let vector = Some(vec![0.1, 0.2, 0.3]); - let payload = Some(test_payload("Test")); + let payload = Some(test_payload("Test", meta(json!({"age": 12})))); storage.insert_point(id, vector.clone(), payload).unwrap(); @@ -32,7 +38,7 @@ fn test_insert_and_get_vector() { fn test_insert_and_get_payload() { let storage = create_test_storage(); let id = Uuid::new_v4(); - let payload = Some(test_payload("Test")); + let payload = Some(test_payload("Test", meta(json!({"age": 12})))); storage.insert_point(id, None, payload.clone()).unwrap(); @@ -47,7 +53,11 @@ fn test_contains_and_delete_point() { assert!(!storage.contains_point(id).unwrap()); storage - .insert_point(id, Some(vec![0.4, 0.5, 0.6]), Some(test_payload("Test"))) + .insert_point( + id, + Some(vec![0.4, 0.5, 0.6]), + Some(test_payload("Test", meta(json!({"age": 12})))), + ) .unwrap(); assert!(storage.contains_point(id).unwrap()); @@ -68,16 +78,41 @@ fn test_list_vectors_respects_offset_limit_and_skips_payload_only_points() { ]; storage - .insert_point(ids[0], Some(vec![1.0, 1.1]), Some(test_payload("one"))) + .insert_point( + ids[0], + Some(vec![1.0, 1.1]), + Some(test_payload("one", meta(json!({"age": 12, "height": 123})))), + ) .unwrap(); storage - .insert_point(ids[1], None, Some(test_payload("payload-only"))) + .insert_point( + ids[1], + None, + Some(test_payload( + "payload-only", + meta(json!({"age": 21, "height": 139})), + )), + ) .unwrap(); storage - .insert_point(ids[2], Some(vec![3.0, 3.1]), Some(test_payload("three"))) + .insert_point( + ids[2], + Some(vec![3.0, 3.1]), + Some(test_payload( + "three", + meta(json!({"age": 49, "height": 154})), + )), + ) .unwrap(); storage - .insert_point(ids[3], Some(vec![4.0, 4.1]), Some(test_payload("four"))) + .insert_point( + ids[3], + Some(vec![4.0, 4.1]), + Some(test_payload( + "four", + meta(json!({"age": 10, "height": 190})), + )), + ) .unwrap(); let (first_page, next_offset) = storage.list_vectors(Uuid::nil(), 2).unwrap().unwrap(); @@ -110,7 +145,10 @@ fn test_create_and_restore_checkpoint() { .insert_point( id_before_checkpoint, Some(vec![0.1, 0.2, 0.3]), - Some(test_payload("before")), + Some(test_payload( + "before", + meta(json!({"age": 10, "height": 190})), + )), ) .unwrap(); let checkpoint = storage.checkpoint_at(temp_dir.path()).unwrap(); @@ -119,7 +157,10 @@ fn test_create_and_restore_checkpoint() { .insert_point( id_after_checkpoint, Some(vec![0.4, 0.5, 0.6]), - Some(test_payload("after")), + Some(test_payload( + "after", + meta(json!({"age": 10, "height": 190})), + )), ) .unwrap(); @@ -129,7 +170,10 @@ fn test_create_and_restore_checkpoint() { assert!(!storage.contains_point(id_after_checkpoint).unwrap()); assert_eq!( storage.get_payload(id_before_checkpoint).unwrap(), - Some(test_payload("before")) + Some(test_payload( + "before", + meta(json!({"age": 10, "height": 190})) + )) ); } @@ -140,7 +184,11 @@ fn test_checkpoint_writes_header() { let id = Uuid::new_v4(); storage - .insert_point(id, Some(vec![0.1, 0.2, 0.3]), Some(test_payload("point"))) + .insert_point( + id, + Some(vec![0.1, 0.2, 0.3]), + Some(test_payload("point", meta(json!({"age": 12})))), + ) .unwrap(); let checkpoint = storage.checkpoint_at(temp_dir.path()).unwrap(); diff --git a/crates/storage/src/rocks_db/tests.rs b/crates/storage/src/rocks_db/tests.rs index 4673f2f..a62718c 100644 --- a/crates/storage/src/rocks_db/tests.rs +++ b/crates/storage/src/rocks_db/tests.rs @@ -1,6 +1,7 @@ use super::*; use defs::ContentType; +use serde_json::{Value, json}; use uuid::Uuid; use tempfile::{TempDir, tempdir}; @@ -12,6 +13,10 @@ fn create_test_db() -> (RocksDbStorage, TempDir) { (db, temp_dir) } +fn meta(v: serde_json::Value) -> Metadata { + v.as_object().unwrap().clone() +} + #[test] fn test_new_rocksdb_storage() { let (db, temp_dir) = create_test_db(); @@ -26,6 +31,7 @@ fn test_insert_and_get_vector() { let payload = Some(Payload { content_type: ContentType::Text, content: "Test".to_string(), + metadata: meta(json!({"age": 12})), }); assert!(db.insert_point(id, vector.clone(), payload).is_ok()); @@ -40,6 +46,7 @@ fn test_insert_and_get_payload() { let payload = Some(Payload { content_type: ContentType::Text, content: "Test".to_string(), + metadata: meta(json!({"age": 12})), }); let vector = None; @@ -49,6 +56,7 @@ fn test_insert_and_get_payload() { let expected = Some(Payload { content_type: ContentType::Text, content: "Test".to_string(), + metadata: meta(json!({"age": 12})), }); assert_eq!(result, expected); } @@ -60,6 +68,7 @@ fn test_contains_point() { let payload = Some(Payload { content_type: ContentType::Text, content: "Test".to_string(), + metadata: meta(json!({"age": 12})), }); assert!(!db.contains_point(id).unwrap()); @@ -77,6 +86,7 @@ fn test_delete_point() { let payload = Some(Payload { content_type: ContentType::Text, content: "Test".to_string(), + metadata: meta(json!({"age": 12})), }); let vector = vec![0.7, 0.8, 0.9]; @@ -138,6 +148,7 @@ fn test_create_and_load_checkpoint() { let payload = Some(Payload { content_type: ContentType::Text, content: "Test".to_string(), + metadata: meta(json!({"age": 12})), }); assert!(