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::{
41    FieldVectorMeta, INDEX_META_FILENAME, IndexMetadata, SegmentMetaInfo, VectorIndexState,
42};
43
44#[cfg(feature = "native")]
45mod helpers;
46#[cfg(feature = "native")]
47pub use helpers::{
48    IndexingStats, SchemaConfig, SchemaFieldConfig, create_index_at_path, create_index_from_sdl,
49    index_documents_from_reader, index_json_document, parse_schema,
50};
51
52/// Default file name for the slice cache
53pub const SLICE_CACHE_FILENAME: &str = "index.slicecache";
54
55/// Index configuration
56#[derive(Debug, Clone)]
57pub struct IndexConfig {
58    /// Number of threads for CPU-intensive tasks (search parallelism)
59    pub num_threads: usize,
60    /// Number of parallel segment builders (documents distributed round-robin)
61    pub num_indexing_threads: usize,
62    /// Number of threads for parallel block compression within each segment
63    pub num_compression_threads: usize,
64    /// Block cache size for term dictionary per segment
65    pub term_cache_blocks: usize,
66    /// Block cache size for document store per segment
67    pub store_cache_blocks: usize,
68    /// Max memory (bytes) across all builders before auto-commit (global limit)
69    pub max_indexing_memory_bytes: usize,
70    /// Merge policy for background segment merging
71    pub merge_policy: Box<dyn crate::merge::MergePolicy>,
72    /// Index optimization mode (adaptive, size-optimized, performance-optimized)
73    pub optimization: crate::structures::IndexOptimization,
74    /// Reload interval in milliseconds for IndexReader (how often to check for new segments)
75    pub reload_interval_ms: u64,
76    /// Maximum number of concurrent background merges per index (default: 4)
77    pub max_concurrent_merges: usize,
78    /// Application-wide background merge gate shared by clones of this
79    /// config. The per-index limit alone multiplied large merge working sets
80    /// by the number of active indexes.
81    #[cfg(feature = "native")]
82    pub background_merge_permits: Arc<tokio::sync::Semaphore>,
83    /// Wall-clock budget for merge-time BP reorder per field (only applies
84    /// when the index has `reorder_on_merge`). A truncated pass still writes
85    /// a valid, better-ordered segment; it is marked `bp_converged = false`
86    /// and the background optimizer deepens it later (warm-started).
87    /// `None` = unbudgeted (BP runs to full depth inside the merge, which can
88    /// hold a merge slot for 10-30+ minutes on 10M+ doc outputs).
89    pub merge_bp_time_budget: Option<std::time::Duration>,
90    /// Memory budget (bytes) for the BP forward index during reorder passes
91    /// (merge-time and background). When a large segment's forward index
92    /// would exceed this, the highest-df dims are dropped from BP's input
93    /// (logged loudly) — clustering quality degrades gracefully. Production
94    /// evidence: 18M-doc merges exceeded the former 2 GB default and dropped
95    /// ~10% of eligible dims; hosts with less headroom may lower this.
96    pub bp_memory_budget_bytes: usize,
97    /// Hard limit on simultaneous whole-segment BP rewrites. This is shared
98    /// by all indexes opened from clones of this config and applies to
99    /// optimizer, merge-time, and manual reorder passes. It is deliberately
100    /// separate from the Rayon pool width: one pass can already use every
101    /// background CPU thread and consume the full BP memory budget.
102    #[cfg(feature = "native")]
103    pub background_reorder_permits: Arc<tokio::sync::Semaphore>,
104    /// Optional process/application-owned Rayon pool for BP work. Supplying
105    /// one lets every index and the optimizer share the same worker threads;
106    /// `None` lazily uses one process-wide cores/2 fallback pool.
107    #[cfg(feature = "native")]
108    pub background_reorder_pool: Option<Arc<rayon::ThreadPool>>,
109}
110
111impl Default for IndexConfig {
112    fn default() -> Self {
113        #[cfg(feature = "native")]
114        let indexing_threads = crate::default_indexing_threads();
115        #[cfg(not(feature = "native"))]
116        let indexing_threads = 1;
117
118        #[cfg(feature = "native")]
119        let compression_threads = crate::default_compression_threads();
120        #[cfg(not(feature = "native"))]
121        let compression_threads = 1;
122
123        Self {
124            num_threads: indexing_threads,
125            num_indexing_threads: 1, // Increase to 2+ for production to avoid stalls during segment build
126            num_compression_threads: compression_threads,
127            term_cache_blocks: 256,
128            store_cache_blocks: 32,
129            max_indexing_memory_bytes: 256 * 1024 * 1024, // 256 MB default
130            // large_scale: wide fan-in + budget/scored selection. Safe for
131            // small indexes too (tier floors only shape *when* segments
132            // merge); merge-time BP is wall-clock budgeted, so giant merges
133            // cannot hold slots indefinitely.
134            merge_policy: Box::new(crate::merge::TieredMergePolicy::large_scale()),
135            optimization: crate::structures::IndexOptimization::default(),
136            reload_interval_ms: 1000, // 1 second default
137            max_concurrent_merges: 4,
138            #[cfg(feature = "native")]
139            background_merge_permits: Arc::new(tokio::sync::Semaphore::new(4)),
140            merge_bp_time_budget: Some(std::time::Duration::from_secs(600)),
141            // 24 GB — mirrors segment::reorder::DEFAULT_MEMORY_BUDGET (that
142            // module is native-only; IndexConfig also compiles for wasm).
143            // A cap, not an allocation: usage is proportional to the segment
144            // being reordered (~4 B/posting + ~32 B/doc). Sized from prod
145            // evidence: a 58M-doc/5B-posting pass estimated 20.1 GB, which
146            // 8/16 GB budgets trimmed by dropping highest-df dims.
147            // 24 GB overflows 32-bit usize (wasm32) — reorder never runs
148            // there, so any large value works; use usize::MAX.
149            #[cfg(target_pointer_width = "64")]
150            bp_memory_budget_bytes: 24 * 1024 * 1024 * 1024,
151            #[cfg(not(target_pointer_width = "64"))]
152            bp_memory_budget_bytes: usize::MAX,
153            #[cfg(feature = "native")]
154            background_reorder_permits: Arc::new(tokio::sync::Semaphore::new(2)),
155            #[cfg(feature = "native")]
156            background_reorder_pool: None,
157        }
158    }
159}
160
161/// Multi-segment async Index
162///
163/// The central concept for search. Owns segment lifecycle and provides:
164/// - `Index::create()` / `Index::open()` - create or open an index
165/// - `index.writer()` - get an IndexWriter for adding documents
166/// - `index.reader()` - get an IndexReader for searching with reload policy
167///
168/// All segment management is delegated to SegmentManager.
169#[cfg(feature = "native")]
170pub struct Index<D: crate::directories::DirectoryWriter + 'static> {
171    directory: Arc<D>,
172    schema: Arc<Schema>,
173    config: IndexConfig,
174    /// Segment manager - owns segments, tracker, metadata, and trained structures
175    segment_manager: Arc<crate::merge::SegmentManager<D>>,
176    /// Cached reader (created lazily, reused across calls)
177    cached_reader: tokio::sync::OnceCell<IndexReader<D>>,
178}
179
180#[cfg(feature = "native")]
181impl<D: crate::directories::DirectoryWriter + 'static> Index<D> {
182    /// Create a new index in the directory
183    pub async fn create(directory: D, schema: Schema, config: IndexConfig) -> Result<Self> {
184        let directory = Arc::new(directory);
185        let schema = Arc::new(schema);
186        // Directory-layer metrics (cold writes, lazy reads) carry the index label
187        directory.set_index_label(schema.index_label());
188        let metadata = IndexMetadata::new((*schema).clone());
189
190        let segment_manager = Arc::new(crate::merge::SegmentManager::new(
191            Arc::clone(&directory),
192            Arc::clone(&schema),
193            metadata,
194            config.merge_policy.clone_box(),
195            config.term_cache_blocks,
196            config.max_concurrent_merges,
197            Arc::clone(&config.background_merge_permits),
198            config.merge_bp_time_budget,
199            config.bp_memory_budget_bytes,
200            Arc::clone(&config.background_reorder_permits),
201            config.background_reorder_pool.clone(),
202        ));
203
204        // Save initial metadata
205        segment_manager.update_metadata(|_| {}).await?;
206
207        Ok(Self {
208            directory,
209            schema,
210            config,
211            segment_manager,
212            cached_reader: tokio::sync::OnceCell::new(),
213        })
214    }
215
216    /// Open an existing index from a directory
217    pub async fn open(directory: D, config: IndexConfig) -> Result<Self> {
218        let directory = Arc::new(directory);
219
220        // Load metadata (includes schema)
221        let metadata = IndexMetadata::load(directory.as_ref()).await?;
222        let schema = Arc::new(metadata.schema.clone());
223        // Directory-layer metrics (cold writes, lazy reads) carry the index label
224        directory.set_index_label(schema.index_label());
225
226        let segment_manager = Arc::new(crate::merge::SegmentManager::new(
227            Arc::clone(&directory),
228            Arc::clone(&schema),
229            metadata,
230            config.merge_policy.clone_box(),
231            config.term_cache_blocks,
232            config.max_concurrent_merges,
233            Arc::clone(&config.background_merge_permits),
234            config.merge_bp_time_budget,
235            config.bp_memory_budget_bytes,
236            Arc::clone(&config.background_reorder_permits),
237            config.background_reorder_pool.clone(),
238        ));
239
240        // Load trained structures into SegmentManager's ArcSwap
241        segment_manager.load_and_publish_trained().await;
242
243        Ok(Self {
244            directory,
245            schema,
246            config,
247            segment_manager,
248            cached_reader: tokio::sync::OnceCell::new(),
249        })
250    }
251
252    /// Get the schema
253    pub fn schema(&self) -> &Schema {
254        &self.schema
255    }
256
257    /// Get the schema as an Arc reference (avoids clone when Arc is needed)
258    pub fn schema_arc(&self) -> &Arc<Schema> {
259        &self.schema
260    }
261
262    /// Get a reference to the underlying directory
263    pub fn directory(&self) -> &D {
264        &self.directory
265    }
266
267    /// Get the segment manager
268    pub fn segment_manager(&self) -> &Arc<crate::merge::SegmentManager<D>> {
269        &self.segment_manager
270    }
271
272    /// Get an IndexReader for searching (with reload policy)
273    ///
274    /// The reader is cached and reused across calls. The reader's internal
275    /// searcher will reload segments based on its reload interval (configurable via IndexConfig).
276    pub async fn reader(&self) -> Result<&IndexReader<D>> {
277        self.cached_reader
278            .get_or_try_init(|| async {
279                IndexReader::from_segment_manager(
280                    Arc::clone(&self.schema),
281                    Arc::clone(&self.segment_manager),
282                    self.config.term_cache_blocks,
283                    self.config.reload_interval_ms,
284                )
285                .await
286            })
287            .await
288    }
289
290    /// Get the config
291    pub fn config(&self) -> &IndexConfig {
292        &self.config
293    }
294
295    /// Get segment readers for query execution (convenience method)
296    pub async fn segment_readers(&self) -> Result<Vec<Arc<crate::segment::SegmentReader>>> {
297        let reader = self.reader().await?;
298        let searcher = reader.searcher().await?;
299        Ok(searcher.segment_readers().to_vec())
300    }
301
302    /// Total number of documents across all segments
303    pub async fn num_docs(&self) -> Result<u32> {
304        let reader = self.reader().await?;
305        let searcher = reader.searcher().await?;
306        Ok(searcher.num_docs())
307    }
308
309    /// Get default fields for search
310    pub fn default_fields(&self) -> Vec<crate::Field> {
311        if !self.schema.default_fields().is_empty() {
312            self.schema.default_fields().to_vec()
313        } else {
314            self.schema
315                .fields()
316                .filter(|(_, entry)| {
317                    entry.indexed && entry.field_type == crate::dsl::FieldType::Text
318                })
319                .map(|(field, _)| field)
320                .collect()
321        }
322    }
323
324    /// Get tokenizer registry
325    pub fn tokenizers(&self) -> Arc<crate::tokenizer::TokenizerRegistry> {
326        Arc::new(crate::tokenizer::TokenizerRegistry::default())
327    }
328
329    /// Create a query parser for this index
330    pub fn query_parser(&self) -> crate::dsl::QueryLanguageParser {
331        let default_fields = self.default_fields();
332        let tokenizers = self.tokenizers();
333
334        let query_routers = self.schema.query_routers();
335        if !query_routers.is_empty()
336            && let Ok(router) = crate::dsl::QueryFieldRouter::from_rules(query_routers)
337        {
338            return crate::dsl::QueryLanguageParser::with_router(
339                Arc::clone(&self.schema),
340                default_fields,
341                tokenizers,
342                router,
343            );
344        }
345
346        crate::dsl::QueryLanguageParser::new(Arc::clone(&self.schema), default_fields, tokenizers)
347    }
348
349    /// Parse and search using a query string
350    pub async fn query(
351        &self,
352        query_str: &str,
353        limit: usize,
354    ) -> Result<crate::query::SearchResponse> {
355        self.query_offset(query_str, limit, 0).await
356    }
357
358    /// Query with offset for pagination
359    pub async fn query_offset(
360        &self,
361        query_str: &str,
362        limit: usize,
363        offset: usize,
364    ) -> Result<crate::query::SearchResponse> {
365        let parser = self.query_parser();
366        let query = parser
367            .parse(query_str)
368            .map_err(crate::error::Error::Query)?;
369        self.search_offset(query.as_ref(), limit, offset).await
370    }
371
372    /// Search and return results
373    pub async fn search(
374        &self,
375        query: &dyn crate::query::Query,
376        limit: usize,
377    ) -> Result<crate::query::SearchResponse> {
378        self.search_offset(query, limit, 0).await
379    }
380
381    /// Search with offset for pagination
382    pub async fn search_offset(
383        &self,
384        query: &dyn crate::query::Query,
385        limit: usize,
386        offset: usize,
387    ) -> Result<crate::query::SearchResponse> {
388        let reader = self.reader().await?;
389        let searcher = reader.searcher().await?;
390
391        #[cfg(feature = "sync")]
392        let (results, total_seen) = {
393            // Sync search: rayon handles segment parallelism internally.
394            // On multi-threaded tokio, use block_in_place to yield the worker;
395            // on single-threaded (tests), call directly.
396            let runtime_flavor = tokio::runtime::Handle::current().runtime_flavor();
397            if runtime_flavor == tokio::runtime::RuntimeFlavor::MultiThread {
398                tokio::task::block_in_place(|| {
399                    searcher.search_with_offset_and_count_sync(query, limit, offset)
400                })?
401            } else {
402                searcher.search_with_offset_and_count_sync(query, limit, offset)?
403            }
404        };
405
406        #[cfg(not(feature = "sync"))]
407        let (results, total_seen) = {
408            searcher
409                .search_with_offset_and_count(query, limit, offset)
410                .await?
411        };
412
413        let total_hits = total_seen;
414        let hits: Vec<crate::query::SearchHit> = results
415            .into_iter()
416            .map(|result| crate::query::SearchHit {
417                address: crate::query::DocAddress::new(result.segment_id, result.doc_id),
418                score: result.score,
419                matched_fields: result.extract_ordinals(),
420            })
421            .collect();
422
423        Ok(crate::query::SearchResponse { hits, total_hits })
424    }
425
426    /// Get a document by its unique address
427    pub async fn get_document(
428        &self,
429        address: &crate::query::DocAddress,
430    ) -> Result<Option<crate::dsl::Document>> {
431        let reader = self.reader().await?;
432        let searcher = reader.searcher().await?;
433        searcher.get_document(address).await
434    }
435
436    /// Get posting lists for a term across all segments
437    pub async fn get_postings(
438        &self,
439        field: crate::Field,
440        term: &[u8],
441    ) -> Result<
442        Vec<(
443            Arc<crate::segment::SegmentReader>,
444            crate::structures::BlockPostingList,
445        )>,
446    > {
447        let segments = self.segment_readers().await?;
448        let mut results = Vec::new();
449
450        for segment in segments {
451            if let Some(postings) = segment.get_postings(field, term).await? {
452                results.push((segment, postings));
453            }
454        }
455
456        Ok(results)
457    }
458}
459
460/// Native-only methods for Index
461#[cfg(feature = "native")]
462impl<D: crate::directories::DirectoryWriter + 'static> Index<D> {
463    /// Get an IndexWriter for adding documents
464    pub fn writer(&self) -> writer::IndexWriter<D> {
465        writer::IndexWriter::from_index(self)
466    }
467}
468
469#[cfg(test)]
470mod tests;
471
472// (tests moved to index/tests/ module)
hermes_core/index/mod.rs

hermes_core/index/
mod.rs
