1use 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#[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
65pub struct Searcher<D: Directory + 'static> {
70 #[cfg(feature = "native")]
72 _snapshot: SegmentSnapshot,
73 _phantom: std::marker::PhantomData<D>,
75 segments: Vec<Arc<SegmentReader>>,
77 schema: Arc<Schema>,
79 default_fields: Vec<crate::Field>,
81 tokenizers: Arc<crate::tokenizer::TokenizerRegistry>,
83 trained_vectors: Arc<TrainedVectorStructures>,
85 global_stats: Arc<LazyGlobalStats>,
87 segment_map: FxHashMap<u128, usize>,
89 total_docs: u32,
91 #[cfg(feature = "sync")]
93 search_pool: Arc<rayon::ThreadPool>,
94 #[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 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 #[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 #[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 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; 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 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 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 let existing_map: FxHashMap<u128, Arc<SegmentReader>> = existing_segments
316 .iter()
317 .map(|seg| (seg.meta().id, Arc::clone(seg)))
318 .collect();
319
320 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 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 let store_cache_directory_namespace = Arc::as_ptr(directory) as usize;
358
359 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 let mut loaded: Vec<(usize, Arc<SegmentReader>)> = Vec::with_capacity(valid_segments.len());
385
386 loaded.extend(reused);
388
389 for ((idx, sid), result) in to_load.into_iter().zip(results) {
391 match result {
392 Ok(mut reader) => {
393 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 loaded.sort_by_key(|(idx, _)| *idx);
408
409 let segments: Vec<Arc<SegmentReader>> = loaded.into_iter().map(|(_, seg)| seg).collect();
410
411 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 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 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 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 pub fn schema(&self) -> &Schema {
522 &self.schema
523 }
524
525 pub fn segment_readers(&self) -> &[Arc<SegmentReader>] {
527 &self.segments
528 }
529
530 pub fn default_fields(&self) -> &[crate::Field] {
532 &self.default_fields
533 }
534
535 pub fn tokenizers(&self) -> &crate::tokenizer::TokenizerRegistry {
537 &self.tokenizers
538 }
539
540 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 pub fn global_stats(&self) -> &Arc<LazyGlobalStats> {
553 &self.global_stats
554 }
555
556 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 pub fn num_docs(&self) -> u32 {
569 self.total_docs
570 }
571
572 pub fn segment_map(&self) -> &FxHashMap<u128, usize> {
574 &self.segment_map
575 }
576
577 #[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 #[cfg(not(feature = "sync"))]
587 pub(crate) fn install_search_cpu<R>(&self, operation: impl FnOnce() -> R) -> R {
588 operation()
589 }
590
591 pub fn num_segments(&self) -> usize {
593 self.segments.len()
594 }
595
596 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 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 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 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 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 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 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 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 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 #[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 const MAX_ASYNC_SEGMENT_SEARCHES: usize = 8;
925 use futures::StreamExt;
926 use futures::TryStreamExt;
927 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 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 #[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 #[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 #[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 #[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 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 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 #[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 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 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 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 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 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 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 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
1398fn 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#[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
1508fn 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#[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 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
1600fn process_rss_bytes() -> u64 {
1602 #[cfg(target_os = "linux")]
1603 {
1604 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 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}