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

hermes_core/index/
mod.rs
