reddb_server/storage/unified/
index.rs

1//! Integrated Index System
2//!
3//! Combines multiple index types for optimal multi-modal queries:
4//! - HNSW: Vector similarity search
5//! - B-tree: Metadata range queries (via MetadataStorage)
6//! - Inverted: Full-text search on content
7//!
8//! # Architecture
9//!
10//! ```text
11//! ┌─────────────────────────────────────────────────────────────┐
12//! │                   IntegratedIndexManager                     │
13//! ├─────────────────────────────────────────────────────────────┤
14//! │  ┌───────────┐  ┌───────────────┐  ┌─────────────────────┐  │
15//! │  │   HNSW    │  │  InvertedIndex│  │  MetadataIndex      │  │
16//! │  │  (vector) │  │  (full-text)  │  │  (B-tree ranges)    │  │
17//! │  └─────┬─────┘  └───────┬───────┘  └──────────┬──────────┘  │
18//! │        │                │                      │             │
19//! │        └────────────────┼──────────────────────┘             │
20//! │                         │                                    │
21//! │              ┌──────────▼──────────┐                        │
22//! │              │   Unified Query     │                        │
23//! │              │   Executor          │                        │
24//! │              └─────────────────────┘                        │
25//! └─────────────────────────────────────────────────────────────┘
26//! ```
27
28use std::collections::{BTreeMap, HashMap, HashSet};
29use std::time::{SystemTime, UNIX_EPOCH};
30
31use parking_lot::RwLock;
32
33use super::entity::EntityId;
34use super::metadata::{Metadata, MetadataStorage, MetadataValue};
35
36// ============================================================================
37// Graph Adjacency Index
38// ============================================================================
39
40/// Edge direction for adjacency lookups
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
42pub enum EdgeDirection {
43    Outgoing,
44    Incoming,
45    Both,
46}
47
48/// An adjacency entry representing a single edge
49#[derive(Debug, Clone)]
50pub struct AdjacencyEntry {
51    /// The edge entity ID
52    pub edge_id: EntityId,
53    /// The target node (for outgoing) or source node (for incoming)
54    pub neighbor_id: EntityId,
55    /// Edge label/type
56    pub label: String,
57    /// Edge weight (optional, default 1.0)
58    pub weight: f32,
59}
60
61/// Graph adjacency index for fast edge/neighbor lookups
62///
63/// Supports O(1) lookups for:
64/// - All outgoing edges from a node
65/// - All incoming edges to a node
66/// - Edges filtered by label
67pub struct GraphAdjacencyIndex {
68    /// Node → outgoing edges: HashMap<source_id, Vec<AdjacencyEntry>>
69    outgoing: RwLock<HashMap<EntityId, Vec<AdjacencyEntry>>>,
70    /// Node → incoming edges: HashMap<target_id, Vec<AdjacencyEntry>>
71    incoming: RwLock<HashMap<EntityId, Vec<AdjacencyEntry>>>,
72    /// Label → edge IDs for label-based filtering
73    by_label: RwLock<HashMap<String, HashSet<EntityId>>>,
74    /// Edge count for stats
75    edge_count: RwLock<usize>,
76    /// Node count (unique nodes with edges)
77    node_count: RwLock<usize>,
78}
79
80impl GraphAdjacencyIndex {
81    /// Create a new empty adjacency index
82    pub fn new() -> Self {
83        Self {
84            outgoing: RwLock::new(HashMap::new()),
85            incoming: RwLock::new(HashMap::new()),
86            by_label: RwLock::new(HashMap::new()),
87            edge_count: RwLock::new(0),
88            node_count: RwLock::new(0),
89        }
90    }
91
92    /// Index an edge for fast lookups
93    pub fn index_edge(
94        &self,
95        edge_id: EntityId,
96        source_id: EntityId,
97        target_id: EntityId,
98        label: &str,
99        weight: f32,
100    ) {
101        // Add to outgoing adjacency
102        {
103            let mut outgoing = self.outgoing.write();
104            let entry = AdjacencyEntry {
105                edge_id,
106                neighbor_id: target_id,
107                label: label.to_string(),
108                weight,
109            };
110            outgoing.entry(source_id).or_default().push(entry);
111        }
112
113        // Add to incoming adjacency
114        {
115            let mut incoming = self.incoming.write();
116            let entry = AdjacencyEntry {
117                edge_id,
118                neighbor_id: source_id,
119                label: label.to_string(),
120                weight,
121            };
122            incoming.entry(target_id).or_default().push(entry);
123        }
124
125        // Add to label index
126        {
127            let mut by_label = self.by_label.write();
128            by_label
129                .entry(label.to_string())
130                .or_default()
131                .insert(edge_id);
132        }
133
134        // Update counts
135        {
136            let mut count = self.edge_count.write();
137            *count += 1;
138        }
139
140        // Update node count (track unique nodes)
141        self.update_node_count();
142    }
143
144    /// Remove an edge from the index
145    pub fn remove_edge(&self, edge_id: EntityId) {
146        // Remove from outgoing
147        {
148            let mut outgoing = self.outgoing.write();
149            for entries in outgoing.values_mut() {
150                entries.retain(|e| e.edge_id != edge_id);
151            }
152        }
153
154        // Remove from incoming
155        {
156            let mut incoming = self.incoming.write();
157            for entries in incoming.values_mut() {
158                entries.retain(|e| e.edge_id != edge_id);
159            }
160        }
161
162        // Remove from label index
163        {
164            let mut by_label = self.by_label.write();
165            for edges in by_label.values_mut() {
166                edges.remove(&edge_id);
167            }
168        }
169
170        // Update counts
171        {
172            let mut count = self.edge_count.write();
173            *count = count.saturating_sub(1);
174        }
175    }
176
177    /// Get neighbors in a direction (optionally filtered by label)
178    pub fn get_neighbors(
179        &self,
180        node_id: EntityId,
181        direction: EdgeDirection,
182        label_filter: Option<&str>,
183    ) -> Vec<AdjacencyEntry> {
184        let mut results = Vec::new();
185
186        if matches!(direction, EdgeDirection::Outgoing | EdgeDirection::Both) {
187            let outgoing = self.outgoing.read();
188            if let Some(entries) = outgoing.get(&node_id) {
189                for entry in entries {
190                    if label_filter.is_none_or(|l| entry.label == l) {
191                        results.push(entry.clone());
192                    }
193                }
194            }
195        }
196
197        if matches!(direction, EdgeDirection::Incoming | EdgeDirection::Both) {
198            let incoming = self.incoming.read();
199            if let Some(entries) = incoming.get(&node_id) {
200                for entry in entries {
201                    if label_filter.is_none_or(|l| entry.label == l) {
202                        results.push(entry.clone());
203                    }
204                }
205            }
206        }
207
208        results
209    }
210
211    /// Get all edges with a specific label
212    pub fn get_edges_by_label(&self, label: &str) -> Vec<EntityId> {
213        let idx = self.by_label.read();
214        idx.get(label)
215            .map(|s| s.iter().copied().collect())
216            .unwrap_or_default()
217    }
218
219    /// Get outgoing degree of a node
220    pub fn out_degree(&self, node_id: EntityId) -> usize {
221        self.outgoing
222            .read()
223            .get(&node_id)
224            .map(|v| v.len())
225            .unwrap_or(0)
226    }
227
228    /// Get incoming degree of a node
229    pub fn in_degree(&self, node_id: EntityId) -> usize {
230        self.incoming
231            .read()
232            .get(&node_id)
233            .map(|v| v.len())
234            .unwrap_or(0)
235    }
236
237    /// Get total degree of a node
238    pub fn degree(&self, node_id: EntityId) -> usize {
239        self.out_degree(node_id) + self.in_degree(node_id)
240    }
241
242    /// Get edge count
243    pub fn edge_count(&self) -> usize {
244        *self.edge_count.read()
245    }
246
247    /// Get node count
248    pub fn node_count(&self) -> usize {
249        *self.node_count.read()
250    }
251
252    /// Clear the entire index
253    pub fn clear(&self) {
254        self.outgoing.write().clear();
255        self.incoming.write().clear();
256        self.by_label.write().clear();
257        *self.edge_count.write() = 0;
258        *self.node_count.write() = 0;
259    }
260
261    fn update_node_count(&self) {
262        let out_nodes: HashSet<_> = self.outgoing.read().keys().copied().collect();
263        let in_nodes: HashSet<_> = self.incoming.read().keys().copied().collect();
264
265        let total: HashSet<_> = out_nodes.union(&in_nodes).collect();
266        *self.node_count.write() = total.len();
267    }
268}
269
270impl Default for GraphAdjacencyIndex {
271    fn default() -> Self {
272        Self::new()
273    }
274}
275
276// ============================================================================
277// Inverted Index for Full-Text Search
278// ============================================================================
279
280/// Token with position information for phrase queries
281#[derive(Debug, Clone)]
282pub struct TokenPosition {
283    pub entity_id: EntityId,
284    pub field: String,
285    pub position: u32,
286}
287
288/// A posting list entry
289#[derive(Debug, Clone)]
290pub struct PostingEntry {
291    pub entity_id: EntityId,
292    pub collection: String,
293    pub field: String,
294    pub positions: Vec<u32>,
295    pub term_frequency: f32,
296}
297
298/// Inverted index for full-text search
299pub struct InvertedIndex {
300    /// Term → Posting list
301    index: RwLock<BTreeMap<String, Vec<PostingEntry>>>,
302    /// Document frequencies for TF-IDF
303    doc_count: RwLock<usize>,
304    /// Indexed fields per collection
305    indexed_fields: RwLock<HashMap<String, HashSet<String>>>,
306}
307
308impl InvertedIndex {
309    /// Create a new empty inverted index
310    pub fn new() -> Self {
311        Self {
312            index: RwLock::new(BTreeMap::new()),
313            doc_count: RwLock::new(0),
314            indexed_fields: RwLock::new(HashMap::new()),
315        }
316    }
317
318    /// Configure which fields to index for a collection
319    pub fn add_indexed_field(&self, collection: &str, field: &str) {
320        self.indexed_fields
321            .write()
322            .entry(collection.to_string())
323            .or_default()
324            .insert(field.to_string());
325    }
326
327    /// Index a document's text content
328    pub fn index_document(
329        &self,
330        collection: &str,
331        entity_id: EntityId,
332        field: &str,
333        content: &str,
334    ) {
335        let tokens = self.tokenize(content);
336        let term_count = tokens.len() as f32;
337
338        // Count term frequencies
339        let mut term_freqs: HashMap<String, Vec<u32>> = HashMap::new();
340        for (position, token) in tokens.iter().enumerate() {
341            term_freqs
342                .entry(token.clone())
343                .or_default()
344                .push(position as u32);
345        }
346
347        // Add to index
348        {
349            let mut index = self.index.write();
350            for (term, positions) in term_freqs {
351                let tf = positions.len() as f32 / term_count.max(1.0);
352
353                let entry = PostingEntry {
354                    entity_id,
355                    collection: collection.to_string(),
356                    field: field.to_string(),
357                    positions,
358                    term_frequency: tf,
359                };
360
361                index.entry(term).or_default().push(entry);
362            }
363        }
364
365        // Update doc count
366        *self.doc_count.write() += 1;
367    }
368
369    /// Remove a document from the index
370    pub fn remove_document(&self, entity_id: EntityId) {
371        let mut index = self.index.write();
372        for postings in index.values_mut() {
373            postings.retain(|p| p.entity_id != entity_id);
374        }
375    }
376
377    /// Search for documents containing all terms (AND query)
378    pub fn search(&self, query: &str, limit: usize) -> Vec<TextSearchResult> {
379        let terms = self.tokenize(query);
380        if terms.is_empty() {
381            return Vec::new();
382        }
383
384        let index = self.index.read();
385
386        let doc_count = *self.doc_count.read();
387
388        // Get posting lists for all terms
389        let mut term_postings: Vec<&Vec<PostingEntry>> = Vec::new();
390        for term in &terms {
391            if let Some(postings) = index.get(term) {
392                term_postings.push(postings);
393            } else {
394                // Term not found, AND query fails
395                return Vec::new();
396            }
397        }
398
399        // Find documents containing all terms
400        let mut scores: HashMap<EntityId, f32> = HashMap::new();
401
402        // Start with first term's documents
403        if let Some(first_postings) = term_postings.first() {
404            for posting in *first_postings {
405                let idf = ((doc_count as f32) / (first_postings.len() as f32 + 1.0)).ln();
406                scores.insert(posting.entity_id, posting.term_frequency * idf);
407            }
408        }
409
410        // Intersect with remaining terms
411        for postings in term_postings.iter().skip(1) {
412            let idf = ((doc_count as f32) / (postings.len() as f32 + 1.0)).ln();
413            let entities_in_term: HashSet<EntityId> =
414                postings.iter().map(|p| p.entity_id).collect();
415
416            // Keep only documents that have this term
417            scores.retain(|id, _| entities_in_term.contains(id));
418
419            // Add TF-IDF score
420            for posting in *postings {
421                if let Some(score) = scores.get_mut(&posting.entity_id) {
422                    *score += posting.term_frequency * idf;
423                }
424            }
425        }
426
427        // Convert to results and sort
428        let mut results: Vec<TextSearchResult> = scores
429            .into_iter()
430            .map(|(entity_id, score)| {
431                // Get collection from first posting
432                let collection = term_postings
433                    .first()
434                    .and_then(|p| p.iter().find(|e| e.entity_id == entity_id))
435                    .map(|p| p.collection.clone())
436                    .unwrap_or_default();
437
438                TextSearchResult {
439                    entity_id,
440                    collection,
441                    score,
442                    matched_terms: terms.clone(),
443                }
444            })
445            .collect();
446
447        results.sort_by(|a, b| {
448            b.score
449                .partial_cmp(&a.score)
450                .unwrap_or(std::cmp::Ordering::Equal)
451        });
452        results.truncate(limit);
453        results
454    }
455
456    /// Search with prefix matching (for autocomplete)
457    pub fn search_prefix(&self, prefix: &str, limit: usize) -> Vec<String> {
458        let prefix_lower = prefix.to_lowercase();
459
460        let index = self.index.read();
461
462        index
463            .range(prefix_lower.clone()..)
464            .take_while(|(term, _)| term.starts_with(&prefix_lower))
465            .take(limit)
466            .map(|(term, _)| term.clone())
467            .collect()
468    }
469
470    /// Simple tokenization - splits on whitespace and punctuation
471    fn tokenize(&self, text: &str) -> Vec<String> {
472        text.to_lowercase()
473            .split(|c: char| !c.is_alphanumeric())
474            .filter(|s| s.len() >= 2)
475            .map(|s| s.to_string())
476            .collect()
477    }
478}
479
480impl Default for InvertedIndex {
481    fn default() -> Self {
482        Self::new()
483    }
484}
485
486/// Result from a text search
487#[derive(Debug, Clone)]
488pub struct TextSearchResult {
489    pub entity_id: EntityId,
490    pub collection: String,
491    pub score: f32,
492    pub matched_terms: Vec<String>,
493}
494
495// ============================================================================
496// Integrated Index Manager
497// ============================================================================
498
499/// Configuration for the integrated index system
500#[derive(Debug, Clone)]
501pub struct IntegratedIndexConfig {
502    /// Enable HNSW indexing for vectors
503    pub enable_hnsw: bool,
504    /// Enable full-text indexing
505    pub enable_fulltext: bool,
506    /// Enable metadata indexing
507    pub enable_metadata: bool,
508    /// Enable graph adjacency indexing
509    pub enable_graph: bool,
510    /// HNSW M parameter (connections per node)
511    pub hnsw_m: usize,
512    /// HNSW ef_construction parameter
513    pub hnsw_ef_construction: usize,
514    /// HNSW ef_search parameter
515    pub hnsw_ef_search: usize,
516}
517
518impl Default for IntegratedIndexConfig {
519    fn default() -> Self {
520        Self {
521            enable_hnsw: true,
522            enable_fulltext: true,
523            enable_metadata: true,
524            enable_graph: true,
525            hnsw_m: 16,
526            hnsw_ef_construction: 100,
527            hnsw_ef_search: 50,
528        }
529    }
530}
531
532/// Statistics about the index system
533#[derive(Debug, Clone, Default)]
534pub struct IndexStats {
535    /// Number of indexed vectors
536    pub vector_count: usize,
537    /// Number of indexed documents (for full-text)
538    pub document_count: usize,
539    /// Number of unique terms
540    pub term_count: usize,
541    /// Number of metadata entries
542    pub metadata_entries: usize,
543    /// Number of graph nodes
544    pub graph_node_count: usize,
545    /// Number of graph edges
546    pub graph_edge_count: usize,
547    /// Index creation timestamp
548    pub created_at: u64,
549    /// Last update timestamp
550    pub updated_at: u64,
551}
552
553/// Types of indices that can be managed
554#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
555pub enum IndexType {
556    /// HNSW vector index
557    Hnsw,
558    /// Inverted full-text index
559    Fulltext,
560    /// B-tree metadata index
561    Metadata,
562    /// Graph adjacency index
563    Graph,
564}
565
566/// Status of an index operation
567#[derive(Debug, Clone)]
568pub enum IndexStatus {
569    /// Index is ready for use
570    Ready,
571    /// Index is being built
572    Building { progress: f32 },
573    /// Index needs rebuild
574    Stale,
575    /// Index is disabled
576    Disabled,
577    /// Index encountered an error
578    Error(String),
579}
580
581/// Index lifecycle event for tracking
582#[derive(Debug, Clone)]
583pub struct IndexEvent {
584    pub index_type: IndexType,
585    pub collection: Option<String>,
586    pub event: IndexEventKind,
587    pub timestamp: u64,
588}
589
590#[derive(Debug, Clone)]
591pub enum IndexEventKind {
592    Created,
593    Dropped,
594    Rebuilt,
595    Updated { entries_affected: usize },
596}
597
598/// Integrated Index Manager combining all index types
599pub struct IntegratedIndexManager {
600    /// Configuration
601    config: IntegratedIndexConfig,
602    /// Inverted index for full-text search
603    text_index: InvertedIndex,
604    /// Metadata storage (provides B-tree indexing)
605    metadata_index: RwLock<MetadataStorage>,
606    /// Collection-specific HNSW indices
607    /// Key: collection name, Value: (dimension, index data)
608    hnsw_indices: RwLock<HashMap<String, HnswIndexInfo>>,
609    /// Graph adjacency index
610    graph_index: GraphAdjacencyIndex,
611    /// Index status tracking
612    index_status: RwLock<HashMap<(IndexType, Option<String>), IndexStatus>>,
613    /// Lifecycle event history
614    event_history: RwLock<Vec<IndexEvent>>,
615    /// Creation timestamp
616    created_at: u64,
617}
618
619/// Information about an HNSW index for a collection
620struct HnswIndexInfo {
621    /// Vector dimension
622    dimension: usize,
623    /// Vectors stored (id → vector)
624    vectors: HashMap<EntityId, Vec<f32>>,
625    /// HNSW graph layers (simplified representation)
626    /// In production, would use the full HNSW implementation
627    entry_point: Option<EntityId>,
628}
629
630mod manager_impl;
631impl Default for IntegratedIndexManager {
632    fn default() -> Self {
633        Self::new()
634    }
635}
636
637/// Result from a vector similarity search
638#[derive(Debug, Clone)]
639pub struct VectorSearchResult {
640    pub entity_id: EntityId,
641    pub collection: String,
642    pub similarity: f32,
643}
644
645/// Filter for metadata queries
646#[derive(Debug, Clone)]
647pub enum MetadataQueryFilter {
648    /// Exact match
649    Equals(MetadataValue),
650    /// Range query (inclusive)
651    Range {
652        min: Option<MetadataValue>,
653        max: Option<MetadataValue>,
654    },
655    /// String contains
656    Contains(String),
657    /// Value in set
658    In(Vec<MetadataValue>),
659}
660
661// ============================================================================
662// Utility Functions
663// ============================================================================
664
665/// Compute cosine similarity between two vectors
666fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
667    if a.len() != b.len() || a.is_empty() {
668        return 0.0;
669    }
670
671    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
672    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
673    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
674
675    if norm_a == 0.0 || norm_b == 0.0 {
676        return 0.0;
677    }
678
679    dot / (norm_a * norm_b)
680}
681
682// ============================================================================
683// Tests
684// ============================================================================
685
686#[cfg(test)]
687mod tests {
688    use super::*;
689
690    #[test]
691    fn test_inverted_index_basic() {
692        let index = InvertedIndex::new();
693
694        index.index_document(
695            "docs",
696            EntityId(1),
697            "content",
698            "hello world rust programming",
699        );
700        index.index_document("docs", EntityId(2), "content", "rust is fast and safe");
701        index.index_document("docs", EntityId(3), "content", "python is easy to learn");
702
703        let results = index.search("rust", 10);
704        assert_eq!(results.len(), 2);
705        assert!(results.iter().any(|r| r.entity_id == EntityId(1)));
706        assert!(results.iter().any(|r| r.entity_id == EntityId(2)));
707    }
708
709    #[test]
710    fn test_inverted_index_and_query() {
711        let index = InvertedIndex::new();
712
713        index.index_document("docs", EntityId(1), "content", "rust programming language");
714        index.index_document("docs", EntityId(2), "content", "rust systems programming");
715        index.index_document(
716            "docs",
717            EntityId(3),
718            "content",
719            "python programming language",
720        );
721
722        // AND query: both "rust" and "programming"
723        let results = index.search("rust programming", 10);
724        assert_eq!(results.len(), 2);
725
726        // "language" appears in docs 1 and 3, "programming" in all three
727        // AND of "language programming" gives docs 1 and 3
728        let results = index.search("language programming", 10);
729        assert_eq!(results.len(), 2);
730    }
731
732    #[test]
733    fn test_prefix_search() {
734        let index = InvertedIndex::new();
735
736        index.index_document("docs", EntityId(1), "content", "programming rust rustacean");
737
738        let suggestions = index.search_prefix("rust", 10);
739        assert!(suggestions.contains(&"rust".to_string()));
740        assert!(suggestions.contains(&"rustacean".to_string()));
741    }
742
743    #[test]
744    fn test_vector_search() {
745        let manager = IntegratedIndexManager::new();
746
747        manager.index_vector("embeddings", EntityId(1), &[1.0, 0.0, 0.0]);
748        manager.index_vector("embeddings", EntityId(2), &[0.9, 0.1, 0.0]);
749        manager.index_vector("embeddings", EntityId(3), &[0.0, 1.0, 0.0]);
750
751        let results = manager.search_similar("embeddings", &[1.0, 0.0, 0.0], 2);
752        assert_eq!(results.len(), 2);
753        assert_eq!(results[0].entity_id, EntityId(1));
754        assert!(results[0].similarity > 0.99);
755    }
756
757    #[test]
758    fn test_cosine_similarity() {
759        let a = [1.0, 0.0, 0.0];
760        let b = [1.0, 0.0, 0.0];
761        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);
762
763        let c = [0.0, 1.0, 0.0];
764        assert!(cosine_similarity(&a, &c).abs() < 0.001);
765    }
766
767    // =========================================================================
768    // Graph Adjacency Index Tests
769    // =========================================================================
770
771    #[test]
772    fn test_graph_adjacency_basic() {
773        let index = GraphAdjacencyIndex::new();
774
775        // Create nodes: 1 -> 2 -> 3
776        index.index_edge(EntityId(100), EntityId(1), EntityId(2), "KNOWS", 1.0);
777        index.index_edge(EntityId(101), EntityId(2), EntityId(3), "KNOWS", 1.0);
778
779        // Check outgoing from node 1
780        let neighbors = index.get_neighbors(EntityId(1), EdgeDirection::Outgoing, None);
781        assert_eq!(neighbors.len(), 1);
782        assert_eq!(neighbors[0].neighbor_id, EntityId(2));
783
784        // Check incoming to node 2
785        let neighbors = index.get_neighbors(EntityId(2), EdgeDirection::Incoming, None);
786        assert_eq!(neighbors.len(), 1);
787        assert_eq!(neighbors[0].neighbor_id, EntityId(1));
788
789        // Check both directions for node 2
790        let neighbors = index.get_neighbors(EntityId(2), EdgeDirection::Both, None);
791        assert_eq!(neighbors.len(), 2);
792    }
793
794    #[test]
795    fn test_graph_adjacency_label_filter() {
796        let index = GraphAdjacencyIndex::new();
797
798        // Create edges with different labels
799        index.index_edge(EntityId(100), EntityId(1), EntityId(2), "KNOWS", 1.0);
800        index.index_edge(EntityId(101), EntityId(1), EntityId(3), "WORKS_WITH", 1.0);
801        index.index_edge(EntityId(102), EntityId(1), EntityId(4), "KNOWS", 1.0);
802
803        // Filter by KNOWS label
804        let neighbors = index.get_neighbors(EntityId(1), EdgeDirection::Outgoing, Some("KNOWS"));
805        assert_eq!(neighbors.len(), 2);
806
807        // Filter by WORKS_WITH label
808        let neighbors =
809            index.get_neighbors(EntityId(1), EdgeDirection::Outgoing, Some("WORKS_WITH"));
810        assert_eq!(neighbors.len(), 1);
811        assert_eq!(neighbors[0].neighbor_id, EntityId(3));
812    }
813
814    #[test]
815    fn test_graph_adjacency_degree() {
816        let index = GraphAdjacencyIndex::new();
817
818        // Create a star graph: 1 -> [2, 3, 4, 5]
819        index.index_edge(EntityId(100), EntityId(1), EntityId(2), "LINK", 1.0);
820        index.index_edge(EntityId(101), EntityId(1), EntityId(3), "LINK", 1.0);
821        index.index_edge(EntityId(102), EntityId(1), EntityId(4), "LINK", 1.0);
822        index.index_edge(EntityId(103), EntityId(1), EntityId(5), "LINK", 1.0);
823
824        assert_eq!(index.out_degree(EntityId(1)), 4);
825        assert_eq!(index.in_degree(EntityId(1)), 0);
826        assert_eq!(index.degree(EntityId(1)), 4);
827
828        // Leaf nodes have in-degree 1
829        assert_eq!(index.in_degree(EntityId(2)), 1);
830        assert_eq!(index.out_degree(EntityId(2)), 0);
831    }
832
833    #[test]
834    fn test_graph_adjacency_edge_by_label() {
835        let index = GraphAdjacencyIndex::new();
836
837        index.index_edge(EntityId(100), EntityId(1), EntityId(2), "A", 1.0);
838        index.index_edge(EntityId(101), EntityId(2), EntityId(3), "B", 1.0);
839        index.index_edge(EntityId(102), EntityId(3), EntityId(4), "A", 1.0);
840
841        let edges_a = index.get_edges_by_label("A");
842        assert_eq!(edges_a.len(), 2);
843        assert!(edges_a.contains(&EntityId(100)));
844        assert!(edges_a.contains(&EntityId(102)));
845
846        let edges_b = index.get_edges_by_label("B");
847        assert_eq!(edges_b.len(), 1);
848        assert!(edges_b.contains(&EntityId(101)));
849    }
850
851    #[test]
852    fn test_graph_adjacency_remove() {
853        let index = GraphAdjacencyIndex::new();
854
855        index.index_edge(EntityId(100), EntityId(1), EntityId(2), "LINK", 1.0);
856        index.index_edge(EntityId(101), EntityId(1), EntityId(3), "LINK", 1.0);
857
858        assert_eq!(index.edge_count(), 2);
859
860        // Remove one edge
861        index.remove_edge(EntityId(100));
862
863        assert_eq!(index.edge_count(), 1);
864        let neighbors = index.get_neighbors(EntityId(1), EdgeDirection::Outgoing, None);
865        assert_eq!(neighbors.len(), 1);
866        assert_eq!(neighbors[0].neighbor_id, EntityId(3));
867    }
868
869    // =========================================================================
870    // Index Lifecycle Tests
871    // =========================================================================
872
873    #[test]
874    fn test_index_lifecycle_create_drop() {
875        let manager = IntegratedIndexManager::new();
876
877        // Create a new HNSW index for a specific collection
878        let result = manager.create_index(IndexType::Hnsw, Some("my_collection"));
879        assert!(result.is_ok());
880
881        // Check status
882        let status = manager.index_status(IndexType::Hnsw, Some("my_collection"));
883        assert!(matches!(status, IndexStatus::Ready));
884
885        // Drop the index
886        let result = manager.drop_index(IndexType::Hnsw, Some("my_collection"));
887        assert!(result.is_ok());
888    }
889
890    #[test]
891    fn test_index_lifecycle_rebuild() {
892        let manager = IntegratedIndexManager::new();
893
894        // Index some vectors
895        manager.index_vector("test", EntityId(1), &[1.0, 0.0, 0.0]);
896        manager.index_vector("test", EntityId(2), &[0.0, 1.0, 0.0]);
897
898        // Rebuild the index
899        let result = manager.rebuild_index(IndexType::Hnsw, Some("test"));
900        assert!(result.is_ok());
901
902        // Check status is ready
903        let status = manager.index_status(IndexType::Hnsw, Some("test"));
904        assert!(matches!(status, IndexStatus::Ready));
905    }
906
907    #[test]
908    fn test_index_stats_with_graph() {
909        let manager = IntegratedIndexManager::new();
910
911        // Add some edges
912        manager.index_edge(EntityId(100), EntityId(1), EntityId(2), "LINK", 1.0);
913        manager.index_edge(EntityId(101), EntityId(2), EntityId(3), "LINK", 1.0);
914
915        let stats = manager.stats();
916        assert_eq!(stats.graph_edge_count, 2);
917        assert!(stats.graph_node_count >= 2); // At least source nodes
918    }
919
920    #[test]
921    fn test_integrated_manager_graph_operations() {
922        let manager = IntegratedIndexManager::new();
923
924        // Index edges through the manager
925        manager.index_edge(EntityId(100), EntityId(1), EntityId(2), "KNOWS", 1.0);
926        manager.index_edge(EntityId(101), EntityId(2), EntityId(3), "KNOWS", 0.5);
927
928        // Query neighbors through the manager
929        let neighbors = manager.get_neighbors(EntityId(1), EdgeDirection::Outgoing, None);
930        assert_eq!(neighbors.len(), 1);
931        assert_eq!(neighbors[0].neighbor_id, EntityId(2));
932        assert_eq!(neighbors[0].weight, 1.0);
933
934        // Check degree
935        assert_eq!(manager.node_degree(EntityId(2), EdgeDirection::Both), 2);
936    }
937}
reddb_server/storage/unified/index.rs

reddb_server/storage/unified/
index.rs
