Skip to main content

hermes_core/index/
searcher.rs

1//! Searcher - read-only search over pre-built segments
2//!
3//! This module provides `Searcher` for read-only search access to indexes.
4//! It can be used standalone (for wasm/read-only) or via `IndexReader` (for native).
5
6use std::sync::Arc;
7
8use rustc_hash::FxHashMap;
9
10use crate::directories::Directory;
11use crate::dsl::Schema;
12use crate::error::Result;
13use crate::query::LazyGlobalStats;
14use crate::segment::{SegmentId, SegmentReader, TrainedVectorStructures};
15#[cfg(feature = "native")]
16use crate::segment::{SegmentSnapshot, SegmentTracker};
17
18/// Immutable resources that must stay identical across `IndexReader` reloads.
19/// The search pool exists only when synchronous scoring is compiled in; native
20/// async-only builds retain the cache policy without spawning unused threads.
21#[cfg(feature = "native")]
22#[derive(Clone)]
23pub(crate) struct SearcherResources {
24    pub(crate) term_cache_blocks: usize,
25    pub(crate) store_cache: Arc<crate::segment::SharedStoreCache>,
26    pub(crate) bmp_io_gate: Arc<super::BmpIoGate>,
27    pub(crate) bmp_io_concurrency: usize,
28    #[cfg(feature = "sync")]
29    pub(crate) search_pool: Arc<rayon::ThreadPool>,
30}
31
32#[cfg(feature = "native")]
33impl SearcherResources {
34    pub(crate) fn new(
35        term_cache_blocks: usize,
36        store_cache_budget_bytes: usize,
37        num_threads: usize,
38        bmp_io_concurrency: usize,
39    ) -> Result<Self> {
40        if num_threads == 0 {
41            return Err(crate::Error::Internal(
42                "IndexConfig.num_threads must be greater than zero".into(),
43            ));
44        }
45        if bmp_io_concurrency == 0 {
46            return Err(crate::Error::Internal(
47                "IndexConfig.bmp_io_concurrency must be greater than zero".into(),
48            ));
49        }
50
51        #[cfg(feature = "sync")]
52        let search_pool = super::shared_search_pool(num_threads)?;
53
54        Ok(Self {
55            term_cache_blocks,
56            store_cache: super::shared_store_cache(store_cache_budget_bytes),
57            bmp_io_gate: super::shared_bmp_io_gate(bmp_io_concurrency),
58            bmp_io_concurrency,
59            #[cfg(feature = "sync")]
60            search_pool,
61        })
62    }
63}
64
65/// Searcher - provides search over loaded segments
66///
67/// For wasm/read-only use, create via `Searcher::open()`.
68/// For native use with Index, this is created via `IndexReader`.
69pub struct Searcher<D: Directory + 'static> {
70    /// Segment snapshot holding refs - prevents deletion during native use
71    #[cfg(feature = "native")]
72    _snapshot: SegmentSnapshot,
73    /// PhantomData for the directory generic
74    _phantom: std::marker::PhantomData<D>,
75    /// Loaded segment readers
76    segments: Vec<Arc<SegmentReader>>,
77    /// Schema
78    schema: Arc<Schema>,
79    /// Default fields for search
80    default_fields: Vec<crate::Field>,
81    /// Tokenizers
82    tokenizers: Arc<crate::tokenizer::TokenizerRegistry>,
83    /// One immutable generation of all index-global ANN artifacts.
84    trained_vectors: Arc<TrainedVectorStructures>,
85    /// Lazy global statistics for cross-segment IDF computation
86    global_stats: Arc<LazyGlobalStats>,
87    /// O(1) segment lookup by segment_id
88    segment_map: FxHashMap<u128, usize>,
89    /// Total document count across all segments
90    total_docs: u32,
91    /// Bounded process-wide-by-width pool for the complete nested search tree.
92    #[cfg(feature = "sync")]
93    search_pool: Arc<rayon::ThreadPool>,
94    /// Shared random-I/O gate and per-query wave width for BMP.
95    #[cfg(feature = "native")]
96    bmp_io_gate: Arc<super::BmpIoGate>,
97    #[cfg(feature = "native")]
98    bmp_io_concurrency: usize,
99}
100
101impl<D: Directory + 'static> Searcher<D> {
102    /// Create a Searcher directly from segment IDs
103    ///
104    /// This is a simpler initialization path that doesn't require SegmentManager.
105    /// Use this for read-only access to pre-built indexes.
106    pub async fn open(
107        directory: Arc<D>,
108        schema: Arc<Schema>,
109        segment_ids: &[String],
110        term_cache_blocks: usize,
111    ) -> Result<Self> {
112        const STANDALONE_STORE_CACHE_BYTES: usize = 32 * 1024 * 1024;
113        #[cfg(feature = "native")]
114        let store_cache = super::shared_store_cache(STANDALONE_STORE_CACHE_BYTES);
115        #[cfg(not(feature = "native"))]
116        let store_cache = Arc::new(crate::segment::SharedStoreCache::new(
117            STANDALONE_STORE_CACHE_BYTES,
118        ));
119        Self::create(
120            directory,
121            schema,
122            segment_ids,
123            Arc::new(TrainedVectorStructures::default()),
124            term_cache_blocks,
125            store_cache,
126        )
127        .await
128    }
129
130    /// Create from a snapshot (for native IndexReader use)
131    #[cfg(feature = "native")]
132    pub(crate) async fn from_snapshot(
133        directory: Arc<D>,
134        schema: Arc<Schema>,
135        snapshot: SegmentSnapshot,
136        trained_vectors: Arc<TrainedVectorStructures>,
137        resources: SearcherResources,
138    ) -> Result<Self> {
139        let (segments, default_fields, global_stats, segment_map, total_docs) = Self::load_common(
140            &directory,
141            &schema,
142            snapshot.segment_ids(),
143            &trained_vectors,
144            resources.term_cache_blocks,
145            Arc::clone(&resources.store_cache),
146            &[],
147        )
148        .await?;
149
150        Ok(Self {
151            _snapshot: snapshot,
152            _phantom: std::marker::PhantomData,
153            segments,
154            schema,
155            default_fields,
156            tokenizers: Arc::new(crate::tokenizer::TokenizerRegistry::default()),
157            trained_vectors,
158            global_stats,
159            segment_map,
160            total_docs,
161            #[cfg(feature = "sync")]
162            search_pool: resources.search_pool,
163            bmp_io_gate: resources.bmp_io_gate,
164            bmp_io_concurrency: resources.bmp_io_concurrency,
165        })
166    }
167
168    /// Create from a snapshot, reusing existing segment readers for unchanged segments.
169    /// This avoids re-opening mmaps, fast fields, sparse indexes, etc. for segments
170    /// that weren't touched by merge.
171    #[cfg(feature = "native")]
172    pub(crate) async fn from_snapshot_reuse(
173        directory: Arc<D>,
174        schema: Arc<Schema>,
175        snapshot: SegmentSnapshot,
176        trained_vectors: Arc<TrainedVectorStructures>,
177        resources: SearcherResources,
178        existing_segments: &[Arc<SegmentReader>],
179    ) -> Result<Self> {
180        let (segments, default_fields, global_stats, segment_map, total_docs) = Self::load_common(
181            &directory,
182            &schema,
183            snapshot.segment_ids(),
184            &trained_vectors,
185            resources.term_cache_blocks,
186            Arc::clone(&resources.store_cache),
187            existing_segments,
188        )
189        .await?;
190
191        Ok(Self {
192            _snapshot: snapshot,
193            _phantom: std::marker::PhantomData,
194            segments,
195            schema,
196            default_fields,
197            tokenizers: Arc::new(crate::tokenizer::TokenizerRegistry::default()),
198            trained_vectors,
199            global_stats,
200            segment_map,
201            total_docs,
202            #[cfg(feature = "sync")]
203            search_pool: resources.search_pool,
204            bmp_io_gate: resources.bmp_io_gate,
205            bmp_io_concurrency: resources.bmp_io_concurrency,
206        })
207    }
208
209    /// Internal create method
210    async fn create(
211        directory: Arc<D>,
212        schema: Arc<Schema>,
213        segment_ids: &[String],
214        trained_vectors: Arc<TrainedVectorStructures>,
215        term_cache_blocks: usize,
216        store_cache: Arc<crate::segment::SharedStoreCache>,
217    ) -> Result<Self> {
218        let (segments, default_fields, global_stats, segment_map, total_docs) = Self::load_common(
219            &directory,
220            &schema,
221            segment_ids,
222            &trained_vectors,
223            term_cache_blocks,
224            store_cache,
225            &[],
226        )
227        .await?;
228
229        #[cfg(feature = "native")]
230        let _snapshot = {
231            let tracker = Arc::new(SegmentTracker::new());
232            SegmentSnapshot::new(tracker, segment_ids.to_vec())
233        };
234
235        #[cfg(feature = "sync")]
236        let search_pool = super::shared_search_pool(crate::default_search_threads())?;
237        #[cfg(feature = "native")]
238        let bmp_io_concurrency = 4;
239        #[cfg(feature = "native")]
240        let bmp_io_gate = super::shared_bmp_io_gate(bmp_io_concurrency);
241
242        let _ = directory; // suppress unused warning on wasm
243        Ok(Self {
244            #[cfg(feature = "native")]
245            _snapshot,
246            _phantom: std::marker::PhantomData,
247            segments,
248            schema,
249            default_fields,
250            tokenizers: Arc::new(crate::tokenizer::TokenizerRegistry::default()),
251            trained_vectors,
252            global_stats,
253            segment_map,
254            total_docs,
255            #[cfg(feature = "sync")]
256            search_pool,
257            #[cfg(feature = "native")]
258            bmp_io_gate,
259            #[cfg(feature = "native")]
260            bmp_io_concurrency,
261        })
262    }
263
264    /// Common loading logic shared by create and from_snapshot
265    async fn load_common(
266        directory: &Arc<D>,
267        schema: &Arc<Schema>,
268        segment_ids: &[String],
269        trained_vectors: &Arc<TrainedVectorStructures>,
270        term_cache_blocks: usize,
271        store_cache: Arc<crate::segment::SharedStoreCache>,
272        existing_segments: &[Arc<SegmentReader>],
273    ) -> Result<(
274        Vec<Arc<SegmentReader>>,
275        Vec<crate::Field>,
276        Arc<LazyGlobalStats>,
277        FxHashMap<u128, usize>,
278        u32,
279    )> {
280        let segments = Self::load_segments(
281            directory,
282            schema,
283            segment_ids,
284            trained_vectors,
285            term_cache_blocks,
286            store_cache,
287            existing_segments,
288        )
289        .await?;
290        let default_fields = Self::build_default_fields(schema);
291        let global_stats = Arc::new(LazyGlobalStats::new(segments.clone()));
292        let (segment_map, total_docs) = Self::build_lookup_tables(&segments);
293        Ok((
294            segments,
295            default_fields,
296            global_stats,
297            segment_map,
298            total_docs,
299        ))
300    }
301
302    /// Load segment readers from IDs (parallel loading for performance).
303    /// Reuses existing segment readers for unchanged segments when `existing_segments`
304    /// is non-empty — avoids re-opening mmaps, fast fields, sparse indexes, etc.
305    async fn load_segments(
306        directory: &Arc<D>,
307        schema: &Arc<Schema>,
308        segment_ids: &[String],
309        trained_vectors: &Arc<TrainedVectorStructures>,
310        term_cache_blocks: usize,
311        store_cache: Arc<crate::segment::SharedStoreCache>,
312        existing_segments: &[Arc<SegmentReader>],
313    ) -> Result<Vec<Arc<SegmentReader>>> {
314        // Build lookup from existing segment readers for reuse
315        let existing_map: FxHashMap<u128, Arc<SegmentReader>> = existing_segments
316            .iter()
317            .map(|seg| (seg.meta().id, Arc::clone(seg)))
318            .collect();
319
320        // Parse segment IDs from metadata. A key that fails to parse means the
321        // metadata is corrupt; fail loud instead of silently serving results
322        // without that segment's documents (the merge path errors with
323        // Corruption on the same input — search must not disagree).
324        let mut valid_segments: Vec<(usize, SegmentId)> = Vec::with_capacity(segment_ids.len());
325        for (idx, id_str) in segment_ids.iter().enumerate() {
326            let sid = SegmentId::from_hex(id_str).ok_or_else(|| {
327                crate::error::Error::Corruption(format!(
328                    "Invalid segment ID in metadata: {id_str:?}"
329                ))
330            })?;
331            valid_segments.push((idx, sid));
332        }
333
334        // Separate into reusable and new segments
335        let mut reused: Vec<(usize, Arc<SegmentReader>)> = Vec::new();
336        let mut to_load: Vec<(usize, SegmentId)> = Vec::new();
337        for (idx, sid) in &valid_segments {
338            if let Some(existing) = existing_map.get(&sid.0) {
339                reused.push((*idx, Arc::clone(existing)));
340            } else {
341                to_load.push((*idx, *sid));
342            }
343        }
344
345        if !existing_segments.is_empty() {
346            log::info!(
347                "[searcher] index={} reusing {} segment readers, loading {} new",
348                schema.index_label(),
349                reused.len(),
350                to_load.len(),
351            );
352        }
353
354        // Include the live directory allocation in document-cache keys.
355        // Segment IDs can legitimately repeat when two independent indexes
356        // are copied/opened in the same process.
357        let store_cache_directory_namespace = Arc::as_ptr(directory) as usize;
358
359        // Load only NEW segments in parallel
360        let futures: Vec<_> = to_load
361            .iter()
362            .map(|(_, segment_id)| {
363                let dir = Arc::clone(directory);
364                let sch = Arc::clone(schema);
365                let store_cache = Arc::clone(&store_cache);
366                let sid = *segment_id;
367                async move {
368                    SegmentReader::open_with_store_cache(
369                        dir.as_ref(),
370                        sid,
371                        sch,
372                        term_cache_blocks,
373                        store_cache_directory_namespace,
374                        store_cache,
375                    )
376                    .await
377                }
378            })
379            .collect();
380
381        let results = futures::future::join_all(futures).await;
382
383        // Collect newly loaded results — fail fast if any segment fails to open
384        let mut loaded: Vec<(usize, Arc<SegmentReader>)> = Vec::with_capacity(valid_segments.len());
385
386        // Add reused segments
387        loaded.extend(reused);
388
389        // Add newly loaded segments
390        for ((idx, sid), result) in to_load.into_iter().zip(results) {
391            match result {
392                Ok(mut reader) => {
393                    // Inject the single immutable index-level artifact generation.
394                    reader.set_trained_vectors(Arc::clone(trained_vectors));
395                    loaded.push((idx, Arc::new(reader)));
396                }
397                Err(e) => {
398                    return Err(crate::error::Error::Internal(format!(
399                        "Failed to open segment {:016x}: {:?}",
400                        sid.0, e
401                    )));
402                }
403            }
404        }
405
406        // Sort by original index to maintain deterministic ordering
407        loaded.sort_by_key(|(idx, _)| *idx);
408
409        let segments: Vec<Arc<SegmentReader>> = loaded.into_iter().map(|(_, seg)| seg).collect();
410
411        // Keep heap, file-backed address space, and pinned residency separate.
412        // Mapped bytes are not necessarily resident; process RSS is the
413        // authoritative whole-process residency measurement.
414        let total_docs: u64 = segments.iter().map(|s| s.meta().num_docs as u64).sum();
415        let mut total_heap = 0usize;
416        let mut total_file_backed = 0u64;
417        let mut total_pinned = 0u64;
418        let mut total_pin_intended = 0u64;
419        for seg in &segments {
420            let stats = seg.memory_stats();
421            let heap = stats.estimated_heap_bytes();
422            let file_backed = stats.file_backed_bytes();
423            total_heap = total_heap.saturating_add(heap);
424            total_file_backed = total_file_backed.saturating_add(file_backed);
425            total_pinned = total_pinned.saturating_add(stats.pinned_metadata_bytes);
426            total_pin_intended = total_pin_intended.saturating_add(stats.pin_intended_bytes);
427            log::info!(
428                "[searcher] index={} segment {:016x}: docs={}, heap_estimate={} \
429                 (term_cache={}, store_cache={}, sparse_vectors={}, dense_vectors={}), \
430                 file_backed={} (term_bloom={}, sparse_vectors={}, dense_vectors={}), \
431                 pinned_metadata={} of {} eligible \
432                 (sparse_vectors={} of {}, dense_vectors={} of {})",
433                schema.index_label(),
434                stats.segment_id,
435                stats.num_docs,
436                crate::format_bytes(heap as u64),
437                crate::format_bytes(stats.term_dict_cache_bytes as u64),
438                crate::format_bytes(stats.store_cache_bytes as u64),
439                crate::format_bytes(stats.sparse_heap_bytes as u64),
440                crate::format_bytes(stats.dense_heap_bytes as u64),
441                crate::format_bytes(file_backed),
442                crate::format_bytes(stats.term_bloom_file_bytes),
443                crate::format_bytes(stats.sparse_file_backed_bytes),
444                crate::format_bytes(stats.dense_file_backed_bytes),
445                crate::format_bytes(stats.pinned_metadata_bytes),
446                crate::format_bytes(stats.pin_intended_bytes),
447                crate::format_bytes(stats.sparse_pinned_metadata_bytes),
448                crate::format_bytes(stats.sparse_pin_intended_bytes),
449                crate::format_bytes(stats.dense_pinned_metadata_bytes),
450                crate::format_bytes(stats.dense_pin_intended_bytes),
451            );
452        }
453        // Log process RSS if available (helps diagnose OOM)
454        let rss_bytes = process_rss_bytes();
455        log::info!(
456            "[searcher] index={} loaded {} segments: total_docs={}, heap_estimate={}, \
457             file_backed={}, pinned_metadata={} of {} eligible, \
458             shared_store_cache={} in {} blocks, process_rss={}",
459            schema.index_label(),
460            segments.len(),
461            total_docs,
462            crate::format_bytes(total_heap as u64),
463            crate::format_bytes(total_file_backed),
464            crate::format_bytes(total_pinned),
465            crate::format_bytes(total_pin_intended),
466            crate::format_bytes(store_cache.total_bytes() as u64),
467            store_cache.total_blocks(),
468            crate::format_bytes(rss_bytes),
469        );
470
471        // One ANN-health line per dense field across the whole index, so an
472        // operator reads leaf skew and extent fragmentation from N_fields
473        // lines instead of N_segments × N_fields open-time lines.
474        let mut ann_per_field: std::collections::BTreeMap<u32, (u64, u64, u32, u32, u64)> =
475            std::collections::BTreeMap::new();
476        for segment in segments.iter() {
477            for &field_id in segment.vector_indexes().keys() {
478                if let Some(health) = segment.ann_health(crate::Field(field_id)) {
479                    let entry = ann_per_field.entry(field_id).or_default();
480                    entry.0 += health.vectors;
481                    entry.1 += health.payload_bytes;
482                    entry.2 += health.runs;
483                    entry.3 += health.clusters_nonempty;
484                    entry.4 = entry.4.max(health.largest_cluster_vectors);
485                }
486            }
487        }
488        for (field_id, (vectors, payload, runs, clusters, largest)) in ann_per_field {
489            log::info!(
490                "[ann_health] index={} field={field_id} aggregate: vectors={vectors} \
491                 payload={} runs={runs} fragmentation={:.2} worst_leaf_vectors={largest}",
492                schema.index_label(),
493                crate::format_bytes(payload),
494                if clusters == 0 {
495                    0.0
496                } else {
497                    f64::from(runs) / f64::from(clusters)
498                },
499            );
500        }
501
502        Ok(segments)
503    }
504
505    /// Build default fields from schema
506    fn build_default_fields(schema: &Schema) -> Vec<crate::Field> {
507        if !schema.default_fields().is_empty() {
508            schema.default_fields().to_vec()
509        } else {
510            schema
511                .fields()
512                .filter(|(_, entry)| {
513                    entry.indexed && entry.field_type == crate::dsl::FieldType::Text
514                })
515                .map(|(field, _)| field)
516                .collect()
517        }
518    }
519
520    /// Get the schema
521    pub fn schema(&self) -> &Schema {
522        &self.schema
523    }
524
525    /// Get segment readers
526    pub fn segment_readers(&self) -> &[Arc<SegmentReader>] {
527        &self.segments
528    }
529
530    /// Get default fields for search
531    pub fn default_fields(&self) -> &[crate::Field] {
532        &self.default_fields
533    }
534
535    /// Get tokenizer registry
536    pub fn tokenizers(&self) -> &crate::tokenizer::TokenizerRegistry {
537        &self.tokenizers
538    }
539
540    /// Get trained centroids
541    pub fn trained_centroids(&self) -> &FxHashMap<u32, Arc<crate::structures::CoarseCentroids>> {
542        &self.trained_vectors.centroids
543    }
544
545    pub fn trained_binary_quantizers(
546        &self,
547    ) -> &FxHashMap<u32, Arc<crate::structures::BinaryCoarseQuantizer>> {
548        &self.trained_vectors.binary_quantizers
549    }
550
551    /// Get lazy global statistics for cross-segment IDF computation
552    pub fn global_stats(&self) -> &Arc<LazyGlobalStats> {
553        &self.global_stats
554    }
555
556    /// Build O(1) lookup tables from loaded segments
557    fn build_lookup_tables(segments: &[Arc<SegmentReader>]) -> (FxHashMap<u128, usize>, u32) {
558        let mut segment_map = FxHashMap::default();
559        let mut total = 0u32;
560        for (i, seg) in segments.iter().enumerate() {
561            segment_map.insert(seg.meta().id, i);
562            total = total.saturating_add(seg.meta().num_docs);
563        }
564        (segment_map, total)
565    }
566
567    /// Get total document count across all segments
568    pub fn num_docs(&self) -> u32 {
569        self.total_docs
570    }
571
572    /// Get O(1) segment_id → index map (used by reranker)
573    pub fn segment_map(&self) -> &FxHashMap<u128, usize> {
574        &self.segment_map
575    }
576
577    /// Run a bounded piece of CPU work inside this index's shared search pool.
578    #[cfg(feature = "sync")]
579    pub(crate) fn install_search_cpu<R: Send>(&self, operation: impl FnOnce() -> R + Send) -> R {
580        self.search_pool.install(operation)
581    }
582
583    /// Async-only/WASM builds execute inline because Rayon is not available.
584    /// Keeping this overload free of `Send` bounds allows browser-backed file
585    /// handles, whose callbacks are deliberately thread-local, to be scored.
586    #[cfg(not(feature = "sync"))]
587    pub(crate) fn install_search_cpu<R>(&self, operation: impl FnOnce() -> R) -> R {
588        operation()
589    }
590
591    /// Get number of segments
592    pub fn num_segments(&self) -> usize {
593        self.segments.len()
594    }
595
596    /// Get a document by (segment_id, local_doc_id)
597    pub async fn doc(&self, segment_id: u128, doc_id: u32) -> Result<Option<crate::dsl::Document>> {
598        if let Some(&idx) = self.segment_map.get(&segment_id) {
599            return self.segments[idx].doc(doc_id).await;
600        }
601        Ok(None)
602    }
603
604    /// Search across all segments and return aggregated results
605    pub async fn search(
606        &self,
607        query: &dyn crate::query::Query,
608        limit: usize,
609    ) -> Result<Vec<crate::query::SearchResult>> {
610        let (results, _) = self.search_with_count(query, limit).await?;
611        Ok(results)
612    }
613
614    /// Search across all segments and return (results, total_seen)
615    /// total_seen is the number of documents that were scored across all segments
616    pub async fn search_with_count(
617        &self,
618        query: &dyn crate::query::Query,
619        limit: usize,
620    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
621        self.search_with_offset_and_count(query, limit, 0).await
622    }
623
624    /// Search with offset for pagination
625    pub async fn search_with_offset(
626        &self,
627        query: &dyn crate::query::Query,
628        limit: usize,
629        offset: usize,
630    ) -> Result<Vec<crate::query::SearchResult>> {
631        let (results, _) = self
632            .search_with_offset_and_count(query, limit, offset)
633            .await?;
634        Ok(results)
635    }
636
637    /// Search with offset and return (results, total_seen)
638    pub async fn search_with_offset_and_count(
639        &self,
640        query: &dyn crate::query::Query,
641        limit: usize,
642        offset: usize,
643    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
644        self.search_internal(query, limit, offset, false).await
645    }
646
647    /// Search with positions (ordinal tracking) and return (results, total_seen)
648    ///
649    /// Use this when you need per-ordinal scores for multi-valued fields.
650    pub async fn search_with_positions(
651        &self,
652        query: &dyn crate::query::Query,
653        limit: usize,
654    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
655        self.search_internal(query, limit, 0, true).await
656    }
657
658    /// Build the paper's single query-level top-γ superblock set, then project
659    /// it back onto segment-local plans.
660    ///
661    /// Treating every immutable segment as an independent LSP index would
662    /// multiply work by the segment count. The prepass retains one global γ
663    /// while preserving Hermes's streaming segment architecture.
664    fn prepare_global_lsp(
665        &self,
666        query: &dyn crate::query::Query,
667        retrieval_depth: usize,
668        parallel: bool,
669    ) -> Result<Vec<Option<std::sync::Arc<crate::query::bmp::LspSegmentPlan>>>> {
670        let total_start = crate::observe::WallTimer::start();
671        let empty = || vec![None; self.segments.len()];
672        if retrieval_depth == 0 {
673            return Ok(empty());
674        }
675        let crate::query::QueryDecomposition::SparseTerms(infos) = query.decompose() else {
676            return Ok(empty());
677        };
678        let Some(&first) = infos.first() else {
679            return Ok(empty());
680        };
681        if infos
682            .iter()
683            .any(|info| info.field != first.field || info.lsp_gamma != first.lsp_gamma)
684        {
685            return Ok(empty());
686        }
687        let field = first.field;
688        let field_label = self.schema.get_field_name(field).unwrap_or("?");
689        let (total_superblocks, total_coarse_groups, planning_depth) = self
690            .segments
691            .iter()
692            .filter_map(|segment| segment.bmp_index(field))
693            .fold(
694                (0usize, 0usize, retrieval_depth),
695                |(total, coarse, depth), bmp| {
696                    (
697                        total.saturating_add(bmp.num_superblocks as usize),
698                        coarse.saturating_add(bmp.num_coarse_groups as usize),
699                        depth.max(crate::query::bmp_executor_limit(
700                            retrieval_depth,
701                            first.over_fetch_factor,
702                            bmp,
703                        )),
704                    )
705                },
706            );
707        let Some(reference_bmp) = self
708            .segments
709            .iter()
710            .find_map(|segment| segment.bmp_index(field))
711        else {
712            return Ok(empty());
713        };
714        if !infos.iter().any(|info| info.candidate) {
715            return Ok(empty());
716        }
717        let prepare_start = crate::observe::WallTimer::start();
718        let Some(prepared_query) =
719            crate::query::bmp::prepare_bmp_query_infos(reference_bmp.dims(), &infos)?
720        else {
721            return Ok(empty());
722        };
723        let infos: std::sync::Arc<[crate::query::SparseTermQueryInfo]> = infos.into();
724        let prepared_query = std::sync::Arc::new(prepared_query);
725        let prepare_secs = prepare_start.secs();
726
727        let local_plans = || {
728            let plan = std::sync::Arc::new(crate::query::bmp::LspSegmentPlan {
729                infos: std::sync::Arc::clone(&infos),
730                prepared_query: std::sync::Arc::clone(&prepared_query),
731                selection: None,
732            });
733            self.segments
734                .iter()
735                .map(|segment| {
736                    segment
737                        .bmp_index(field)
738                        .map(|_| std::sync::Arc::clone(&plan))
739                })
740                .collect()
741        };
742        let gamma = first
743            .lsp_gamma
744            .unwrap_or_else(|| crate::query::bmp::recommended_lsp_gamma(planning_depth));
745        if gamma == 0 || gamma >= total_superblocks {
746            // A cap covering the whole index is exhaustive. Let each segment
747            // compute and traverse its local order once instead of building a
748            // query-global heap and retaining an all-superblock selection.
749            crate::observe::bmp_lsp(
750                self.schema.index_label(),
751                field_label,
752                total_start.secs(),
753                prepare_secs,
754                0.0,
755                0.0,
756                total_superblocks,
757                gamma,
758                total_coarse_groups,
759                0,
760                0,
761            );
762            return Ok(local_plans());
763        }
764        let hierarchy_scan_start = crate::observe::WallTimer::start();
765        let prepare = |segment: &std::sync::Arc<crate::segment::SegmentReader>| {
766            segment
767                .bmp_index(field)
768                .map(|bmp| crate::query::bmp::prepare_lsp_coarse_ubs(bmp, &prepared_query))
769                .transpose()
770        };
771
772        #[cfg(feature = "sync")]
773        let coarse_bounds: Vec<Option<Vec<f32>>> = if parallel {
774            use rayon::prelude::*;
775            self.search_pool.install(|| {
776                self.segments
777                    .par_iter()
778                    .map(prepare)
779                    .collect::<Result<Vec<_>>>()
780            })?
781        } else {
782            self.segments
783                .iter()
784                .map(prepare)
785                .collect::<Result<Vec<_>>>()?
786        };
787        #[cfg(not(feature = "sync"))]
788        let coarse_bounds: Vec<Option<Vec<f32>>> = {
789            let _ = parallel;
790            self.segments
791                .iter()
792                .map(prepare)
793                .collect::<Result<Vec<_>>>()?
794        };
795        let hierarchy_scan_secs = hierarchy_scan_start.secs();
796
797        let select_start = crate::observe::WallTimer::start();
798        let selection =
799            select_global_lsp_hierarchical(&coarse_bounds, gamma, |segment, group, out| {
800                let bmp = self.segments[segment].bmp_index(field).ok_or_else(|| {
801                    crate::Error::Internal(
802                        "BMP coarse plan references a segment without the sparse field".into(),
803                    )
804                })?;
805                crate::query::bmp::expand_lsp_coarse_group(bmp, &prepared_query, group, out)
806            })?;
807        let mut plans = Vec::with_capacity(self.segments.len());
808        for (segment, selected) in selection.selected.into_iter().enumerate() {
809            if self.segments[segment].bmp_index(field).is_none() {
810                plans.push(None);
811                continue;
812            }
813            let (selected_superblocks, selected_bounds): (Vec<_>, Vec<_>) =
814                selected.into_iter().unzip();
815            plans.push(Some(std::sync::Arc::new(
816                crate::query::bmp::LspSegmentPlan {
817                    infos: std::sync::Arc::clone(&infos),
818                    prepared_query: std::sync::Arc::clone(&prepared_query),
819                    selection: Some(crate::query::bmp::LspSelection {
820                        sb_ubs: selected_bounds,
821                        sb_order: selected_superblocks,
822                    }),
823                },
824            )));
825        }
826        let select_secs = select_start.secs();
827        crate::observe::bmp_lsp(
828            self.schema.index_label(),
829            field_label,
830            total_start.secs(),
831            prepare_secs,
832            hierarchy_scan_secs,
833            select_secs,
834            total_superblocks,
835            gamma,
836            total_coarse_groups,
837            selection.expanded_groups,
838            selection.evaluated_superblocks,
839        );
840        log::debug!(
841            "[searcher] BMP hierarchical LSP: index={}, field={}, coarse_groups={}/{}, E_superblocks={}/{}, gamma={}",
842            self.schema.index_label(),
843            field_label,
844            selection.expanded_groups,
845            coarse_bounds
846                .iter()
847                .filter_map(Option::as_ref)
848                .map(Vec::len)
849                .sum::<usize>(),
850            selection.evaluated_superblocks,
851            total_superblocks,
852            gamma,
853        );
854        Ok(plans)
855    }
856
857    fn ordered_lsp_segments(
858        &self,
859        plans: &[Option<std::sync::Arc<crate::query::bmp::LspSegmentPlan>>],
860    ) -> Vec<usize> {
861        let mut order: Vec<usize> = (0..self.segments.len()).collect();
862        order.sort_unstable_by(|&left, &right| {
863            let left_priority = plans[left]
864                .as_ref()
865                .map_or(f32::NEG_INFINITY, |plan| plan.priority());
866            let right_priority = plans[right]
867                .as_ref()
868                .map_or(f32::NEG_INFINITY, |plan| plan.priority());
869            right_priority
870                .total_cmp(&left_priority)
871                .then_with(|| {
872                    self.segments[right]
873                        .num_docs()
874                        .cmp(&self.segments[left].num_docs())
875                })
876                .then_with(|| {
877                    self.segments[left]
878                        .meta()
879                        .id
880                        .cmp(&self.segments[right].meta().id)
881                })
882        });
883        order
884    }
885
886    #[inline]
887    fn bmp_wave_width(&self) -> usize {
888        #[cfg(feature = "native")]
889        {
890            self.bmp_io_concurrency
891        }
892        #[cfg(not(feature = "native"))]
893        {
894            4
895        }
896    }
897
898    /// Internal search implementation
899    async fn search_internal(
900        &self,
901        query: &dyn crate::query::Query,
902        limit: usize,
903        offset: usize,
904        collect_positions: bool,
905    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
906        let fetch_limit = checked_search_window(limit, offset)?;
907
908        // Use rayon + block_in_place for CPU-bound scoring (sync feature required).
909        // Offloads the scoring loop from tokio workers so search doesn't starve
910        // other async tasks. Works for any segment count (rayon degrades gracefully
911        // to inline execution for a single segment).
912        // Only works on multi-threaded tokio runtime (block_in_place panics on current_thread).
913        #[cfg(feature = "sync")]
914        if !self.segments.is_empty()
915            && tokio::runtime::Handle::current().runtime_flavor()
916                == tokio::runtime::RuntimeFlavor::MultiThread
917        {
918            return self.search_internal_parallel(query, fetch_limit, offset, collect_positions);
919        }
920
921        // No segments, no sync feature, or current_thread runtime: use an
922        // explicitly bounded async stream. Starting every segment at once can
923        // retain `segments × top_k` results while the slowest I/O completes.
924        const MAX_ASYNC_SEGMENT_SEARCHES: usize = 8;
925        use futures::StreamExt;
926        use futures::TryStreamExt;
927        // Cross-segment top-k floor (see search_internal_sync). Concurrent
928        // segments share it via an atomic; ordering is best-effort.
929        let shared = crate::query::SharedThreshold::new();
930        let lsp_plans = self.prepare_global_lsp(query, fetch_limit, false)?;
931        let mut total_seen: u32 = 0;
932        let mut merged = Vec::new();
933        let mut merge_scratch = Vec::new();
934        let bmp_planned = lsp_plans.iter().any(Option::is_some);
935        let order = if bmp_planned {
936            self.ordered_lsp_segments(&lsp_plans)
937        } else {
938            (0..self.segments.len()).collect()
939        };
940        #[cfg(feature = "native")]
941        let bmp_io_gate = Arc::clone(&self.bmp_io_gate);
942        let run_segment = |segment_index: usize| {
943            let segment = Arc::clone(&self.segments[segment_index]);
944            let lsp_plan = lsp_plans[segment_index].clone();
945            let shared = shared.clone();
946            #[cfg(feature = "native")]
947            let bmp_io_gate = Arc::clone(&bmp_io_gate);
948            async move {
949                if lsp_plan.as_ref().is_some_and(|plan| !plan.has_work()) {
950                    return Ok((Vec::new(), 0u32));
951                }
952                #[cfg(feature = "native")]
953                let _io_permit = if lsp_plan.is_some() {
954                    Some(bmp_io_gate.acquire_async().await)
955                } else {
956                    None
957                };
958                let sid = segment.meta().id;
959                let (mut results, segment_seen) = crate::query::search_segment_shared_planned(
960                    segment.as_ref(),
961                    query,
962                    fetch_limit,
963                    collect_positions,
964                    shared.clone(),
965                    lsp_plan,
966                )
967                .await?;
968                if fetch_limit > 0 && results.len() >= fetch_limit {
969                    shared.raise(results[fetch_limit - 1].score);
970                }
971                for result in &mut results {
972                    result.segment_id = sid;
973                }
974                Ok::<_, crate::error::Error>((results, segment_seen))
975            }
976        };
977
978        let mut remainder = order.as_slice();
979        if bmp_planned {
980            // Match the synchronous policy: score the highest-bound pilot
981            // first so lower-bound async work starts with a useful theta.
982            if let Some((&pilot, rest)) = order.split_first() {
983                let (batch, segment_seen) = run_segment(pilot).await?;
984                total_seen = total_seen.saturating_add(segment_seen);
985                merge_ranked_reuse(&mut merged, batch, fetch_limit, &mut merge_scratch);
986                remainder = rest;
987            }
988        }
989        let concurrency = if bmp_planned {
990            self.bmp_wave_width()
991        } else {
992            MAX_ASYNC_SEGMENT_SEARCHES
993        };
994        let searches = futures::stream::iter(remainder.iter().copied().map(run_segment))
995            .buffer_unordered(concurrency);
996        futures::pin_mut!(searches);
997        while let Some((batch, segment_seen)) = searches.try_next().await? {
998            total_seen = total_seen.saturating_add(segment_seen);
999            merge_ranked_reuse(&mut merged, batch, fetch_limit, &mut merge_scratch);
1000        }
1001
1002        let results = apply_result_offset(merged, fetch_limit, offset);
1003        Ok((results, total_seen))
1004    }
1005
1006    /// Multi-segment parallel search using rayon (CPU-bound scoring on thread pool).
1007    ///
1008    /// `block_in_place` tells tokio this worker is occupied so it can steal tasks.
1009    /// `rayon::par_iter` distributes segment scoring across the rayon thread pool.
1010    #[cfg(feature = "sync")]
1011    fn search_internal_parallel(
1012        &self,
1013        query: &dyn crate::query::Query,
1014        fetch_limit: usize,
1015        offset: usize,
1016        collect_positions: bool,
1017    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1018        tokio::task::block_in_place(|| {
1019            self.search_internal_sync(query, fetch_limit, offset, collect_positions)
1020        })
1021    }
1022
1023    /// Sync body of the parallel search: rayon par_iter over segments.
1024    /// Callers must already be off the async reactor (block_in_place or a
1025    /// rayon/blocking thread) — safe to nest inside another par_iter
1026    /// (rayon work-stealing composes).
1027    #[cfg(feature = "sync")]
1028    fn search_internal_sync(
1029        &self,
1030        query: &dyn crate::query::Query,
1031        fetch_limit: usize,
1032        offset: usize,
1033        collect_positions: bool,
1034    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1035        let (merged, total_seen) =
1036            self.search_segments_sync(query, fetch_limit, collect_positions)?;
1037
1038        let results = apply_result_offset(merged, fetch_limit, offset);
1039        Ok((results, total_seen))
1040    }
1041
1042    /// Score all segments with one shared threshold.
1043    ///
1044    /// Ordinary queries retain full CPU parallelism. BMP uses a highest-bound
1045    /// pilot followed by bounded waves, and every active BMP segment also
1046    /// holds a process-wide random-I/O permit. This lets theta mature before
1047    /// lower-bound segments touch pageable D/payload pages and prevents
1048    /// concurrent queries from multiplying the wave width.
1049    #[cfg(feature = "sync")]
1050    fn search_segments_sync(
1051        &self,
1052        query: &dyn crate::query::Query,
1053        fetch_limit: usize,
1054        collect_positions: bool,
1055    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1056        use rayon::prelude::*;
1057
1058        let lsp_plans = self.prepare_global_lsp(query, fetch_limit, true)?;
1059        let shared = crate::query::SharedThreshold::new();
1060        let run_segment = |segment_index: &usize| {
1061            let segment = &self.segments[*segment_index];
1062            let lsp_plan = lsp_plans[*segment_index].clone();
1063            if lsp_plan.as_ref().is_some_and(|plan| !plan.has_work()) {
1064                return Ok((Vec::new(), 0u32));
1065            }
1066            let _io_permit = lsp_plan.as_ref().map(|_| self.bmp_io_gate.acquire());
1067            let sid = segment.meta().id;
1068            let (mut results, segment_seen) = crate::query::search_segment_shared_sync_planned(
1069                segment.as_ref(),
1070                query,
1071                fetch_limit,
1072                collect_positions,
1073                shared.clone(),
1074                lsp_plan,
1075            )?;
1076            if fetch_limit > 0 && results.len() >= fetch_limit {
1077                shared.raise(results[fetch_limit - 1].score);
1078            }
1079            for result in &mut results {
1080                result.segment_id = sid;
1081            }
1082            Ok::<_, crate::Error>((results, segment_seen))
1083        };
1084
1085        if !lsp_plans.iter().any(Option::is_some) {
1086            return self.search_pool.install(|| {
1087                (0..self.segments.len())
1088                    .into_par_iter()
1089                    .map(|segment| run_segment(&segment))
1090                    .try_reduce(
1091                        || (Vec::new(), 0u32),
1092                        |(left, left_seen), (right, right_seen)| {
1093                            Ok((
1094                                merge_two_ranked(left, right, fetch_limit),
1095                                left_seen.saturating_add(right_seen),
1096                            ))
1097                        },
1098                    )
1099            });
1100        }
1101
1102        let order = self.ordered_lsp_segments(&lsp_plans);
1103
1104        let mut merged = Vec::new();
1105        let mut merge_scratch = Vec::new();
1106        let mut total_seen = 0u32;
1107        if let Some((&pilot, rest)) = order.split_first() {
1108            let (pilot_results, pilot_seen) = self.search_pool.install(|| run_segment(&pilot))?;
1109            merge_ranked_reuse(&mut merged, pilot_results, fetch_limit, &mut merge_scratch);
1110            total_seen = total_seen.saturating_add(pilot_seen);
1111
1112            for wave in rest.chunks(self.bmp_wave_width()) {
1113                let batches = self
1114                    .search_pool
1115                    .install(|| wave.par_iter().map(run_segment).collect::<Result<Vec<_>>>())?;
1116                for (results, seen) in batches {
1117                    merge_ranked_reuse(&mut merged, results, fetch_limit, &mut merge_scratch);
1118                    total_seen = total_seen.saturating_add(seen);
1119                }
1120            }
1121        }
1122        Ok((merged, total_seen))
1123    }
1124
1125    /// Synchronous search across all segments using rayon for parallelism.
1126    ///
1127    /// This is the async-free boundary — no tokio involvement from here down.
1128    #[cfg(feature = "sync")]
1129    pub fn search_with_offset_and_count_sync(
1130        &self,
1131        query: &dyn crate::query::Query,
1132        limit: usize,
1133        offset: usize,
1134    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1135        let fetch_limit = checked_search_window(limit, offset)?;
1136        let (merged, total_seen) = self.search_segments_sync(query, fetch_limit, false)?;
1137
1138        let results = apply_result_offset(merged, fetch_limit, offset);
1139        Ok((results, total_seen))
1140    }
1141
1142    /// Hybrid search: run several queries independently and fuse their
1143    /// ranked lists (union) into a single top-`limit` result.
1144    ///
1145    /// Unlike [`Self::search_and_rerank`] — which can only re-score
1146    /// documents the first-stage query already found — fusion keeps
1147    /// documents found by *any* of the queries. Typical use is sparse
1148    /// (BM25/SPLADE) + dense vector hybrid retrieval with
1149    /// `FusionMethod::Rrf { k: 60.0 }`.
1150    ///
1151    /// Fusion happens at **chunk granularity**: per-ordinal scores are
1152    /// collected from each sub-query, fused per `(doc, ordinal)` key, then
1153    /// combined into a doc score with `combiner`
1154    /// (`MultiValueCombiner::Max` recommended — same-chunk corroboration
1155    /// across verticals compounds, scattered noise does not). Fused results
1156    /// carry per-chunk `positions`.
1157    ///
1158    /// Each query is paired with a weight scaling its contribution.
1159    /// `fetch_limit` is the per-query candidate depth. Request-facing adapters
1160    /// default to at most [`crate::query::MAX_CANDIDATE_OVERSUBSCRIPTION`] times
1161    /// the result window and never multiply an existing rerank pool again.
1162    pub async fn search_fused(
1163        &self,
1164        queries: &[(&dyn crate::query::Query, f32)],
1165        fetch_limit: usize,
1166        limit: usize,
1167        method: crate::query::FusionMethod,
1168        combiner: crate::query::MultiValueCombiner,
1169    ) -> Result<Vec<crate::query::SearchResult>> {
1170        let (results, _) = self
1171            .search_fused_with_count(queries, fetch_limit, limit, method, combiner)
1172            .await?;
1173        Ok(results)
1174    }
1175
1176    /// Fusion variant that also returns the aggregate number of documents
1177    /// scored by all sub-queries. This lets request-facing callers use the
1178    /// parallel fusion path without rerunning sub-queries for observability.
1179    pub async fn search_fused_with_count(
1180        &self,
1181        queries: &[(&dyn crate::query::Query, f32)],
1182        fetch_limit: usize,
1183        limit: usize,
1184        method: crate::query::FusionMethod,
1185        combiner: crate::query::MultiValueCombiner,
1186    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1187        if queries.is_empty() {
1188            return Err(crate::Error::Query(
1189                "fusion requires at least one sub-query".to_string(),
1190            ));
1191        }
1192        if queries.len() > crate::query::MAX_FUSION_SUB_QUERIES {
1193            return Err(crate::Error::Query(format!(
1194                "fusion supports at most {} sub-queries, got {}",
1195                crate::query::MAX_FUSION_SUB_QUERIES,
1196                queries.len()
1197            )));
1198        }
1199        if fetch_limit == 0 {
1200            return Err(crate::Error::Query(
1201                "fusion fetch_limit must be greater than zero".to_string(),
1202            ));
1203        }
1204        let candidate_slots = fetch_limit
1205            .checked_mul(queries.len())
1206            .ok_or_else(|| crate::Error::Query("fusion candidate budget overflow".to_string()))?;
1207        if candidate_slots > crate::query::MAX_FUSION_CANDIDATE_SLOTS {
1208            return Err(crate::Error::Query(format!(
1209                "fusion candidate budget must not exceed {}, got {candidate_slots}",
1210                crate::query::MAX_FUSION_CANDIDATE_SLOTS
1211            )));
1212        }
1213        for (index, &(_, weight)) in queries.iter().enumerate() {
1214            if !weight.is_finite() || weight < 0.0 {
1215                return Err(crate::Error::Query(format!(
1216                    "fusion query weight at index {index} must be finite and non-negative, \
1217                     got {weight}"
1218                )));
1219            }
1220        }
1221        if let crate::query::FusionMethod::Rrf { k } = method
1222            && (!k.is_finite() || k < 0.0)
1223        {
1224            return Err(crate::Error::Query(format!(
1225                "fusion RRF k must be finite and non-negative, got {k}"
1226            )));
1227        }
1228        combiner.validate().map_err(crate::Error::Query)?;
1229
1230        // Each sub-query already fans out across every segment. Keep fusion
1231        // sequential at the outer level so queries do not contend for the
1232        // same rayon pool, mmap pages, and memory bandwidth, and so each
1233        // query's shared threshold converges as early as possible.
1234        #[cfg(feature = "sync")]
1235        if !self.segments.is_empty()
1236            && tokio::runtime::Handle::current().runtime_flavor()
1237                == tokio::runtime::RuntimeFlavor::MultiThread
1238        {
1239            let lists: Vec<(Vec<crate::query::SearchResult>, f32, u32)> =
1240                tokio::task::block_in_place(|| {
1241                    queries
1242                        .iter()
1243                        .map(|&(query, weight)| {
1244                            let (results, seen) =
1245                                self.search_internal_sync(query, fetch_limit, 0, true)?;
1246                            Ok((results, weight, seen))
1247                        })
1248                        .collect::<Result<Vec<_>>>()
1249                })?;
1250            let mut total_seen = 0u32;
1251            let ranked_lists = lists
1252                .into_iter()
1253                .map(|(results, weight, seen)| {
1254                    total_seen = total_seen.saturating_add(seen);
1255                    (results, weight)
1256                })
1257                .collect();
1258            let fused =
1259                crate::query::try_fuse_ranked_lists_chunked(ranked_lists, method, combiner, limit)
1260                    .map_err(crate::Error::Query)?;
1261            return Ok((fused, total_seen));
1262        }
1263
1264        // Async/current-thread fallback uses the same outer execution shape
1265        // and preserves input list order for deterministic rank ties.
1266        let mut lists = Vec::with_capacity(queries.len());
1267        for &(query, weight) in queries {
1268            let (results, seen) = self.search_with_positions(query, fetch_limit).await?;
1269            lists.push((results, weight, seen));
1270        }
1271        let mut total_seen = 0u32;
1272        let ranked_lists = lists
1273            .into_iter()
1274            .map(|(results, weight, seen)| {
1275                total_seen = total_seen.saturating_add(seen);
1276                (results, weight)
1277            })
1278            .collect();
1279        let fused =
1280            crate::query::try_fuse_ranked_lists_chunked(ranked_lists, method, combiner, limit)
1281                .map_err(crate::Error::Query)?;
1282        Ok((fused, total_seen))
1283    }
1284
1285    /// Two-stage search: L1 retrieval + L2 dense vector reranking
1286    ///
1287    /// Runs the query to get `l1_limit` candidates, then reranks by exact
1288    /// dense vector distance and returns the top `final_limit` results.
1289    pub async fn search_and_rerank(
1290        &self,
1291        query: &dyn crate::query::Query,
1292        l1_limit: usize,
1293        final_limit: usize,
1294        config: &crate::query::RerankerConfig,
1295    ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1296        let (candidates, total_seen) = self.search_with_count(query, l1_limit).await?;
1297        let reranked = crate::query::rerank(self, &candidates, config, final_limit).await?;
1298        Ok((reranked, total_seen))
1299    }
1300
1301    /// Parse query string and search (convenience method)
1302    pub async fn query(
1303        &self,
1304        query_str: &str,
1305        limit: usize,
1306    ) -> Result<crate::query::SearchResponse> {
1307        self.query_offset(query_str, limit, 0).await
1308    }
1309
1310    /// Parse query string and search with offset (convenience method)
1311    pub async fn query_offset(
1312        &self,
1313        query_str: &str,
1314        limit: usize,
1315        offset: usize,
1316    ) -> Result<crate::query::SearchResponse> {
1317        let parser = self.query_parser();
1318        let query = parser
1319            .parse(query_str)
1320            .map_err(crate::error::Error::Query)?;
1321
1322        let (results, _total_seen) = self
1323            .search_internal(query.as_ref(), limit, offset, false)
1324            .await?;
1325
1326        let total_hits = results.len() as u32;
1327        let hits: Vec<crate::query::SearchHit> = results
1328            .into_iter()
1329            .map(|result| crate::query::SearchHit {
1330                address: crate::query::DocAddress::new(result.segment_id, result.doc_id),
1331                score: result.score,
1332                matched_fields: result.extract_ordinals(),
1333            })
1334            .collect();
1335
1336        Ok(crate::query::SearchResponse { hits, total_hits })
1337    }
1338
1339    /// Get query parser for this searcher
1340    pub fn query_parser(&self) -> crate::dsl::QueryLanguageParser {
1341        let query_routers = self.schema.query_routers();
1342        if !query_routers.is_empty()
1343            && let Ok(router) = crate::dsl::QueryFieldRouter::from_rules(query_routers)
1344        {
1345            return crate::dsl::QueryLanguageParser::with_router(
1346                Arc::clone(&self.schema),
1347                self.default_fields.clone(),
1348                Arc::clone(&self.tokenizers),
1349                router,
1350            );
1351        }
1352
1353        crate::dsl::QueryLanguageParser::new(
1354            Arc::clone(&self.schema),
1355            self.default_fields.clone(),
1356            Arc::clone(&self.tokenizers),
1357        )
1358    }
1359
1360    /// Get a document by address (segment_id + local doc_id)
1361    pub async fn get_document(
1362        &self,
1363        address: &crate::query::DocAddress,
1364    ) -> Result<Option<crate::dsl::Document>> {
1365        self.get_document_with_fields(address, None).await
1366    }
1367
1368    /// Get a document by address, hydrating only the specified field IDs.
1369    ///
1370    /// If `fields` is `None`, all fields are hydrated (including dense vectors).
1371    /// If `fields` is `Some(set)`, only dense vector fields in the set are read
1372    /// from flat storage — skipping expensive mmap reads for unrequested vectors.
1373    pub async fn get_document_with_fields(
1374        &self,
1375        address: &crate::query::DocAddress,
1376        fields: Option<&rustc_hash::FxHashSet<u32>>,
1377    ) -> Result<Option<crate::dsl::Document>> {
1378        let segment_id = address.segment_id_u128().ok_or_else(|| {
1379            crate::error::Error::Query(format!("Invalid segment ID: {}", address.segment_id()))
1380        })?;
1381
1382        if let Some(&idx) = self.segment_map.get(&segment_id) {
1383            return self.segments[idx]
1384                .doc_with_fields(address.doc_id, fields)
1385                .await;
1386        }
1387
1388        Ok(None)
1389    }
1390}
1391
1392struct HierarchicalLspSelection {
1393    selected: Vec<Vec<(u32, f32)>>,
1394    expanded_groups: usize,
1395    evaluated_superblocks: usize,
1396}
1397
1398/// Select the exact global top-gamma E cells through safe H upper bounds.
1399///
1400/// Best-first expansion stops only when the next H bound is strictly below
1401/// the current gamma-th E bound. Equality must still expand because the
1402/// deterministic `(score, segment, superblock)` tie-break can change
1403/// membership. Memory is O(number of H cells + gamma), never O(all E cells).
1404fn select_global_lsp_hierarchical(
1405    coarse_bounds: &[Option<Vec<f32>>],
1406    gamma: usize,
1407    mut expand: impl FnMut(usize, u32, &mut Vec<(u32, f32)>) -> Result<()>,
1408) -> Result<HierarchicalLspSelection> {
1409    if gamma == 0 {
1410        return Err(crate::Error::Internal(
1411            "hierarchical LSP selection requires a positive gamma".into(),
1412        ));
1413    }
1414    let mut frontier = std::collections::BinaryHeap::<(u32, usize, u32)>::new();
1415    for (segment, bounds) in coarse_bounds.iter().enumerate() {
1416        let Some(bounds) = bounds else {
1417            continue;
1418        };
1419        for (coarse_group, &bound) in bounds.iter().enumerate() {
1420            debug_assert!(bound.is_finite() && bound >= 0.0);
1421            if bound > 0.0 {
1422                frontier.push((bound.to_bits(), segment, coarse_group as u32));
1423            }
1424        }
1425    }
1426    let mut top =
1427        std::collections::BinaryHeap::<std::cmp::Reverse<(u32, usize, u32)>>::with_capacity(
1428            gamma.min(65_536),
1429        );
1430    let mut expanded =
1431        Vec::with_capacity(crate::segment::reader::bmp::BMP_COARSE_SUPERBLOCKS as usize);
1432    let mut expanded_groups = 0usize;
1433    let mut evaluated_superblocks = 0usize;
1434    while let Some((coarse_bound, segment, coarse_group)) = frontier.pop() {
1435        if top.len() == gamma && top.peek().is_some_and(|minimum| coarse_bound < minimum.0.0) {
1436            break;
1437        }
1438        expand(segment, coarse_group, &mut expanded)?;
1439        expanded_groups += 1;
1440        evaluated_superblocks = evaluated_superblocks.saturating_add(expanded.len());
1441        for &(superblock, bound) in &expanded {
1442            debug_assert!(bound.is_finite() && bound >= 0.0);
1443            if bound <= 0.0 {
1444                continue;
1445            }
1446            let candidate = (bound.to_bits(), segment, superblock);
1447            if top.len() < gamma {
1448                top.push(std::cmp::Reverse(candidate));
1449            } else if top.peek().is_some_and(|minimum| candidate > minimum.0) {
1450                top.pop();
1451                top.push(std::cmp::Reverse(candidate));
1452            }
1453        }
1454    }
1455
1456    let mut selected = vec![Vec::<(u32, f32)>::new(); coarse_bounds.len()];
1457    for std::cmp::Reverse((bound, segment, superblock)) in top {
1458        selected[segment].push((superblock, f32::from_bits(bound)));
1459    }
1460    for segment in &mut selected {
1461        segment.sort_unstable_by(|&(left_sb, left_bound), &(right_sb, right_bound)| {
1462            right_bound
1463                .total_cmp(&left_bound)
1464                .then_with(|| left_sb.cmp(&right_sb))
1465        });
1466    }
1467    Ok(HierarchicalLspSelection {
1468        selected,
1469        expanded_groups,
1470        evaluated_superblocks,
1471    })
1472}
1473
1474/// Select one query-level top-γ set without allocating one tuple per
1475/// superblock. Prepared BMP bounds are finite and non-negative, so their f32
1476/// bit patterns have the same order as their numeric values.
1477#[cfg(test)]
1478fn select_global_lsp_superblocks(bounds: &[Option<Vec<f32>>], gamma: usize) -> Vec<Vec<u32>> {
1479    let mut top =
1480        std::collections::BinaryHeap::<std::cmp::Reverse<(u32, usize, u32)>>::with_capacity(
1481            gamma.min(65_536),
1482        );
1483    for (segment, segment_bounds) in bounds.iter().enumerate() {
1484        let Some(segment_bounds) = segment_bounds else {
1485            continue;
1486        };
1487        for (superblock, &bound) in segment_bounds.iter().enumerate() {
1488            if bound <= 0.0 {
1489                continue;
1490            }
1491            let candidate = (bound.to_bits(), segment, superblock as u32);
1492            if top.len() < gamma {
1493                top.push(std::cmp::Reverse(candidate));
1494            } else if top.peek().is_some_and(|minimum| candidate > minimum.0) {
1495                top.pop();
1496                top.push(std::cmp::Reverse(candidate));
1497            }
1498        }
1499    }
1500
1501    let mut selected = vec![Vec::<u32>::new(); bounds.len()];
1502    for std::cmp::Reverse((_, segment, superblock)) in top {
1503        selected[segment].push(superblock);
1504    }
1505    selected
1506}
1507
1508/// Merge one segment batch into the running top-k while alternating two
1509/// retained buffers. Search results are moved, not cloned, and segment fan-out
1510/// no longer allocates a fresh `limit`-sized vector for every merge.
1511fn merge_ranked_reuse(
1512    merged: &mut Vec<crate::query::SearchResult>,
1513    batch: Vec<crate::query::SearchResult>,
1514    limit: usize,
1515    scratch: &mut Vec<crate::query::SearchResult>,
1516) {
1517    scratch.clear();
1518    let output_len = limit.min(merged.len().saturating_add(batch.len()));
1519    if scratch.capacity() < output_len {
1520        scratch.reserve_exact(output_len);
1521    }
1522    {
1523        let mut left = merged.drain(..).peekable();
1524        let mut right = batch.into_iter().peekable();
1525        while scratch.len() < output_len {
1526            let take_left = match (left.peek(), right.peek()) {
1527                (Some(left), Some(right)) => {
1528                    !crate::query::compare_search_results_desc(left, right).is_gt()
1529                }
1530                (Some(_), None) => true,
1531                (None, Some(_)) => false,
1532                (None, None) => break,
1533            };
1534            if take_left {
1535                scratch.push(left.next().expect("peeked left result"));
1536            } else {
1537                scratch.push(right.next().expect("peeked right result"));
1538            }
1539        }
1540    }
1541    std::mem::swap(merged, scratch);
1542}
1543
1544/// Merge two canonically sorted batches while moving (not cloning) hits.
1545/// Synchronous parallel reductions use this eagerly, so retained cross-segment
1546/// results stay O(k) instead of O(number_of_segments × k).
1547#[cfg(any(feature = "sync", test))]
1548fn merge_two_ranked(
1549    left: Vec<crate::query::SearchResult>,
1550    right: Vec<crate::query::SearchResult>,
1551    limit: usize,
1552) -> Vec<crate::query::SearchResult> {
1553    let mut left = left.into_iter().peekable();
1554    let mut right = right.into_iter().peekable();
1555    let mut merged = Vec::with_capacity(limit.min(left.len().saturating_add(right.len())));
1556
1557    while merged.len() < limit {
1558        let take_left = match (left.peek(), right.peek()) {
1559            (Some(left), Some(right)) => {
1560                !crate::query::compare_search_results_desc(left, right).is_gt()
1561            }
1562            (Some(_), None) => true,
1563            (None, Some(_)) => false,
1564            (None, None) => break,
1565        };
1566        if take_left {
1567            merged.push(left.next().expect("peeked left result"));
1568        } else {
1569            merged.push(right.next().expect("peeked right result"));
1570        }
1571    }
1572    merged
1573}
1574
1575fn apply_result_offset(
1576    mut results: Vec<crate::query::SearchResult>,
1577    fetch_limit: usize,
1578    offset: usize,
1579) -> Vec<crate::query::SearchResult> {
1580    if offset == 0 {
1581        results.truncate(fetch_limit);
1582        return results;
1583    }
1584    // Pagination usually returns a small window from a much larger fetch.
1585    // Allocate that small result rather than retaining the fetch-sized backing
1586    // allocation through response serialization.
1587    results
1588        .into_iter()
1589        .skip(offset)
1590        .take(fetch_limit.saturating_sub(offset))
1591        .collect()
1592}
1593
1594fn checked_search_window(limit: usize, offset: usize) -> Result<usize> {
1595    offset
1596        .checked_add(limit)
1597        .ok_or_else(|| crate::Error::Query("search offset + limit overflow".into()))
1598}
1599
1600/// Get current process RSS in bytes (best-effort, returns zero on failure).
1601fn process_rss_bytes() -> u64 {
1602    #[cfg(target_os = "linux")]
1603    {
1604        // Read from /proc/self/status — VmRSS line
1605        if let Ok(status) = std::fs::read_to_string("/proc/self/status") {
1606            for line in status.lines() {
1607                if let Some(rest) = line.strip_prefix("VmRSS:") {
1608                    let kib: u64 = rest
1609                        .trim()
1610                        .trim_end_matches("kB")
1611                        .trim()
1612                        .parse()
1613                        .unwrap_or(0);
1614                    return kib.saturating_mul(1024);
1615                }
1616            }
1617        }
1618        0
1619    }
1620    #[cfg(target_os = "macos")]
1621    {
1622        // Use mach_task_self / task_info via raw syscall
1623        use std::mem;
1624        #[repr(C)]
1625        struct TaskBasicInfo {
1626            virtual_size: u64,
1627            resident_size: u64,
1628            resident_size_max: u64,
1629            user_time: [u32; 2],
1630            system_time: [u32; 2],
1631            policy: i32,
1632            suspend_count: i32,
1633        }
1634        unsafe extern "C" {
1635            fn mach_task_self() -> u32;
1636            fn task_info(task: u32, flavor: u32, info: *mut TaskBasicInfo, count: *mut u32) -> i32;
1637        }
1638        const MACH_TASK_BASIC_INFO: u32 = 20;
1639        let mut info: TaskBasicInfo = unsafe { mem::zeroed() };
1640        let mut count = (mem::size_of::<TaskBasicInfo>() / mem::size_of::<u32>()) as u32;
1641        let ret = unsafe {
1642            task_info(
1643                mach_task_self(),
1644                MACH_TASK_BASIC_INFO,
1645                &mut info,
1646                &mut count,
1647            )
1648        };
1649        if ret == 0 { info.resident_size } else { 0 }
1650    }
1651    #[cfg(not(any(target_os = "linux", target_os = "macos")))]
1652    {
1653        0
1654    }
1655}
1656
1657#[cfg(test)]
1658mod load_segments_tests {
1659    use super::*;
1660
1661    #[tokio::test]
1662    async fn searcher_open_fails_loud_on_corrupt_metadata_segment_id() {
1663        let directory = Arc::new(crate::directories::RamDirectory::new());
1664        let schema = Arc::new(crate::dsl::SchemaBuilder::default().build());
1665
1666        let result =
1667            Searcher::open(directory, schema, &["not-a-hex-segment-id".to_string()], 8).await;
1668
1669        match result {
1670            Ok(searcher) => panic!(
1671                "corrupt segment ID must fail loud instead of silently serving {} segments",
1672                searcher.segment_readers().len()
1673            ),
1674            Err(crate::error::Error::Corruption(message)) => {
1675                assert!(message.contains("not-a-hex-segment-id"), "{message}");
1676            }
1677            Err(other) => panic!("expected Corruption error for invalid segment ID, got: {other}"),
1678        }
1679    }
1680}
1681
1682#[cfg(test)]
1683mod search_window_tests {
1684    use super::{
1685        apply_result_offset, checked_search_window, merge_ranked_reuse, merge_two_ranked,
1686        select_global_lsp_hierarchical, select_global_lsp_superblocks,
1687    };
1688    use crate::query::SearchResult;
1689
1690    fn result(segment_id: u128, doc_id: u32, score: f32) -> SearchResult {
1691        SearchResult {
1692            doc_id,
1693            score,
1694            segment_id,
1695            positions: Vec::new(),
1696        }
1697    }
1698
1699    #[test]
1700    fn search_window_is_checked() {
1701        assert_eq!(checked_search_window(7, 5).unwrap(), 12);
1702        assert!(checked_search_window(1, usize::MAX).is_err());
1703    }
1704
1705    #[test]
1706    fn bounded_merge_preserves_canonical_order_and_ties() {
1707        let left = vec![result(2, 9, 10.0), result(2, 3, 7.0)];
1708        let right = vec![result(1, 8, 10.0), result(1, 2, 7.0)];
1709
1710        let expected = merge_two_ranked(left.clone(), right.clone(), 3);
1711        let mut merged = left;
1712        let mut scratch = Vec::new();
1713        merge_ranked_reuse(&mut merged, right, 3, &mut scratch);
1714        assert_eq!(merged, expected);
1715        assert!(scratch.is_empty());
1716        let keys: Vec<_> = merged
1717            .iter()
1718            .map(|result| (result.score, result.segment_id, result.doc_id))
1719            .collect();
1720        assert_eq!(keys, vec![(10.0, 1, 8), (10.0, 2, 9), (7.0, 1, 2)]);
1721    }
1722
1723    #[test]
1724    fn result_offset_returns_only_the_requested_window() {
1725        let results = (0..8)
1726            .map(|doc_id| result(1, doc_id, 8.0 - doc_id as f32))
1727            .collect();
1728
1729        let page = apply_result_offset(results, 5, 2);
1730        assert_eq!(
1731            page.iter().map(|result| result.doc_id).collect::<Vec<_>>(),
1732            vec![2, 3, 4]
1733        );
1734    }
1735
1736    #[test]
1737    fn lsp_gamma_is_global_not_per_segment() {
1738        let bounds = vec![
1739            Some(vec![9.0, 1.0, 8.0]),
1740            Some(vec![7.0, 6.0, 0.0]),
1741            None,
1742            Some(vec![5.0, 4.0]),
1743        ];
1744        let mut selected = select_global_lsp_superblocks(&bounds, 4);
1745        for segment in &mut selected {
1746            segment.sort_unstable();
1747        }
1748        assert_eq!(selected.iter().map(Vec::len).sum::<usize>(), 4);
1749        assert_eq!(selected[0], vec![0, 2]);
1750        assert_eq!(selected[1], vec![0, 1]);
1751        assert!(selected[2].is_empty());
1752        assert!(selected[3].is_empty());
1753    }
1754
1755    #[test]
1756    fn hierarchical_lsp_matches_full_e_scan_exactly() {
1757        const GROUP: usize = crate::segment::reader::bmp::BMP_COARSE_SUPERBLOCKS as usize;
1758        let full = vec![
1759            Some(
1760                (0..700)
1761                    .map(|index| 1_000.0 - index as f32)
1762                    .collect::<Vec<_>>(),
1763            ),
1764            Some(
1765                (0..530)
1766                    .map(|index| 995.0 - index as f32 * 1.25)
1767                    .collect::<Vec<_>>(),
1768            ),
1769            None,
1770        ];
1771        let coarse: Vec<Option<Vec<f32>>> = full
1772            .iter()
1773            .map(|segment| {
1774                segment.as_ref().map(|bounds| {
1775                    bounds
1776                        .chunks(GROUP)
1777                        .map(|group| group.iter().copied().fold(0.0, f32::max))
1778                        .collect()
1779                })
1780            })
1781            .collect();
1782        let gamma = 37;
1783        let hierarchy = select_global_lsp_hierarchical(&coarse, gamma, |segment, group, output| {
1784            output.clear();
1785            let Some(bounds) = &full[segment] else {
1786                return Ok(());
1787            };
1788            let start = group as usize * GROUP;
1789            output.extend(
1790                bounds[start..(start + GROUP).min(bounds.len())]
1791                    .iter()
1792                    .enumerate()
1793                    .map(|(within, &bound)| ((start + within) as u32, bound)),
1794            );
1795            Ok(())
1796        })
1797        .unwrap();
1798        let mut expected = select_global_lsp_superblocks(&full, gamma);
1799        for segment in &mut expected {
1800            segment.sort_unstable();
1801        }
1802        let mut actual: Vec<Vec<u32>> = hierarchy
1803            .selected
1804            .iter()
1805            .map(|segment| {
1806                let mut ids: Vec<_> = segment.iter().map(|&(id, _)| id).collect();
1807                ids.sort_unstable();
1808                ids
1809            })
1810            .collect();
1811        for segment in &mut actual {
1812            segment.sort_unstable();
1813        }
1814        assert_eq!(actual, expected);
1815        assert_eq!(actual.iter().map(Vec::len).sum::<usize>(), gamma);
1816        assert!(
1817            hierarchy.evaluated_superblocks
1818                < full.iter().filter_map(Option::as_ref).map(Vec::len).sum()
1819        );
1820    }
1821}
1822
1823#[cfg(test)]
1824mod fusion_parallelism_tests {
1825    #[test]
1826    fn fusion_keeps_parallelism_at_the_segment_level() {
1827        let source = include_str!("searcher.rs");
1828        let body = source
1829            .split("pub async fn search_fused_with_count")
1830            .nth(1)
1831            .and_then(|tail| tail.split("/// Two-stage search").next())
1832            .expect("bounded fusion search implementation");
1833
1834        assert!(
1835            !body.contains(".par_iter()"),
1836            "sub-query parallelism nests over segment parallelism"
1837        );
1838        assert!(
1839            !body.contains(".buffered("),
1840            "async fusion fallback must preserve the same bounded execution shape"
1841        );
1842    }
1843}