zeppelin_embed/lifecycle/
rescored_scan.rs

1//! Explicit quantized candidate selection followed by exact selected-row scores.
2
3/// Controls the candidate frontier of each active or sealed segment.
4///
5/// The frontier is `min(eligible_rows, k * oversample)`, plus all coarse-score
6/// boundary ties. The limit bounds exact row reads per segment and vector
7/// producer invocation, including ties; exceeding it is a typed error.
8#[derive(Clone, Copy, Debug, Eq, PartialEq)]
9pub struct ScanRescoreOptions {
10    oversample: usize,
11    max_candidates_per_segment: usize,
12}
13
14impl ScanRescoreOptions {
15    /// Creates explicit controls; neither value may be zero.
16    pub const fn new(
17        oversample: usize,
18        max_candidates_per_segment: usize,
19    ) -> Result<Self, crate::quant::RescoreError> {
20        if oversample == 0 {
21            return Err(crate::quant::RescoreError::ZeroOversample);
22        }
23        if max_candidates_per_segment == 0 {
24            return Err(crate::quant::RescoreError::ZeroCandidateLimit);
25        }
26        Ok(Self {
27            oversample,
28            max_candidates_per_segment,
29        })
30    }
31
32    /// Requested coarse frontier multiplier before boundary ties.
33    #[must_use]
34    pub const fn oversample(self) -> usize {
35        self.oversample
36    }
37
38    /// Hard exact-row-read limit per segment and vector producer invocation.
39    #[must_use]
40    pub const fn max_candidates_per_segment(self) -> usize {
41        self.max_candidates_per_segment
42    }
43
44    pub(crate) fn frontier(self, k: usize, eligible: u64) -> Result<usize, super::QueryError> {
45        let requested = k
46            .checked_mul(self.oversample)
47            .ok_or_else(|| rescore_error(crate::quant::RescoreError::ArithmeticOverflow))?;
48        let eligible = usize::try_from(eligible)
49            .map_err(|_| rescore_error(crate::quant::RescoreError::ArithmeticOverflow))?;
50        let frontier = requested.min(eligible);
51        self.check_count(frontier)?;
52        Ok(frontier)
53    }
54
55    fn check_count(self, requested: usize) -> Result<(), super::QueryError> {
56        if requested > self.max_candidates_per_segment {
57            return Err(rescore_error(
58                crate::quant::RescoreError::CandidateLimitExceeded {
59                    requested,
60                    maximum: self.max_candidates_per_segment,
61                },
62            ));
63        }
64        Ok(())
65    }
66}
67
68fn rescore_error(error: crate::quant::RescoreError) -> super::QueryError {
69    super::QueryError::Scan(crate::scan::ScanError::Rescore(error))
70}
71
72pub(crate) fn accumulate(
73    total: &mut crate::diag::ScanRescoreCounters,
74    next: crate::diag::ScanRescoreCounters,
75) -> Result<(), super::QueryError> {
76    let add = |a: u64, b: u64| {
77        a.checked_add(b).ok_or(super::QueryError::Scan(
78            crate::scan::ScanError::ArithmeticOverflow,
79        ))
80    };
81    total.coarse_rows = add(total.coarse_rows, next.coarse_rows)?;
82    total.eligible_rows = add(total.eligible_rows, next.eligible_rows)?;
83    total.candidates_rescored = add(total.candidates_rescored, next.candidates_rescored)?;
84    total.coarse_bytes = add(total.coarse_bytes, next.coarse_bytes)?;
85    total.rescore_bytes = add(total.rescore_bytes, next.rescore_bytes)?;
86    Ok(())
87}
88
89/// Local candidate storage accepted by the shared selected-row scorer.
90pub(crate) trait RescoreCandidate {
91    fn row_id(&self) -> usize;
92    fn score(&self) -> f32;
93    fn set_score(&mut self, score: f32);
94}
95
96impl RescoreCandidate for crate::scan::ScanCandidate {
97    fn row_id(&self) -> usize {
98        self.row_id
99    }
100    fn score(&self) -> f32 {
101        self.score
102    }
103    fn set_score(&mut self, score: f32) {
104        self.score = score;
105    }
106}
107
108/// Reuses the shared exact scorer on only the selected, validated row slices.
109/// Candidate storage is updated in place; the caller owns its charge through
110/// the subsequent identity join and merge.
111#[allow(clippy::too_many_arguments)]
112pub(crate) fn rescore<C: RescoreCandidate>(
113    candidates: &mut Vec<C>,
114    stats: &mut crate::scan::ScanStats,
115    vectors: &[f32],
116    query: &[f32],
117    k: usize,
118    eligible: u64,
119    options: ScanRescoreOptions,
120    cancellation: &super::QueryCancellation<'_>,
121    accounting: &std::sync::Arc<super::stats::Accounting>,
122    memory: &mut super::ExactScanMemory,
123) -> Result<(crate::diag::ScanRescoreCounters, bool, Option<f32>), super::QueryError> {
124    use super::{
125        QueryError, StoreError,
126        stats::{Accounted, AllocationComponent},
127    };
128    use crate::scan::ScanError;
129    cancellation.check_graph().map_err(super::map_scan_error)?;
130    let dimension = query.len();
131    if dimension == 0 {
132        return Err(QueryError::Scan(ScanError::ZeroDimension));
133    }
134    if !vectors.len().is_multiple_of(dimension) {
135        return Err(QueryError::Scan(ScanError::RowDataLength {
136            dimension,
137            actual: vectors.len(),
138        }));
139    }
140    let count = candidates.len();
141    options.check_count(count)?;
142    memory
143        .candidates
144        .set(
145            candidates
146                .capacity()
147                .checked_mul(std::mem::size_of::<C>())
148                .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?,
149        )
150        .map_err(QueryError::Store)?;
151    let mut rows = Accounted::try_with_capacity(accounting, count, AllocationComponent::Temporary)
152        .map_err(QueryError::Store)?;
153    for (position, candidate) in candidates.iter().enumerate() {
154        if position.is_multiple_of(64) {
155            cancellation.check_graph().map_err(super::map_scan_error)?;
156        }
157        let start = candidate
158            .row_id()
159            .checked_mul(dimension)
160            .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
161        let end = start
162            .checked_add(dimension)
163            .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
164        let vector = vectors.get(start..end).ok_or_else(|| {
165            rescore_error(crate::quant::RescoreError::CandidateRowOutOfRange {
166                position,
167                row_index: candidate.row_id(),
168                row_count: vectors.len() / dimension,
169            })
170        })?;
171        rows.push((candidate.row_id(), vector))
172            .map_err(QueryError::Store)?;
173    }
174    let mut position = 0;
175    let mut worst: Option<f32> = None;
176    let mut narrowing: Option<(usize, f64)> = None;
177    let scored = crate::quant::exact_squared_l2_with_sink(
178        query,
179        rows.iter().copied(),
180        |_, checkpoint| {
181            if checkpoint {
182                cancellation.check_graph().map_err(super::map_scan_error)?;
183            }
184            Ok(())
185        },
186        |row_id, exact| {
187            let candidate = candidates
188                .get_mut(position)
189                .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
190            position += 1;
191            let score = exact as f32;
192            if !score.is_finite() {
193                if narrowing.is_none_or(|(row, previous)| {
194                    exact.total_cmp(&previous).is_gt()
195                        || (exact.total_cmp(&previous).is_eq() && row_id < row)
196                }) {
197                    narrowing = Some((row_id, exact));
198                }
199            } else {
200                candidate.set_score(score);
201                if worst.is_none_or(|previous| score.total_cmp(&previous).is_lt()) {
202                    worst = Some(score);
203                }
204            }
205            Ok(())
206        },
207    )
208    .map_err(|error| match error {
209        crate::quant::RescoreCheckError::Rescore(error) => super::map_l2_rescore_error(error),
210        crate::quant::RescoreCheckError::Check(error) => error,
211    })?;
212    cancellation.check_graph().map_err(super::map_scan_error)?;
213    if let Some((row_id, _)) = narrowing {
214        return Err(QueryError::Scan(ScanError::NonFiniteScore { row_id }));
215    }
216    let scored =
217        u64::try_from(scored).map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?;
218    let coordinates = scored
219        .checked_mul(dimension as u64)
220        .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
221    let bytes = coordinates
222        .checked_mul(4)
223        .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
224    let counters = crate::diag::ScanRescoreCounters {
225        coarse_rows: stats.dims_touched / dimension as u64,
226        eligible_rows: eligible,
227        candidates_rescored: scored,
228        coarse_bytes: stats.bytes_read,
229        rescore_bytes: bytes,
230    };
231    stats.dims_touched = stats
232        .dims_touched
233        .checked_add(coordinates)
234        .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
235    stats.bytes_read = stats
236        .bytes_read
237        .checked_add(bytes)
238        .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
239    let caller = std::thread::current().id();
240    if scored > 0 && !stats.worker_thread_ids.contains(&caller) {
241        let needed = stats
242            .worker_thread_ids
243            .len()
244            .checked_add(1)
245            .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
246        memory
247            .worker_ids
248            .set(
249                needed
250                    .max(stats.worker_thread_ids.capacity())
251                    .checked_mul(std::mem::size_of::<std::thread::ThreadId>())
252                    .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?,
253            )
254            .map_err(QueryError::Store)?;
255        stats.worker_thread_ids.try_reserve_exact(1).map_err(|_| {
256            QueryError::Store(StoreError::AllocationFailed {
257                needed: memory.worker_ids.bytes(),
258                component: "rescore worker IDs",
259            })
260        })?;
261        stats.worker_thread_ids.push(caller);
262    }
263    stats.threads_used = stats.worker_thread_ids.len();
264    let exhaustive = scored == eligible;
265
266    candidates.sort_unstable_by(|a, b| {
267        b.score()
268            .total_cmp(&a.score())
269            .then(a.row_id().cmp(&b.row_id()))
270    });
271    crate::scan::truncate_to_k_with_score_ties(candidates, k, |candidate| candidate.score());
272    Ok((counters, exhaustive, if exhaustive { worst } else { None }))
273}
zeppelin_embed/lifecycle/rescored_scan.rs

zeppelin_embed/lifecycle/
rescored_scan.rs
