diff --git a/README.md b/README.md index 9382c92f..4b53cb48 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ a low-latency, memory-efficient, distributed cache written in Rust. designed to - **key commands** — DEL, EXISTS, EXPIRE, TTL, PEXPIRE, PTTL, PERSIST, TYPE, SCAN, KEYS, RENAME - **server commands** — PING, ECHO, INFO, DBSIZE, FLUSHDB, BGSAVE, BGREWRITEAOF, AUTH, QUIT - **pub/sub** — SUBSCRIBE, UNSUBSCRIBE, PSUBSCRIBE, PUNSUBSCRIBE, PUBLISH, plus PUBSUB introspection +- **vector similarity search** — HNSW-backed approximate nearest neighbor search with cosine, L2, and inner product metrics (compile with `--features vector`) - **protobuf storage** — schema-validated protobuf values with field-level access (compile with `--features protobuf`) - **authentication** — `--requirepass` for redis-compatible AUTH (legacy and username/password forms) - **tls support** — redis-compatible TLS on a separate port, with optional mTLS for client certificates @@ -144,6 +145,46 @@ cargo build --release --features protobuf field-level operations decode/mutate/re-encode on the server, so clients don't need protobuf libraries for simple reads and writes. nested paths use dot notation (e.g., `address.city`). complex types (repeated, map, nested messages) require `PROTO.GET`/`PROTO.SET` for full replacement. +## vector similarity search + +ember supports HNSW-backed approximate nearest neighbor search for building recommendation systems, semantic search, and RAG pipelines. compile with `--features vector` to enable. + +```bash +# build with vector support +cargo build --release --features vector +``` + +**commands**: + +| command | description | +|---------|-------------| +| `VADD key element f32 [f32 ...] [METRIC COSINE\|L2\|IP] [QUANT F32\|F16\|Q8] [M n] [EF n]` | add a vector to the set | +| `VSIM key f32 [f32 ...] COUNT k [EF n] [WITHSCORES]` | k nearest neighbors | +| `VREM key element` | remove a vector | +| `VGET key element` | retrieve stored vector values | +| `VCARD key` | number of vectors in the set | +| `VDIM key` | dimensionality of the vector set | +| `VINFO key` | metadata: dim, count, metric, quantization, M, ef | + +index configuration (METRIC, QUANT, M, EF) is set on the first VADD and locked after that. dimension is inferred from the first vector's length. each key owns its own independent HNSW index. + +```bash +# store some embeddings +VADD docs doc1 0.1 0.2 0.3 METRIC COSINE +VADD docs doc2 0.9 0.1 0.0 +VADD docs doc3 0.0 0.8 0.2 + +# find 2 nearest neighbors +VSIM docs 0.1 0.3 0.2 COUNT 2 WITHSCORES +# => 1) "doc1" 2) "0.05" 3) "doc3" 4) "0.12" + +VCARD docs # => (integer) 3 +VDIM docs # => (integer) 3 +VINFO docs # => metric, quantization, dim, count, M, ef +``` + +distance metrics: **COSINE** (default), **L2** (squared euclidean), **IP** (inner product). quantization: **F32** (default), **F16** (half precision), **Q8** (8-bit integer). lower precision uses less memory at a small accuracy cost. + ## configuration | flag | default | description | @@ -257,7 +298,7 @@ contributions welcome — see [CONTRIBUTING.md](CONTRIBUTING.md). | 4 | clustering (raft, gossip, slots, migration) | ✅ complete | | 5 | developer experience (observability, CLI, clients) | 🚧 in progress | -**current**: 94 commands, 989 tests, ~21k lines of code (excluding tests) +**current**: 101 commands, 796+ tests, ~22k lines of code (excluding tests) ## security diff --git a/crates/ember-core/src/keyspace.rs b/crates/ember-core/src/keyspace.rs index cf6fedd0..d3e04983 100644 --- a/crates/ember-core/src/keyspace.rs +++ b/crates/ember-core/src/keyspace.rs @@ -1926,17 +1926,17 @@ impl Keyspace { self.entries.insert(key.to_owned(), Entry::new(value, None)); } - let entry = self - .entries - .get_mut(key) - .expect("just inserted or verified"); + let entry = match self.entries.get_mut(key) { + Some(e) => e, + None => return Err(VectorWriteError::IndexError("entry missing".into())), + }; let old_entry_size = memory::entry_size(key, &entry.value); let added = match entry.value { Value::Vector(ref mut vs) => vs .add(element.clone(), &vector) .map_err(|e| VectorWriteError::IndexError(e.to_string()))?, - _ => unreachable!(), + _ => return Err(VectorWriteError::WrongType), }; entry.touch(); diff --git a/crates/ember-core/src/types/vector.rs b/crates/ember-core/src/types/vector.rs index 0d7ac18d..48fdcc51 100644 --- a/crates/ember-core/src/types/vector.rs +++ b/crates/ember-core/src/types/vector.rs @@ -152,6 +152,9 @@ pub enum VectorError { #[error("dimension mismatch: index has {expected}, got {got}")] DimensionMismatch { expected: usize, got: usize }, + #[error("vector contains non-finite value (NaN or infinity)")] + NonFinite, + #[error("usearch error: {0}")] Index(String), } @@ -227,7 +230,7 @@ impl VectorSet { /// Adds or replaces a vector for the given element name. /// /// Returns `true` if a new element was added, `false` if an existing - /// element was updated. + /// element was updated. Rejects vectors containing NaN or infinity. pub fn add(&mut self, element: String, vector: &[f32]) -> Result { if vector.len() != self.dim { return Err(VectorError::DimensionMismatch { @@ -236,6 +239,10 @@ impl VectorSet { }); } + if vector.iter().any(|v| !v.is_finite()) { + return Err(VectorError::NonFinite); + } + // ensure capacity (double when full, amortized O(1)) if self.index.size() >= self.index.capacity() { let new_cap = (self.index.capacity() * 2).max(64); @@ -253,7 +260,7 @@ impl VectorSet { false } else { let key = self.next_key; - self.next_key += 1; + self.next_key = self.next_key.wrapping_add(1); self.index .add(key, vector) .map_err(|e| VectorError::Index(e.to_string()))?; @@ -280,10 +287,13 @@ impl VectorSet { } } + /// Default search beam width. Used when the caller doesn't specify one. + const DEFAULT_EF_SEARCH: usize = 64; + /// Searches for the k nearest neighbors of the given query vector. /// /// `ef_search` controls the search beam width (higher = more accurate, - /// slower). Pass 0 to use usearch's default. + /// slower). Pass 0 to use the default (64). pub fn search( &self, query: &[f32], @@ -297,14 +307,22 @@ impl VectorSet { }); } + if query.iter().any(|v| !v.is_finite()) { + return Err(VectorError::NonFinite); + } + if self.elements.is_empty() { return Ok(Vec::new()); } - // temporarily adjust search expansion if requested - if ef_search > 0 { - self.index.change_expansion_search(ef_search); - } + // always set expansion_search explicitly so a previous call's value + // doesn't leak into this search + let ef = if ef_search > 0 { + ef_search + } else { + Self::DEFAULT_EF_SEARCH + }; + self.index.change_expansion_search(ef); let matches = self .index @@ -427,26 +445,53 @@ impl fmt::Debug for VectorSet { } } -impl Clone for VectorSet { - fn clone(&self) -> Self { - // rebuild the index from scratch — usearch Index doesn't implement Clone +impl VectorSet { + /// Fallible clone — rebuilds the HNSW index from scratch. + /// + /// Usearch's `Index` doesn't implement Clone, so we create a new index + /// with the same config and re-insert every vector. Returns an error + /// if the index can't be created or a vector fails to insert. + pub fn try_clone(&self) -> Result { let mut new = Self::new( self.dim, self.metric, self.quantization, self.connectivity, self.expansion_add, - ) - .expect("clone: failed to create index with same config"); + )?; for (name, &key) in &self.elements { let mut buffer = vec![0.0f32; self.dim]; if self.index.get(key, &mut buffer).is_ok() { - let _ = new.add(name.clone(), &buffer); + new.add(name.clone(), &buffer)?; } } - new + Ok(new) + } +} + +impl Clone for VectorSet { + fn clone(&self) -> Self { + // try_clone rebuilds the index from scratch. if it fails (e.g. OOM + // inside usearch), return an empty set with the same config rather + // than panicking. data loss is preferable to a server crash. + self.try_clone().unwrap_or_else(|_| { + Self::new( + self.dim, + self.metric, + self.quantization, + self.connectivity, + self.expansion_add, + ) + .unwrap_or_else(|_| { + // if we can't even create an empty index with the same + // config, fall back to the smallest possible valid set. + // this only happens under extreme memory pressure. + Self::new(self.dim, self.metric, self.quantization, 2, 2) + .expect("failed to allocate even a minimal index — system is out of memory") + }) + }) } } @@ -639,4 +684,79 @@ mod tests { let vec = vs.get("a").unwrap(); assert!((vec[0] - 1.0).abs() < 0.01); } + + #[test] + fn reject_nan_in_add() { + let mut vs = make_set(3); + let err = vs.add("bad".into(), &[1.0, f32::NAN, 0.0]).unwrap_err(); + assert!(matches!(err, VectorError::NonFinite)); + assert_eq!(vs.len(), 0); + } + + #[test] + fn reject_infinity_in_add() { + let mut vs = make_set(3); + let err = vs + .add("bad".into(), &[f32::INFINITY, 0.0, 0.0]) + .unwrap_err(); + assert!(matches!(err, VectorError::NonFinite)); + } + + #[test] + fn reject_nan_in_search() { + let mut vs = make_set(3); + vs.add("a".into(), &[1.0, 0.0, 0.0]).unwrap(); + let err = vs.search(&[f32::NAN, 0.0, 0.0], 1, 0).unwrap_err(); + assert!(matches!(err, VectorError::NonFinite)); + } + + #[test] + fn search_ef_does_not_leak_between_calls() { + let mut vs = make_set(3); + vs.add("a".into(), &[1.0, 0.0, 0.0]).unwrap(); + vs.add("b".into(), &[0.0, 1.0, 0.0]).unwrap(); + + // first search with high ef + let r1 = vs.search(&[0.9, 0.1, 0.0], 1, 200).unwrap(); + assert_eq!(r1.len(), 1); + + // second search with default ef should still work correctly + let r2 = vs.search(&[0.9, 0.1, 0.0], 1, 0).unwrap(); + assert_eq!(r2.len(), 1); + assert_eq!(r2[0].element, r1[0].element); + } + + #[test] + fn try_clone_returns_ok() { + let mut vs = make_set(3); + vs.add("a".into(), &[1.0, 2.0, 3.0]).unwrap(); + + let cloned = vs.try_clone().unwrap(); + assert_eq!(cloned.len(), 1); + assert_eq!(cloned.get("a").unwrap(), vec![1.0, 2.0, 3.0]); + } + + #[test] + fn u8_round_trip_metric() { + for metric in [ + DistanceMetric::Cosine, + DistanceMetric::L2, + DistanceMetric::InnerProduct, + ] { + let byte: u8 = metric.into(); + assert_eq!(DistanceMetric::from_u8(byte), metric); + } + } + + #[test] + fn u8_round_trip_quantization() { + for quant in [ + QuantizationType::F32, + QuantizationType::F16, + QuantizationType::I8, + ] { + let byte: u8 = quant.into(); + assert_eq!(QuantizationType::from_u8(byte), quant); + } + } } diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 232e886c..794e93f0 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -1531,6 +1531,52 @@ mod tests { assert_eq!(rec, decoded); } + #[cfg(feature = "vector")] + #[test] + fn record_round_trip_vadd() { + let rec = AofRecord::VAdd { + key: "embeddings".into(), + element: "doc1".into(), + vector: vec![0.1, 0.2, 0.3], + metric: 0, // cosine + quantization: 0, // f32 + connectivity: 16, + expansion_add: 64, + }; + let bytes = rec.to_bytes().unwrap(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + + #[cfg(feature = "vector")] + #[test] + fn record_round_trip_vadd_high_dim() { + let rec = AofRecord::VAdd { + key: "vecs".into(), + element: "e".into(), + vector: vec![0.0; 1536], // typical embedding dimension + metric: 1, // l2 + quantization: 1, // f16 + connectivity: 32, + expansion_add: 128, + }; + let bytes = rec.to_bytes().unwrap(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + + #[cfg(feature = "vector")] + #[test] + fn record_round_trip_vrem() { + let rec = AofRecord::VRem { + key: "embeddings".into(), + element: "doc1".into(), + }; + let bytes = rec.to_bytes().unwrap(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + #[cfg(feature = "encryption")] mod encrypted { use super::*; diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index ca845e04..baf0890d 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -958,6 +958,137 @@ mod tests { assert!(result.entries[0].ttl.is_some()); } + #[cfg(feature = "vector")] + #[test] + fn vector_snapshot_recovery() { + let dir = temp_dir(); + let path = snapshot::snapshot_path(dir.path(), 0); + + { + let mut writer = SnapshotWriter::create(&path, 0).unwrap(); + writer + .write_entry(&SnapEntry { + key: "embeddings".into(), + value: SnapValue::Vector { + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + dim: 3, + elements: vec![ + ("doc1".into(), vec![1.0, 0.0, 0.0]), + ("doc2".into(), vec![0.0, 1.0, 0.0]), + ], + }, + expire_ms: -1, + }) + .unwrap(); + writer.finish().unwrap(); + } + + let result = recover_shard(dir.path(), 0); + assert!(result.loaded_snapshot); + assert_eq!(result.entries.len(), 1); + match &result.entries[0].value { + RecoveredValue::Vector { + metric, + quantization, + elements, + .. + } => { + assert_eq!(*metric, 0); + assert_eq!(*quantization, 0); + assert_eq!(elements.len(), 2); + // dim is inferred from the vector length + assert_eq!(elements[0].1.len(), 3); + } + other => panic!("expected Vector, got {other:?}"), + } + } + + #[cfg(feature = "vector")] + #[test] + fn vector_aof_replay() { + let dir = temp_dir(); + let path = aof::aof_path(dir.path(), 0); + + { + let mut writer = AofWriter::open(&path).unwrap(); + writer + .write_record(&AofRecord::VAdd { + key: "vecs".into(), + element: "a".into(), + vector: vec![1.0, 0.0, 0.0], + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + }) + .unwrap(); + writer + .write_record(&AofRecord::VAdd { + key: "vecs".into(), + element: "b".into(), + vector: vec![0.0, 1.0, 0.0], + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + }) + .unwrap(); + writer + .write_record(&AofRecord::VRem { + key: "vecs".into(), + element: "a".into(), + }) + .unwrap(); + writer.sync().unwrap(); + } + + let result = recover_shard(dir.path(), 0); + assert!(result.replayed_aof); + assert_eq!(result.entries.len(), 1); + match &result.entries[0].value { + RecoveredValue::Vector { elements, .. } => { + assert_eq!(elements.len(), 1); + assert_eq!(elements[0].0, "b"); + } + other => panic!("expected Vector, got {other:?}"), + } + } + + #[cfg(feature = "vector")] + #[test] + fn vector_vrem_auto_deletes_empty() { + let dir = temp_dir(); + let path = aof::aof_path(dir.path(), 0); + + { + let mut writer = AofWriter::open(&path).unwrap(); + writer + .write_record(&AofRecord::VAdd { + key: "vecs".into(), + element: "only".into(), + vector: vec![1.0, 2.0], + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + }) + .unwrap(); + writer + .write_record(&AofRecord::VRem { + key: "vecs".into(), + element: "only".into(), + }) + .unwrap(); + writer.sync().unwrap(); + } + + let result = recover_shard(dir.path(), 0); + assert!(result.entries.is_empty()); + } + #[cfg(feature = "protobuf")] #[test] fn proto_schemas_recovered_from_aof() { diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index 41a923c9..9c840421 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -1021,6 +1021,78 @@ mod tests { ); } + #[cfg(feature = "vector")] + #[test] + fn vector_entries_round_trip() { + let dir = temp_dir(); + let path = dir.path().join("vec.snap"); + + let entries = vec![SnapEntry { + key: "embeddings".into(), + value: SnapValue::Vector { + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + dim: 3, + elements: vec![ + ("doc1".into(), vec![0.1, 0.2, 0.3]), + ("doc2".into(), vec![0.4, 0.5, 0.6]), + ], + }, + expire_ms: -1, + }]; + + { + let mut writer = SnapshotWriter::create(&path, 0).unwrap(); + for entry in &entries { + writer.write_entry(entry).unwrap(); + } + writer.finish().unwrap(); + } + + let mut reader = SnapshotReader::open(&path).unwrap(); + let mut got = Vec::new(); + while let Some(entry) = reader.read_entry().unwrap() { + got.push(entry); + } + assert_eq!(entries, got); + reader.verify_footer().unwrap(); + } + + #[cfg(feature = "vector")] + #[test] + fn vector_empty_set_round_trip() { + let dir = temp_dir(); + let path = dir.path().join("vec_empty.snap"); + + let entries = vec![SnapEntry { + key: "empty_vecs".into(), + value: SnapValue::Vector { + metric: 2, // inner product + quantization: 2, + connectivity: 8, + expansion_add: 32, + dim: 128, + elements: vec![], + }, + expire_ms: 5000, + }]; + + { + let mut writer = SnapshotWriter::create(&path, 0).unwrap(); + for entry in &entries { + writer.write_entry(entry).unwrap(); + } + writer.finish().unwrap(); + } + + let mut reader = SnapshotReader::open(&path).unwrap(); + let got = reader.read_entry().unwrap().unwrap(); + assert_eq!(entries[0], got); + reader.verify_footer().unwrap(); + } + #[cfg(feature = "encryption")] mod encrypted { use super::*;