hermes_core/
index.rs

1//! Index - multi-segment async search index
2//!
3//! Components:
4//! - Index: main entry point for searching
5//! - IndexWriter: for adding documents and committing segments
6//! - Supports multiple segments with merge
7
8use std::path::Path;
9use std::sync::Arc;
10
11use parking_lot::RwLock;
12#[cfg(feature = "native")]
13use rustc_hash::FxHashMap;
14
15use crate::DocId;
16#[cfg(feature = "native")]
17use crate::directories::DirectoryWriter;
18use crate::directories::{Directory, SliceCachingDirectory};
19use crate::dsl::{Document, Field, Schema};
20use crate::error::{Error, Result};
21#[cfg(feature = "native")]
22use crate::segment::{SegmentBuilder, SegmentBuilderConfig, SegmentMerger};
23use crate::segment::{SegmentId, SegmentReader};
24use crate::structures::BlockPostingList;
25#[cfg(feature = "native")]
26use crate::tokenizer::BoxedTokenizer;
27
28#[cfg(feature = "native")]
29use tokio::sync::Mutex as AsyncMutex;
30
31/// Default file name for the slice cache
32pub const SLICE_CACHE_FILENAME: &str = "index.slicecache";
33
34/// Index configuration
35#[derive(Debug, Clone)]
36pub struct IndexConfig {
37    /// Number of threads for CPU-intensive tasks (search parallelism)
38    pub num_threads: usize,
39    /// Number of parallel segment builders (documents distributed round-robin)
40    pub num_indexing_threads: usize,
41    /// Number of threads for parallel block compression within each segment
42    pub num_compression_threads: usize,
43    /// Block cache size for term dictionary per segment
44    pub term_cache_blocks: usize,
45    /// Block cache size for document store per segment
46    pub store_cache_blocks: usize,
47    /// Max documents per segment before auto-commit
48    pub max_docs_per_segment: u32,
49    /// Min segments to trigger merge
50    pub merge_threshold: usize,
51    /// Index optimization mode (adaptive, size-optimized, performance-optimized)
52    pub optimization: crate::structures::IndexOptimization,
53}
54
55impl Default for IndexConfig {
56    fn default() -> Self {
57        #[cfg(feature = "native")]
58        let cpus = num_cpus::get().max(1);
59        #[cfg(not(feature = "native"))]
60        let cpus = 1;
61
62        Self {
63            num_threads: cpus,
64            num_indexing_threads: 1,
65            num_compression_threads: cpus,
66            term_cache_blocks: 256,
67            store_cache_blocks: 32,
68            max_docs_per_segment: 100_000,
69            merge_threshold: 5,
70            optimization: crate::structures::IndexOptimization::default(),
71        }
72    }
73}
74
75/// Multi-segment async Index
76///
77/// The main entry point for searching. Manages multiple segments
78/// and provides unified search across all of them.
79pub struct Index<D: Directory> {
80    directory: Arc<D>,
81    schema: Arc<Schema>,
82    config: IndexConfig,
83    segments: RwLock<Vec<Arc<SegmentReader>>>,
84    default_fields: Vec<crate::Field>,
85    tokenizers: Arc<crate::tokenizer::TokenizerRegistry>,
86    #[cfg(feature = "native")]
87    thread_pool: Arc<rayon::ThreadPool>,
88}
89
90impl<D: Directory> Index<D> {
91    /// Open an existing index from a directory
92    pub async fn open(directory: D, config: IndexConfig) -> Result<Self> {
93        let directory = Arc::new(directory);
94
95        // Read schema
96        let schema_slice = directory.open_read(Path::new("schema.json")).await?;
97        let schema_bytes = schema_slice.read_bytes().await?;
98        let schema: Schema = serde_json::from_slice(schema_bytes.as_slice())
99            .map_err(|e| Error::Serialization(e.to_string()))?;
100        let schema = Arc::new(schema);
101
102        // Read segment list
103        let segments = Self::load_segments(&directory, &schema, &config).await?;
104
105        #[cfg(feature = "native")]
106        let thread_pool = {
107            let pool = rayon::ThreadPoolBuilder::new()
108                .num_threads(config.num_threads)
109                .build()
110                .map_err(|e| Error::Io(std::io::Error::other(e)))?;
111            Arc::new(pool)
112        };
113
114        // Use schema's default_fields if specified, otherwise fall back to all indexed text fields
115        let default_fields: Vec<crate::Field> = if !schema.default_fields().is_empty() {
116            schema.default_fields().to_vec()
117        } else {
118            schema
119                .fields()
120                .filter(|(_, entry)| {
121                    entry.indexed && entry.field_type == crate::dsl::FieldType::Text
122                })
123                .map(|(field, _)| field)
124                .collect()
125        };
126
127        Ok(Self {
128            directory,
129            schema,
130            config,
131            segments: RwLock::new(segments),
132            default_fields,
133            tokenizers: Arc::new(crate::tokenizer::TokenizerRegistry::default()),
134            #[cfg(feature = "native")]
135            thread_pool,
136        })
137    }
138
139    async fn load_segments(
140        directory: &Arc<D>,
141        schema: &Arc<Schema>,
142        config: &IndexConfig,
143    ) -> Result<Vec<Arc<SegmentReader>>> {
144        // Read segments.json which lists all segment IDs
145        let segments_path = Path::new("segments.json");
146        if !directory.exists(segments_path).await? {
147            return Ok(Vec::new());
148        }
149
150        let segments_slice = directory.open_read(segments_path).await?;
151        let segments_bytes = segments_slice.read_bytes().await?;
152        let segment_ids: Vec<String> = serde_json::from_slice(segments_bytes.as_slice())
153            .map_err(|e| Error::Serialization(e.to_string()))?;
154
155        let mut segments = Vec::new();
156        let mut doc_id_offset = 0u32;
157
158        for id_str in segment_ids {
159            let segment_id = SegmentId::from_hex(&id_str)
160                .ok_or_else(|| Error::Corruption(format!("Invalid segment ID: {}", id_str)))?;
161            let reader = SegmentReader::open(
162                directory.as_ref(),
163                segment_id,
164                Arc::clone(schema),
165                doc_id_offset,
166                config.term_cache_blocks,
167            )
168            .await?;
169
170            doc_id_offset += reader.meta().num_docs;
171            segments.push(Arc::new(reader));
172        }
173
174        Ok(segments)
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    /// Total number of documents across all segments
188    pub fn num_docs(&self) -> u32 {
189        self.segments.read().iter().map(|s| s.num_docs()).sum()
190    }
191
192    /// Get a document by global doc_id (async)
193    pub async fn doc(&self, doc_id: DocId) -> Result<Option<Document>> {
194        let segments = self.segments.read().clone();
195
196        let mut offset = 0u32;
197        for segment in segments.iter() {
198            let segment_docs = segment.meta().num_docs;
199            if doc_id < offset + segment_docs {
200                let local_doc_id = doc_id - offset;
201                return segment.doc(local_doc_id).await;
202            }
203            offset += segment_docs;
204        }
205
206        Ok(None)
207    }
208
209    /// Get posting lists for a term across all segments (async)
210    pub async fn get_postings(
211        &self,
212        field: Field,
213        term: &[u8],
214    ) -> Result<Vec<(Arc<SegmentReader>, BlockPostingList)>> {
215        let segments = self.segments.read().clone();
216        let mut results = Vec::new();
217
218        for segment in segments.iter() {
219            if let Some(postings) = segment.get_postings(field, term).await? {
220                results.push((Arc::clone(segment), postings));
221            }
222        }
223
224        Ok(results)
225    }
226
227    /// Execute CPU-intensive work on thread pool (native only)
228    #[cfg(feature = "native")]
229    pub async fn spawn_blocking<F, R>(&self, f: F) -> R
230    where
231        F: FnOnce() -> R + Send + 'static,
232        R: Send + 'static,
233    {
234        let (tx, rx) = tokio::sync::oneshot::channel();
235        self.thread_pool.spawn(move || {
236            let result = f();
237            let _ = tx.send(result);
238        });
239        rx.await.expect("Thread pool task panicked")
240    }
241
242    /// Get segment readers for query execution
243    pub fn segment_readers(&self) -> Vec<Arc<SegmentReader>> {
244        self.segments.read().clone()
245    }
246
247    /// Reload segments from directory (after new segments added)
248    pub async fn reload(&self) -> Result<()> {
249        let new_segments = Self::load_segments(&self.directory, &self.schema, &self.config).await?;
250        *self.segments.write() = new_segments;
251        Ok(())
252    }
253
254    /// Search across all segments
255    pub async fn search(
256        &self,
257        query: &dyn crate::query::Query,
258        limit: usize,
259    ) -> Result<Vec<crate::query::SearchResult>> {
260        let segments = self.segments.read().clone();
261        let mut all_results = Vec::new();
262
263        for segment in &segments {
264            let results = crate::query::search_segment(segment.as_ref(), query, limit).await?;
265            all_results.extend(results);
266        }
267
268        // Sort by score descending
269        all_results.sort_by(|a, b| {
270            b.score
271                .partial_cmp(&a.score)
272                .unwrap_or(std::cmp::Ordering::Equal)
273        });
274        all_results.truncate(limit);
275
276        Ok(all_results)
277    }
278
279    /// Search and return results with document addresses (no document content)
280    pub async fn search_with_addresses(
281        &self,
282        query: &dyn crate::query::Query,
283        limit: usize,
284    ) -> Result<crate::query::SearchResponse> {
285        let segments = self.segments.read().clone();
286        let mut all_results: Vec<(u128, crate::query::SearchResult)> = Vec::new();
287
288        for segment in &segments {
289            let segment_id = segment.meta().id;
290            let results = crate::query::search_segment(segment.as_ref(), query, limit).await?;
291            for result in results {
292                all_results.push((segment_id, result));
293            }
294        }
295
296        // Sort by score descending
297        all_results.sort_by(|a, b| {
298            b.1.score
299                .partial_cmp(&a.1.score)
300                .unwrap_or(std::cmp::Ordering::Equal)
301        });
302        all_results.truncate(limit);
303
304        let total_hits = all_results.len() as u32;
305        let hits: Vec<crate::query::SearchHit> = all_results
306            .into_iter()
307            .map(|(segment_id, result)| crate::query::SearchHit {
308                address: crate::query::DocAddress::new(segment_id, result.doc_id),
309                score: result.score,
310            })
311            .collect();
312
313        Ok(crate::query::SearchResponse { hits, total_hits })
314    }
315
316    /// Get a document by its unique address (segment_id + local doc_id)
317    pub async fn get_document(
318        &self,
319        address: &crate::query::DocAddress,
320    ) -> Result<Option<Document>> {
321        let segment_id = address
322            .segment_id_u128()
323            .ok_or_else(|| Error::Query(format!("Invalid segment ID: {}", address.segment_id)))?;
324
325        let segments = self.segments.read().clone();
326        for segment in &segments {
327            if segment.meta().id == segment_id {
328                return segment.doc(address.doc_id).await;
329            }
330        }
331
332        Ok(None)
333    }
334
335    /// Get the default fields for this index
336    pub fn default_fields(&self) -> &[crate::Field] {
337        &self.default_fields
338    }
339
340    /// Set the default fields for query parsing
341    pub fn set_default_fields(&mut self, fields: Vec<crate::Field>) {
342        self.default_fields = fields;
343    }
344
345    /// Get the tokenizer registry
346    pub fn tokenizers(&self) -> &Arc<crate::tokenizer::TokenizerRegistry> {
347        &self.tokenizers
348    }
349
350    /// Create a query parser for this index
351    ///
352    /// If the schema contains query router rules, they will be used to route
353    /// queries to specific fields based on regex patterns.
354    pub fn query_parser(&self) -> crate::dsl::QueryLanguageParser {
355        // Check if schema has query routers
356        let query_routers = self.schema.query_routers();
357        if !query_routers.is_empty() {
358            // Try to create a router from the schema's rules
359            if let Ok(router) = crate::dsl::QueryFieldRouter::from_rules(query_routers) {
360                return crate::dsl::QueryLanguageParser::with_router(
361                    Arc::clone(&self.schema),
362                    self.default_fields.clone(),
363                    Arc::clone(&self.tokenizers),
364                    router,
365                );
366            }
367        }
368
369        // Fall back to parser without router
370        crate::dsl::QueryLanguageParser::new(
371            Arc::clone(&self.schema),
372            self.default_fields.clone(),
373            Arc::clone(&self.tokenizers),
374        )
375    }
376
377    /// Parse and search using a query string
378    ///
379    /// Accepts both query language syntax (field:term, AND, OR, NOT, grouping)
380    /// and simple text (tokenized and searched across default fields).
381    /// Returns document addresses (segment_id + doc_id) without document content.
382    pub async fn query(
383        &self,
384        query_str: &str,
385        limit: usize,
386    ) -> Result<crate::query::SearchResponse> {
387        let parser = self.query_parser();
388        let query = parser.parse(query_str).map_err(Error::Query)?;
389        self.search_with_addresses(query.as_ref(), limit).await
390    }
391}
392
393/// Methods for opening index with slice caching
394impl<D: Directory> Index<SliceCachingDirectory<D>> {
395    /// Open an index with slice caching, automatically loading the cache file if present
396    ///
397    /// This wraps the directory in a SliceCachingDirectory and attempts to load
398    /// any existing slice cache file to prefill the cache with hot data.
399    pub async fn open_with_cache(
400        directory: D,
401        config: IndexConfig,
402        cache_max_bytes: usize,
403    ) -> Result<Self> {
404        let caching_dir = SliceCachingDirectory::new(directory, cache_max_bytes);
405
406        // Try to load existing slice cache
407        let cache_path = Path::new(SLICE_CACHE_FILENAME);
408        if let Ok(true) = caching_dir.inner().exists(cache_path).await
409            && let Ok(slice) = caching_dir.inner().open_read(cache_path).await
410            && let Ok(bytes) = slice.read_bytes().await
411        {
412            let _ = caching_dir.deserialize(bytes.as_slice());
413        }
414
415        Self::open(caching_dir, config).await
416    }
417
418    /// Serialize the current slice cache to the index directory
419    ///
420    /// This saves all cached slices to a single file that can be loaded
421    /// on subsequent index opens for faster startup.
422    #[cfg(feature = "native")]
423    pub async fn save_slice_cache(&self) -> Result<()>
424    where
425        D: DirectoryWriter,
426    {
427        let cache_data = self.directory.serialize();
428        let cache_path = Path::new(SLICE_CACHE_FILENAME);
429        self.directory
430            .inner()
431            .write(cache_path, &cache_data)
432            .await?;
433        Ok(())
434    }
435
436    /// Get slice cache statistics
437    pub fn slice_cache_stats(&self) -> crate::directories::SliceCacheStats {
438        self.directory.stats()
439    }
440}
441
442/// Warm up the slice cache by opening an index and performing typical read operations
443///
444/// This function opens an index using a SliceCachingDirectory, performs operations
445/// that would typically be done during search (reading term dictionaries, posting lists),
446/// and then serializes the cache to a file for future use.
447///
448/// The resulting cache file contains all the "hot" data that was read during warmup,
449/// allowing subsequent index opens to prefill the cache and avoid cold-start latency.
450#[cfg(feature = "native")]
451pub async fn warmup_and_save_slice_cache<D: DirectoryWriter>(
452    directory: D,
453    config: IndexConfig,
454    cache_max_bytes: usize,
455) -> Result<()> {
456    let caching_dir = SliceCachingDirectory::new(directory, cache_max_bytes);
457    let index = Index::open(caching_dir, config).await?;
458
459    // Warm up by loading segment metadata and term dictionaries
460    // The SegmentReader::open already reads essential metadata
461    // Additional warmup can be done by iterating terms or doing sample queries
462
463    // Save the cache
464    index.save_slice_cache().await?;
465
466    Ok(())
467}
468
469#[cfg(feature = "native")]
470impl<D: Directory> Clone for Index<D> {
471    fn clone(&self) -> Self {
472        Self {
473            directory: Arc::clone(&self.directory),
474            schema: Arc::clone(&self.schema),
475            config: self.config.clone(),
476            segments: RwLock::new(self.segments.read().clone()),
477            default_fields: self.default_fields.clone(),
478            tokenizers: Arc::clone(&self.tokenizers),
479            thread_pool: Arc::clone(&self.thread_pool),
480        }
481    }
482}
483
484/// Async IndexWriter for adding documents and committing segments
485///
486/// Features:
487/// - Streams documents to disk immediately (no in-memory document storage)
488/// - Uses string interning for terms (reduced allocations)
489/// - Uses hashbrown HashMap (faster than BTreeMap)
490/// - Spills large posting lists to disk (bounded memory)
491/// - Uses memory-mapped files for intermediate data
492#[cfg(feature = "native")]
493pub struct IndexWriter<D: DirectoryWriter> {
494    directory: Arc<D>,
495    schema: Arc<Schema>,
496    config: IndexConfig,
497    builder_config: SegmentBuilderConfig,
498    tokenizers: FxHashMap<Field, BoxedTokenizer>,
499    /// Segment builder
500    builder: AsyncMutex<Option<SegmentBuilder>>,
501    /// List of committed segment IDs (hex strings)
502    segment_ids: AsyncMutex<Vec<String>>,
503}
504
505#[cfg(feature = "native")]
506impl<D: DirectoryWriter> IndexWriter<D> {
507    /// Create a new index in the directory
508    pub async fn create(directory: D, schema: Schema, config: IndexConfig) -> Result<Self> {
509        Self::create_with_config(directory, schema, config, SegmentBuilderConfig::default()).await
510    }
511
512    /// Create a new index with custom builder config
513    pub async fn create_with_config(
514        directory: D,
515        schema: Schema,
516        config: IndexConfig,
517        builder_config: SegmentBuilderConfig,
518    ) -> Result<Self> {
519        let directory = Arc::new(directory);
520        let schema = Arc::new(schema);
521
522        // Write schema
523        let schema_bytes =
524            serde_json::to_vec(&*schema).map_err(|e| Error::Serialization(e.to_string()))?;
525        directory
526            .write(Path::new("schema.json"), &schema_bytes)
527            .await?;
528
529        // Write empty segments list
530        let segments_bytes = serde_json::to_vec(&Vec::<String>::new())
531            .map_err(|e| Error::Serialization(e.to_string()))?;
532        directory
533            .write(Path::new("segments.json"), &segments_bytes)
534            .await?;
535
536        Ok(Self {
537            directory,
538            schema,
539            config,
540            builder_config,
541            tokenizers: FxHashMap::default(),
542            builder: AsyncMutex::new(None),
543            segment_ids: AsyncMutex::new(Vec::new()),
544        })
545    }
546
547    /// Open an existing index for writing
548    pub async fn open(directory: D, config: IndexConfig) -> Result<Self> {
549        Self::open_with_config(directory, config, SegmentBuilderConfig::default()).await
550    }
551
552    /// Open an existing index with custom builder config
553    pub async fn open_with_config(
554        directory: D,
555        config: IndexConfig,
556        builder_config: SegmentBuilderConfig,
557    ) -> Result<Self> {
558        let directory = Arc::new(directory);
559
560        // Read schema
561        let schema_slice = directory.open_read(Path::new("schema.json")).await?;
562        let schema_bytes = schema_slice.read_bytes().await?;
563        let schema: Schema = serde_json::from_slice(schema_bytes.as_slice())
564            .map_err(|e| Error::Serialization(e.to_string()))?;
565        let schema = Arc::new(schema);
566
567        // Read existing segment IDs (hex strings)
568        let segments_slice = directory.open_read(Path::new("segments.json")).await?;
569        let segments_bytes = segments_slice.read_bytes().await?;
570        let segment_ids: Vec<String> = serde_json::from_slice(segments_bytes.as_slice())
571            .map_err(|e| Error::Serialization(e.to_string()))?;
572
573        Ok(Self {
574            directory,
575            schema,
576            config,
577            builder_config,
578            tokenizers: FxHashMap::default(),
579            builder: AsyncMutex::new(None),
580            segment_ids: AsyncMutex::new(segment_ids),
581        })
582    }
583
584    /// Get the schema
585    pub fn schema(&self) -> &Schema {
586        &self.schema
587    }
588
589    /// Set tokenizer for a field
590    pub fn set_tokenizer<T: crate::tokenizer::Tokenizer>(&mut self, field: Field, tokenizer: T) {
591        self.tokenizers.insert(field, Box::new(tokenizer));
592    }
593
594    /// Add a document
595    ///
596    /// Documents are streamed to disk immediately, keeping memory usage bounded.
597    /// When the builder reaches `max_docs_per_segment`, it is committed and a new one starts.
598    pub async fn add_document(&self, doc: Document) -> Result<DocId> {
599        let mut builder_guard = self.builder.lock().await;
600
601        // Initialize builder if needed
602        if builder_guard.is_none() {
603            let mut builder =
604                SegmentBuilder::new((*self.schema).clone(), self.builder_config.clone())?;
605            for (field, tokenizer) in &self.tokenizers {
606                builder.set_tokenizer(*field, tokenizer.clone_box());
607            }
608            *builder_guard = Some(builder);
609        }
610
611        let builder = builder_guard.as_mut().unwrap();
612        let doc_id = builder.add_document(doc)?;
613
614        // Check if we need to commit
615        if builder.num_docs() >= self.config.max_docs_per_segment {
616            let full_builder = builder_guard.take().unwrap();
617            drop(builder_guard); // Release lock before async operation
618            self.commit_builder(full_builder).await?;
619        }
620
621        Ok(doc_id)
622    }
623
624    /// Commit all pending segments to disk
625    pub async fn commit(&self) -> Result<()> {
626        let mut builder_guard = self.builder.lock().await;
627
628        if let Some(builder) = builder_guard.take()
629            && builder.num_docs() > 0
630        {
631            drop(builder_guard);
632            self.commit_builder(builder).await?;
633        }
634
635        Ok(())
636    }
637
638    async fn commit_builder(&self, builder: SegmentBuilder) -> Result<()> {
639        let segment_id = SegmentId::new();
640        builder.build(self.directory.as_ref(), segment_id).await?;
641
642        // Update segment list
643        let mut segment_ids = self.segment_ids.lock().await;
644        segment_ids.push(segment_id.to_hex());
645
646        let segments_bytes =
647            serde_json::to_vec(&*segment_ids).map_err(|e| Error::Serialization(e.to_string()))?;
648        self.directory
649            .write(Path::new("segments.json"), &segments_bytes)
650            .await?;
651
652        // Check if merge is needed
653        if segment_ids.len() >= self.config.merge_threshold {
654            drop(segment_ids); // Release lock
655            self.maybe_merge().await?;
656        }
657
658        Ok(())
659    }
660
661    /// Merge segments if threshold is reached
662    async fn maybe_merge(&self) -> Result<()> {
663        let segment_ids = self.segment_ids.lock().await;
664
665        if segment_ids.len() < self.config.merge_threshold {
666            return Ok(());
667        }
668
669        let ids_to_merge: Vec<String> = segment_ids.clone();
670        drop(segment_ids);
671
672        // Load segment readers
673        let mut readers = Vec::new();
674        let mut doc_offset = 0u32;
675
676        for id_str in &ids_to_merge {
677            let segment_id = SegmentId::from_hex(id_str)
678                .ok_or_else(|| Error::Corruption(format!("Invalid segment ID: {}", id_str)))?;
679            let reader = SegmentReader::open(
680                self.directory.as_ref(),
681                segment_id,
682                Arc::clone(&self.schema),
683                doc_offset,
684                self.config.term_cache_blocks,
685            )
686            .await?;
687            doc_offset += reader.meta().num_docs;
688            readers.push(reader);
689        }
690
691        // Merge into new segment
692        let merger = SegmentMerger::new(Arc::clone(&self.schema));
693        let new_segment_id = SegmentId::new();
694        merger
695            .merge(self.directory.as_ref(), &readers, new_segment_id)
696            .await?;
697
698        // Update segment list
699        let mut segment_ids = self.segment_ids.lock().await;
700        segment_ids.clear();
701        segment_ids.push(new_segment_id.to_hex());
702
703        let segments_bytes =
704            serde_json::to_vec(&*segment_ids).map_err(|e| Error::Serialization(e.to_string()))?;
705        self.directory
706            .write(Path::new("segments.json"), &segments_bytes)
707            .await?;
708
709        // Delete old segments
710        for id_str in ids_to_merge {
711            if let Some(segment_id) = SegmentId::from_hex(&id_str) {
712                let _ = crate::segment::delete_segment(self.directory.as_ref(), segment_id).await;
713            }
714        }
715
716        Ok(())
717    }
718
719    /// Force merge all segments into one
720    pub async fn force_merge(&self) -> Result<()> {
721        // First commit any pending documents
722        self.commit().await?;
723
724        let segment_ids = self.segment_ids.lock().await;
725        if segment_ids.len() <= 1 {
726            return Ok(());
727        }
728
729        let ids_to_merge: Vec<String> = segment_ids.clone();
730        drop(segment_ids);
731
732        // Load all segments
733        let mut readers = Vec::new();
734        let mut doc_offset = 0u32;
735
736        for id_str in &ids_to_merge {
737            let segment_id = SegmentId::from_hex(id_str)
738                .ok_or_else(|| Error::Corruption(format!("Invalid segment ID: {}", id_str)))?;
739            let reader = SegmentReader::open(
740                self.directory.as_ref(),
741                segment_id,
742                Arc::clone(&self.schema),
743                doc_offset,
744                self.config.term_cache_blocks,
745            )
746            .await?;
747            doc_offset += reader.meta().num_docs;
748            readers.push(reader);
749        }
750
751        // Merge
752        let merger = SegmentMerger::new(Arc::clone(&self.schema));
753        let new_segment_id = SegmentId::new();
754        merger
755            .merge(self.directory.as_ref(), &readers, new_segment_id)
756            .await?;
757
758        // Update segment list
759        let mut segment_ids = self.segment_ids.lock().await;
760        segment_ids.clear();
761        segment_ids.push(new_segment_id.to_hex());
762
763        let segments_bytes =
764            serde_json::to_vec(&*segment_ids).map_err(|e| Error::Serialization(e.to_string()))?;
765        self.directory
766            .write(Path::new("segments.json"), &segments_bytes)
767            .await?;
768
769        // Delete old segments
770        for id_str in ids_to_merge {
771            if let Some(segment_id) = SegmentId::from_hex(&id_str) {
772                let _ = crate::segment::delete_segment(self.directory.as_ref(), segment_id).await;
773            }
774        }
775
776        Ok(())
777    }
778}
779
780#[cfg(test)]
781mod tests {
782    use super::*;
783    use crate::directories::RamDirectory;
784    use crate::dsl::SchemaBuilder;
785
786    #[tokio::test]
787    async fn test_index_create_and_search() {
788        let mut schema_builder = SchemaBuilder::default();
789        let title = schema_builder.add_text_field("title", true, true);
790        let body = schema_builder.add_text_field("body", true, true);
791        let schema = schema_builder.build();
792
793        let dir = RamDirectory::new();
794        let config = IndexConfig::default();
795
796        // Create index and add documents
797        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
798            .await
799            .unwrap();
800
801        let mut doc1 = Document::new();
802        doc1.add_text(title, "Hello World");
803        doc1.add_text(body, "This is the first document");
804        writer.add_document(doc1).await.unwrap();
805
806        let mut doc2 = Document::new();
807        doc2.add_text(title, "Goodbye World");
808        doc2.add_text(body, "This is the second document");
809        writer.add_document(doc2).await.unwrap();
810
811        writer.commit().await.unwrap();
812
813        // Open for reading
814        let index = Index::open(dir, config).await.unwrap();
815        assert_eq!(index.num_docs(), 2);
816
817        // Check postings
818        let postings = index.get_postings(title, b"world").await.unwrap();
819        assert_eq!(postings.len(), 1); // One segment
820        assert_eq!(postings[0].1.doc_count(), 2); // Two docs with "world"
821
822        // Retrieve document
823        let doc = index.doc(0).await.unwrap().unwrap();
824        assert_eq!(doc.get_first(title).unwrap().as_text(), Some("Hello World"));
825    }
826
827    #[tokio::test]
828    async fn test_multiple_segments() {
829        let mut schema_builder = SchemaBuilder::default();
830        let title = schema_builder.add_text_field("title", true, true);
831        let schema = schema_builder.build();
832
833        let dir = RamDirectory::new();
834        let config = IndexConfig {
835            max_docs_per_segment: 5, // Small segments for testing
836            merge_threshold: 10,     // Don't auto-merge
837            ..Default::default()
838        };
839
840        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
841            .await
842            .unwrap();
843
844        // Add documents in batches to create multiple segments
845        for batch in 0..3 {
846            for i in 0..5 {
847                let mut doc = Document::new();
848                doc.add_text(title, format!("Document {} batch {}", i, batch));
849                writer.add_document(doc).await.unwrap();
850            }
851            writer.commit().await.unwrap();
852        }
853
854        // Open and check
855        let index = Index::open(dir, config).await.unwrap();
856        assert_eq!(index.num_docs(), 15);
857        assert_eq!(index.segment_readers().len(), 3);
858    }
859
860    #[tokio::test]
861    async fn test_segment_merge() {
862        let mut schema_builder = SchemaBuilder::default();
863        let title = schema_builder.add_text_field("title", true, true);
864        let schema = schema_builder.build();
865
866        let dir = RamDirectory::new();
867        let config = IndexConfig {
868            max_docs_per_segment: 3,
869            merge_threshold: 100, // Don't auto-merge
870            ..Default::default()
871        };
872
873        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
874            .await
875            .unwrap();
876
877        // Create multiple segments
878        for i in 0..9 {
879            let mut doc = Document::new();
880            doc.add_text(title, format!("Document {}", i));
881            writer.add_document(doc).await.unwrap();
882        }
883        writer.commit().await.unwrap();
884
885        // Should have 3 segments
886        let index = Index::open(dir.clone(), config.clone()).await.unwrap();
887        assert_eq!(index.segment_readers().len(), 3);
888
889        // Force merge
890        let writer = IndexWriter::open(dir.clone(), config.clone())
891            .await
892            .unwrap();
893        writer.force_merge().await.unwrap();
894
895        // Should have 1 segment now
896        let index = Index::open(dir, config).await.unwrap();
897        assert_eq!(index.segment_readers().len(), 1);
898        assert_eq!(index.num_docs(), 9);
899
900        // Verify all documents accessible
901        for i in 0..9 {
902            let doc = index.doc(i).await.unwrap().unwrap();
903            assert_eq!(
904                doc.get_first(title).unwrap().as_text(),
905                Some(format!("Document {}", i).as_str())
906            );
907        }
908    }
909
910    #[tokio::test]
911    async fn test_match_query() {
912        let mut schema_builder = SchemaBuilder::default();
913        let title = schema_builder.add_text_field("title", true, true);
914        let body = schema_builder.add_text_field("body", true, true);
915        let schema = schema_builder.build();
916
917        let dir = RamDirectory::new();
918        let config = IndexConfig::default();
919
920        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
921            .await
922            .unwrap();
923
924        let mut doc1 = Document::new();
925        doc1.add_text(title, "rust programming");
926        doc1.add_text(body, "Learn rust language");
927        writer.add_document(doc1).await.unwrap();
928
929        let mut doc2 = Document::new();
930        doc2.add_text(title, "python programming");
931        doc2.add_text(body, "Learn python language");
932        writer.add_document(doc2).await.unwrap();
933
934        writer.commit().await.unwrap();
935
936        let index = Index::open(dir, config).await.unwrap();
937
938        // Test match query with multiple default fields
939        let results = index.query("rust", 10).await.unwrap();
940        assert_eq!(results.hits.len(), 1);
941
942        // Test match query with multiple tokens
943        let results = index.query("rust programming", 10).await.unwrap();
944        assert!(!results.hits.is_empty());
945
946        // Verify hit has address (segment_id + doc_id)
947        let hit = &results.hits[0];
948        assert!(!hit.address.segment_id.is_empty(), "Should have segment_id");
949
950        // Verify document retrieval by address
951        let doc = index.get_document(&hit.address).await.unwrap().unwrap();
952        assert!(
953            !doc.field_values().is_empty(),
954            "Doc should have field values"
955        );
956
957        // Also verify doc retrieval directly by global doc_id
958        let doc = index.doc(0).await.unwrap().unwrap();
959        assert!(
960            !doc.field_values().is_empty(),
961            "Doc should have field values"
962        );
963    }
964
965    #[tokio::test]
966    async fn test_slice_cache_warmup_and_load() {
967        use crate::directories::SliceCachingDirectory;
968
969        let mut schema_builder = SchemaBuilder::default();
970        let title = schema_builder.add_text_field("title", true, true);
971        let body = schema_builder.add_text_field("body", true, true);
972        let schema = schema_builder.build();
973
974        let dir = RamDirectory::new();
975        let config = IndexConfig::default();
976
977        // Create index with some documents
978        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
979            .await
980            .unwrap();
981
982        for i in 0..10 {
983            let mut doc = Document::new();
984            doc.add_text(title, format!("Document {} about rust", i));
985            doc.add_text(body, format!("This is body text number {}", i));
986            writer.add_document(doc).await.unwrap();
987        }
988        writer.commit().await.unwrap();
989
990        // Open with slice caching and perform some operations to warm up cache
991        let caching_dir = SliceCachingDirectory::new(dir.clone(), 1024 * 1024);
992        let index = Index::open(caching_dir, config.clone()).await.unwrap();
993
994        // Perform a search to warm up the cache
995        let results = index.query("rust", 10).await.unwrap();
996        assert!(!results.hits.is_empty());
997
998        // Check cache stats - should have cached some data
999        let stats = index.slice_cache_stats();
1000        assert!(stats.total_bytes > 0, "Cache should have data after search");
1001
1002        // Save the cache
1003        index.save_slice_cache().await.unwrap();
1004
1005        // Verify cache file was written
1006        assert!(
1007            dir.exists(Path::new(super::SLICE_CACHE_FILENAME))
1008                .await
1009                .unwrap()
1010        );
1011
1012        // Now open with cache loading
1013        let index2 = Index::open_with_cache(dir.clone(), config.clone(), 1024 * 1024)
1014            .await
1015            .unwrap();
1016
1017        // Cache should be prefilled
1018        let stats2 = index2.slice_cache_stats();
1019        assert!(
1020            stats2.total_bytes > 0,
1021            "Cache should be prefilled from file"
1022        );
1023
1024        // Search should still work
1025        let results2 = index2.query("rust", 10).await.unwrap();
1026        assert_eq!(results.hits.len(), results2.hits.len());
1027    }
1028
1029    #[tokio::test]
1030    async fn test_multivalue_field_indexing_and_search() {
1031        let mut schema_builder = SchemaBuilder::default();
1032        let uris = schema_builder.add_text_field("uris", true, true);
1033        let title = schema_builder.add_text_field("title", true, true);
1034        let schema = schema_builder.build();
1035
1036        let dir = RamDirectory::new();
1037        let config = IndexConfig::default();
1038
1039        // Create index and add document with multi-value field
1040        let writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
1041            .await
1042            .unwrap();
1043
1044        let mut doc = Document::new();
1045        doc.add_text(uris, "one");
1046        doc.add_text(uris, "two");
1047        doc.add_text(title, "Test Document");
1048        writer.add_document(doc).await.unwrap();
1049
1050        // Add another document with different uris
1051        let mut doc2 = Document::new();
1052        doc2.add_text(uris, "three");
1053        doc2.add_text(title, "Another Document");
1054        writer.add_document(doc2).await.unwrap();
1055
1056        writer.commit().await.unwrap();
1057
1058        // Open for reading
1059        let index = Index::open(dir, config).await.unwrap();
1060        assert_eq!(index.num_docs(), 2);
1061
1062        // Verify document retrieval preserves all values
1063        let doc = index.doc(0).await.unwrap().unwrap();
1064        let all_uris: Vec<_> = doc.get_all(uris).collect();
1065        assert_eq!(all_uris.len(), 2, "Should have 2 uris values");
1066        assert_eq!(all_uris[0].as_text(), Some("one"));
1067        assert_eq!(all_uris[1].as_text(), Some("two"));
1068
1069        // Verify to_json returns array for multi-value field
1070        let json = doc.to_json(index.schema());
1071        let uris_json = json.get("uris").unwrap();
1072        assert!(uris_json.is_array(), "Multi-value field should be an array");
1073        let uris_arr = uris_json.as_array().unwrap();
1074        assert_eq!(uris_arr.len(), 2);
1075        assert_eq!(uris_arr[0].as_str(), Some("one"));
1076        assert_eq!(uris_arr[1].as_str(), Some("two"));
1077
1078        // Verify both values are searchable
1079        let results = index.query("uris:one", 10).await.unwrap();
1080        assert_eq!(results.hits.len(), 1, "Should find doc with 'one'");
1081        assert_eq!(results.hits[0].address.doc_id, 0);
1082
1083        let results = index.query("uris:two", 10).await.unwrap();
1084        assert_eq!(results.hits.len(), 1, "Should find doc with 'two'");
1085        assert_eq!(results.hits[0].address.doc_id, 0);
1086
1087        let results = index.query("uris:three", 10).await.unwrap();
1088        assert_eq!(results.hits.len(), 1, "Should find doc with 'three'");
1089        assert_eq!(results.hits[0].address.doc_id, 1);
1090
1091        // Verify searching for non-existent value returns no results
1092        let results = index.query("uris:nonexistent", 10).await.unwrap();
1093        assert_eq!(results.hits.len(), 0, "Should not find non-existent value");
1094    }
1095}
hermes_core/index.rs

hermes_core/
index.rs
