Skip to main content

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 std::sync::Arc;
16
17mod searcher;
18pub use searcher::Searcher;
19
20#[cfg(feature = "native")]
21mod primary_key;
22#[cfg(feature = "native")]
23mod reader;
24#[cfg(feature = "native")]
25mod vector_builder;
26#[cfg(all(feature = "wasm", not(feature = "native")))]
27mod wasm_writer;
28#[cfg(feature = "native")]
29mod writer;
30#[cfg(feature = "native")]
31pub use primary_key::PrimaryKeyIndex;
32#[cfg(feature = "native")]
33pub use reader::IndexReader;
34#[cfg(all(feature = "wasm", not(feature = "native")))]
35pub use wasm_writer::IndexWriter as WasmIndexWriter;
36#[cfg(feature = "native")]
37pub use writer::{IndexWriter, PreparedCommit};
38
39mod metadata;
40pub use metadata::{FieldVectorMeta, INDEX_META_FILENAME, IndexMetadata, VectorIndexState};
41
42#[cfg(feature = "native")]
43mod helpers;
44#[cfg(feature = "native")]
45pub use helpers::{
46    IndexingStats, SchemaConfig, SchemaFieldConfig, create_index_at_path, create_index_from_sdl,
47    index_documents_from_reader, index_json_document, parse_schema,
48};
49
50/// Default file name for the slice cache
51pub const SLICE_CACHE_FILENAME: &str = "index.slicecache";
52
53/// Index configuration
54#[derive(Debug, Clone)]
55pub struct IndexConfig {
56    /// Number of threads for CPU-intensive tasks (search parallelism)
57    pub num_threads: usize,
58    /// Number of parallel segment builders (documents distributed round-robin)
59    pub num_indexing_threads: usize,
60    /// Number of threads for parallel block compression within each segment
61    pub num_compression_threads: usize,
62    /// Block cache size for term dictionary per segment
63    pub term_cache_blocks: usize,
64    /// Block cache size for document store per segment
65    pub store_cache_blocks: usize,
66    /// Max memory (bytes) across all builders before auto-commit (global limit)
67    pub max_indexing_memory_bytes: usize,
68    /// Merge policy for background segment merging
69    pub merge_policy: Box<dyn crate::merge::MergePolicy>,
70    /// Index optimization mode (adaptive, size-optimized, performance-optimized)
71    pub optimization: crate::structures::IndexOptimization,
72    /// Reload interval in milliseconds for IndexReader (how often to check for new segments)
73    pub reload_interval_ms: u64,
74    /// Maximum number of concurrent background merges (default: 4)
75    pub max_concurrent_merges: usize,
76}
77
78impl Default for IndexConfig {
79    fn default() -> Self {
80        #[cfg(feature = "native")]
81        let indexing_threads = crate::default_indexing_threads();
82        #[cfg(not(feature = "native"))]
83        let indexing_threads = 1;
84
85        #[cfg(feature = "native")]
86        let compression_threads = crate::default_compression_threads();
87        #[cfg(not(feature = "native"))]
88        let compression_threads = 1;
89
90        Self {
91            num_threads: indexing_threads,
92            num_indexing_threads: 1, // Increase to 2+ for production to avoid stalls during segment build
93            num_compression_threads: compression_threads,
94            term_cache_blocks: 256,
95            store_cache_blocks: 32,
96            max_indexing_memory_bytes: 256 * 1024 * 1024, // 256 MB default
97            merge_policy: Box::new(crate::merge::TieredMergePolicy::default()),
98            optimization: crate::structures::IndexOptimization::default(),
99            reload_interval_ms: 1000, // 1 second default
100            max_concurrent_merges: 4,
101        }
102    }
103}
104
105/// Multi-segment async Index
106///
107/// The central concept for search. Owns segment lifecycle and provides:
108/// - `Index::create()` / `Index::open()` - create or open an index
109/// - `index.writer()` - get an IndexWriter for adding documents
110/// - `index.reader()` - get an IndexReader for searching with reload policy
111///
112/// All segment management is delegated to SegmentManager.
113#[cfg(feature = "native")]
114pub struct Index<D: crate::directories::DirectoryWriter + 'static> {
115    directory: Arc<D>,
116    schema: Arc<Schema>,
117    config: IndexConfig,
118    /// Segment manager - owns segments, tracker, metadata, and trained structures
119    segment_manager: Arc<crate::merge::SegmentManager<D>>,
120    /// Cached reader (created lazily, reused across calls)
121    cached_reader: tokio::sync::OnceCell<IndexReader<D>>,
122}
123
124#[cfg(feature = "native")]
125impl<D: crate::directories::DirectoryWriter + 'static> Index<D> {
126    /// Create a new index in the directory
127    pub async fn create(directory: D, schema: Schema, config: IndexConfig) -> Result<Self> {
128        let directory = Arc::new(directory);
129        let schema = Arc::new(schema);
130        // Directory-layer metrics (cold writes, lazy reads) carry the index label
131        directory.set_index_label(schema.index_label());
132        let metadata = IndexMetadata::new((*schema).clone());
133
134        let segment_manager = Arc::new(crate::merge::SegmentManager::new(
135            Arc::clone(&directory),
136            Arc::clone(&schema),
137            metadata,
138            config.merge_policy.clone_box(),
139            config.term_cache_blocks,
140            config.max_concurrent_merges,
141        ));
142
143        // Save initial metadata
144        segment_manager.update_metadata(|_| {}).await?;
145
146        Ok(Self {
147            directory,
148            schema,
149            config,
150            segment_manager,
151            cached_reader: tokio::sync::OnceCell::new(),
152        })
153    }
154
155    /// Open an existing index from a directory
156    pub async fn open(directory: D, config: IndexConfig) -> Result<Self> {
157        let directory = Arc::new(directory);
158
159        // Load metadata (includes schema)
160        let metadata = IndexMetadata::load(directory.as_ref()).await?;
161        let schema = Arc::new(metadata.schema.clone());
162        // Directory-layer metrics (cold writes, lazy reads) carry the index label
163        directory.set_index_label(schema.index_label());
164
165        let segment_manager = Arc::new(crate::merge::SegmentManager::new(
166            Arc::clone(&directory),
167            Arc::clone(&schema),
168            metadata,
169            config.merge_policy.clone_box(),
170            config.term_cache_blocks,
171            config.max_concurrent_merges,
172        ));
173
174        // Load trained structures into SegmentManager's ArcSwap
175        segment_manager.load_and_publish_trained().await;
176
177        Ok(Self {
178            directory,
179            schema,
180            config,
181            segment_manager,
182            cached_reader: tokio::sync::OnceCell::new(),
183        })
184    }
185
186    /// Get the schema
187    pub fn schema(&self) -> &Schema {
188        &self.schema
189    }
190
191    /// Get the schema as an Arc reference (avoids clone when Arc is needed)
192    pub fn schema_arc(&self) -> &Arc<Schema> {
193        &self.schema
194    }
195
196    /// Get a reference to the underlying directory
197    pub fn directory(&self) -> &D {
198        &self.directory
199    }
200
201    /// Get the segment manager
202    pub fn segment_manager(&self) -> &Arc<crate::merge::SegmentManager<D>> {
203        &self.segment_manager
204    }
205
206    /// Get an IndexReader for searching (with reload policy)
207    ///
208    /// The reader is cached and reused across calls. The reader's internal
209    /// searcher will reload segments based on its reload interval (configurable via IndexConfig).
210    pub async fn reader(&self) -> Result<&IndexReader<D>> {
211        self.cached_reader
212            .get_or_try_init(|| async {
213                IndexReader::from_segment_manager(
214                    Arc::clone(&self.schema),
215                    Arc::clone(&self.segment_manager),
216                    self.config.term_cache_blocks,
217                    self.config.reload_interval_ms,
218                )
219                .await
220            })
221            .await
222    }
223
224    /// Get the config
225    pub fn config(&self) -> &IndexConfig {
226        &self.config
227    }
228
229    /// Get segment readers for query execution (convenience method)
230    pub async fn segment_readers(&self) -> Result<Vec<Arc<crate::segment::SegmentReader>>> {
231        let reader = self.reader().await?;
232        let searcher = reader.searcher().await?;
233        Ok(searcher.segment_readers().to_vec())
234    }
235
236    /// Total number of documents across all segments
237    pub async fn num_docs(&self) -> Result<u32> {
238        let reader = self.reader().await?;
239        let searcher = reader.searcher().await?;
240        Ok(searcher.num_docs())
241    }
242
243    /// Get default fields for search
244    pub fn default_fields(&self) -> Vec<crate::Field> {
245        if !self.schema.default_fields().is_empty() {
246            self.schema.default_fields().to_vec()
247        } else {
248            self.schema
249                .fields()
250                .filter(|(_, entry)| {
251                    entry.indexed && entry.field_type == crate::dsl::FieldType::Text
252                })
253                .map(|(field, _)| field)
254                .collect()
255        }
256    }
257
258    /// Get tokenizer registry
259    pub fn tokenizers(&self) -> Arc<crate::tokenizer::TokenizerRegistry> {
260        Arc::new(crate::tokenizer::TokenizerRegistry::default())
261    }
262
263    /// Create a query parser for this index
264    pub fn query_parser(&self) -> crate::dsl::QueryLanguageParser {
265        let default_fields = self.default_fields();
266        let tokenizers = self.tokenizers();
267
268        let query_routers = self.schema.query_routers();
269        if !query_routers.is_empty()
270            && let Ok(router) = crate::dsl::QueryFieldRouter::from_rules(query_routers)
271        {
272            return crate::dsl::QueryLanguageParser::with_router(
273                Arc::clone(&self.schema),
274                default_fields,
275                tokenizers,
276                router,
277            );
278        }
279
280        crate::dsl::QueryLanguageParser::new(Arc::clone(&self.schema), default_fields, tokenizers)
281    }
282
283    /// Parse and search using a query string
284    pub async fn query(
285        &self,
286        query_str: &str,
287        limit: usize,
288    ) -> Result<crate::query::SearchResponse> {
289        self.query_offset(query_str, limit, 0).await
290    }
291
292    /// Query with offset for pagination
293    pub async fn query_offset(
294        &self,
295        query_str: &str,
296        limit: usize,
297        offset: usize,
298    ) -> Result<crate::query::SearchResponse> {
299        let parser = self.query_parser();
300        let query = parser
301            .parse(query_str)
302            .map_err(crate::error::Error::Query)?;
303        self.search_offset(query.as_ref(), limit, offset).await
304    }
305
306    /// Search and return results
307    pub async fn search(
308        &self,
309        query: &dyn crate::query::Query,
310        limit: usize,
311    ) -> Result<crate::query::SearchResponse> {
312        self.search_offset(query, limit, 0).await
313    }
314
315    /// Search with offset for pagination
316    pub async fn search_offset(
317        &self,
318        query: &dyn crate::query::Query,
319        limit: usize,
320        offset: usize,
321    ) -> Result<crate::query::SearchResponse> {
322        let reader = self.reader().await?;
323        let searcher = reader.searcher().await?;
324
325        #[cfg(feature = "sync")]
326        let (results, total_seen) = {
327            // Sync search: rayon handles segment parallelism internally.
328            // On multi-threaded tokio, use block_in_place to yield the worker;
329            // on single-threaded (tests), call directly.
330            let runtime_flavor = tokio::runtime::Handle::current().runtime_flavor();
331            if runtime_flavor == tokio::runtime::RuntimeFlavor::MultiThread {
332                tokio::task::block_in_place(|| {
333                    searcher.search_with_offset_and_count_sync(query, limit, offset)
334                })?
335            } else {
336                searcher.search_with_offset_and_count_sync(query, limit, offset)?
337            }
338        };
339
340        #[cfg(not(feature = "sync"))]
341        let (results, total_seen) = {
342            searcher
343                .search_with_offset_and_count(query, limit, offset)
344                .await?
345        };
346
347        let total_hits = total_seen;
348        let hits: Vec<crate::query::SearchHit> = results
349            .into_iter()
350            .map(|result| crate::query::SearchHit {
351                address: crate::query::DocAddress::new(result.segment_id, result.doc_id),
352                score: result.score,
353                matched_fields: result.extract_ordinals(),
354            })
355            .collect();
356
357        Ok(crate::query::SearchResponse { hits, total_hits })
358    }
359
360    /// Get a document by its unique address
361    pub async fn get_document(
362        &self,
363        address: &crate::query::DocAddress,
364    ) -> Result<Option<crate::dsl::Document>> {
365        let reader = self.reader().await?;
366        let searcher = reader.searcher().await?;
367        searcher.get_document(address).await
368    }
369
370    /// Get posting lists for a term across all segments
371    pub async fn get_postings(
372        &self,
373        field: crate::Field,
374        term: &[u8],
375    ) -> Result<
376        Vec<(
377            Arc<crate::segment::SegmentReader>,
378            crate::structures::BlockPostingList,
379        )>,
380    > {
381        let segments = self.segment_readers().await?;
382        let mut results = Vec::new();
383
384        for segment in segments {
385            if let Some(postings) = segment.get_postings(field, term).await? {
386                results.push((segment, postings));
387            }
388        }
389
390        Ok(results)
391    }
392}
393
394/// Native-only methods for Index
395#[cfg(feature = "native")]
396impl<D: crate::directories::DirectoryWriter + 'static> Index<D> {
397    /// Get an IndexWriter for adding documents
398    pub fn writer(&self) -> writer::IndexWriter<D> {
399        writer::IndexWriter::from_index(self)
400    }
401}
402
403#[cfg(test)]
404mod tests;
405
406// (tests moved to index/tests/ module)