11use std:: sync:: Arc ;
22
3- use async_trait:: async_trait;
43use chroma_blockstore:: { BlockfileFlusher , BlockfileReader , BlockfileWriter } ;
54use chroma_error:: { ChromaError , ErrorCodes } ;
65use chroma_types:: {
@@ -14,40 +13,6 @@ use uuid::Uuid;
1413
1514use crate :: sparse:: types:: { decode_u32, encode_u32, Score , TopKHeap } ;
1615
17- // ── Two-phase re-scoring ────────────────────────────────────────────
18-
19- /// Batch-oriented rescorer for two-phase retrieval. MaxScore generates
20- /// an oversampled candidate set with approximate (quantized) scores;
21- /// the rescorer computes exact scores so the caller can pick the true
22- /// top-k.
23- #[ async_trait]
24- pub trait SparseRescorer : Send + Sync {
25- async fn rescore_batch ( & self , doc_ids : & [ u32 ] , query : & [ ( u32 , f32 ) ] ) -> Vec < f32 > ;
26- }
27-
28- /// Re-score an oversampled candidate set and return the final top-k
29- /// by exact score.
30- pub async fn rescore_and_select (
31- candidates : Vec < Score > ,
32- k : usize ,
33- query : & [ ( u32 , f32 ) ] ,
34- rescorer : & dyn SparseRescorer ,
35- ) -> Vec < Score > {
36- if candidates. is_empty ( ) || k == 0 {
37- return vec ! [ ] ;
38- }
39-
40- let doc_ids: Vec < u32 > = candidates. iter ( ) . map ( |s| s. offset ) . collect ( ) ;
41- let exact_scores = rescorer. rescore_batch ( & doc_ids, query) . await ;
42-
43- let mut heap = TopKHeap :: new ( k) ;
44- for ( i, & score) in exact_scores. iter ( ) . enumerate ( ) {
45- heap. push ( score, doc_ids[ i] ) ;
46- }
47-
48- heap. into_sorted_vec ( )
49- }
50-
5116const DEFAULT_BLOCK_SIZE : u32 = 1024 ;
5217
5318pub const SPARSE_POSTING_BLOCK_SIZE_BYTES : usize = 1024 * 1024 ;
0 commit comments