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] reusing {} segment readers, loading {} new",
348 reused.len(),
349 to_load.len(),
350 );
351 }
352
353 let store_cache_directory_namespace = Arc::as_ptr(directory) as usize;
357
358 let futures: Vec<_> = to_load
360 .iter()
361 .map(|(_, segment_id)| {
362 let dir = Arc::clone(directory);
363 let sch = Arc::clone(schema);
364 let store_cache = Arc::clone(&store_cache);
365 let sid = *segment_id;
366 async move {
367 SegmentReader::open_with_store_cache(
368 dir.as_ref(),
369 sid,
370 sch,
371 term_cache_blocks,
372 store_cache_directory_namespace,
373 store_cache,
374 )
375 .await
376 }
377 })
378 .collect();
379
380 let results = futures::future::join_all(futures).await;
381
382 let mut loaded: Vec<(usize, Arc<SegmentReader>)> = Vec::with_capacity(valid_segments.len());
384
385 loaded.extend(reused);
387
388 for ((idx, sid), result) in to_load.into_iter().zip(results) {
390 match result {
391 Ok(mut reader) => {
392 reader.set_trained_vectors(Arc::clone(trained_vectors));
394 loaded.push((idx, Arc::new(reader)));
395 }
396 Err(e) => {
397 return Err(crate::error::Error::Internal(format!(
398 "Failed to open segment {:016x}: {:?}",
399 sid.0, e
400 )));
401 }
402 }
403 }
404
405 loaded.sort_by_key(|(idx, _)| *idx);
407
408 let segments: Vec<Arc<SegmentReader>> = loaded.into_iter().map(|(_, seg)| seg).collect();
409
410 let total_docs: u64 = segments.iter().map(|s| s.meta().num_docs as u64).sum();
414 let mut total_heap = 0usize;
415 let mut total_file_backed = 0u64;
416 let mut total_pinned = 0u64;
417 let mut total_pin_intended = 0u64;
418 for seg in &segments {
419 let stats = seg.memory_stats();
420 let heap = stats.estimated_heap_bytes();
421 let file_backed = stats.file_backed_bytes();
422 total_heap = total_heap.saturating_add(heap);
423 total_file_backed = total_file_backed.saturating_add(file_backed);
424 total_pinned = total_pinned.saturating_add(stats.pinned_metadata_bytes);
425 total_pin_intended = total_pin_intended.saturating_add(stats.pin_intended_bytes);
426 log::info!(
427 "[searcher] segment {:016x}: docs={}, heap_estimate={} \
428 (term_cache={}, store_cache={}, sparse_vectors={}, dense_vectors={}), \
429 file_backed={} (term_bloom={}, sparse_vectors={}, dense_vectors={}), \
430 pinned_metadata={} of {} eligible \
431 (sparse_vectors={} of {}, dense_vectors={} of {})",
432 stats.segment_id,
433 stats.num_docs,
434 crate::format_bytes(heap as u64),
435 crate::format_bytes(stats.term_dict_cache_bytes as u64),
436 crate::format_bytes(stats.store_cache_bytes as u64),
437 crate::format_bytes(stats.sparse_heap_bytes as u64),
438 crate::format_bytes(stats.dense_heap_bytes as u64),
439 crate::format_bytes(file_backed),
440 crate::format_bytes(stats.term_bloom_file_bytes),
441 crate::format_bytes(stats.sparse_file_backed_bytes),
442 crate::format_bytes(stats.dense_file_backed_bytes),
443 crate::format_bytes(stats.pinned_metadata_bytes),
444 crate::format_bytes(stats.pin_intended_bytes),
445 crate::format_bytes(stats.sparse_pinned_metadata_bytes),
446 crate::format_bytes(stats.sparse_pin_intended_bytes),
447 crate::format_bytes(stats.dense_pinned_metadata_bytes),
448 crate::format_bytes(stats.dense_pin_intended_bytes),
449 );
450 }
451 let rss_bytes = process_rss_bytes();
453 log::info!(
454 "[searcher] loaded {} segments: total_docs={}, heap_estimate={}, \
455 file_backed={}, pinned_metadata={} of {} eligible, \
456 shared_store_cache={} in {} blocks, process_rss={}",
457 segments.len(),
458 total_docs,
459 crate::format_bytes(total_heap as u64),
460 crate::format_bytes(total_file_backed),
461 crate::format_bytes(total_pinned),
462 crate::format_bytes(total_pin_intended),
463 crate::format_bytes(store_cache.total_bytes() as u64),
464 store_cache.total_blocks(),
465 crate::format_bytes(rss_bytes),
466 );
467
468 Ok(segments)
469 }
470
471 fn build_default_fields(schema: &Schema) -> Vec<crate::Field> {
473 if !schema.default_fields().is_empty() {
474 schema.default_fields().to_vec()
475 } else {
476 schema
477 .fields()
478 .filter(|(_, entry)| {
479 entry.indexed && entry.field_type == crate::dsl::FieldType::Text
480 })
481 .map(|(field, _)| field)
482 .collect()
483 }
484 }
485
486 pub fn schema(&self) -> &Schema {
488 &self.schema
489 }
490
491 pub fn segment_readers(&self) -> &[Arc<SegmentReader>] {
493 &self.segments
494 }
495
496 pub fn default_fields(&self) -> &[crate::Field] {
498 &self.default_fields
499 }
500
501 pub fn tokenizers(&self) -> &crate::tokenizer::TokenizerRegistry {
503 &self.tokenizers
504 }
505
506 pub fn trained_centroids(&self) -> &FxHashMap<u32, Arc<crate::structures::CoarseCentroids>> {
508 &self.trained_vectors.centroids
509 }
510
511 pub fn trained_binary_quantizers(
512 &self,
513 ) -> &FxHashMap<u32, Arc<crate::structures::BinaryCoarseQuantizer>> {
514 &self.trained_vectors.binary_quantizers
515 }
516
517 pub fn global_stats(&self) -> &Arc<LazyGlobalStats> {
519 &self.global_stats
520 }
521
522 fn build_lookup_tables(segments: &[Arc<SegmentReader>]) -> (FxHashMap<u128, usize>, u32) {
524 let mut segment_map = FxHashMap::default();
525 let mut total = 0u32;
526 for (i, seg) in segments.iter().enumerate() {
527 segment_map.insert(seg.meta().id, i);
528 total = total.saturating_add(seg.meta().num_docs);
529 }
530 (segment_map, total)
531 }
532
533 pub fn num_docs(&self) -> u32 {
535 self.total_docs
536 }
537
538 pub fn segment_map(&self) -> &FxHashMap<u128, usize> {
540 &self.segment_map
541 }
542
543 #[cfg(feature = "sync")]
545 pub(crate) fn install_search_cpu<R: Send>(&self, operation: impl FnOnce() -> R + Send) -> R {
546 self.search_pool.install(operation)
547 }
548
549 #[cfg(not(feature = "sync"))]
553 pub(crate) fn install_search_cpu<R>(&self, operation: impl FnOnce() -> R) -> R {
554 operation()
555 }
556
557 pub fn num_segments(&self) -> usize {
559 self.segments.len()
560 }
561
562 pub async fn doc(&self, segment_id: u128, doc_id: u32) -> Result<Option<crate::dsl::Document>> {
564 if let Some(&idx) = self.segment_map.get(&segment_id) {
565 return self.segments[idx].doc(doc_id).await;
566 }
567 Ok(None)
568 }
569
570 pub async fn search(
572 &self,
573 query: &dyn crate::query::Query,
574 limit: usize,
575 ) -> Result<Vec<crate::query::SearchResult>> {
576 let (results, _) = self.search_with_count(query, limit).await?;
577 Ok(results)
578 }
579
580 pub async fn search_with_count(
583 &self,
584 query: &dyn crate::query::Query,
585 limit: usize,
586 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
587 self.search_with_offset_and_count(query, limit, 0).await
588 }
589
590 pub async fn search_with_offset(
592 &self,
593 query: &dyn crate::query::Query,
594 limit: usize,
595 offset: usize,
596 ) -> Result<Vec<crate::query::SearchResult>> {
597 let (results, _) = self
598 .search_with_offset_and_count(query, limit, offset)
599 .await?;
600 Ok(results)
601 }
602
603 pub async fn search_with_offset_and_count(
605 &self,
606 query: &dyn crate::query::Query,
607 limit: usize,
608 offset: usize,
609 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
610 self.search_internal(query, limit, offset, false).await
611 }
612
613 pub async fn search_with_positions(
617 &self,
618 query: &dyn crate::query::Query,
619 limit: usize,
620 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
621 self.search_internal(query, limit, 0, true).await
622 }
623
624 fn prepare_global_lsp(
631 &self,
632 query: &dyn crate::query::Query,
633 retrieval_depth: usize,
634 parallel: bool,
635 ) -> Result<Vec<Option<std::sync::Arc<crate::query::bmp::LspSegmentPlan>>>> {
636 let total_start = std::time::Instant::now();
637 let empty = || vec![None; self.segments.len()];
638 if retrieval_depth == 0 {
639 return Ok(empty());
640 }
641 let crate::query::QueryDecomposition::SparseTerms(infos) = query.decompose() else {
642 return Ok(empty());
643 };
644 let Some(&first) = infos.first() else {
645 return Ok(empty());
646 };
647 if infos
648 .iter()
649 .any(|info| info.field != first.field || info.lsp_gamma != first.lsp_gamma)
650 {
651 return Ok(empty());
652 }
653 let field = first.field;
654 let field_label = self.schema.get_field_name(field).unwrap_or("?");
655 let (total_superblocks, total_coarse_groups, planning_depth) = self
656 .segments
657 .iter()
658 .filter_map(|segment| segment.bmp_index(field))
659 .fold(
660 (0usize, 0usize, retrieval_depth),
661 |(total, coarse, depth), bmp| {
662 (
663 total.saturating_add(bmp.num_superblocks as usize),
664 coarse.saturating_add(bmp.num_coarse_groups as usize),
665 depth.max(crate::query::bmp_executor_limit(
666 retrieval_depth,
667 first.over_fetch_factor,
668 bmp,
669 )),
670 )
671 },
672 );
673 let Some(reference_bmp) = self
674 .segments
675 .iter()
676 .find_map(|segment| segment.bmp_index(field))
677 else {
678 return Ok(empty());
679 };
680 if !infos.iter().any(|info| info.candidate) {
681 return Ok(empty());
682 }
683 let prepare_start = std::time::Instant::now();
684 let Some(prepared_query) =
685 crate::query::bmp::prepare_bmp_query_infos(reference_bmp.dims(), &infos)?
686 else {
687 return Ok(empty());
688 };
689 let infos: std::sync::Arc<[crate::query::SparseTermQueryInfo]> = infos.into();
690 let prepared_query = std::sync::Arc::new(prepared_query);
691 let prepare_secs = prepare_start.elapsed().as_secs_f64();
692
693 let local_plans = || {
694 let plan = std::sync::Arc::new(crate::query::bmp::LspSegmentPlan {
695 infos: std::sync::Arc::clone(&infos),
696 prepared_query: std::sync::Arc::clone(&prepared_query),
697 selection: None,
698 });
699 self.segments
700 .iter()
701 .map(|segment| {
702 segment
703 .bmp_index(field)
704 .map(|_| std::sync::Arc::clone(&plan))
705 })
706 .collect()
707 };
708 let gamma = first
709 .lsp_gamma
710 .unwrap_or_else(|| crate::query::bmp::recommended_lsp_gamma(planning_depth));
711 if gamma == 0 || gamma >= total_superblocks {
712 crate::observe::bmp_lsp(
716 self.schema.index_label(),
717 field_label,
718 total_start.elapsed().as_secs_f64(),
719 prepare_secs,
720 0.0,
721 0.0,
722 total_superblocks,
723 gamma,
724 total_coarse_groups,
725 0,
726 0,
727 );
728 return Ok(local_plans());
729 }
730 let hierarchy_scan_start = std::time::Instant::now();
731 let prepare = |segment: &std::sync::Arc<crate::segment::SegmentReader>| {
732 segment
733 .bmp_index(field)
734 .map(|bmp| crate::query::bmp::prepare_lsp_coarse_ubs(bmp, &prepared_query))
735 .transpose()
736 };
737
738 #[cfg(feature = "sync")]
739 let coarse_bounds: Vec<Option<Vec<f32>>> = if parallel {
740 use rayon::prelude::*;
741 self.search_pool.install(|| {
742 self.segments
743 .par_iter()
744 .map(prepare)
745 .collect::<Result<Vec<_>>>()
746 })?
747 } else {
748 self.segments
749 .iter()
750 .map(prepare)
751 .collect::<Result<Vec<_>>>()?
752 };
753 #[cfg(not(feature = "sync"))]
754 let coarse_bounds: Vec<Option<Vec<f32>>> = {
755 let _ = parallel;
756 self.segments
757 .iter()
758 .map(prepare)
759 .collect::<Result<Vec<_>>>()?
760 };
761 let hierarchy_scan_secs = hierarchy_scan_start.elapsed().as_secs_f64();
762
763 let select_start = std::time::Instant::now();
764 let selection =
765 select_global_lsp_hierarchical(&coarse_bounds, gamma, |segment, group, out| {
766 let bmp = self.segments[segment].bmp_index(field).ok_or_else(|| {
767 crate::Error::Internal(
768 "BMP coarse plan references a segment without the sparse field".into(),
769 )
770 })?;
771 crate::query::bmp::expand_lsp_coarse_group(bmp, &prepared_query, group, out)
772 })?;
773 let mut plans = Vec::with_capacity(self.segments.len());
774 for (segment, selected) in selection.selected.into_iter().enumerate() {
775 if self.segments[segment].bmp_index(field).is_none() {
776 plans.push(None);
777 continue;
778 }
779 let (selected_superblocks, selected_bounds): (Vec<_>, Vec<_>) =
780 selected.into_iter().unzip();
781 plans.push(Some(std::sync::Arc::new(
782 crate::query::bmp::LspSegmentPlan {
783 infos: std::sync::Arc::clone(&infos),
784 prepared_query: std::sync::Arc::clone(&prepared_query),
785 selection: Some(crate::query::bmp::LspSelection {
786 sb_ubs: selected_bounds,
787 sb_order: selected_superblocks,
788 }),
789 },
790 )));
791 }
792 let select_secs = select_start.elapsed().as_secs_f64();
793 crate::observe::bmp_lsp(
794 self.schema.index_label(),
795 field_label,
796 total_start.elapsed().as_secs_f64(),
797 prepare_secs,
798 hierarchy_scan_secs,
799 select_secs,
800 total_superblocks,
801 gamma,
802 total_coarse_groups,
803 selection.expanded_groups,
804 selection.evaluated_superblocks,
805 );
806 log::debug!(
807 "BMP hierarchical LSP: field={}, coarse_groups={}/{}, E_superblocks={}/{}, gamma={}",
808 field_label,
809 selection.expanded_groups,
810 coarse_bounds
811 .iter()
812 .filter_map(Option::as_ref)
813 .map(Vec::len)
814 .sum::<usize>(),
815 selection.evaluated_superblocks,
816 total_superblocks,
817 gamma,
818 );
819 Ok(plans)
820 }
821
822 fn ordered_lsp_segments(
823 &self,
824 plans: &[Option<std::sync::Arc<crate::query::bmp::LspSegmentPlan>>],
825 ) -> Vec<usize> {
826 let mut order: Vec<usize> = (0..self.segments.len()).collect();
827 order.sort_unstable_by(|&left, &right| {
828 let left_priority = plans[left]
829 .as_ref()
830 .map_or(f32::NEG_INFINITY, |plan| plan.priority());
831 let right_priority = plans[right]
832 .as_ref()
833 .map_or(f32::NEG_INFINITY, |plan| plan.priority());
834 right_priority
835 .total_cmp(&left_priority)
836 .then_with(|| {
837 self.segments[right]
838 .num_docs()
839 .cmp(&self.segments[left].num_docs())
840 })
841 .then_with(|| {
842 self.segments[left]
843 .meta()
844 .id
845 .cmp(&self.segments[right].meta().id)
846 })
847 });
848 order
849 }
850
851 #[inline]
852 fn bmp_wave_width(&self) -> usize {
853 #[cfg(feature = "native")]
854 {
855 self.bmp_io_concurrency
856 }
857 #[cfg(not(feature = "native"))]
858 {
859 4
860 }
861 }
862
863 async fn search_internal(
865 &self,
866 query: &dyn crate::query::Query,
867 limit: usize,
868 offset: usize,
869 collect_positions: bool,
870 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
871 let fetch_limit = checked_search_window(limit, offset)?;
872
873 #[cfg(feature = "sync")]
879 if !self.segments.is_empty()
880 && tokio::runtime::Handle::current().runtime_flavor()
881 == tokio::runtime::RuntimeFlavor::MultiThread
882 {
883 return self.search_internal_parallel(query, fetch_limit, offset, collect_positions);
884 }
885
886 const MAX_ASYNC_SEGMENT_SEARCHES: usize = 8;
890 use futures::StreamExt;
891 use futures::TryStreamExt;
892 let shared = crate::query::SharedThreshold::new();
895 let lsp_plans = self.prepare_global_lsp(query, fetch_limit, false)?;
896 let mut total_seen: u32 = 0;
897 let mut merged = Vec::new();
898 let mut merge_scratch = Vec::new();
899 let bmp_planned = lsp_plans.iter().any(Option::is_some);
900 let order = if bmp_planned {
901 self.ordered_lsp_segments(&lsp_plans)
902 } else {
903 (0..self.segments.len()).collect()
904 };
905 #[cfg(feature = "native")]
906 let bmp_io_gate = Arc::clone(&self.bmp_io_gate);
907 let run_segment = |segment_index: usize| {
908 let segment = Arc::clone(&self.segments[segment_index]);
909 let lsp_plan = lsp_plans[segment_index].clone();
910 let shared = shared.clone();
911 #[cfg(feature = "native")]
912 let bmp_io_gate = Arc::clone(&bmp_io_gate);
913 async move {
914 if lsp_plan.as_ref().is_some_and(|plan| !plan.has_work()) {
915 return Ok((Vec::new(), 0u32));
916 }
917 #[cfg(feature = "native")]
918 let _io_permit = if lsp_plan.is_some() {
919 Some(bmp_io_gate.acquire_async().await)
920 } else {
921 None
922 };
923 let sid = segment.meta().id;
924 let (mut results, segment_seen) = crate::query::search_segment_shared_planned(
925 segment.as_ref(),
926 query,
927 fetch_limit,
928 collect_positions,
929 shared.clone(),
930 lsp_plan,
931 )
932 .await?;
933 if fetch_limit > 0 && results.len() >= fetch_limit {
934 shared.raise(results[fetch_limit - 1].score);
935 }
936 for result in &mut results {
937 result.segment_id = sid;
938 }
939 Ok::<_, crate::error::Error>((results, segment_seen))
940 }
941 };
942
943 let mut remainder = order.as_slice();
944 if bmp_planned {
945 if let Some((&pilot, rest)) = order.split_first() {
948 let (batch, segment_seen) = run_segment(pilot).await?;
949 total_seen = total_seen.saturating_add(segment_seen);
950 merge_ranked_reuse(&mut merged, batch, fetch_limit, &mut merge_scratch);
951 remainder = rest;
952 }
953 }
954 let concurrency = if bmp_planned {
955 self.bmp_wave_width()
956 } else {
957 MAX_ASYNC_SEGMENT_SEARCHES
958 };
959 let searches = futures::stream::iter(remainder.iter().copied().map(run_segment))
960 .buffer_unordered(concurrency);
961 futures::pin_mut!(searches);
962 while let Some((batch, segment_seen)) = searches.try_next().await? {
963 total_seen = total_seen.saturating_add(segment_seen);
964 merge_ranked_reuse(&mut merged, batch, fetch_limit, &mut merge_scratch);
965 }
966
967 let results = apply_result_offset(merged, fetch_limit, offset);
968 Ok((results, total_seen))
969 }
970
971 #[cfg(feature = "sync")]
976 fn search_internal_parallel(
977 &self,
978 query: &dyn crate::query::Query,
979 fetch_limit: usize,
980 offset: usize,
981 collect_positions: bool,
982 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
983 tokio::task::block_in_place(|| {
984 self.search_internal_sync(query, fetch_limit, offset, collect_positions)
985 })
986 }
987
988 #[cfg(feature = "sync")]
993 fn search_internal_sync(
994 &self,
995 query: &dyn crate::query::Query,
996 fetch_limit: usize,
997 offset: usize,
998 collect_positions: bool,
999 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1000 let (merged, total_seen) =
1001 self.search_segments_sync(query, fetch_limit, collect_positions)?;
1002
1003 let results = apply_result_offset(merged, fetch_limit, offset);
1004 Ok((results, total_seen))
1005 }
1006
1007 #[cfg(feature = "sync")]
1015 fn search_segments_sync(
1016 &self,
1017 query: &dyn crate::query::Query,
1018 fetch_limit: usize,
1019 collect_positions: bool,
1020 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1021 use rayon::prelude::*;
1022
1023 let lsp_plans = self.prepare_global_lsp(query, fetch_limit, true)?;
1024 let shared = crate::query::SharedThreshold::new();
1025 let run_segment = |segment_index: &usize| {
1026 let segment = &self.segments[*segment_index];
1027 let lsp_plan = lsp_plans[*segment_index].clone();
1028 if lsp_plan.as_ref().is_some_and(|plan| !plan.has_work()) {
1029 return Ok((Vec::new(), 0u32));
1030 }
1031 let _io_permit = lsp_plan.as_ref().map(|_| self.bmp_io_gate.acquire());
1032 let sid = segment.meta().id;
1033 let (mut results, segment_seen) = crate::query::search_segment_shared_sync_planned(
1034 segment.as_ref(),
1035 query,
1036 fetch_limit,
1037 collect_positions,
1038 shared.clone(),
1039 lsp_plan,
1040 )?;
1041 if fetch_limit > 0 && results.len() >= fetch_limit {
1042 shared.raise(results[fetch_limit - 1].score);
1043 }
1044 for result in &mut results {
1045 result.segment_id = sid;
1046 }
1047 Ok::<_, crate::Error>((results, segment_seen))
1048 };
1049
1050 if !lsp_plans.iter().any(Option::is_some) {
1051 return self.search_pool.install(|| {
1052 (0..self.segments.len())
1053 .into_par_iter()
1054 .map(|segment| run_segment(&segment))
1055 .try_reduce(
1056 || (Vec::new(), 0u32),
1057 |(left, left_seen), (right, right_seen)| {
1058 Ok((
1059 merge_two_ranked(left, right, fetch_limit),
1060 left_seen.saturating_add(right_seen),
1061 ))
1062 },
1063 )
1064 });
1065 }
1066
1067 let order = self.ordered_lsp_segments(&lsp_plans);
1068
1069 let mut merged = Vec::new();
1070 let mut merge_scratch = Vec::new();
1071 let mut total_seen = 0u32;
1072 if let Some((&pilot, rest)) = order.split_first() {
1073 let (pilot_results, pilot_seen) = self.search_pool.install(|| run_segment(&pilot))?;
1074 merge_ranked_reuse(&mut merged, pilot_results, fetch_limit, &mut merge_scratch);
1075 total_seen = total_seen.saturating_add(pilot_seen);
1076
1077 for wave in rest.chunks(self.bmp_wave_width()) {
1078 let batches = self
1079 .search_pool
1080 .install(|| wave.par_iter().map(run_segment).collect::<Result<Vec<_>>>())?;
1081 for (results, seen) in batches {
1082 merge_ranked_reuse(&mut merged, results, fetch_limit, &mut merge_scratch);
1083 total_seen = total_seen.saturating_add(seen);
1084 }
1085 }
1086 }
1087 Ok((merged, total_seen))
1088 }
1089
1090 #[cfg(feature = "sync")]
1094 pub fn search_with_offset_and_count_sync(
1095 &self,
1096 query: &dyn crate::query::Query,
1097 limit: usize,
1098 offset: usize,
1099 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1100 let fetch_limit = checked_search_window(limit, offset)?;
1101 let (merged, total_seen) = self.search_segments_sync(query, fetch_limit, false)?;
1102
1103 let results = apply_result_offset(merged, fetch_limit, offset);
1104 Ok((results, total_seen))
1105 }
1106
1107 pub async fn search_fused(
1128 &self,
1129 queries: &[(&dyn crate::query::Query, f32)],
1130 fetch_limit: usize,
1131 limit: usize,
1132 method: crate::query::FusionMethod,
1133 combiner: crate::query::MultiValueCombiner,
1134 ) -> Result<Vec<crate::query::SearchResult>> {
1135 let (results, _) = self
1136 .search_fused_with_count(queries, fetch_limit, limit, method, combiner)
1137 .await?;
1138 Ok(results)
1139 }
1140
1141 pub async fn search_fused_with_count(
1145 &self,
1146 queries: &[(&dyn crate::query::Query, f32)],
1147 fetch_limit: usize,
1148 limit: usize,
1149 method: crate::query::FusionMethod,
1150 combiner: crate::query::MultiValueCombiner,
1151 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1152 if queries.is_empty() {
1153 return Err(crate::Error::Query(
1154 "fusion requires at least one sub-query".to_string(),
1155 ));
1156 }
1157 if queries.len() > crate::query::MAX_FUSION_SUB_QUERIES {
1158 return Err(crate::Error::Query(format!(
1159 "fusion supports at most {} sub-queries, got {}",
1160 crate::query::MAX_FUSION_SUB_QUERIES,
1161 queries.len()
1162 )));
1163 }
1164 if fetch_limit == 0 {
1165 return Err(crate::Error::Query(
1166 "fusion fetch_limit must be greater than zero".to_string(),
1167 ));
1168 }
1169 let candidate_slots = fetch_limit
1170 .checked_mul(queries.len())
1171 .ok_or_else(|| crate::Error::Query("fusion candidate budget overflow".to_string()))?;
1172 if candidate_slots > crate::query::MAX_FUSION_CANDIDATE_SLOTS {
1173 return Err(crate::Error::Query(format!(
1174 "fusion candidate budget must not exceed {}, got {candidate_slots}",
1175 crate::query::MAX_FUSION_CANDIDATE_SLOTS
1176 )));
1177 }
1178 for (index, &(_, weight)) in queries.iter().enumerate() {
1179 if !weight.is_finite() || weight < 0.0 {
1180 return Err(crate::Error::Query(format!(
1181 "fusion query weight at index {index} must be finite and non-negative, \
1182 got {weight}"
1183 )));
1184 }
1185 }
1186 if let crate::query::FusionMethod::Rrf { k } = method
1187 && (!k.is_finite() || k < 0.0)
1188 {
1189 return Err(crate::Error::Query(format!(
1190 "fusion RRF k must be finite and non-negative, got {k}"
1191 )));
1192 }
1193 combiner.validate().map_err(crate::Error::Query)?;
1194
1195 #[cfg(feature = "sync")]
1200 if !self.segments.is_empty()
1201 && tokio::runtime::Handle::current().runtime_flavor()
1202 == tokio::runtime::RuntimeFlavor::MultiThread
1203 {
1204 let lists: Vec<(Vec<crate::query::SearchResult>, f32, u32)> =
1205 tokio::task::block_in_place(|| {
1206 queries
1207 .iter()
1208 .map(|&(query, weight)| {
1209 let (results, seen) =
1210 self.search_internal_sync(query, fetch_limit, 0, true)?;
1211 Ok((results, weight, seen))
1212 })
1213 .collect::<Result<Vec<_>>>()
1214 })?;
1215 let mut total_seen = 0u32;
1216 let ranked_lists = lists
1217 .into_iter()
1218 .map(|(results, weight, seen)| {
1219 total_seen = total_seen.saturating_add(seen);
1220 (results, weight)
1221 })
1222 .collect();
1223 let fused =
1224 crate::query::try_fuse_ranked_lists_chunked(ranked_lists, method, combiner, limit)
1225 .map_err(crate::Error::Query)?;
1226 return Ok((fused, total_seen));
1227 }
1228
1229 let mut lists = Vec::with_capacity(queries.len());
1232 for &(query, weight) in queries {
1233 let (results, seen) = self.search_with_positions(query, fetch_limit).await?;
1234 lists.push((results, weight, seen));
1235 }
1236 let mut total_seen = 0u32;
1237 let ranked_lists = lists
1238 .into_iter()
1239 .map(|(results, weight, seen)| {
1240 total_seen = total_seen.saturating_add(seen);
1241 (results, weight)
1242 })
1243 .collect();
1244 let fused =
1245 crate::query::try_fuse_ranked_lists_chunked(ranked_lists, method, combiner, limit)
1246 .map_err(crate::Error::Query)?;
1247 Ok((fused, total_seen))
1248 }
1249
1250 pub async fn search_and_rerank(
1255 &self,
1256 query: &dyn crate::query::Query,
1257 l1_limit: usize,
1258 final_limit: usize,
1259 config: &crate::query::RerankerConfig,
1260 ) -> Result<(Vec<crate::query::SearchResult>, u32)> {
1261 let (candidates, total_seen) = self.search_with_count(query, l1_limit).await?;
1262 let reranked = crate::query::rerank(self, &candidates, config, final_limit).await?;
1263 Ok((reranked, total_seen))
1264 }
1265
1266 pub async fn query(
1268 &self,
1269 query_str: &str,
1270 limit: usize,
1271 ) -> Result<crate::query::SearchResponse> {
1272 self.query_offset(query_str, limit, 0).await
1273 }
1274
1275 pub async fn query_offset(
1277 &self,
1278 query_str: &str,
1279 limit: usize,
1280 offset: usize,
1281 ) -> Result<crate::query::SearchResponse> {
1282 let parser = self.query_parser();
1283 let query = parser
1284 .parse(query_str)
1285 .map_err(crate::error::Error::Query)?;
1286
1287 let (results, _total_seen) = self
1288 .search_internal(query.as_ref(), limit, offset, false)
1289 .await?;
1290
1291 let total_hits = results.len() as u32;
1292 let hits: Vec<crate::query::SearchHit> = results
1293 .into_iter()
1294 .map(|result| crate::query::SearchHit {
1295 address: crate::query::DocAddress::new(result.segment_id, result.doc_id),
1296 score: result.score,
1297 matched_fields: result.extract_ordinals(),
1298 })
1299 .collect();
1300
1301 Ok(crate::query::SearchResponse { hits, total_hits })
1302 }
1303
1304 pub fn query_parser(&self) -> crate::dsl::QueryLanguageParser {
1306 let query_routers = self.schema.query_routers();
1307 if !query_routers.is_empty()
1308 && let Ok(router) = crate::dsl::QueryFieldRouter::from_rules(query_routers)
1309 {
1310 return crate::dsl::QueryLanguageParser::with_router(
1311 Arc::clone(&self.schema),
1312 self.default_fields.clone(),
1313 Arc::clone(&self.tokenizers),
1314 router,
1315 );
1316 }
1317
1318 crate::dsl::QueryLanguageParser::new(
1319 Arc::clone(&self.schema),
1320 self.default_fields.clone(),
1321 Arc::clone(&self.tokenizers),
1322 )
1323 }
1324
1325 pub async fn get_document(
1327 &self,
1328 address: &crate::query::DocAddress,
1329 ) -> Result<Option<crate::dsl::Document>> {
1330 self.get_document_with_fields(address, None).await
1331 }
1332
1333 pub async fn get_document_with_fields(
1339 &self,
1340 address: &crate::query::DocAddress,
1341 fields: Option<&rustc_hash::FxHashSet<u32>>,
1342 ) -> Result<Option<crate::dsl::Document>> {
1343 let segment_id = address.segment_id_u128().ok_or_else(|| {
1344 crate::error::Error::Query(format!("Invalid segment ID: {}", address.segment_id()))
1345 })?;
1346
1347 if let Some(&idx) = self.segment_map.get(&segment_id) {
1348 return self.segments[idx]
1349 .doc_with_fields(address.doc_id, fields)
1350 .await;
1351 }
1352
1353 Ok(None)
1354 }
1355}
1356
1357struct HierarchicalLspSelection {
1358 selected: Vec<Vec<(u32, f32)>>,
1359 expanded_groups: usize,
1360 evaluated_superblocks: usize,
1361}
1362
1363fn select_global_lsp_hierarchical(
1370 coarse_bounds: &[Option<Vec<f32>>],
1371 gamma: usize,
1372 mut expand: impl FnMut(usize, u32, &mut Vec<(u32, f32)>) -> Result<()>,
1373) -> Result<HierarchicalLspSelection> {
1374 if gamma == 0 {
1375 return Err(crate::Error::Internal(
1376 "hierarchical LSP selection requires a positive gamma".into(),
1377 ));
1378 }
1379 let mut frontier = std::collections::BinaryHeap::<(u32, usize, u32)>::new();
1380 for (segment, bounds) in coarse_bounds.iter().enumerate() {
1381 let Some(bounds) = bounds else {
1382 continue;
1383 };
1384 for (coarse_group, &bound) in bounds.iter().enumerate() {
1385 debug_assert!(bound.is_finite() && bound >= 0.0);
1386 if bound > 0.0 {
1387 frontier.push((bound.to_bits(), segment, coarse_group as u32));
1388 }
1389 }
1390 }
1391 let mut top =
1392 std::collections::BinaryHeap::<std::cmp::Reverse<(u32, usize, u32)>>::with_capacity(
1393 gamma.min(65_536),
1394 );
1395 let mut expanded =
1396 Vec::with_capacity(crate::segment::reader::bmp::BMP_COARSE_SUPERBLOCKS as usize);
1397 let mut expanded_groups = 0usize;
1398 let mut evaluated_superblocks = 0usize;
1399 while let Some((coarse_bound, segment, coarse_group)) = frontier.pop() {
1400 if top.len() == gamma && top.peek().is_some_and(|minimum| coarse_bound < minimum.0.0) {
1401 break;
1402 }
1403 expand(segment, coarse_group, &mut expanded)?;
1404 expanded_groups += 1;
1405 evaluated_superblocks = evaluated_superblocks.saturating_add(expanded.len());
1406 for &(superblock, bound) in &expanded {
1407 debug_assert!(bound.is_finite() && bound >= 0.0);
1408 if bound <= 0.0 {
1409 continue;
1410 }
1411 let candidate = (bound.to_bits(), segment, superblock);
1412 if top.len() < gamma {
1413 top.push(std::cmp::Reverse(candidate));
1414 } else if top.peek().is_some_and(|minimum| candidate > minimum.0) {
1415 top.pop();
1416 top.push(std::cmp::Reverse(candidate));
1417 }
1418 }
1419 }
1420
1421 let mut selected = vec![Vec::<(u32, f32)>::new(); coarse_bounds.len()];
1422 for std::cmp::Reverse((bound, segment, superblock)) in top {
1423 selected[segment].push((superblock, f32::from_bits(bound)));
1424 }
1425 for segment in &mut selected {
1426 segment.sort_unstable_by(|&(left_sb, left_bound), &(right_sb, right_bound)| {
1427 right_bound
1428 .total_cmp(&left_bound)
1429 .then_with(|| left_sb.cmp(&right_sb))
1430 });
1431 }
1432 Ok(HierarchicalLspSelection {
1433 selected,
1434 expanded_groups,
1435 evaluated_superblocks,
1436 })
1437}
1438
1439#[cfg(test)]
1443fn select_global_lsp_superblocks(bounds: &[Option<Vec<f32>>], gamma: usize) -> Vec<Vec<u32>> {
1444 let mut top =
1445 std::collections::BinaryHeap::<std::cmp::Reverse<(u32, usize, u32)>>::with_capacity(
1446 gamma.min(65_536),
1447 );
1448 for (segment, segment_bounds) in bounds.iter().enumerate() {
1449 let Some(segment_bounds) = segment_bounds else {
1450 continue;
1451 };
1452 for (superblock, &bound) in segment_bounds.iter().enumerate() {
1453 if bound <= 0.0 {
1454 continue;
1455 }
1456 let candidate = (bound.to_bits(), segment, superblock as u32);
1457 if top.len() < gamma {
1458 top.push(std::cmp::Reverse(candidate));
1459 } else if top.peek().is_some_and(|minimum| candidate > minimum.0) {
1460 top.pop();
1461 top.push(std::cmp::Reverse(candidate));
1462 }
1463 }
1464 }
1465
1466 let mut selected = vec![Vec::<u32>::new(); bounds.len()];
1467 for std::cmp::Reverse((_, segment, superblock)) in top {
1468 selected[segment].push(superblock);
1469 }
1470 selected
1471}
1472
1473fn merge_ranked_reuse(
1477 merged: &mut Vec<crate::query::SearchResult>,
1478 batch: Vec<crate::query::SearchResult>,
1479 limit: usize,
1480 scratch: &mut Vec<crate::query::SearchResult>,
1481) {
1482 scratch.clear();
1483 let output_len = limit.min(merged.len().saturating_add(batch.len()));
1484 if scratch.capacity() < output_len {
1485 scratch.reserve_exact(output_len);
1486 }
1487 {
1488 let mut left = merged.drain(..).peekable();
1489 let mut right = batch.into_iter().peekable();
1490 while scratch.len() < output_len {
1491 let take_left = match (left.peek(), right.peek()) {
1492 (Some(left), Some(right)) => {
1493 !crate::query::compare_search_results_desc(left, right).is_gt()
1494 }
1495 (Some(_), None) => true,
1496 (None, Some(_)) => false,
1497 (None, None) => break,
1498 };
1499 if take_left {
1500 scratch.push(left.next().expect("peeked left result"));
1501 } else {
1502 scratch.push(right.next().expect("peeked right result"));
1503 }
1504 }
1505 }
1506 std::mem::swap(merged, scratch);
1507}
1508
1509#[cfg(any(feature = "native", test))]
1513fn merge_two_ranked(
1514 left: Vec<crate::query::SearchResult>,
1515 right: Vec<crate::query::SearchResult>,
1516 limit: usize,
1517) -> Vec<crate::query::SearchResult> {
1518 let mut left = left.into_iter().peekable();
1519 let mut right = right.into_iter().peekable();
1520 let mut merged = Vec::with_capacity(limit.min(left.len().saturating_add(right.len())));
1521
1522 while merged.len() < limit {
1523 let take_left = match (left.peek(), right.peek()) {
1524 (Some(left), Some(right)) => {
1525 !crate::query::compare_search_results_desc(left, right).is_gt()
1526 }
1527 (Some(_), None) => true,
1528 (None, Some(_)) => false,
1529 (None, None) => break,
1530 };
1531 if take_left {
1532 merged.push(left.next().expect("peeked left result"));
1533 } else {
1534 merged.push(right.next().expect("peeked right result"));
1535 }
1536 }
1537 merged
1538}
1539
1540fn apply_result_offset(
1541 mut results: Vec<crate::query::SearchResult>,
1542 fetch_limit: usize,
1543 offset: usize,
1544) -> Vec<crate::query::SearchResult> {
1545 if offset == 0 {
1546 results.truncate(fetch_limit);
1547 return results;
1548 }
1549 results
1553 .into_iter()
1554 .skip(offset)
1555 .take(fetch_limit.saturating_sub(offset))
1556 .collect()
1557}
1558
1559fn checked_search_window(limit: usize, offset: usize) -> Result<usize> {
1560 offset
1561 .checked_add(limit)
1562 .ok_or_else(|| crate::Error::Query("search offset + limit overflow".into()))
1563}
1564
1565fn process_rss_bytes() -> u64 {
1567 #[cfg(target_os = "linux")]
1568 {
1569 if let Ok(status) = std::fs::read_to_string("/proc/self/status") {
1571 for line in status.lines() {
1572 if let Some(rest) = line.strip_prefix("VmRSS:") {
1573 let kib: u64 = rest
1574 .trim()
1575 .trim_end_matches("kB")
1576 .trim()
1577 .parse()
1578 .unwrap_or(0);
1579 return kib.saturating_mul(1024);
1580 }
1581 }
1582 }
1583 0
1584 }
1585 #[cfg(target_os = "macos")]
1586 {
1587 use std::mem;
1589 #[repr(C)]
1590 struct TaskBasicInfo {
1591 virtual_size: u64,
1592 resident_size: u64,
1593 resident_size_max: u64,
1594 user_time: [u32; 2],
1595 system_time: [u32; 2],
1596 policy: i32,
1597 suspend_count: i32,
1598 }
1599 unsafe extern "C" {
1600 fn mach_task_self() -> u32;
1601 fn task_info(task: u32, flavor: u32, info: *mut TaskBasicInfo, count: *mut u32) -> i32;
1602 }
1603 const MACH_TASK_BASIC_INFO: u32 = 20;
1604 let mut info: TaskBasicInfo = unsafe { mem::zeroed() };
1605 let mut count = (mem::size_of::<TaskBasicInfo>() / mem::size_of::<u32>()) as u32;
1606 let ret = unsafe {
1607 task_info(
1608 mach_task_self(),
1609 MACH_TASK_BASIC_INFO,
1610 &mut info,
1611 &mut count,
1612 )
1613 };
1614 if ret == 0 { info.resident_size } else { 0 }
1615 }
1616 #[cfg(not(any(target_os = "linux", target_os = "macos")))]
1617 {
1618 0
1619 }
1620}
1621
1622#[cfg(test)]
1623mod load_segments_tests {
1624 use super::*;
1625
1626 #[tokio::test]
1627 async fn searcher_open_fails_loud_on_corrupt_metadata_segment_id() {
1628 let directory = Arc::new(crate::directories::RamDirectory::new());
1629 let schema = Arc::new(crate::dsl::SchemaBuilder::default().build());
1630
1631 let result =
1632 Searcher::open(directory, schema, &["not-a-hex-segment-id".to_string()], 8).await;
1633
1634 match result {
1635 Ok(searcher) => panic!(
1636 "corrupt segment ID must fail loud instead of silently serving {} segments",
1637 searcher.segment_readers().len()
1638 ),
1639 Err(crate::error::Error::Corruption(message)) => {
1640 assert!(message.contains("not-a-hex-segment-id"), "{message}");
1641 }
1642 Err(other) => panic!("expected Corruption error for invalid segment ID, got: {other}"),
1643 }
1644 }
1645}
1646
1647#[cfg(test)]
1648mod search_window_tests {
1649 use super::{
1650 apply_result_offset, checked_search_window, merge_ranked_reuse, merge_two_ranked,
1651 select_global_lsp_hierarchical, select_global_lsp_superblocks,
1652 };
1653 use crate::query::SearchResult;
1654
1655 fn result(segment_id: u128, doc_id: u32, score: f32) -> SearchResult {
1656 SearchResult {
1657 doc_id,
1658 score,
1659 segment_id,
1660 positions: Vec::new(),
1661 }
1662 }
1663
1664 #[test]
1665 fn search_window_is_checked() {
1666 assert_eq!(checked_search_window(7, 5).unwrap(), 12);
1667 assert!(checked_search_window(1, usize::MAX).is_err());
1668 }
1669
1670 #[test]
1671 fn bounded_merge_preserves_canonical_order_and_ties() {
1672 let left = vec![result(2, 9, 10.0), result(2, 3, 7.0)];
1673 let right = vec![result(1, 8, 10.0), result(1, 2, 7.0)];
1674
1675 let expected = merge_two_ranked(left.clone(), right.clone(), 3);
1676 let mut merged = left;
1677 let mut scratch = Vec::new();
1678 merge_ranked_reuse(&mut merged, right, 3, &mut scratch);
1679 assert_eq!(merged, expected);
1680 assert!(scratch.is_empty());
1681 let keys: Vec<_> = merged
1682 .iter()
1683 .map(|result| (result.score, result.segment_id, result.doc_id))
1684 .collect();
1685 assert_eq!(keys, vec![(10.0, 1, 8), (10.0, 2, 9), (7.0, 1, 2)]);
1686 }
1687
1688 #[test]
1689 fn result_offset_returns_only_the_requested_window() {
1690 let results = (0..8)
1691 .map(|doc_id| result(1, doc_id, 8.0 - doc_id as f32))
1692 .collect();
1693
1694 let page = apply_result_offset(results, 5, 2);
1695 assert_eq!(
1696 page.iter().map(|result| result.doc_id).collect::<Vec<_>>(),
1697 vec![2, 3, 4]
1698 );
1699 }
1700
1701 #[test]
1702 fn lsp_gamma_is_global_not_per_segment() {
1703 let bounds = vec![
1704 Some(vec![9.0, 1.0, 8.0]),
1705 Some(vec![7.0, 6.0, 0.0]),
1706 None,
1707 Some(vec![5.0, 4.0]),
1708 ];
1709 let mut selected = select_global_lsp_superblocks(&bounds, 4);
1710 for segment in &mut selected {
1711 segment.sort_unstable();
1712 }
1713 assert_eq!(selected.iter().map(Vec::len).sum::<usize>(), 4);
1714 assert_eq!(selected[0], vec![0, 2]);
1715 assert_eq!(selected[1], vec![0, 1]);
1716 assert!(selected[2].is_empty());
1717 assert!(selected[3].is_empty());
1718 }
1719
1720 #[test]
1721 fn hierarchical_lsp_matches_full_e_scan_exactly() {
1722 const GROUP: usize = crate::segment::reader::bmp::BMP_COARSE_SUPERBLOCKS as usize;
1723 let full = vec![
1724 Some(
1725 (0..700)
1726 .map(|index| 1_000.0 - index as f32)
1727 .collect::<Vec<_>>(),
1728 ),
1729 Some(
1730 (0..530)
1731 .map(|index| 995.0 - index as f32 * 1.25)
1732 .collect::<Vec<_>>(),
1733 ),
1734 None,
1735 ];
1736 let coarse: Vec<Option<Vec<f32>>> = full
1737 .iter()
1738 .map(|segment| {
1739 segment.as_ref().map(|bounds| {
1740 bounds
1741 .chunks(GROUP)
1742 .map(|group| group.iter().copied().fold(0.0, f32::max))
1743 .collect()
1744 })
1745 })
1746 .collect();
1747 let gamma = 37;
1748 let hierarchy = select_global_lsp_hierarchical(&coarse, gamma, |segment, group, output| {
1749 output.clear();
1750 let Some(bounds) = &full[segment] else {
1751 return Ok(());
1752 };
1753 let start = group as usize * GROUP;
1754 output.extend(
1755 bounds[start..(start + GROUP).min(bounds.len())]
1756 .iter()
1757 .enumerate()
1758 .map(|(within, &bound)| ((start + within) as u32, bound)),
1759 );
1760 Ok(())
1761 })
1762 .unwrap();
1763 let mut expected = select_global_lsp_superblocks(&full, gamma);
1764 for segment in &mut expected {
1765 segment.sort_unstable();
1766 }
1767 let mut actual: Vec<Vec<u32>> = hierarchy
1768 .selected
1769 .iter()
1770 .map(|segment| {
1771 let mut ids: Vec<_> = segment.iter().map(|&(id, _)| id).collect();
1772 ids.sort_unstable();
1773 ids
1774 })
1775 .collect();
1776 for segment in &mut actual {
1777 segment.sort_unstable();
1778 }
1779 assert_eq!(actual, expected);
1780 assert_eq!(actual.iter().map(Vec::len).sum::<usize>(), gamma);
1781 assert!(
1782 hierarchy.evaluated_superblocks
1783 < full.iter().filter_map(Option::as_ref).map(Vec::len).sum()
1784 );
1785 }
1786}
1787
1788#[cfg(test)]
1789mod fusion_parallelism_tests {
1790 #[test]
1791 fn fusion_keeps_parallelism_at_the_segment_level() {
1792 let source = include_str!("searcher.rs");
1793 let body = source
1794 .split("pub async fn search_fused_with_count")
1795 .nth(1)
1796 .and_then(|tail| tail.split("/// Two-stage search").next())
1797 .expect("bounded fusion search implementation");
1798
1799 assert!(
1800 !body.contains(".par_iter()"),
1801 "sub-query parallelism nests over segment parallelism"
1802 );
1803 assert!(
1804 !body.contains(".buffered("),
1805 "async fusion fallback must preserve the same bounded execution shape"
1806 );
1807 }
1808}