hermes_core/index/
mod.rs

1//! Index - multi-segment async search index
2//!
3//! The `Index` is the central concept that provides:
4//! - `Index::create()` / `Index::open()` - create or open an index
5//! - `index.writer()` - get an IndexWriter for adding documents
6//! - `index.reader()` - get an IndexReader for searching (with reload policy)
7//!
8//! The Index owns the SegmentManager which handles segment lifecycle and tracking.
9
10#[cfg(feature = "native")]
11use crate::dsl::Schema;
12#[cfg(feature = "native")]
13use crate::error::Result;
14#[cfg(feature = "native")]
15use crate::structures::{CoarseCentroids, PQCodebook};
16#[cfg(feature = "native")]
17use rustc_hash::FxHashMap;
18#[cfg(feature = "native")]
19use std::sync::Arc;
20
21mod searcher;
22pub use searcher::Searcher;
23
24#[cfg(feature = "native")]
25mod reader;
26#[cfg(feature = "native")]
27mod vector_builder;
28#[cfg(feature = "native")]
29mod writer;
30#[cfg(feature = "native")]
31pub use reader::IndexReader;
32#[cfg(feature = "native")]
33pub use writer::IndexWriter;
34
35mod metadata;
36pub use metadata::{FieldVectorMeta, INDEX_META_FILENAME, IndexMetadata, VectorIndexState};
37
38#[cfg(feature = "native")]
39mod helpers;
40#[cfg(feature = "native")]
41pub use helpers::{
42    IndexingStats, SchemaConfig, SchemaFieldConfig, create_index_at_path, create_index_from_sdl,
43    index_documents_from_reader, index_json_document, parse_schema,
44};
45
46/// Default file name for the slice cache
47pub const SLICE_CACHE_FILENAME: &str = "index.slicecache";
48
49/// Index configuration
50#[derive(Debug, Clone)]
51pub struct IndexConfig {
52    /// Number of threads for CPU-intensive tasks (search parallelism)
53    pub num_threads: usize,
54    /// Number of parallel segment builders (documents distributed round-robin)
55    pub num_indexing_threads: usize,
56    /// Number of threads for parallel block compression within each segment
57    pub num_compression_threads: usize,
58    /// Block cache size for term dictionary per segment
59    pub term_cache_blocks: usize,
60    /// Block cache size for document store per segment
61    pub store_cache_blocks: usize,
62    /// Max memory (bytes) across all builders before auto-commit (global limit)
63    pub max_indexing_memory_bytes: usize,
64    /// Merge policy for background segment merging
65    pub merge_policy: Box<dyn crate::merge::MergePolicy>,
66    /// Index optimization mode (adaptive, size-optimized, performance-optimized)
67    pub optimization: crate::structures::IndexOptimization,
68    /// Reload interval in milliseconds for IndexReader (how often to check for new segments)
69    pub reload_interval_ms: u64,
70}
71
72impl Default for IndexConfig {
73    fn default() -> Self {
74        #[cfg(feature = "native")]
75        let cpus = num_cpus::get().max(1);
76        #[cfg(not(feature = "native"))]
77        let cpus = 1;
78
79        Self {
80            num_threads: cpus,
81            num_indexing_threads: 1,
82            num_compression_threads: cpus,
83            term_cache_blocks: 256,
84            store_cache_blocks: 32,
85            max_indexing_memory_bytes: 2 * 1024 * 1024 * 1024, // 256 MB default
86            merge_policy: Box::new(crate::merge::TieredMergePolicy::default()),
87            optimization: crate::structures::IndexOptimization::default(),
88            reload_interval_ms: 1000, // 1 second default
89        }
90    }
91}
92
93/// Multi-segment async Index
94///
95/// The central concept for search. Owns segment lifecycle and provides:
96/// - `Index::create()` / `Index::open()` - create or open an index
97/// - `index.writer()` - get an IndexWriter for adding documents
98/// - `index.reader()` - get an IndexReader for searching with reload policy
99///
100/// All segment management is delegated to SegmentManager.
101#[cfg(feature = "native")]
102pub struct Index<D: crate::directories::DirectoryWriter + 'static> {
103    directory: Arc<D>,
104    schema: Arc<Schema>,
105    config: IndexConfig,
106    /// Segment manager - owns segments, tracker, and metadata
107    segment_manager: Arc<crate::merge::SegmentManager<D>>,
108    /// Trained centroids for vector search
109    trained_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
110    /// Trained codebooks for vector search
111    trained_codebooks: FxHashMap<u32, Arc<PQCodebook>>,
112    /// Cached reader (created lazily, reused across calls)
113    cached_reader: tokio::sync::OnceCell<IndexReader<D>>,
114}
115
116#[cfg(feature = "native")]
117impl<D: crate::directories::DirectoryWriter + 'static> Index<D> {
118    /// Create a new index in the directory
119    pub async fn create(directory: D, schema: Schema, config: IndexConfig) -> Result<Self> {
120        let directory = Arc::new(directory);
121        let schema = Arc::new(schema);
122        let metadata = IndexMetadata::new((*schema).clone());
123
124        let segment_manager = Arc::new(crate::merge::SegmentManager::new(
125            Arc::clone(&directory),
126            Arc::clone(&schema),
127            metadata,
128            config.merge_policy.clone_box(),
129            config.term_cache_blocks,
130        ));
131
132        // Save initial metadata
133        segment_manager.update_metadata(|_| {}).await?;
134
135        Ok(Self {
136            directory,
137            schema,
138            config,
139            segment_manager,
140            trained_centroids: FxHashMap::default(),
141            trained_codebooks: FxHashMap::default(),
142            cached_reader: tokio::sync::OnceCell::new(),
143        })
144    }
145
146    /// Open an existing index from a directory
147    pub async fn open(directory: D, config: IndexConfig) -> Result<Self> {
148        let directory = Arc::new(directory);
149
150        // Load metadata (includes schema)
151        let metadata = IndexMetadata::load(directory.as_ref()).await?;
152        let schema = Arc::new(metadata.schema.clone());
153
154        // Load trained structures
155        let (trained_centroids, trained_codebooks) =
156            metadata.load_trained_structures(directory.as_ref()).await;
157
158        let segment_manager = Arc::new(crate::merge::SegmentManager::new(
159            Arc::clone(&directory),
160            Arc::clone(&schema),
161            metadata,
162            config.merge_policy.clone_box(),
163            config.term_cache_blocks,
164        ));
165
166        Ok(Self {
167            directory,
168            schema,
169            config,
170            segment_manager,
171            trained_centroids,
172            trained_codebooks,
173            cached_reader: tokio::sync::OnceCell::new(),
174        })
175    }
176
177    /// Get the schema
178    pub fn schema(&self) -> &Schema {
179        &self.schema
180    }
181
182    /// Get a reference to the underlying directory
183    pub fn directory(&self) -> &D {
184        &self.directory
185    }
186
187    /// Get the segment manager
188    pub fn segment_manager(&self) -> &Arc<crate::merge::SegmentManager<D>> {
189        &self.segment_manager
190    }
191
192    /// Get an IndexReader for searching (with reload policy)
193    ///
194    /// The reader is cached and reused across calls. The reader's internal
195    /// searcher will reload segments based on its reload interval (configurable via IndexConfig).
196    pub async fn reader(&self) -> Result<&IndexReader<D>> {
197        self.cached_reader
198            .get_or_try_init(|| async {
199                IndexReader::from_segment_manager_with_reload_interval(
200                    Arc::clone(&self.schema),
201                    Arc::clone(&self.segment_manager),
202                    self.trained_centroids.clone(),
203                    self.trained_codebooks.clone(),
204                    self.config.term_cache_blocks,
205                    self.config.reload_interval_ms,
206                )
207                .await
208            })
209            .await
210    }
211
212    /// Get the config
213    pub fn config(&self) -> &IndexConfig {
214        &self.config
215    }
216
217    /// Get trained centroids
218    pub fn trained_centroids(&self) -> &FxHashMap<u32, Arc<CoarseCentroids>> {
219        &self.trained_centroids
220    }
221
222    /// Get trained codebooks
223    pub fn trained_codebooks(&self) -> &FxHashMap<u32, Arc<PQCodebook>> {
224        &self.trained_codebooks
225    }
226
227    /// Get segment readers for query execution (convenience method)
228    pub async fn segment_readers(&self) -> Result<Vec<Arc<crate::segment::SegmentReader>>> {
229        let reader = self.reader().await?;
230        let searcher = reader.searcher().await?;
231        Ok(searcher.segment_readers().to_vec())
232    }
233
234    /// Total number of documents across all segments
235    pub async fn num_docs(&self) -> Result<u32> {
236        let reader = self.reader().await?;
237        let searcher = reader.searcher().await?;
238        Ok(searcher.num_docs())
239    }
240
241    /// Get a document by global doc_id
242    pub async fn doc(&self, doc_id: crate::DocId) -> Result<Option<crate::dsl::Document>> {
243        let reader = self.reader().await?;
244        let searcher = reader.searcher().await?;
245        searcher.doc(doc_id).await
246    }
247
248    /// Get default fields for search
249    pub fn default_fields(&self) -> Vec<crate::Field> {
250        if !self.schema.default_fields().is_empty() {
251            self.schema.default_fields().to_vec()
252        } else {
253            self.schema
254                .fields()
255                .filter(|(_, entry)| {
256                    entry.indexed && entry.field_type == crate::dsl::FieldType::Text
257                })
258                .map(|(field, _)| field)
259                .collect()
260        }
261    }
262
263    /// Get tokenizer registry
264    pub fn tokenizers(&self) -> Arc<crate::tokenizer::TokenizerRegistry> {
265        Arc::new(crate::tokenizer::TokenizerRegistry::default())
266    }
267
268    /// Create a query parser for this index
269    pub fn query_parser(&self) -> crate::dsl::QueryLanguageParser {
270        let default_fields = self.default_fields();
271        let tokenizers = self.tokenizers();
272
273        let query_routers = self.schema.query_routers();
274        if !query_routers.is_empty()
275            && let Ok(router) = crate::dsl::QueryFieldRouter::from_rules(query_routers)
276        {
277            return crate::dsl::QueryLanguageParser::with_router(
278                Arc::clone(&self.schema),
279                default_fields,
280                tokenizers,
281                router,
282            );
283        }
284
285        crate::dsl::QueryLanguageParser::new(Arc::clone(&self.schema), default_fields, tokenizers)
286    }
287
288    /// Parse and search using a query string
289    pub async fn query(
290        &self,
291        query_str: &str,
292        limit: usize,
293    ) -> Result<crate::query::SearchResponse> {
294        self.query_offset(query_str, limit, 0).await
295    }
296
297    /// Query with offset for pagination
298    pub async fn query_offset(
299        &self,
300        query_str: &str,
301        limit: usize,
302        offset: usize,
303    ) -> Result<crate::query::SearchResponse> {
304        let parser = self.query_parser();
305        let query = parser
306            .parse(query_str)
307            .map_err(crate::error::Error::Query)?;
308        self.search_offset(query.as_ref(), limit, offset).await
309    }
310
311    /// Search and return results
312    pub async fn search(
313        &self,
314        query: &dyn crate::query::Query,
315        limit: usize,
316    ) -> Result<crate::query::SearchResponse> {
317        self.search_offset(query, limit, 0).await
318    }
319
320    /// Search with offset for pagination
321    pub async fn search_offset(
322        &self,
323        query: &dyn crate::query::Query,
324        limit: usize,
325        offset: usize,
326    ) -> Result<crate::query::SearchResponse> {
327        let reader = self.reader().await?;
328        let searcher = reader.searcher().await?;
329        let segments = searcher.segment_readers();
330
331        let mut all_results: Vec<(u128, crate::query::SearchResult)> = Vec::new();
332        let fetch_limit = offset + limit;
333
334        for segment in segments {
335            let segment_id = segment.meta().id;
336            let results =
337                crate::query::search_segment(segment.as_ref(), query, fetch_limit).await?;
338            for result in results {
339                all_results.push((segment_id, result));
340            }
341        }
342
343        all_results.sort_by(|a, b| {
344            b.1.score
345                .partial_cmp(&a.1.score)
346                .unwrap_or(std::cmp::Ordering::Equal)
347        });
348
349        let total_hits = all_results.len() as u32;
350
351        let hits: Vec<crate::query::SearchHit> = all_results
352            .into_iter()
353            .skip(offset)
354            .take(limit)
355            .map(|(segment_id, result)| crate::query::SearchHit {
356                address: crate::query::DocAddress::new(segment_id, result.doc_id),
357                score: result.score,
358                matched_fields: result.extract_ordinals(),
359            })
360            .collect();
361
362        Ok(crate::query::SearchResponse { hits, total_hits })
363    }
364
365    /// Get a document by its unique address
366    pub async fn get_document(
367        &self,
368        address: &crate::query::DocAddress,
369    ) -> Result<Option<crate::dsl::Document>> {
370        let segment_id = address.segment_id_u128().ok_or_else(|| {
371            crate::error::Error::Query(format!("Invalid segment ID: {}", address.segment_id))
372        })?;
373
374        let reader = self.reader().await?;
375        let searcher = reader.searcher().await?;
376
377        for segment in searcher.segment_readers() {
378            if segment.meta().id == segment_id {
379                return segment.doc(address.doc_id).await;
380            }
381        }
382
383        Ok(None)
384    }
385
386    /// Reload is no longer needed - reader handles this automatically
387    pub async fn reload(&self) -> Result<()> {
388        // No-op - reader reloads automatically based on policy
389        Ok(())
390    }
391
392    /// Get posting lists for a term across all segments
393    pub async fn get_postings(
394        &self,
395        field: crate::Field,
396        term: &[u8],
397    ) -> Result<
398        Vec<(
399            Arc<crate::segment::SegmentReader>,
400            crate::structures::BlockPostingList,
401        )>,
402    > {
403        let segments = self.segment_readers().await?;
404        let mut results = Vec::new();
405
406        for segment in segments {
407            if let Some(postings) = segment.get_postings(field, term).await? {
408                results.push((segment, postings));
409            }
410        }
411
412        Ok(results)
413    }
414}
415
416/// Native-only methods for Index
417#[cfg(feature = "native")]
418impl<D: crate::directories::DirectoryWriter + 'static> Index<D> {
419    /// Get an IndexWriter for adding documents
420    pub fn writer(&self) -> writer::IndexWriter<D> {
421        writer::IndexWriter::from_index(self)
422    }
423}
424
425// TODO: Add back warmup_and_save_slice_cache when slice caching is re-integrated
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use crate::directories::RamDirectory;
431    use crate::dsl::{Document, SchemaBuilder};
432
433    #[tokio::test]
434    async fn test_index_create_and_search() {
435        let mut schema_builder = SchemaBuilder::default();
436        let title = schema_builder.add_text_field("title", true, true);
437        let body = schema_builder.add_text_field("body", true, true);
438        let schema = schema_builder.build();
439
440        let dir = RamDirectory::new();
441        let config = IndexConfig::default();
442
443        // Create index and add documents
444        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
445            .await
446            .unwrap();
447
448        let mut doc1 = Document::new();
449        doc1.add_text(title, "Hello World");
450        doc1.add_text(body, "This is the first document");
451        writer.add_document(doc1).unwrap();
452
453        let mut doc2 = Document::new();
454        doc2.add_text(title, "Goodbye World");
455        doc2.add_text(body, "This is the second document");
456        writer.add_document(doc2).unwrap();
457
458        writer.commit().await.unwrap();
459
460        // Open for reading
461        let index = Index::open(dir, config).await.unwrap();
462        assert_eq!(index.num_docs().await.unwrap(), 2);
463
464        // Check postings
465        let postings = index.get_postings(title, b"world").await.unwrap();
466        assert_eq!(postings.len(), 1); // One segment
467        assert_eq!(postings[0].1.doc_count(), 2); // Two docs with "world"
468
469        // Retrieve document
470        let doc = index.doc(0).await.unwrap().unwrap();
471        assert_eq!(doc.get_first(title).unwrap().as_text(), Some("Hello World"));
472    }
473
474    #[tokio::test]
475    async fn test_multiple_segments() {
476        let mut schema_builder = SchemaBuilder::default();
477        let title = schema_builder.add_text_field("title", true, true);
478        let schema = schema_builder.build();
479
480        let dir = RamDirectory::new();
481        let config = IndexConfig {
482            max_indexing_memory_bytes: 1024, // Very small to trigger frequent flushes
483            ..Default::default()
484        };
485
486        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
487            .await
488            .unwrap();
489
490        // Add documents in batches to create multiple segments
491        for batch in 0..3 {
492            for i in 0..5 {
493                let mut doc = Document::new();
494                doc.add_text(title, format!("Document {} batch {}", i, batch));
495                writer.add_document(doc).unwrap();
496            }
497            writer.commit().await.unwrap();
498        }
499
500        // Open and check
501        let index = Index::open(dir, config).await.unwrap();
502        assert_eq!(index.num_docs().await.unwrap(), 15);
503        // With queue-based indexing, exact segment count varies
504        assert!(
505            index.segment_readers().await.unwrap().len() >= 2,
506            "Expected multiple segments"
507        );
508    }
509
510    #[tokio::test]
511    async fn test_segment_merge() {
512        let mut schema_builder = SchemaBuilder::default();
513        let title = schema_builder.add_text_field("title", true, true);
514        let schema = schema_builder.build();
515
516        let dir = RamDirectory::new();
517        let config = IndexConfig {
518            max_indexing_memory_bytes: 512, // Very small to trigger frequent flushes
519            ..Default::default()
520        };
521
522        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
523            .await
524            .unwrap();
525
526        // Create multiple segments by flushing between batches
527        for batch in 0..3 {
528            for i in 0..3 {
529                let mut doc = Document::new();
530                doc.add_text(title, format!("Document {} batch {}", i, batch));
531                writer.add_document(doc).unwrap();
532            }
533            writer.flush().await.unwrap();
534        }
535        writer.commit().await.unwrap();
536
537        // Should have multiple segments (at least 2, one per flush with docs)
538        let index = Index::open(dir.clone(), config.clone()).await.unwrap();
539        assert!(
540            index.segment_readers().await.unwrap().len() >= 2,
541            "Expected multiple segments"
542        );
543
544        // Force merge
545        let writer = IndexWriter::open(dir.clone(), config.clone())
546            .await
547            .unwrap();
548        writer.force_merge().await.unwrap();
549
550        // Should have 1 segment now
551        let index = Index::open(dir, config).await.unwrap();
552        assert_eq!(index.segment_readers().await.unwrap().len(), 1);
553        assert_eq!(index.num_docs().await.unwrap(), 9);
554
555        // Verify all documents accessible (order may vary with queue-based indexing)
556        let mut found_docs = 0;
557        for i in 0..9 {
558            if index.doc(i).await.unwrap().is_some() {
559                found_docs += 1;
560            }
561        }
562        assert_eq!(found_docs, 9);
563    }
564
565    #[tokio::test]
566    async fn test_match_query() {
567        let mut schema_builder = SchemaBuilder::default();
568        let title = schema_builder.add_text_field("title", true, true);
569        let body = schema_builder.add_text_field("body", true, true);
570        let schema = schema_builder.build();
571
572        let dir = RamDirectory::new();
573        let config = IndexConfig::default();
574
575        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
576            .await
577            .unwrap();
578
579        let mut doc1 = Document::new();
580        doc1.add_text(title, "rust programming");
581        doc1.add_text(body, "Learn rust language");
582        writer.add_document(doc1).unwrap();
583
584        let mut doc2 = Document::new();
585        doc2.add_text(title, "python programming");
586        doc2.add_text(body, "Learn python language");
587        writer.add_document(doc2).unwrap();
588
589        writer.commit().await.unwrap();
590
591        let index = Index::open(dir, config).await.unwrap();
592
593        // Test match query with multiple default fields
594        let results = index.query("rust", 10).await.unwrap();
595        assert_eq!(results.hits.len(), 1);
596
597        // Test match query with multiple tokens
598        let results = index.query("rust programming", 10).await.unwrap();
599        assert!(!results.hits.is_empty());
600
601        // Verify hit has address (segment_id + doc_id)
602        let hit = &results.hits[0];
603        assert!(!hit.address.segment_id.is_empty(), "Should have segment_id");
604
605        // Verify document retrieval by address
606        let doc = index.get_document(&hit.address).await.unwrap().unwrap();
607        assert!(
608            !doc.field_values().is_empty(),
609            "Doc should have field values"
610        );
611
612        // Also verify doc retrieval directly by global doc_id
613        let doc = index.doc(0).await.unwrap().unwrap();
614        assert!(
615            !doc.field_values().is_empty(),
616            "Doc should have field values"
617        );
618    }
619
620    #[tokio::test]
621    async fn test_slice_cache_warmup_and_load() {
622        use crate::directories::SliceCachingDirectory;
623
624        let mut schema_builder = SchemaBuilder::default();
625        let title = schema_builder.add_text_field("title", true, true);
626        let body = schema_builder.add_text_field("body", true, true);
627        let schema = schema_builder.build();
628
629        let dir = RamDirectory::new();
630        let config = IndexConfig::default();
631
632        // Create index with some documents
633        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
634            .await
635            .unwrap();
636
637        for i in 0..10 {
638            let mut doc = Document::new();
639            doc.add_text(title, format!("Document {} about rust", i));
640            doc.add_text(body, format!("This is body text number {}", i));
641            writer.add_document(doc).unwrap();
642        }
643        writer.commit().await.unwrap();
644
645        // Open with slice caching and perform some operations to warm up cache
646        let caching_dir = SliceCachingDirectory::new(dir.clone(), 1024 * 1024);
647        let index = Index::open(caching_dir, config.clone()).await.unwrap();
648
649        // Perform a search to warm up the cache
650        let results = index.query("rust", 10).await.unwrap();
651        assert!(!results.hits.is_empty());
652
653        // Check cache stats - should have cached some data
654        let stats = index.directory.stats();
655        assert!(stats.total_bytes > 0, "Cache should have data after search");
656    }
657
658    #[tokio::test]
659    async fn test_multivalue_field_indexing_and_search() {
660        let mut schema_builder = SchemaBuilder::default();
661        let uris = schema_builder.add_text_field("uris", true, true);
662        let title = schema_builder.add_text_field("title", true, true);
663        let schema = schema_builder.build();
664
665        let dir = RamDirectory::new();
666        let config = IndexConfig::default();
667
668        // Create index and add document with multi-value field
669        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
670            .await
671            .unwrap();
672
673        let mut doc = Document::new();
674        doc.add_text(uris, "one");
675        doc.add_text(uris, "two");
676        doc.add_text(title, "Test Document");
677        writer.add_document(doc).unwrap();
678
679        // Add another document with different uris
680        let mut doc2 = Document::new();
681        doc2.add_text(uris, "three");
682        doc2.add_text(title, "Another Document");
683        writer.add_document(doc2).unwrap();
684
685        writer.commit().await.unwrap();
686
687        // Open for reading
688        let index = Index::open(dir, config).await.unwrap();
689        assert_eq!(index.num_docs().await.unwrap(), 2);
690
691        // Verify document retrieval preserves all values
692        let doc = index.doc(0).await.unwrap().unwrap();
693        let all_uris: Vec<_> = doc.get_all(uris).collect();
694        assert_eq!(all_uris.len(), 2, "Should have 2 uris values");
695        assert_eq!(all_uris[0].as_text(), Some("one"));
696        assert_eq!(all_uris[1].as_text(), Some("two"));
697
698        // Verify to_json returns array for multi-value field
699        let json = doc.to_json(index.schema());
700        let uris_json = json.get("uris").unwrap();
701        assert!(uris_json.is_array(), "Multi-value field should be an array");
702        let uris_arr = uris_json.as_array().unwrap();
703        assert_eq!(uris_arr.len(), 2);
704        assert_eq!(uris_arr[0].as_str(), Some("one"));
705        assert_eq!(uris_arr[1].as_str(), Some("two"));
706
707        // Verify both values are searchable
708        let results = index.query("uris:one", 10).await.unwrap();
709        assert_eq!(results.hits.len(), 1, "Should find doc with 'one'");
710        assert_eq!(results.hits[0].address.doc_id, 0);
711
712        let results = index.query("uris:two", 10).await.unwrap();
713        assert_eq!(results.hits.len(), 1, "Should find doc with 'two'");
714        assert_eq!(results.hits[0].address.doc_id, 0);
715
716        let results = index.query("uris:three", 10).await.unwrap();
717        assert_eq!(results.hits.len(), 1, "Should find doc with 'three'");
718        assert_eq!(results.hits[0].address.doc_id, 1);
719
720        // Verify searching for non-existent value returns no results
721        let results = index.query("uris:nonexistent", 10).await.unwrap();
722        assert_eq!(results.hits.len(), 0, "Should not find non-existent value");
723    }
724
725    /// Comprehensive test for WAND optimization in BooleanQuery OR queries
726    ///
727    /// This test verifies that:
728    /// 1. BooleanQuery with multiple SHOULD term queries uses WAND automatically
729    /// 2. Search results are correct regardless of WAND optimization
730    /// 3. Scores are reasonable for matching documents
731    #[tokio::test]
732    async fn test_wand_optimization_for_or_queries() {
733        use crate::query::{BooleanQuery, TermQuery};
734
735        let mut schema_builder = SchemaBuilder::default();
736        let content = schema_builder.add_text_field("content", true, true);
737        let schema = schema_builder.build();
738
739        let dir = RamDirectory::new();
740        let config = IndexConfig::default();
741
742        // Create index with documents containing various terms
743        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
744            .await
745            .unwrap();
746
747        // Doc 0: contains "rust" and "programming"
748        let mut doc = Document::new();
749        doc.add_text(content, "rust programming language is fast");
750        writer.add_document(doc).unwrap();
751
752        // Doc 1: contains "rust" only
753        let mut doc = Document::new();
754        doc.add_text(content, "rust is a systems language");
755        writer.add_document(doc).unwrap();
756
757        // Doc 2: contains "programming" only
758        let mut doc = Document::new();
759        doc.add_text(content, "programming is fun");
760        writer.add_document(doc).unwrap();
761
762        // Doc 3: contains "python" (neither rust nor programming)
763        let mut doc = Document::new();
764        doc.add_text(content, "python is easy to learn");
765        writer.add_document(doc).unwrap();
766
767        // Doc 4: contains both "rust" and "programming" multiple times
768        let mut doc = Document::new();
769        doc.add_text(content, "rust rust programming programming systems");
770        writer.add_document(doc).unwrap();
771
772        writer.commit().await.unwrap();
773
774        // Open for reading
775        let index = Index::open(dir.clone(), config.clone()).await.unwrap();
776
777        // Test 1: Pure OR query with multiple terms (should use WAND automatically)
778        let or_query = BooleanQuery::new()
779            .should(TermQuery::text(content, "rust"))
780            .should(TermQuery::text(content, "programming"));
781
782        let results = index.search(&or_query, 10).await.unwrap();
783
784        // Should find docs 0, 1, 2, 4 (all that contain "rust" OR "programming")
785        assert_eq!(results.hits.len(), 4, "Should find exactly 4 documents");
786
787        let doc_ids: Vec<u32> = results.hits.iter().map(|h| h.address.doc_id).collect();
788        assert!(doc_ids.contains(&0), "Should find doc 0");
789        assert!(doc_ids.contains(&1), "Should find doc 1");
790        assert!(doc_ids.contains(&2), "Should find doc 2");
791        assert!(doc_ids.contains(&4), "Should find doc 4");
792        assert!(
793            !doc_ids.contains(&3),
794            "Should NOT find doc 3 (only has 'python')"
795        );
796
797        // Test 2: Single term query (should NOT use WAND, but still work)
798        let single_query = BooleanQuery::new().should(TermQuery::text(content, "rust"));
799
800        let results = index.search(&single_query, 10).await.unwrap();
801        assert_eq!(results.hits.len(), 3, "Should find 3 documents with 'rust'");
802
803        // Test 3: Query with MUST (should NOT use WAND)
804        let must_query = BooleanQuery::new()
805            .must(TermQuery::text(content, "rust"))
806            .should(TermQuery::text(content, "programming"));
807
808        let results = index.search(&must_query, 10).await.unwrap();
809        // Must have "rust", optionally "programming"
810        assert_eq!(results.hits.len(), 3, "Should find 3 documents with 'rust'");
811
812        // Test 4: Query with MUST_NOT (should NOT use WAND)
813        let must_not_query = BooleanQuery::new()
814            .should(TermQuery::text(content, "rust"))
815            .should(TermQuery::text(content, "programming"))
816            .must_not(TermQuery::text(content, "systems"));
817
818        let results = index.search(&must_not_query, 10).await.unwrap();
819        // Should exclude docs with "systems" (doc 1 and 4)
820        let doc_ids: Vec<u32> = results.hits.iter().map(|h| h.address.doc_id).collect();
821        assert!(
822            !doc_ids.contains(&1),
823            "Should NOT find doc 1 (has 'systems')"
824        );
825        assert!(
826            !doc_ids.contains(&4),
827            "Should NOT find doc 4 (has 'systems')"
828        );
829
830        // Test 5: Verify top-k limit works correctly with WAND
831        let or_query = BooleanQuery::new()
832            .should(TermQuery::text(content, "rust"))
833            .should(TermQuery::text(content, "programming"));
834
835        let results = index.search(&or_query, 2).await.unwrap();
836        assert_eq!(results.hits.len(), 2, "Should return only top 2 results");
837
838        // Top results should be docs that match both terms (higher scores)
839        // Doc 0 and 4 contain both "rust" and "programming"
840    }
841
842    /// Test that WAND optimization produces same results as non-WAND for correctness
843    #[tokio::test]
844    async fn test_wand_results_match_standard_boolean() {
845        use crate::query::{BooleanQuery, TermQuery, WandOrQuery};
846
847        let mut schema_builder = SchemaBuilder::default();
848        let content = schema_builder.add_text_field("content", true, true);
849        let schema = schema_builder.build();
850
851        let dir = RamDirectory::new();
852        let config = IndexConfig::default();
853
854        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
855            .await
856            .unwrap();
857
858        // Add several documents
859        for i in 0..10 {
860            let mut doc = Document::new();
861            let text = match i % 4 {
862                0 => "apple banana cherry",
863                1 => "apple orange",
864                2 => "banana grape",
865                _ => "cherry date",
866            };
867            doc.add_text(content, text);
868            writer.add_document(doc).unwrap();
869        }
870
871        writer.commit().await.unwrap();
872        let index = Index::open(dir.clone(), config.clone()).await.unwrap();
873
874        // Compare explicit WandOrQuery with auto-optimized BooleanQuery
875        let wand_query = WandOrQuery::new(content).term("apple").term("banana");
876
877        let bool_query = BooleanQuery::new()
878            .should(TermQuery::text(content, "apple"))
879            .should(TermQuery::text(content, "banana"));
880
881        let wand_results = index.search(&wand_query, 10).await.unwrap();
882        let bool_results = index.search(&bool_query, 10).await.unwrap();
883
884        // Both should find the same documents
885        assert_eq!(
886            wand_results.hits.len(),
887            bool_results.hits.len(),
888            "WAND and Boolean should find same number of docs"
889        );
890
891        let wand_docs: std::collections::HashSet<u32> =
892            wand_results.hits.iter().map(|h| h.address.doc_id).collect();
893        let bool_docs: std::collections::HashSet<u32> =
894            bool_results.hits.iter().map(|h| h.address.doc_id).collect();
895
896        assert_eq!(
897            wand_docs, bool_docs,
898            "WAND and Boolean should find same documents"
899        );
900    }
901
902    #[tokio::test]
903    async fn test_vector_index_threshold_switch() {
904        use crate::dsl::{DenseVectorConfig, VectorIndexType};
905
906        // Create schema with dense vector field configured for IVF-RaBitQ
907        let mut schema_builder = SchemaBuilder::default();
908        let title = schema_builder.add_text_field("title", true, true);
909        let embedding = schema_builder.add_dense_vector_field_with_config(
910            "embedding",
911            true, // indexed
912            true, // stored
913            DenseVectorConfig {
914                dim: 8,
915                index_type: VectorIndexType::IvfRaBitQ,
916                store_raw: true,
917                num_clusters: Some(4), // Small for test
918                nprobe: 2,
919                mrl_dim: None,
920                build_threshold: Some(50), // Build when we have 50+ vectors
921            },
922        );
923        let schema = schema_builder.build();
924
925        let dir = RamDirectory::new();
926        let config = IndexConfig::default();
927
928        // Phase 1: Add vectors below threshold (should use Flat index)
929        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
930            .await
931            .unwrap();
932
933        // Add 30 documents (below threshold of 50)
934        for i in 0..30 {
935            let mut doc = Document::new();
936            doc.add_text(title, format!("Document {}", i));
937            // Simple embedding: [i, i, i, i, i, i, i, i] normalized
938            let vec: Vec<f32> = (0..8).map(|_| (i as f32) / 30.0).collect();
939            doc.add_dense_vector(embedding, vec);
940            writer.add_document(doc).unwrap();
941        }
942        writer.commit().await.unwrap();
943
944        // Open index and verify it's using Flat (not built yet)
945        let index = Index::open(dir.clone(), config.clone()).await.unwrap();
946        assert!(
947            index.trained_centroids.is_empty(),
948            "Should not have trained centroids below threshold"
949        );
950
951        // Search should work with Flat index
952        let query_vec: Vec<f32> = vec![0.5; 8];
953        let segments = index.segment_readers().await.unwrap();
954        assert!(!segments.is_empty());
955
956        let results = segments[0]
957            .search_dense_vector(
958                embedding,
959                &query_vec,
960                5,
961                1,
962                crate::query::MultiValueCombiner::Max,
963            )
964            .unwrap();
965        assert!(!results.is_empty(), "Flat search should return results");
966
967        // Phase 2: Add more vectors to cross threshold
968        let writer = IndexWriter::open(dir.clone(), config.clone())
969            .await
970            .unwrap();
971
972        // Add 30 more documents (total 60, above threshold of 50)
973        for i in 30..60 {
974            let mut doc = Document::new();
975            doc.add_text(title, format!("Document {}", i));
976            let vec: Vec<f32> = (0..8).map(|_| (i as f32) / 60.0).collect();
977            doc.add_dense_vector(embedding, vec);
978            writer.add_document(doc).unwrap();
979        }
980        // Commit auto-triggers vector index build when threshold is crossed
981        writer.commit().await.unwrap();
982
983        // Verify centroids were trained (auto-triggered)
984        assert!(
985            writer.is_vector_index_built(embedding).await,
986            "Vector index should be built after crossing threshold"
987        );
988
989        // Reopen index and verify trained structures are loaded
990        let index = Index::open(dir.clone(), config.clone()).await.unwrap();
991        assert!(
992            index.trained_centroids.contains_key(&embedding.0),
993            "Should have loaded trained centroids for embedding field"
994        );
995
996        // Search should still work
997        let segments = index.segment_readers().await.unwrap();
998        let results = segments[0]
999            .search_dense_vector(
1000                embedding,
1001                &query_vec,
1002                5,
1003                1,
1004                crate::query::MultiValueCombiner::Max,
1005            )
1006            .unwrap();
1007        assert!(
1008            !results.is_empty(),
1009            "Search should return results after build"
1010        );
1011
1012        // Phase 3: Verify calling build_vector_index again is a no-op
1013        let writer = IndexWriter::open(dir.clone(), config.clone())
1014            .await
1015            .unwrap();
1016        writer.build_vector_index().await.unwrap(); // Should skip training
1017
1018        // Still built
1019        assert!(writer.is_vector_index_built(embedding).await);
1020    }
1021}
hermes_core/index/mod.rs

hermes_core/index/
mod.rs
