1pub(crate) mod bmp;
4pub(crate) mod loader;
5mod types;
6
7pub use bmp::BmpIndex;
8#[cfg(feature = "diagnostics")]
9pub use types::DimRawData;
10pub use types::{SparseIndex, VectorIndex, VectorSearchResult};
11
12#[derive(Debug, Clone, Default)]
14pub struct SegmentMemoryStats {
15 pub segment_id: u128,
17 pub num_docs: u32,
19 pub term_dict_cache_bytes: usize,
21 pub store_cache_bytes: usize,
23 pub sparse_index_bytes: usize,
25 pub dense_index_bytes: usize,
27 pub bloom_filter_bytes: usize,
29 pub pinned_metadata_bytes: u64,
31 pub pin_intended_bytes: u64,
34}
35
36impl SegmentMemoryStats {
37 pub fn total_bytes(&self) -> usize {
39 self.term_dict_cache_bytes
40 + self.store_cache_bytes
41 + self.sparse_index_bytes
42 + self.dense_index_bytes
43 + self.bloom_filter_bytes
44 }
45}
46
47use std::sync::Arc;
48
49use rustc_hash::FxHashMap;
50
51use super::vector_data::LazyFlatVectorData;
52use crate::directories::{Directory, FileHandle};
53use crate::dsl::{DenseVectorQuantization, Document, Field, Schema};
54use crate::structures::{
55 AsyncSSTableReader, BlockPostingList, CoarseCentroids, IVFPQIndex, IVFRaBitQIndex, PQCodebook,
56 RaBitQIndex, SSTableStats, TermInfo,
57};
58use crate::{DocId, Error, Result};
59
60use super::store::{AsyncStoreReader, RawStoreBlock};
61use super::types::{SegmentFiles, SegmentId, SegmentMeta};
62
63pub(crate) fn combine_ordinal_results(
69 raw: impl IntoIterator<Item = (u32, u16, f32)>,
70 combiner: crate::query::MultiValueCombiner,
71 limit: usize,
72) -> Vec<VectorSearchResult> {
73 let collected: Vec<(u32, u16, f32)> = raw.into_iter().collect();
74
75 let num_raw = collected.len();
76 if log::log_enabled!(log::Level::Debug) {
77 let mut ids: Vec<u32> = collected.iter().map(|(d, _, _)| *d).collect();
78 ids.sort_unstable();
79 ids.dedup();
80 log::debug!(
81 "combine_ordinal_results: {} raw entries, {} unique docs, combiner={:?}, limit={}",
82 num_raw,
83 ids.len(),
84 combiner,
85 limit
86 );
87 }
88
89 let all_single = collected.iter().all(|&(_, ord, _)| ord == 0);
91 if all_single {
92 let mut results: Vec<VectorSearchResult> = collected
93 .into_iter()
94 .map(|(doc_id, _, score)| VectorSearchResult::new(doc_id, score, vec![(0, score)]))
95 .collect();
96 results.sort_unstable_by(|a, b| {
97 b.score
98 .partial_cmp(&a.score)
99 .unwrap_or(std::cmp::Ordering::Equal)
100 });
101 results.truncate(limit);
102 return results;
103 }
104
105 let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
107 rustc_hash::FxHashMap::default();
108 for (doc_id, ordinal, score) in collected {
109 doc_ordinals
110 .entry(doc_id as DocId)
111 .or_default()
112 .push((ordinal as u32, score));
113 }
114 let mut results: Vec<VectorSearchResult> = doc_ordinals
115 .into_iter()
116 .map(|(doc_id, ordinals)| {
117 let combined_score = combiner.combine(&ordinals);
118 VectorSearchResult::new(doc_id, combined_score, ordinals)
119 })
120 .collect();
121 results.sort_unstable_by(|a, b| {
122 b.score
123 .partial_cmp(&a.score)
124 .unwrap_or(std::cmp::Ordering::Equal)
125 });
126 results.truncate(limit);
127 results
128}
129
130pub struct SegmentReader {
136 meta: SegmentMeta,
137 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
139 postings_handle: FileHandle,
141 store: Arc<AsyncStoreReader>,
143 schema: Arc<Schema>,
144 vector_indexes: FxHashMap<u32, VectorIndex>,
146 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
148 coarse_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
150 sparse_indexes: FxHashMap<u32, SparseIndex>,
152 bmp_indexes: FxHashMap<u32, BmpIndex>,
154 positions_handle: Option<FileHandle>,
156 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
158 shared_threshold: std::sync::atomic::AtomicU32,
162 #[cfg(feature = "native")]
164 pin_report: crate::segment::pin::PinReport,
165}
166
167impl SegmentReader {
168 pub async fn open<D: Directory>(
170 dir: &D,
171 segment_id: SegmentId,
172 schema: Arc<Schema>,
173 cache_blocks: usize,
174 ) -> Result<Self> {
175 let files = SegmentFiles::new(segment_id.0);
176
177 let meta_slice = dir.open_read(&files.meta).await?;
179 let meta_bytes = meta_slice.read_bytes().await?;
180 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
181 debug_assert_eq!(meta.id, segment_id.0);
182
183 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
185 let term_dict = AsyncSSTableReader::open(term_dict_handle, cache_blocks).await?;
186
187 let postings_handle = dir.open_lazy(&files.postings).await?;
189
190 let store_handle = dir.open_lazy(&files.store).await?;
192 let store = AsyncStoreReader::open(store_handle, cache_blocks).await?;
193
194 let vectors_data = loader::load_vectors_file(dir, &files, &schema).await?;
196 let vector_indexes = vectors_data.indexes;
197 let flat_vectors = vectors_data.flat_vectors;
198
199 #[cfg(feature = "native")]
204 for (field_id, lazy_flat) in &flat_vectors {
205 if vector_indexes.contains_key(field_id) {
206 lazy_flat.advise_random_access();
207 }
208 }
209
210 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
212 let sparse_indexes = sparse_data.maxscore_indexes;
213 let bmp_indexes = sparse_data.bmp_indexes;
214
215 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
217
218 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
220
221 {
223 let mut parts = vec![format!(
224 "[segment] loaded {:016x}: docs={}",
225 segment_id.0, meta.num_docs
226 )];
227 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
228 parts.push(format!(
229 "dense: {} ann + {} flat fields",
230 vector_indexes.len(),
231 flat_vectors.len()
232 ));
233 }
234 for (field_id, idx) in &sparse_indexes {
235 parts.push(format!(
236 "sparse field {}: {} dims, ~{:.1} KB",
237 field_id,
238 idx.num_dimensions(),
239 idx.num_dimensions() as f64 * 24.0 / 1024.0
240 ));
241 }
242 for (field_id, idx) in &bmp_indexes {
243 parts.push(format!(
244 "bmp field {}: {} dims, {} blocks",
245 field_id,
246 idx.dims(),
247 idx.num_blocks
248 ));
249 }
250 if !fast_fields.is_empty() {
251 parts.push(format!("fast: {} fields", fast_fields.len()));
252 }
253 log::debug!("{}", parts.join(", "));
254 }
255
256 #[allow(unused_mut)]
257 let mut reader = Self {
258 meta,
259 term_dict: Arc::new(term_dict),
260 postings_handle,
261 store: Arc::new(store),
262 schema,
263 vector_indexes,
264 flat_vectors,
265 coarse_centroids: FxHashMap::default(),
266 sparse_indexes,
267 bmp_indexes,
268 positions_handle,
269 fast_fields,
270 shared_threshold: std::sync::atomic::AtomicU32::new(0),
271 #[cfg(feature = "native")]
272 pin_report: Default::default(),
273 };
274
275 #[cfg(feature = "native")]
277 reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
278
279 Ok(reader)
280 }
281
282 #[cfg(feature = "native")]
291 pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
292 use crate::segment::pin::PinReport;
293
294 if !policy.is_enabled() {
295 return;
296 }
297 let mut remaining = policy.budget_bytes;
298 let mut report = PinReport::default();
299
300 for bmp in self.bmp_indexes.values_mut() {
302 bmp.pin_block_starts(policy.mode, &mut remaining, &mut report);
303 }
304 for sparse in self.sparse_indexes.values_mut() {
306 sparse.pin_skip_section(policy.mode, &mut remaining, &mut report);
307 }
308 for flat in self.flat_vectors.values_mut() {
310 flat.pin_doc_ids(policy.mode, &mut remaining, &mut report);
311 }
312 for bmp in self.bmp_indexes.values_mut() {
313 bmp.pin_doc_maps(policy.mode, &mut remaining, &mut report);
314 }
315 for bmp in self.bmp_indexes.values_mut() {
317 bmp.pin_sb_grid(policy.mode, &mut remaining, &mut report);
318 }
319
320 if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
321 log::warn!(
322 "[pin] segment {:016x}: pinned {}/{} bytes (budget skipped {}, mlock failed {}) — raise HERMES_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
323 self.meta.id,
324 report.pinned_bytes,
325 report.intended_bytes,
326 report.skipped_budget_bytes,
327 report.failed_bytes,
328 );
329 } else if report.pinned_bytes > 0 {
330 log::info!(
331 "[pin] segment {:016x}: pinned {} bytes of hot metadata ({:?})",
332 self.meta.id,
333 report.pinned_bytes,
334 policy.mode,
335 );
336 }
337 self.pin_report = report;
338 }
339
340 #[inline]
342 pub fn reset_shared_threshold(&self) {
343 self.shared_threshold
344 .store(0, std::sync::atomic::Ordering::Relaxed);
345 }
346
347 #[inline]
349 pub fn shared_threshold_f32(&self) -> f32 {
350 f32::from_bits(
351 self.shared_threshold
352 .load(std::sync::atomic::Ordering::Relaxed),
353 )
354 }
355
356 #[inline]
358 pub fn update_shared_threshold(&self, new_threshold: f32) {
359 use std::sync::atomic::Ordering::Relaxed;
360 let new_bits = new_threshold.to_bits();
361 let mut current_bits = self.shared_threshold.load(Relaxed);
362 while new_threshold > f32::from_bits(current_bits) {
363 match self.shared_threshold.compare_exchange_weak(
364 current_bits,
365 new_bits,
366 Relaxed,
367 Relaxed,
368 ) {
369 Ok(_) => return,
370 Err(actual) => current_bits = actual,
371 }
372 }
373 }
374
375 pub fn meta(&self) -> &SegmentMeta {
376 &self.meta
377 }
378
379 pub fn num_docs(&self) -> u32 {
380 self.meta.num_docs
381 }
382
383 pub fn avg_field_len(&self, field: Field) -> f32 {
385 self.meta.avg_field_len(field)
386 }
387
388 pub fn schema(&self) -> &Schema {
389 &self.schema
390 }
391
392 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
394 &self.sparse_indexes
395 }
396
397 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
399 self.sparse_indexes.get(&field.0)
400 }
401
402 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
404 self.bmp_indexes.get(&field.0)
405 }
406
407 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
409 &self.bmp_indexes
410 }
411
412 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
414 &self.vector_indexes
415 }
416
417 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
419 &self.flat_vectors
420 }
421
422 pub fn fast_field(
424 &self,
425 field_id: u32,
426 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
427 self.fast_fields.get(&field_id)
428 }
429
430 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
432 &self.fast_fields
433 }
434
435 pub fn term_dict_stats(&self) -> SSTableStats {
437 self.term_dict.stats()
438 }
439
440 pub fn memory_stats(&self) -> SegmentMemoryStats {
442 let term_dict_stats = self.term_dict.stats();
443
444 let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
446
447 let store_cache_bytes = self.store.cached_blocks() * 4096;
449
450 let sparse_index_bytes: usize = self
452 .sparse_indexes
453 .values()
454 .map(|s| s.estimated_memory_bytes())
455 .sum::<usize>()
456 + self
457 .bmp_indexes
458 .values()
459 .map(|b| b.estimated_memory_bytes())
460 .sum::<usize>();
461
462 let dense_index_bytes: usize = self
465 .vector_indexes
466 .values()
467 .map(|v| v.estimated_memory_bytes())
468 .sum();
469
470 #[cfg(feature = "native")]
471 let (pinned_metadata_bytes, pin_intended_bytes) =
472 (self.pin_report.pinned_bytes, self.pin_report.intended_bytes);
473 #[cfg(not(feature = "native"))]
474 let (pinned_metadata_bytes, pin_intended_bytes) = (0u64, 0u64);
475
476 SegmentMemoryStats {
477 segment_id: self.meta.id,
478 num_docs: self.meta.num_docs,
479 term_dict_cache_bytes,
480 store_cache_bytes,
481 sparse_index_bytes,
482 dense_index_bytes,
483 bloom_filter_bytes: term_dict_stats.bloom_filter_size,
484 pinned_metadata_bytes,
485 pin_intended_bytes,
486 }
487 }
488
489 pub async fn get_postings(
494 &self,
495 field: Field,
496 term: &[u8],
497 ) -> Result<Option<BlockPostingList>> {
498 log::debug!(
499 "SegmentReader::get_postings field={} term_len={}",
500 field.0,
501 term.len()
502 );
503
504 let mut key = Vec::with_capacity(4 + term.len());
506 key.extend_from_slice(&field.0.to_le_bytes());
507 key.extend_from_slice(term);
508
509 let term_info = match self.term_dict.get(&key).await? {
511 Some(info) => {
512 log::debug!("SegmentReader::get_postings found term_info");
513 info
514 }
515 None => {
516 log::debug!("SegmentReader::get_postings term not found");
517 return Ok(None);
518 }
519 };
520
521 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
523 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
525 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
526 posting_list.push(doc_id, tf);
527 }
528 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
529 return Ok(Some(block_list));
530 }
531
532 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
534 Error::Corruption("TermInfo has neither inline nor external data".to_string())
535 })?;
536
537 let start = posting_offset;
538 let end = start + posting_len;
539
540 if end > self.postings_handle.len() {
541 return Err(Error::Corruption(
542 "Posting offset out of bounds".to_string(),
543 ));
544 }
545
546 let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
547 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
548
549 Ok(Some(block_list))
550 }
551
552 pub async fn get_prefix_postings(
554 &self,
555 field: Field,
556 prefix: &[u8],
557 ) -> Result<Vec<BlockPostingList>> {
558 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
560 key_prefix.extend_from_slice(&field.0.to_le_bytes());
561 key_prefix.extend_from_slice(prefix);
562
563 let entries = self.term_dict.prefix_scan(&key_prefix).await?;
564 let mut results = Vec::with_capacity(entries.len());
565
566 for (_key, term_info) in entries {
567 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
568 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
569 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
570 posting_list.push(doc_id, tf);
571 }
572 results.push(BlockPostingList::from_posting_list(&posting_list)?);
573 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
574 let start = posting_offset;
575 let end = start + posting_len;
576 if end > self.postings_handle.len() {
577 continue;
578 }
579 let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
580 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
581 }
582 }
583
584 Ok(results)
585 }
586
587 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
592 self.doc_with_fields(local_doc_id, None).await
593 }
594
595 pub async fn doc_with_fields(
601 &self,
602 local_doc_id: DocId,
603 fields: Option<&rustc_hash::FxHashSet<u32>>,
604 ) -> Result<Option<Document>> {
605 let mut doc = match fields {
606 Some(set) => {
607 let field_ids: Vec<u32> = set.iter().copied().collect();
608 match self
609 .store
610 .get_fields(local_doc_id, &self.schema, &field_ids)
611 .await
612 {
613 Ok(Some(d)) => d,
614 Ok(None) => return Ok(None),
615 Err(e) => return Err(Error::from(e)),
616 }
617 }
618 None => match self.store.get(local_doc_id, &self.schema).await {
619 Ok(Some(d)) => d,
620 Ok(None) => return Ok(None),
621 Err(e) => return Err(Error::from(e)),
622 },
623 };
624
625 for (&field_id, lazy_flat) in &self.flat_vectors {
627 if let Some(set) = fields
629 && !set.contains(&field_id)
630 {
631 continue;
632 }
633
634 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
635 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
636 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
637 let flat_idx = start + j;
638 if is_binary {
639 let vbs = lazy_flat.vector_byte_size();
640 let mut raw = vec![0u8; vbs];
641 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
642 Ok(()) => {
643 doc.add_binary_dense_vector(Field(field_id), raw);
644 }
645 Err(e) => {
646 log::warn!("Failed to hydrate binary vector field {}: {}", field_id, e);
647 }
648 }
649 } else {
650 match lazy_flat.get_vector(flat_idx).await {
651 Ok(vec) => {
652 doc.add_dense_vector(Field(field_id), vec);
653 }
654 Err(e) => {
655 log::warn!("Failed to hydrate vector field {}: {}", field_id, e);
656 }
657 }
658 }
659 }
660 }
661
662 Ok(Some(doc))
663 }
664
665 pub async fn prefetch_terms(
667 &self,
668 field: Field,
669 start_term: &[u8],
670 end_term: &[u8],
671 ) -> Result<()> {
672 let mut start_key = Vec::with_capacity(4 + start_term.len());
673 start_key.extend_from_slice(&field.0.to_le_bytes());
674 start_key.extend_from_slice(start_term);
675
676 let mut end_key = Vec::with_capacity(4 + end_term.len());
677 end_key.extend_from_slice(&field.0.to_le_bytes());
678 end_key.extend_from_slice(end_term);
679
680 self.term_dict.prefetch_range(&start_key, &end_key).await?;
681 Ok(())
682 }
683
684 pub fn store_has_dict(&self) -> bool {
686 self.store.has_dict()
687 }
688
689 pub fn store(&self) -> &super::store::AsyncStoreReader {
691 &self.store
692 }
693
694 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
696 self.store.raw_blocks()
697 }
698
699 pub fn store_data_slice(&self) -> &FileHandle {
701 self.store.data_slice()
702 }
703
704 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
706 self.term_dict.all_entries().await.map_err(Error::from)
707 }
708
709 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
714 let entries = self.term_dict.all_entries().await?;
715 let mut result = Vec::with_capacity(entries.len());
716
717 for (key, term_info) in entries {
718 if key.len() > 4 {
720 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
721 let term_bytes = &key[4..];
722 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
723 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
724 }
725 }
726 }
727
728 Ok(result)
729 }
730
731 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
733 self.term_dict.iter()
734 }
735
736 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
740 self.term_dict
741 .prefetch_all_data_bulk()
742 .await
743 .map_err(crate::Error::from)
744 }
745
746 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
748 let start = offset;
749 let end = start + len;
750 let bytes = self.postings_handle.read_bytes_range(start..end).await?;
751 Ok(bytes.to_vec())
752 }
753
754 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
756 let handle = match &self.positions_handle {
757 Some(h) => h,
758 None => return Ok(None),
759 };
760 let start = offset;
761 let end = start + len;
762 let bytes = handle.read_bytes_range(start..end).await?;
763 Ok(Some(bytes.to_vec()))
764 }
765
766 pub fn has_positions_file(&self) -> bool {
768 self.positions_handle.is_some()
769 }
770
771 fn score_quantized_batch(
777 query: &[f32],
778 raw: &[u8],
779 quant: crate::dsl::DenseVectorQuantization,
780 dim: usize,
781 scores: &mut [f32],
782 unit_norm: bool,
783 ) {
784 use crate::dsl::DenseVectorQuantization;
785 use crate::structures::simd;
786 match (quant, unit_norm) {
787 (DenseVectorQuantization::F32, false) => {
788 let num_floats = scores.len() * dim;
789 debug_assert!(
790 (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
791 "f32 vector data not 4-byte aligned — vectors file may use legacy format"
792 );
793 let vectors: &[f32] =
794 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
795 simd::batch_cosine_scores(query, vectors, dim, scores);
796 }
797 (DenseVectorQuantization::F32, true) => {
798 let num_floats = scores.len() * dim;
799 debug_assert!(
800 (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
801 "f32 vector data not 4-byte aligned"
802 );
803 let vectors: &[f32] =
804 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
805 simd::batch_dot_scores(query, vectors, dim, scores);
806 }
807 (DenseVectorQuantization::F16, false) => {
808 simd::batch_cosine_scores_f16(query, raw, dim, scores);
809 }
810 (DenseVectorQuantization::F16, true) => {
811 simd::batch_dot_scores_f16(query, raw, dim, scores);
812 }
813 (DenseVectorQuantization::UInt8, false) => {
814 simd::batch_cosine_scores_u8(query, raw, dim, scores);
815 }
816 (DenseVectorQuantization::UInt8, true) => {
817 simd::batch_dot_scores_u8(query, raw, dim, scores);
818 }
819 (DenseVectorQuantization::Binary, _) => {
820 unreachable!("Binary quantization should not reach score_quantized_batch");
822 }
823 }
824 }
825
826 pub async fn search_dense_vector(
832 &self,
833 field: Field,
834 query: &[f32],
835 k: usize,
836 nprobe: usize,
837 rerank_factor: f32,
838 combiner: crate::query::MultiValueCombiner,
839 ) -> Result<Vec<VectorSearchResult>> {
840 let ann_index = self.vector_indexes.get(&field.0);
841 let lazy_flat = self.flat_vectors.get(&field.0);
842
843 if ann_index.is_none() && lazy_flat.is_none() {
845 return Ok(Vec::new());
846 }
847
848 let unit_norm = self
850 .schema
851 .get_field_entry(field)
852 .and_then(|e| e.dense_vector_config.as_ref())
853 .is_some_and(|c| c.unit_norm);
854
855 const BRUTE_FORCE_BATCH: usize = 4096;
857
858 let fetch_k = (k as f32 * rerank_factor.max(1.0)).ceil() as usize;
859
860 let t0 = std::time::Instant::now();
862 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
863 match index {
865 VectorIndex::RaBitQ(lazy) => {
866 let rabitq = lazy.get().ok_or_else(|| {
867 Error::Schema("RaBitQ index deserialization failed".to_string())
868 })?;
869 rabitq
870 .search(query, fetch_k)
871 .into_iter()
872 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
873 .collect()
874 }
875 VectorIndex::IVF(lazy) => {
876 let (index, codebook) = lazy.get().ok_or_else(|| {
877 Error::Schema("IVF index deserialization failed".to_string())
878 })?;
879 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
880 Error::Schema(format!(
881 "IVF index requires coarse centroids for field {}",
882 field.0
883 ))
884 })?;
885 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
886 index
887 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
888 .into_iter()
889 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
890 .collect()
891 }
892 VectorIndex::ScaNN(lazy) => {
893 let (index, codebook) = lazy.get().ok_or_else(|| {
894 Error::Schema("ScaNN index deserialization failed".to_string())
895 })?;
896 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
897 Error::Schema(format!(
898 "ScaNN index requires coarse centroids for field {}",
899 field.0
900 ))
901 })?;
902 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
903 index
904 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
905 .into_iter()
906 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
907 .collect()
908 }
909 VectorIndex::BinaryIvf(_) => {
910 Vec::new()
912 }
913 }
914 } else if let Some(lazy_flat) = lazy_flat {
915 log::debug!(
918 "[search_dense] field {}: brute-force on {} vectors (dim={}, quant={:?})",
919 field.0,
920 lazy_flat.num_vectors,
921 lazy_flat.dim,
922 lazy_flat.quantization
923 );
924 let dim = lazy_flat.dim;
925 let n = lazy_flat.num_vectors;
926 let quant = lazy_flat.quantization;
927 let mut collector = crate::query::ScoreCollector::new(fetch_k);
928 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
929
930 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
931 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
932 let batch_bytes = lazy_flat
933 .read_vectors_batch(batch_start, batch_count)
934 .await
935 .map_err(crate::Error::Io)?;
936 let raw = batch_bytes.as_slice();
937
938 Self::score_quantized_batch(
939 query,
940 raw,
941 quant,
942 dim,
943 &mut scores[..batch_count],
944 unit_norm,
945 );
946
947 for (i, &score) in scores.iter().enumerate().take(batch_count) {
948 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
949 collector.insert_with_ordinal(doc_id, score, ordinal);
950 }
951 }
952
953 collector
954 .into_sorted_results()
955 .into_iter()
956 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
957 .collect()
958 } else {
959 return Ok(Vec::new());
960 };
961 let l1_elapsed = t0.elapsed();
962 {
963 let kind = match ann_index {
964 Some(VectorIndex::RaBitQ(_)) => "rabitq",
965 Some(VectorIndex::IVF(_)) => "ivf_rabitq",
966 Some(VectorIndex::ScaNN(_)) => "scann",
967 Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
968 None => "flat",
969 };
970 crate::observe::dense_l1(field.0, kind, l1_elapsed.as_secs_f64(), results.len());
971 }
972 log::debug!(
973 "[search_dense] field {}: L1 returned {} candidates in {:.1}ms",
974 field.0,
975 results.len(),
976 l1_elapsed.as_secs_f64() * 1000.0
977 );
978
979 if ann_index.is_some()
982 && !results.is_empty()
983 && let Some(lazy_flat) = lazy_flat
984 {
985 let t_rerank = std::time::Instant::now();
986 let dim = lazy_flat.dim;
987 let quant = lazy_flat.quantization;
988 let vbs = lazy_flat.vector_byte_size();
989
990 let mut resolved: Vec<(usize, usize)> = Vec::new(); for (ri, c) in results.iter().enumerate() {
993 let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
994 for (j, &(_, ord)) in entries.iter().enumerate() {
995 if ord == c.1 {
996 resolved.push((ri, start + j));
997 break;
998 }
999 }
1000 }
1001
1002 let t_resolve = t_rerank.elapsed();
1003 if !resolved.is_empty() {
1004 resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
1006
1007 #[cfg(feature = "native")]
1012 lazy_flat.prefetch_vectors(resolved.iter().map(|&(_, flat_idx)| flat_idx));
1013
1014 let t_read = std::time::Instant::now();
1016 let mut raw_buf = vec![0u8; resolved.len() * vbs];
1017 for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
1018 let _ = lazy_flat
1019 .read_vector_raw_into(
1020 flat_idx,
1021 &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
1022 )
1023 .await;
1024 }
1025
1026 let read_elapsed = t_read.elapsed();
1027
1028 let t_score = std::time::Instant::now();
1030 let mut scores = vec![0f32; resolved.len()];
1031 Self::score_quantized_batch(query, &raw_buf, quant, dim, &mut scores, unit_norm);
1032 let score_elapsed = t_score.elapsed();
1033
1034 for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
1036 results[ri].2 = scores[buf_idx];
1037 }
1038
1039 crate::observe::dense_rerank(
1040 field.0,
1041 t_rerank.elapsed().as_secs_f64(),
1042 t_resolve.as_secs_f64(),
1043 read_elapsed.as_secs_f64(),
1044 resolved.len(),
1045 );
1046 log::debug!(
1047 "[search_dense] field {}: rerank {} vectors (dim={}, quant={:?}, {}B/vec): resolve={:.1}ms read={:.1}ms score={:.1}ms",
1048 field.0,
1049 resolved.len(),
1050 dim,
1051 quant,
1052 vbs,
1053 t_resolve.as_secs_f64() * 1000.0,
1054 read_elapsed.as_secs_f64() * 1000.0,
1055 score_elapsed.as_secs_f64() * 1000.0,
1056 );
1057 }
1058
1059 if results.len() > fetch_k {
1060 results.select_nth_unstable_by(fetch_k, |a, b| b.2.total_cmp(&a.2));
1061 results.truncate(fetch_k);
1062 }
1063 results.sort_unstable_by(|a, b| b.2.total_cmp(&a.2));
1064 log::debug!(
1065 "[search_dense] field {}: rerank total={:.1}ms",
1066 field.0,
1067 t_rerank.elapsed().as_secs_f64() * 1000.0
1068 );
1069 }
1070
1071 Ok(combine_ordinal_results(results, combiner, k))
1072 }
1073
1074 fn binary_ivf_nprobe(&self, field: Field) -> Option<usize> {
1076 self.schema
1077 .get_field_entry(field)
1078 .and_then(|e| e.binary_dense_vector_config.as_ref())
1079 .map(|c| c.nprobe)
1080 .filter(|&n| n > 0)
1081 }
1082
1083 pub async fn search_binary_dense_vector(
1087 &self,
1088 field: Field,
1089 query: &[u8],
1090 k: usize,
1091 combiner: crate::query::MultiValueCombiner,
1092 ) -> Result<Vec<VectorSearchResult>> {
1093 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0)
1096 && let Some(ivf) = lazy.get()
1097 {
1098 let nprobe = self.binary_ivf_nprobe(field);
1099 let results = ivf.search(query, k, nprobe);
1100 return Ok(combine_ordinal_results(results, combiner, k));
1101 }
1102
1103 let lazy_flat = match self.flat_vectors.get(&field.0) {
1104 Some(f) => f,
1105 None => return Ok(Vec::new()),
1106 };
1107
1108 const BRUTE_FORCE_BATCH: usize = 8192; let dim_bits = lazy_flat.dim;
1111 let byte_len = lazy_flat.vector_byte_size();
1112 let n = lazy_flat.num_vectors;
1113
1114 if byte_len != query.len() {
1115 return Err(Error::Schema(format!(
1116 "Binary query vector byte length {} != field byte length {}",
1117 query.len(),
1118 byte_len
1119 )));
1120 }
1121
1122 let mut collector = crate::query::ScoreCollector::new(k);
1123 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
1124
1125 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
1126 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
1127 let batch_bytes = lazy_flat
1128 .read_vectors_batch(batch_start, batch_count)
1129 .await
1130 .map_err(crate::Error::Io)?;
1131 let raw = batch_bytes.as_slice();
1132
1133 crate::structures::simd::batch_hamming_scores(
1134 query,
1135 raw,
1136 byte_len,
1137 dim_bits,
1138 &mut scores[..batch_count],
1139 );
1140
1141 let threshold = collector.threshold();
1142 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1143 if score > threshold {
1144 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1145 collector.insert_with_ordinal(doc_id, score, ordinal);
1146 }
1147 }
1148 }
1149
1150 let results: Vec<(u32, u16, f32)> = collector
1151 .into_sorted_results()
1152 .into_iter()
1153 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
1154 .collect();
1155
1156 Ok(combine_ordinal_results(results, combiner, k))
1157 }
1158
1159 pub fn has_dense_vector_index(&self, field: Field) -> bool {
1161 self.vector_indexes.contains_key(&field.0) || self.flat_vectors.contains_key(&field.0)
1162 }
1163
1164 pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
1166 match self.vector_indexes.get(&field.0) {
1167 Some(VectorIndex::RaBitQ(lazy)) => lazy.get().cloned(),
1168 _ => None,
1169 }
1170 }
1171
1172 pub fn get_ivf_vector_index(
1174 &self,
1175 field: Field,
1176 ) -> Option<(Arc<IVFRaBitQIndex>, Arc<crate::structures::RaBitQCodebook>)> {
1177 match self.vector_indexes.get(&field.0) {
1178 Some(VectorIndex::IVF(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
1179 _ => None,
1180 }
1181 }
1182
1183 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
1185 self.coarse_centroids.get(&field_id)
1186 }
1187
1188 pub fn set_coarse_centroids(&mut self, centroids: FxHashMap<u32, Arc<CoarseCentroids>>) {
1190 self.coarse_centroids = centroids;
1191 }
1192
1193 pub fn get_scann_vector_index(
1195 &self,
1196 field: Field,
1197 ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
1198 match self.vector_indexes.get(&field.0) {
1199 Some(VectorIndex::ScaNN(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
1200 _ => None,
1201 }
1202 }
1203
1204 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
1206 self.vector_indexes.get(&field.0)
1207 }
1208
1209 pub async fn get_positions(
1214 &self,
1215 field: Field,
1216 term: &[u8],
1217 ) -> Result<Option<crate::structures::PositionPostingList>> {
1218 let handle = match &self.positions_handle {
1220 Some(h) => h,
1221 None => return Ok(None),
1222 };
1223
1224 let mut key = Vec::with_capacity(4 + term.len());
1226 key.extend_from_slice(&field.0.to_le_bytes());
1227 key.extend_from_slice(term);
1228
1229 let term_info = match self.term_dict.get(&key).await? {
1231 Some(info) => info,
1232 None => return Ok(None),
1233 };
1234
1235 let (offset, length) = match term_info.position_info() {
1237 Some((o, l)) => (o, l),
1238 None => return Ok(None),
1239 };
1240
1241 let slice = handle.slice(offset..offset + length);
1243 let data = slice.read_bytes().await?;
1244
1245 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
1247
1248 Ok(Some(pos_list))
1249 }
1250
1251 pub fn has_positions(&self, field: Field) -> bool {
1253 if let Some(entry) = self.schema.get_field_entry(field) {
1255 entry.positions.is_some()
1256 } else {
1257 false
1258 }
1259 }
1260}
1261
1262#[cfg(feature = "sync")]
1264impl SegmentReader {
1265 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
1267 let mut key = Vec::with_capacity(4 + term.len());
1269 key.extend_from_slice(&field.0.to_le_bytes());
1270 key.extend_from_slice(term);
1271
1272 let term_info = match self.term_dict.get_sync(&key)? {
1274 Some(info) => info,
1275 None => return Ok(None),
1276 };
1277
1278 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1280 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1281 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1282 posting_list.push(doc_id, tf);
1283 }
1284 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
1285 return Ok(Some(block_list));
1286 }
1287
1288 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
1290 Error::Corruption("TermInfo has neither inline nor external data".to_string())
1291 })?;
1292
1293 let start = posting_offset;
1294 let end = start + posting_len;
1295
1296 if end > self.postings_handle.len() {
1297 return Err(Error::Corruption(
1298 "Posting offset out of bounds".to_string(),
1299 ));
1300 }
1301
1302 let posting_bytes = self.postings_handle.read_bytes_range_sync(start..end)?;
1303 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
1304
1305 Ok(Some(block_list))
1306 }
1307
1308 pub fn get_prefix_postings_sync(
1310 &self,
1311 field: Field,
1312 prefix: &[u8],
1313 ) -> Result<Vec<BlockPostingList>> {
1314 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
1315 key_prefix.extend_from_slice(&field.0.to_le_bytes());
1316 key_prefix.extend_from_slice(prefix);
1317
1318 let entries = self.term_dict.prefix_scan_sync(&key_prefix)?;
1319 let mut results = Vec::with_capacity(entries.len());
1320
1321 for (_key, term_info) in entries {
1322 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1323 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1324 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1325 posting_list.push(doc_id, tf);
1326 }
1327 results.push(BlockPostingList::from_posting_list(&posting_list)?);
1328 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
1329 let start = posting_offset;
1330 let end = start + posting_len;
1331 if end > self.postings_handle.len() {
1332 continue;
1333 }
1334 let posting_bytes = self.postings_handle.read_bytes_range_sync(start..end)?;
1335 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
1336 }
1337 }
1338
1339 Ok(results)
1340 }
1341
1342 pub fn get_positions_sync(
1344 &self,
1345 field: Field,
1346 term: &[u8],
1347 ) -> Result<Option<crate::structures::PositionPostingList>> {
1348 let handle = match &self.positions_handle {
1349 Some(h) => h,
1350 None => return Ok(None),
1351 };
1352
1353 let mut key = Vec::with_capacity(4 + term.len());
1355 key.extend_from_slice(&field.0.to_le_bytes());
1356 key.extend_from_slice(term);
1357
1358 let term_info = match self.term_dict.get_sync(&key)? {
1360 Some(info) => info,
1361 None => return Ok(None),
1362 };
1363
1364 let (offset, length) = match term_info.position_info() {
1365 Some((o, l)) => (o, l),
1366 None => return Ok(None),
1367 };
1368
1369 let slice = handle.slice(offset..offset + length);
1370 let data = slice.read_bytes_sync()?;
1371
1372 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
1373 Ok(Some(pos_list))
1374 }
1375
1376 pub fn search_dense_vector_sync(
1379 &self,
1380 field: Field,
1381 query: &[f32],
1382 k: usize,
1383 nprobe: usize,
1384 rerank_factor: f32,
1385 combiner: crate::query::MultiValueCombiner,
1386 ) -> Result<Vec<VectorSearchResult>> {
1387 let ann_index = self.vector_indexes.get(&field.0);
1388 let lazy_flat = self.flat_vectors.get(&field.0);
1389
1390 if ann_index.is_none() && lazy_flat.is_none() {
1391 return Ok(Vec::new());
1392 }
1393
1394 let unit_norm = self
1395 .schema
1396 .get_field_entry(field)
1397 .and_then(|e| e.dense_vector_config.as_ref())
1398 .is_some_and(|c| c.unit_norm);
1399
1400 const BRUTE_FORCE_BATCH: usize = 4096;
1401 let fetch_k = (k as f32 * rerank_factor.max(1.0)).ceil() as usize;
1402
1403 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
1404 match index {
1406 VectorIndex::RaBitQ(lazy) => {
1407 let rabitq = lazy.get().ok_or_else(|| {
1408 Error::Schema("RaBitQ index deserialization failed".to_string())
1409 })?;
1410 rabitq
1411 .search(query, fetch_k)
1412 .into_iter()
1413 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1414 .collect()
1415 }
1416 VectorIndex::IVF(lazy) => {
1417 let (index, codebook) = lazy.get().ok_or_else(|| {
1418 Error::Schema("IVF index deserialization failed".to_string())
1419 })?;
1420 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
1421 Error::Schema(format!(
1422 "IVF index requires coarse centroids for field {}",
1423 field.0
1424 ))
1425 })?;
1426 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
1427 index
1428 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
1429 .into_iter()
1430 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1431 .collect()
1432 }
1433 VectorIndex::ScaNN(lazy) => {
1434 let (index, codebook) = lazy.get().ok_or_else(|| {
1435 Error::Schema("ScaNN index deserialization failed".to_string())
1436 })?;
1437 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
1438 Error::Schema(format!(
1439 "ScaNN index requires coarse centroids for field {}",
1440 field.0
1441 ))
1442 })?;
1443 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
1444 index
1445 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
1446 .into_iter()
1447 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1448 .collect()
1449 }
1450 VectorIndex::BinaryIvf(_) => {
1451 Vec::new()
1453 }
1454 }
1455 } else if let Some(lazy_flat) = lazy_flat {
1456 let dim = lazy_flat.dim;
1458 let n = lazy_flat.num_vectors;
1459 let quant = lazy_flat.quantization;
1460 let mut collector = crate::query::ScoreCollector::new(fetch_k);
1461 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
1462
1463 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
1464 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
1465 let batch_bytes = lazy_flat
1466 .read_vectors_batch_sync(batch_start, batch_count)
1467 .map_err(crate::Error::Io)?;
1468 let raw = batch_bytes.as_slice();
1469
1470 Self::score_quantized_batch(
1471 query,
1472 raw,
1473 quant,
1474 dim,
1475 &mut scores[..batch_count],
1476 unit_norm,
1477 );
1478
1479 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1480 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1481 collector.insert_with_ordinal(doc_id, score, ordinal);
1482 }
1483 }
1484
1485 collector
1486 .into_sorted_results()
1487 .into_iter()
1488 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
1489 .collect()
1490 } else {
1491 return Ok(Vec::new());
1492 };
1493
1494 if ann_index.is_some()
1496 && !results.is_empty()
1497 && let Some(lazy_flat) = lazy_flat
1498 {
1499 let dim = lazy_flat.dim;
1500 let quant = lazy_flat.quantization;
1501 let vbs = lazy_flat.vector_byte_size();
1502
1503 let mut resolved: Vec<(usize, usize)> = Vec::new();
1504 for (ri, c) in results.iter().enumerate() {
1505 let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
1506 for (j, &(_, ord)) in entries.iter().enumerate() {
1507 if ord == c.1 {
1508 resolved.push((ri, start + j));
1509 break;
1510 }
1511 }
1512 }
1513
1514 if !resolved.is_empty() {
1515 resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
1516 let mut raw_buf = vec![0u8; resolved.len() * vbs];
1517 for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
1518 let _ = lazy_flat.read_vector_raw_into_sync(
1519 flat_idx,
1520 &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
1521 );
1522 }
1523
1524 let mut scores = vec![0f32; resolved.len()];
1525 Self::score_quantized_batch(query, &raw_buf, quant, dim, &mut scores, unit_norm);
1526
1527 for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
1528 results[ri].2 = scores[buf_idx];
1529 }
1530 }
1531
1532 if results.len() > fetch_k {
1533 results.select_nth_unstable_by(fetch_k, |a, b| b.2.total_cmp(&a.2));
1534 results.truncate(fetch_k);
1535 }
1536 results.sort_unstable_by(|a, b| b.2.total_cmp(&a.2));
1537 }
1538
1539 Ok(combine_ordinal_results(results, combiner, k))
1540 }
1541
1542 #[cfg(feature = "sync")]
1547 pub fn search_binary_dense_vector_sync(
1548 &self,
1549 field: Field,
1550 query: &[u8],
1551 k: usize,
1552 combiner: crate::query::MultiValueCombiner,
1553 ) -> Result<Vec<VectorSearchResult>> {
1554 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0)
1557 && let Some(ivf) = lazy.get()
1558 {
1559 let nprobe = self.binary_ivf_nprobe(field);
1560 let results = ivf.search(query, k, nprobe);
1561 return Ok(combine_ordinal_results(results, combiner, k));
1562 }
1563
1564 let lazy_flat = match self.flat_vectors.get(&field.0) {
1565 Some(f) => f,
1566 None => return Ok(Vec::new()),
1567 };
1568
1569 const BRUTE_FORCE_BATCH: usize = 8192; let dim_bits = lazy_flat.dim;
1572 let byte_len = lazy_flat.vector_byte_size();
1573 let n = lazy_flat.num_vectors;
1574
1575 if byte_len != query.len() {
1576 return Err(Error::Schema(format!(
1577 "Binary query vector byte length {} != field byte length {}",
1578 query.len(),
1579 byte_len
1580 )));
1581 }
1582
1583 let mut collector = crate::query::ScoreCollector::new(k);
1584 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
1585
1586 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
1587 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
1588 let batch_bytes = lazy_flat
1589 .read_vectors_batch_sync(batch_start, batch_count)
1590 .map_err(crate::Error::Io)?;
1591 let raw = batch_bytes.as_slice();
1592
1593 crate::structures::simd::batch_hamming_scores(
1594 query,
1595 raw,
1596 byte_len,
1597 dim_bits,
1598 &mut scores[..batch_count],
1599 );
1600
1601 let threshold = collector.threshold();
1602 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1603 if score > threshold {
1604 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1605 collector.insert_with_ordinal(doc_id, score, ordinal);
1606 }
1607 }
1608 }
1609
1610 let results: Vec<(u32, u16, f32)> = collector
1611 .into_sorted_results()
1612 .into_iter()
1613 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
1614 .collect();
1615
1616 Ok(combine_ordinal_results(results, combiner, k))
1617 }
1618}