Skip to main content

zeppelin_embed/scan/
mod.rs

1//! Exact flat scanning over row-major vector layouts.
2
3pub mod parallel;
4pub(crate) mod topk;
5#[cfg(any(test, feature = "test-seams"))]
6#[doc(hidden)]
7pub mod vector_fault;
8
9pub use parallel::{ScanOptions, ScanOutcome, ScanStats, physical_thread_capacity};
10
11use crate::kernels;
12use crate::lifecycle::QueryCancellation;
13#[cfg(any(test, feature = "test-seams"))]
14use crate::quant::est_dot_bit4;
15use crate::quant::{
16    Bit4Factors, Bit4Query, Int8Query, Int8Vec, QuantError, QuantScheme, dot_int8_query,
17    est_dot_bit4_batch,
18};
19
20use topk::BoundedTopK;
21
22/// Immutable row-major full-precision rows with cached input validation.
23#[derive(Clone, Debug)]
24pub struct F32Rows {
25    values: Box<[f32]>,
26    first_non_finite: Option<usize>,
27}
28
29impl F32Rows {
30    /// Takes ownership of row-major values and caches their first non-finite scalar.
31    #[must_use]
32    pub fn new(values: Vec<f32>) -> Self {
33        let first_non_finite = values.iter().position(|value| !value.is_finite());
34        Self {
35            values: values.into_boxed_slice(),
36            first_non_finite,
37        }
38    }
39
40    /// Returns the immutable row-major scalar buffer.
41    #[must_use]
42    pub fn values(&self) -> &[f32] {
43        &self.values
44    }
45
46    const fn first_non_finite(&self) -> Option<usize> {
47        self.first_non_finite
48    }
49}
50
51/// One exact scan candidate.
52#[derive(Clone, Copy, Debug, PartialEq)]
53pub struct ScanCandidate {
54    /// Zero-based row id.
55    pub row_id: usize,
56    /// Larger-is-better exact scan score.
57    pub score: f32,
58}
59
60/// Truncates a best-first candidate list without splitting the k-th score tie.
61///
62/// Candidates must already be sorted by descending score.
63pub(crate) fn truncate_to_k_with_score_ties<T>(
64    candidates: &mut Vec<T>,
65    k: usize,
66    score: impl Fn(&T) -> f32,
67) {
68    if k == 0 {
69        candidates.clear();
70        return;
71    }
72    if candidates.len() <= k {
73        return;
74    }
75    let Some(boundary_score) = candidates.get(k.saturating_sub(1)).map(&score) else {
76        candidates.truncate(k);
77        return;
78    };
79    let mut keep = k;
80    for candidate in candidates.iter().skip(k) {
81        if score(candidate).total_cmp(&boundary_score).is_ne() {
82            break;
83        }
84        keep = keep.saturating_add(1);
85    }
86    candidates.truncate(keep);
87}
88
89/// Per-row affine reconstruction factors for signed-byte scan rows.
90#[derive(Clone, Copy, Debug, PartialEq)]
91#[repr(C)]
92pub struct Int8Factors {
93    scale: f32,
94    offset: f32,
95}
96
97impl Int8Factors {
98    /// Constructs validated affine factors from [`crate::quant::quantize_int8`].
99    ///
100    /// # Errors
101    ///
102    /// Returns [`Int8FactorsError`] for a negative or non-finite scale or a
103    /// non-finite offset.
104    pub fn new(scale: f32, offset: f32) -> Result<Self, Int8FactorsError> {
105        if !scale.is_finite() {
106            return Err(Int8FactorsError::NonFiniteScale);
107        }
108        if scale < 0.0 {
109            return Err(Int8FactorsError::NegativeScale);
110        }
111        if !offset.is_finite() {
112            return Err(Int8FactorsError::NonFiniteOffset);
113        }
114        Ok(Self { scale, offset })
115    }
116}
117
118/// Typed failure while constructing signed-byte row factors.
119#[derive(Clone, Copy, Debug, Eq, PartialEq)]
120pub enum Int8FactorsError {
121    /// Scale was NaN or infinite.
122    NonFiniteScale,
123    /// Scale was finite but negative.
124    NegativeScale,
125    /// Offset was NaN or infinite.
126    NonFiniteOffset,
127}
128
129impl std::fmt::Display for Int8FactorsError {
130    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131        match self {
132            Self::NonFiniteScale => formatter.write_str("Int8 row scale must be finite"),
133            Self::NegativeScale => formatter.write_str("Int8 row scale must not be negative"),
134            Self::NonFiniteOffset => formatter.write_str("Int8 row offset must be finite"),
135        }
136    }
137}
138
139impl std::error::Error for Int8FactorsError {}
140
141/// Strongly typed query representation for an exact scan.
142#[derive(Clone, Copy, Debug)]
143pub enum ScanQuery<'a> {
144    /// Full-precision query coordinates.
145    F32(&'a [f32]),
146    /// IEEE-f16 query bit patterns.
147    F16(&'a [u16]),
148    /// Prepared affine signed-byte query.
149    Int8(&'a Int8Query),
150    /// Prepared four-bit Extended-RaBitQ query.
151    Bit4(&'a Bit4Query),
152}
153
154/// Strongly typed row representation for an exact scan.
155#[derive(Clone, Copy, Debug)]
156pub enum ScanRows<'a> {
157    /// Contiguous row-major full-precision rows.
158    F32RowMajor(&'a F32Rows),
159    /// Borrowed contiguous row-major full-precision rows.
160    ///
161    /// Persisted F32 segments use this variant so a query does not copy the
162    /// mapped vector region merely to enter the scan executor.
163    F32BorrowedRowMajor(&'a [f32]),
164    /// Contiguous row-major IEEE-f16 bit patterns.
165    F16RowMajor(&'a [u16]),
166    /// Contiguous row-major signed-byte codes and affine row factors.
167    Int8RowMajor {
168        /// One signed code per coordinate and row.
169        codes: &'a [i8],
170        /// One validated affine factor pair per row.
171        factors: &'a [Int8Factors],
172    },
173    /// Contiguous row-major packed four-bit codes and estimator factors.
174    Bit4RowMajor {
175        /// MSB-first packed four-bit codes.
176        codes: &'a [u8],
177        /// One Extended-RaBitQ factor record per row.
178        factors: &'a [Bit4Factors],
179    },
180}
181
182/// Borrowed inputs for one exact scan.
183#[derive(Clone, Copy, Debug)]
184pub struct ScanRequest<'a> {
185    /// Query representation, whose scheme must match `rows`.
186    pub query: ScanQuery<'a>,
187    /// Candidate row representation.
188    pub rows: ScanRows<'a>,
189    /// Optional allow-list of zero-based row ids.
190    pub row_mask: Option<&'a roaring::RoaringBitmap>,
191}
192
193/// Typed failure from exact scan validation or scoring.
194#[derive(Clone, Debug, Eq, PartialEq)]
195pub enum ScanError {
196    /// The query dimension was zero.
197    ZeroDimension,
198    /// Row data was not an exact multiple of the query dimension.
199    RowDataLength {
200        /// Query dimension.
201        dimension: usize,
202        /// Supplied scalar count.
203        actual: usize,
204    },
205    /// Query and row representations selected different schemes.
206    SchemeMismatch {
207        /// Query scheme.
208        query: QuantScheme,
209        /// Row scheme.
210        rows: QuantScheme,
211    },
212    /// Row-factor count differed from the encoded row count.
213    FactorCount {
214        /// Encoded row count.
215        expected: usize,
216        /// Supplied factor count.
217        actual: usize,
218    },
219    /// An input coordinate was NaN or infinite.
220    NonFiniteInput {
221        /// Zero-based scalar position in the affected input.
222        index: usize,
223    },
224    /// A computed row score was NaN or infinite.
225    NonFiniteScore {
226        /// Zero-based row id.
227        row_id: usize,
228    },
229    /// A quantization scorer rejected encoded input.
230    Quant(QuantError),
231    /// Explicit candidate-rescore controls or selected rows were rejected.
232    Rescore(crate::quant::RescoreError),
233    /// Candidate-window arithmetic overflowed `usize`.
234    ArithmeticOverflow,
235    /// A persistent scan worker panicked before returning a typed result.
236    WorkerPanicked,
237    /// The operating system could not report a usable CPU count.
238    CpuCount(String),
239    /// A store query deadline expired. Partial results are never returned.
240    Timeout {
241        /// Permanently false.
242        partial: bool,
243    },
244    /// A caller-provided token cancelled the store query.
245    Cancelled {
246        /// Permanently false.
247        partial: bool,
248    },
249    /// Store close cancelled an already admitted query.
250    ReadCancelled {
251        /// Permanently false.
252        partial: bool,
253    },
254}
255
256impl std::fmt::Display for ScanError {
257    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
258        match self {
259            Self::ZeroDimension => formatter.write_str("scan dimension must not be zero"),
260            Self::RowDataLength { dimension, actual } => write!(
261                formatter,
262                "scan row data length {actual} is not divisible by dimension {dimension}"
263            ),
264            Self::SchemeMismatch { query, rows } => {
265                write!(
266                    formatter,
267                    "scan scheme mismatch: query={query:?}, rows={rows:?}"
268                )
269            }
270            Self::FactorCount { expected, actual } => write!(
271                formatter,
272                "scan row factor count mismatch: expected {expected}, got {actual}"
273            ),
274            Self::NonFiniteInput { index } => {
275                write!(formatter, "scan input is non-finite at scalar {index}")
276            }
277            Self::NonFiniteScore { row_id } => {
278                write!(formatter, "scan score is non-finite at row {row_id}")
279            }
280            Self::Quant(error) => write!(formatter, "scan quantization error: {error}"),
281            Self::Rescore(error) => write!(formatter, "scan rescore error: {error}"),
282            Self::ArithmeticOverflow => {
283                formatter.write_str("scan candidate-window arithmetic overflowed")
284            }
285            Self::WorkerPanicked => formatter.write_str("a persistent scan worker panicked"),
286            Self::CpuCount(error) => {
287                write!(formatter, "could not determine scan CPU count: {error}")
288            }
289            Self::Timeout { partial } => {
290                write!(formatter, "scan deadline expired (partial={partial})")
291            }
292            Self::Cancelled { partial } => {
293                write!(formatter, "scan was cancelled (partial={partial})")
294            }
295            Self::ReadCancelled { partial } => {
296                write!(formatter, "store close cancelled scan (partial={partial})")
297            }
298        }
299    }
300}
301
302impl std::error::Error for ScanError {}
303
304impl From<QuantError> for ScanError {
305    fn from(error: QuantError) -> Self {
306        Self::Quant(error)
307    }
308}
309
310/// Returns the exact best `k` candidates in permanent parity order.
311///
312/// Scores are ordered descending. Equal scores are always ordered by ascending
313/// row id. This is a permanent public contract shared by row-major and parallel
314/// implementations. Store-internal document-presenting paths retain score ties
315/// until document identity is available.
316///
317/// # Errors
318///
319/// Returns [`ScanError`] for invalid shapes, non-finite inputs, or non-finite
320/// scores.
321pub fn top_k(request: ScanRequest<'_>, k: usize) -> Result<Vec<ScanCandidate>, ScanError> {
322    let mut candidates = top_k_with_ties(request, k)?;
323    candidates.truncate(k);
324    Ok(candidates)
325}
326
327/// Returns the exact best candidates, including every k-th-score boundary tie.
328pub(crate) fn top_k_with_ties(
329    request: ScanRequest<'_>,
330    k: usize,
331) -> Result<Vec<ScanCandidate>, ScanError> {
332    scan_top_k(request, k)
333}
334
335#[cfg(any(test, feature = "test-seams"))]
336pub(crate) fn scan_partition_with_vector_faults(
337    request: ScanRequest<'_>,
338    k: usize,
339    range: std::ops::Range<usize>,
340    cancellation: Option<&QueryCancellation<'_>>,
341    controller: &vector_fault::VectorFaultController,
342    source: vector_fault::VectorRowSource,
343    tier: vector_fault::VectorSearchTier,
344) -> Result<PartitionScan, ScanError> {
345    check_cancellation(cancellation)?;
346    let geometry = scan_geometry(request)?;
347    if range.start > range.end || range.end > geometry.row_count {
348        return Err(ScanError::ArithmeticOverflow);
349    }
350    let (candidates, scored_rows) =
351        scan_faulted_range(request, k, range, controller, source, tier)?;
352    let dimensions = request_dimension(request)?;
353    let dims_touched = u64::try_from(scored_rows)
354        .map_err(|_| ScanError::ArithmeticOverflow)?
355        .checked_mul(u64::try_from(dimensions).map_err(|_| ScanError::ArithmeticOverflow)?)
356        .ok_or(ScanError::ArithmeticOverflow)?;
357    let bytes_read = scored_rows
358        .checked_mul(request_bytes_per_row(request)?)
359        .and_then(|bytes| u64::try_from(bytes).ok())
360        .ok_or(ScanError::ArithmeticOverflow)?;
361    Ok(PartitionScan {
362        candidates,
363        dims_touched,
364        bytes_read,
365        worker_thread_id: std::thread::current().id(),
366    })
367}
368
369#[cfg(any(test, feature = "test-seams"))]
370fn scan_faulted_range(
371    request: ScanRequest<'_>,
372    k: usize,
373    range: std::ops::Range<usize>,
374    controller: &vector_fault::VectorFaultController,
375    source: vector_fault::VectorRowSource,
376    tier: vector_fault::VectorSearchTier,
377) -> Result<(Vec<ScanCandidate>, usize), ScanError> {
378    let geometry = scan_geometry(request)?;
379    if range.start > range.end || range.end > geometry.row_count {
380        return Err(ScanError::ArithmeticOverflow);
381    }
382    let mut selected = BoundedTopK::new(k.min(geometry.row_count));
383    let mut scored_rows = 0_usize;
384    for row in range {
385        if !row_is_allowed(request.row_mask, row) {
386            continue;
387        }
388        let candidate = score_faulted_row(request, row, controller, source, tier)?;
389        scored_rows = scored_rows
390            .checked_add(1)
391            .ok_or(ScanError::ArithmeticOverflow)?;
392        if controller.after_eligible_row(source, tier, row) {
393            return Err(ScanError::Cancelled { partial: false });
394        }
395        selected.push(candidate);
396    }
397    Ok((selected.into_sorted_with_ties(), scored_rows))
398}
399
400#[cfg(any(test, feature = "test-seams"))]
401fn score_faulted_row(
402    request: ScanRequest<'_>,
403    row: usize,
404    controller: &vector_fault::VectorFaultController,
405    source: vector_fault::VectorRowSource,
406    tier: vector_fault::VectorSearchTier,
407) -> Result<ScanCandidate, ScanError> {
408    let score = match (request.query, request.rows) {
409        (ScanQuery::F32(query), ScanRows::F32RowMajor(rows)) => {
410            let values = scalar_row_range(rows.values(), query.len(), row..row + 1)?;
411            kernels::dot_f32(query, values)
412        }
413        (ScanQuery::F32(query), ScanRows::F32BorrowedRowMajor(rows)) => {
414            let values = scalar_row_range(rows, query.len(), row..row + 1)?;
415            kernels::dot_f32(query, values)
416        }
417        (ScanQuery::F16(query), ScanRows::F16RowMajor(rows)) => {
418            let values = scalar_row_range(rows, query.len(), row..row + 1)?;
419            kernels::dot_f16(query, values)
420        }
421        (ScanQuery::Int8(query), ScanRows::Int8RowMajor { codes, factors }) => {
422            let values = scalar_row_range(codes, query.len(), row..row + 1)?;
423            let factor = factors.get(row).ok_or(ScanError::ArithmeticOverflow)?;
424            let mut scale = factor.scale;
425            let mut offset = factor.offset;
426            controller.corrupt_int8(source, tier, row, &mut scale, &mut offset);
427            dot_int8_query(
428                query,
429                Int8Vec {
430                    codes: values,
431                    scale,
432                    offset,
433                },
434            )?
435        }
436        (ScanQuery::Bit4(query), ScanRows::Bit4RowMajor { codes, factors }) => {
437            let values = scalar_row_range(codes, query.len().div_ceil(2), row..row + 1)?;
438            let mut owned_codes = values.to_vec();
439            let mut factor = *factors.get(row).ok_or(ScanError::ArithmeticOverflow)?;
440            controller.corrupt_bit4(source, tier, row, &mut owned_codes, &mut factor);
441            est_dot_bit4(query, &owned_codes, factor)?
442        }
443        (query, rows) => {
444            return Err(ScanError::SchemeMismatch {
445                query: query.scheme(),
446                rows: rows.scheme(),
447            });
448        }
449    };
450    if !score.is_finite() {
451        return Err(ScanError::NonFiniteScore { row_id: row });
452    }
453    Ok(ScanCandidate { row_id: row, score })
454}
455
456#[cfg(any(test, feature = "test-seams"))]
457fn request_dimension(request: ScanRequest<'_>) -> Result<usize, ScanError> {
458    let dimension = match request.query {
459        ScanQuery::F32(query) => query.len(),
460        ScanQuery::F16(query) => query.len(),
461        ScanQuery::Int8(query) => query.len(),
462        ScanQuery::Bit4(query) => query.len(),
463    };
464    if dimension == 0 {
465        Err(ScanError::ZeroDimension)
466    } else {
467        Ok(dimension)
468    }
469}
470
471#[cfg(any(test, feature = "test-seams"))]
472fn request_bytes_per_row(request: ScanRequest<'_>) -> Result<usize, ScanError> {
473    match request.query {
474        ScanQuery::F32(query) => query
475            .len()
476            .checked_mul(std::mem::size_of::<f32>())
477            .ok_or(ScanError::ArithmeticOverflow),
478        ScanQuery::F16(query) => query
479            .len()
480            .checked_mul(std::mem::size_of::<u16>())
481            .ok_or(ScanError::ArithmeticOverflow),
482        ScanQuery::Int8(query) => Ok(query.len()),
483        ScanQuery::Bit4(query) => Ok(query.len().div_ceil(2)),
484    }
485}
486
487/// Forces the planner's gather executor for threshold calibration.
488///
489/// This seam is available only to repository test and benchmark tooling. It
490/// keeps calibration on the shipping gather implementation without making a
491/// forced execution branch part of the product API.
492#[cfg(any(test, feature = "benchmark-kernels"))]
493pub fn calibration_gather_top_k(
494    query: ScanQuery<'_>,
495    rows: ScanRows<'_>,
496    allow_list: &crate::meta::DocBitmap,
497    k: usize,
498) -> Result<Vec<ScanCandidate>, ScanError> {
499    Ok(gather_top_k(
500        ScanRequest {
501            query,
502            rows,
503            row_mask: Some(allow_list.as_roaring()),
504        },
505        k,
506        None,
507    )?
508    .candidates)
509}
510
511/// Forces the planner's masked full-sweep executor for threshold calibration.
512///
513/// This seam is available only to repository test and benchmark tooling.
514#[cfg(any(test, feature = "benchmark-kernels"))]
515pub fn calibration_masked_top_k(
516    query: ScanQuery<'_>,
517    rows: ScanRows<'_>,
518    allow_list: &crate::meta::DocBitmap,
519    k: usize,
520) -> Result<Vec<ScanCandidate>, ScanError> {
521    scan_top_k(
522        ScanRequest {
523            query,
524            rows,
525            row_mask: Some(allow_list.as_roaring()),
526        },
527        k,
528    )
529}
530
531fn scan_top_k(request: ScanRequest<'_>, k: usize) -> Result<Vec<ScanCandidate>, ScanError> {
532    match (request.query, request.rows) {
533        (ScanQuery::F32(query), ScanRows::F32RowMajor(rows)) => {
534            scan_f32(query, rows, request.row_mask, k)
535        }
536        (ScanQuery::F32(query), ScanRows::F32BorrowedRowMajor(rows)) => {
537            validate_f32_slice(query, rows)?;
538            scan_f32_rows(query, rows, request.row_mask, k, 0, None)
539        }
540        (ScanQuery::F16(query), ScanRows::F16RowMajor(rows)) => {
541            scan_f16(query, rows, request.row_mask, k)
542        }
543        (ScanQuery::Int8(query), ScanRows::Int8RowMajor { codes, factors }) => {
544            scan_int8(query, codes, factors, request.row_mask, k)
545        }
546        (ScanQuery::Bit4(query), ScanRows::Bit4RowMajor { codes, factors }) => {
547            scan_bit4(query, codes, factors, request.row_mask, k)
548        }
549        (query, rows) => Err(ScanError::SchemeMismatch {
550            query: query.scheme(),
551            rows: rows.scheme(),
552        }),
553    }
554}
555
556#[derive(Debug)]
557pub(crate) struct PartitionScan {
558    pub(crate) candidates: Vec<ScanCandidate>,
559    pub(crate) dims_touched: u64,
560    pub(crate) bytes_read: u64,
561    pub(crate) worker_thread_id: std::thread::ThreadId,
562}
563
564impl ScanQuery<'_> {
565    const fn scheme(self) -> QuantScheme {
566        match self {
567            Self::F32(_) => QuantScheme::F32,
568            Self::F16(_) => QuantScheme::F16,
569            Self::Int8(_) => QuantScheme::Int8,
570            Self::Bit4(_) => QuantScheme::Bit4,
571        }
572    }
573}
574
575impl ScanRows<'_> {
576    const fn scheme(self) -> QuantScheme {
577        match self {
578            Self::F32RowMajor(_) | Self::F32BorrowedRowMajor(_) => QuantScheme::F32,
579            Self::F16RowMajor(_) => QuantScheme::F16,
580            Self::Int8RowMajor { .. } => QuantScheme::Int8,
581            Self::Bit4RowMajor { .. } => QuantScheme::Bit4,
582        }
583    }
584}
585
586#[derive(Clone, Copy, Debug)]
587pub(crate) struct ScanGeometry {
588    pub(crate) row_count: usize,
589    pub(crate) work_units: usize,
590}
591
592pub(crate) fn scan_geometry(request: ScanRequest<'_>) -> Result<ScanGeometry, ScanError> {
593    let query_scheme = request.query.scheme();
594    let row_scheme = request.rows.scheme();
595    if query_scheme != row_scheme {
596        return Err(ScanError::SchemeMismatch {
597            query: query_scheme,
598            rows: row_scheme,
599        });
600    }
601    let query_dimension = match request.query {
602        ScanQuery::F32(query) => query.len(),
603        ScanQuery::F16(query) => query.len(),
604        ScanQuery::Int8(query) => query.len(),
605        ScanQuery::Bit4(query) => query.len(),
606    };
607    if query_dimension == 0 {
608        return Err(ScanError::ZeroDimension);
609    }
610    let (row_count, work_units) = match request.rows {
611        ScanRows::F32RowMajor(rows) => {
612            let geometry = row_major_geometry(rows.values().len(), query_dimension)?;
613            let ScanQuery::F32(query) = request.query else {
614                return Err(ScanError::ArithmeticOverflow);
615            };
616            validate_f32(query, rows)?;
617            geometry
618        }
619        ScanRows::F32BorrowedRowMajor(rows) => {
620            let geometry = row_major_geometry(rows.len(), query_dimension)?;
621            let ScanQuery::F32(query) = request.query else {
622                return Err(ScanError::ArithmeticOverflow);
623            };
624            validate_f32_slice(query, rows)?;
625            geometry
626        }
627        ScanRows::F16RowMajor(rows) => row_major_geometry(rows.len(), query_dimension)?,
628        ScanRows::Int8RowMajor { codes, factors } => {
629            let geometry = row_major_geometry(codes.len(), query_dimension)?;
630            require_factor_count(geometry.0, factors.len())?;
631            geometry
632        }
633        ScanRows::Bit4RowMajor { codes, factors } => {
634            let row_width = query_dimension.div_ceil(2);
635            let geometry = row_major_geometry(codes.len(), row_width)?;
636            require_factor_count(geometry.0, factors.len())?;
637            geometry
638        }
639    };
640    Ok(ScanGeometry {
641        row_count,
642        work_units,
643    })
644}
645
646fn row_major_geometry(scalar_count: usize, row_width: usize) -> Result<(usize, usize), ScanError> {
647    if !scalar_count.is_multiple_of(row_width) {
648        return Err(ScanError::RowDataLength {
649            dimension: row_width,
650            actual: scalar_count,
651        });
652    }
653    let rows = scalar_count / row_width;
654    Ok((rows, rows))
655}
656
657fn require_factor_count(expected: usize, actual: usize) -> Result<(), ScanError> {
658    if expected != actual {
659        return Err(ScanError::FactorCount { expected, actual });
660    }
661    Ok(())
662}
663
664fn query_dimensions(query: ScanQuery<'_>) -> usize {
665    match query {
666        ScanQuery::F32(query) => query.len(),
667        ScanQuery::F16(query) => query.len(),
668        ScanQuery::Int8(query) => query.len(),
669        ScanQuery::Bit4(query) => query.len(),
670    }
671}
672
673fn query_bytes_per_row(query: ScanQuery<'_>) -> Result<usize, ScanError> {
674    match query {
675        ScanQuery::F32(query) => query
676            .len()
677            .checked_mul(std::mem::size_of::<f32>())
678            .ok_or(ScanError::ArithmeticOverflow),
679        ScanQuery::F16(query) => query
680            .len()
681            .checked_mul(std::mem::size_of::<u16>())
682            .ok_or(ScanError::ArithmeticOverflow),
683        ScanQuery::Int8(query) => Ok(query.len()),
684        ScanQuery::Bit4(query) => Ok(query.len().div_ceil(2)),
685    }
686}
687
688fn partition_stat_counts(
689    scored_rows: usize,
690    query: ScanQuery<'_>,
691) -> Result<(u64, u64), ScanError> {
692    let dimensions =
693        u64::try_from(query_dimensions(query)).map_err(|_| ScanError::ArithmeticOverflow)?;
694    let dims_touched = u64::try_from(scored_rows)
695        .map_err(|_| ScanError::ArithmeticOverflow)?
696        .checked_mul(dimensions)
697        .ok_or(ScanError::ArithmeticOverflow)?;
698    let bytes_per_row = query_bytes_per_row(query)?;
699    let bytes_read = scored_rows
700        .checked_mul(bytes_per_row)
701        .and_then(|bytes| u64::try_from(bytes).ok())
702        .ok_or(ScanError::ArithmeticOverflow)?;
703    Ok((dims_touched, bytes_read))
704}
705
706pub(crate) fn scan_partition(
707    request: ScanRequest<'_>,
708    k: usize,
709    range: std::ops::Range<usize>,
710    cancellation: Option<&QueryCancellation<'_>>,
711) -> Result<PartitionScan, ScanError> {
712    check_cancellation(cancellation)?;
713    let geometry = scan_geometry(request)?;
714    if range.start > range.end || range.end > geometry.row_count {
715        return Err(ScanError::ArithmeticOverflow);
716    }
717    let first_row = range.start;
718    let candidates = match (request.query, request.rows) {
719        (ScanQuery::F32(query), ScanRows::F32RowMajor(rows)) => {
720            let rows = scalar_row_range(rows.values(), query.len(), range.clone())?;
721            scan_f32_rows(query, rows, request.row_mask, k, first_row, cancellation)?
722        }
723        (ScanQuery::F32(query), ScanRows::F32BorrowedRowMajor(rows)) => {
724            let rows = scalar_row_range(rows, query.len(), range.clone())?;
725            scan_f32_rows(query, rows, request.row_mask, k, first_row, cancellation)?
726        }
727        (ScanQuery::F16(query), ScanRows::F16RowMajor(rows)) => {
728            let rows = scalar_row_range(rows, query.len(), range.clone())?;
729            scan_f16_rows(query, rows, request.row_mask, k, first_row, cancellation)?
730        }
731        (ScanQuery::Int8(query), ScanRows::Int8RowMajor { codes, factors }) => {
732            let codes = scalar_row_range(codes, query.len(), range.clone())?;
733            let factors = factors
734                .get(range.clone())
735                .ok_or(ScanError::ArithmeticOverflow)?;
736            scan_int8_rows(
737                query,
738                codes,
739                factors,
740                request.row_mask,
741                k,
742                first_row,
743                cancellation,
744            )?
745        }
746        (ScanQuery::Bit4(query), ScanRows::Bit4RowMajor { codes, factors }) => {
747            let codes = scalar_row_range(codes, query.len().div_ceil(2), range.clone())?;
748            let factors = factors
749                .get(range.clone())
750                .ok_or(ScanError::ArithmeticOverflow)?;
751            scan_bit4_rows(
752                query,
753                codes,
754                factors,
755                request.row_mask,
756                k,
757                first_row,
758                cancellation,
759            )?
760        }
761        (query, rows) => {
762            return Err(ScanError::SchemeMismatch {
763                query: query.scheme(),
764                rows: rows.scheme(),
765            });
766        }
767    };
768    let scored_rows = match request.rows {
769        ScanRows::Bit4RowMajor { .. } => range.end - range.start,
770        _ => allowed_row_count(request.row_mask, range.clone(), cancellation)?,
771    };
772    let (dims_touched, bytes_read) = partition_stat_counts(scored_rows, request.query)?;
773    Ok(PartitionScan {
774        candidates,
775        dims_touched,
776        bytes_read,
777        worker_thread_id: std::thread::current().id(),
778    })
779}
780
781/// Scores only rows enumerated by the request's allow-list.
782///
783/// This is the gather counterpart to [`scan_partition`]. It deliberately
784/// shares the same scheme validation and scoring kernels so planner branch
785/// selection can change work without changing score semantics.
786pub(crate) fn gather_top_k(
787    request: ScanRequest<'_>,
788    k: usize,
789    cancellation: Option<&QueryCancellation<'_>>,
790) -> Result<ScanOutcome, ScanError> {
791    check_cancellation(cancellation)?;
792    let geometry = scan_geometry(request)?;
793    let mut selected = BoundedTopK::new(k.min(geometry.row_count));
794    let mut scored_rows = 0_usize;
795    if let Some(allow_list) = request.row_mask {
796        for (ordinal, row) in allow_list.iter().enumerate() {
797            check_cancellation_at_row(cancellation, ordinal)?;
798            let row = usize::try_from(row).map_err(|_| ScanError::ArithmeticOverflow)?;
799            if row >= geometry.row_count {
800                continue;
801            }
802            selected.push(score_gather_row(request, row)?);
803            scored_rows = scored_rows
804                .checked_add(1)
805                .ok_or(ScanError::ArithmeticOverflow)?;
806        }
807    } else {
808        for row in 0..geometry.row_count {
809            check_cancellation_at_row(cancellation, row)?;
810            selected.push(score_gather_row(request, row)?);
811            scored_rows = scored_rows
812                .checked_add(1)
813                .ok_or(ScanError::ArithmeticOverflow)?;
814        }
815    }
816    let (dims_touched, bytes_read) = partition_stat_counts(scored_rows, request.query)?;
817    let worker_thread_ids = (scored_rows != 0)
818        .then(|| std::thread::current().id())
819        .into_iter()
820        .collect::<Vec<_>>();
821    Ok(ScanOutcome {
822        candidates: selected.into_sorted_with_ties(),
823        worst_score: None,
824        stats: ScanStats {
825            dims_touched,
826            bytes_read,
827            threads_used: worker_thread_ids.len(),
828            worker_thread_ids,
829        },
830    })
831}
832
833fn score_gather_row(request: ScanRequest<'_>, row: usize) -> Result<ScanCandidate, ScanError> {
834    let score = match (request.query, request.rows) {
835        (ScanQuery::F32(query), ScanRows::F32RowMajor(rows)) => {
836            let values = scalar_row_range(rows.values(), query.len(), row..row + 1)?;
837            kernels::dot_f32(query, values)
838        }
839        (ScanQuery::F32(query), ScanRows::F32BorrowedRowMajor(rows)) => {
840            let values = scalar_row_range(rows, query.len(), row..row + 1)?;
841            kernels::dot_f32(query, values)
842        }
843        (ScanQuery::F16(query), ScanRows::F16RowMajor(rows)) => {
844            let values = scalar_row_range(rows, query.len(), row..row + 1)?;
845            kernels::dot_f16(query, values)
846        }
847        (ScanQuery::Int8(query), ScanRows::Int8RowMajor { codes, factors }) => {
848            let values = scalar_row_range(codes, query.len(), row..row + 1)?;
849            let factor = factors.get(row).ok_or(ScanError::ArithmeticOverflow)?;
850            dot_int8_query(
851                query,
852                Int8Vec {
853                    codes: values,
854                    scale: factor.scale,
855                    offset: factor.offset,
856                },
857            )?
858        }
859        (ScanQuery::Bit4(query), ScanRows::Bit4RowMajor { codes, factors }) => {
860            let values = scalar_row_range(codes, query.len().div_ceil(2), row..row + 1)?;
861            let factor = factors.get(row).ok_or(ScanError::ArithmeticOverflow)?;
862            let mut score = [0.0_f32; 1];
863            est_dot_bit4_batch(query, values, std::slice::from_ref(factor), &mut score)?;
864            score[0]
865        }
866        (query, rows) => {
867            return Err(ScanError::SchemeMismatch {
868                query: query.scheme(),
869                rows: rows.scheme(),
870            });
871        }
872    };
873    if !score.is_finite() {
874        return Err(ScanError::NonFiniteScore { row_id: row });
875    }
876    Ok(ScanCandidate { row_id: row, score })
877}
878
879const CANCELLATION_CHECK_ROWS: usize = 64;
880
881fn check_cancellation(cancellation: Option<&QueryCancellation<'_>>) -> Result<(), ScanError> {
882    cancellation.map_or(Ok(()), QueryCancellation::check)
883}
884
885fn check_cancellation_at_row(
886    cancellation: Option<&QueryCancellation<'_>>,
887    local_row: usize,
888) -> Result<(), ScanError> {
889    if local_row.is_multiple_of(CANCELLATION_CHECK_ROWS) {
890        check_cancellation(cancellation)?;
891    }
892    Ok(())
893}
894
895fn row_is_allowed(row_mask: Option<&roaring::RoaringBitmap>, row_id: usize) -> bool {
896    !row_mask.is_some_and(|mask| u32::try_from(row_id).map_or(true, |id| !mask.contains(id)))
897}
898
899fn scalar_row_range<T>(
900    values: &[T],
901    row_width: usize,
902    rows: std::ops::Range<usize>,
903) -> Result<&[T], ScanError> {
904    let start = rows
905        .start
906        .checked_mul(row_width)
907        .ok_or(ScanError::ArithmeticOverflow)?;
908    let end = rows
909        .end
910        .checked_mul(row_width)
911        .ok_or(ScanError::ArithmeticOverflow)?;
912    values.get(start..end).ok_or(ScanError::ArithmeticOverflow)
913}
914
915fn allowed_row_count(
916    row_mask: Option<&roaring::RoaringBitmap>,
917    rows: std::ops::Range<usize>,
918    cancellation: Option<&QueryCancellation<'_>>,
919) -> Result<usize, ScanError> {
920    let mut allowed = 0_usize;
921    for (local_row, row_id) in rows.enumerate() {
922        check_cancellation_at_row(cancellation, local_row)?;
923        if !row_mask.is_some_and(|mask| u32::try_from(row_id).map_or(true, |id| !mask.contains(id)))
924        {
925            allowed = allowed
926                .checked_add(1)
927                .ok_or(ScanError::ArithmeticOverflow)?;
928        }
929    }
930    Ok(allowed)
931}
932
933fn scan_f32(
934    query: &[f32],
935    rows: &F32Rows,
936    row_mask: Option<&roaring::RoaringBitmap>,
937    k: usize,
938) -> Result<Vec<ScanCandidate>, ScanError> {
939    if query.is_empty() {
940        return Err(ScanError::ZeroDimension);
941    }
942    if !rows.values().len().is_multiple_of(query.len()) {
943        return Err(ScanError::RowDataLength {
944            dimension: query.len(),
945            actual: rows.values().len(),
946        });
947    }
948    validate_f32(query, rows)?;
949    scan_f32_rows(query, rows.values(), row_mask, k, 0, None)
950}
951
952fn validate_f32(query: &[f32], rows: &F32Rows) -> Result<(), ScanError> {
953    if let Some(index) = query.iter().position(|value| !value.is_finite()) {
954        return Err(ScanError::NonFiniteInput { index });
955    }
956    if let Some(local_index) = rows.first_non_finite() {
957        let index = query
958            .len()
959            .checked_add(local_index)
960            .ok_or(ScanError::ArithmeticOverflow)?;
961        return Err(ScanError::NonFiniteInput { index });
962    }
963    Ok(())
964}
965
966fn validate_f32_slice(query: &[f32], rows: &[f32]) -> Result<(), ScanError> {
967    if let Some(index) = query.iter().position(|value| !value.is_finite()) {
968        return Err(ScanError::NonFiniteInput { index });
969    }
970    if let Some(local_index) = rows.iter().position(|value| !value.is_finite()) {
971        let index = query
972            .len()
973            .checked_add(local_index)
974            .ok_or(ScanError::ArithmeticOverflow)?;
975        return Err(ScanError::NonFiniteInput { index });
976    }
977    Ok(())
978}
979
980fn scan_f32_rows(
981    query: &[f32],
982    rows: &[f32],
983    row_mask: Option<&roaring::RoaringBitmap>,
984    k: usize,
985    first_row: usize,
986    cancellation: Option<&QueryCancellation<'_>>,
987) -> Result<Vec<ScanCandidate>, ScanError> {
988    let row_count = rows.len() / query.len();
989    let mut selected = BoundedTopK::new(k.min(row_count));
990    for (local_row, row) in rows.chunks_exact(query.len()).enumerate() {
991        check_cancellation_at_row(cancellation, local_row)?;
992        let row_id = first_row
993            .checked_add(local_row)
994            .ok_or(ScanError::ArithmeticOverflow)?;
995        if row_mask.is_some_and(|mask| u32::try_from(row_id).map_or(true, |id| !mask.contains(id)))
996        {
997            continue;
998        }
999        let score = kernels::dot_f32(query, row);
1000        if !score.is_finite() {
1001            return Err(ScanError::NonFiniteScore { row_id });
1002        }
1003        selected.push(ScanCandidate { row_id, score });
1004    }
1005    Ok(selected.into_sorted_with_ties())
1006}
1007
1008fn scan_f16(
1009    query: &[u16],
1010    rows: &[u16],
1011    row_mask: Option<&roaring::RoaringBitmap>,
1012    k: usize,
1013) -> Result<Vec<ScanCandidate>, ScanError> {
1014    scan_f16_rows(query, rows, row_mask, k, 0, None)
1015}
1016
1017fn scan_f16_rows(
1018    query: &[u16],
1019    rows: &[u16],
1020    row_mask: Option<&roaring::RoaringBitmap>,
1021    k: usize,
1022    first_row: usize,
1023    cancellation: Option<&QueryCancellation<'_>>,
1024) -> Result<Vec<ScanCandidate>, ScanError> {
1025    if query.is_empty() {
1026        return Err(ScanError::ZeroDimension);
1027    }
1028    if !rows.len().is_multiple_of(query.len()) {
1029        return Err(ScanError::RowDataLength {
1030            dimension: query.len(),
1031            actual: rows.len(),
1032        });
1033    }
1034    let row_count = rows.len() / query.len();
1035    let mut selected = BoundedTopK::new(k.min(row_count));
1036    for (local_row, row) in rows.chunks_exact(query.len()).enumerate() {
1037        check_cancellation_at_row(cancellation, local_row)?;
1038        let row_id = first_row
1039            .checked_add(local_row)
1040            .ok_or(ScanError::ArithmeticOverflow)?;
1041        if row_mask.is_some_and(|mask| u32::try_from(row_id).map_or(true, |id| !mask.contains(id)))
1042        {
1043            continue;
1044        }
1045        let score = kernels::dot_f16(query, row);
1046        if !score.is_finite() {
1047            return Err(ScanError::NonFiniteScore { row_id });
1048        }
1049        selected.push(ScanCandidate { row_id, score });
1050    }
1051    Ok(selected.into_sorted_with_ties())
1052}
1053
1054fn scan_int8(
1055    query: &Int8Query,
1056    codes: &[i8],
1057    factors: &[Int8Factors],
1058    row_mask: Option<&roaring::RoaringBitmap>,
1059    k: usize,
1060) -> Result<Vec<ScanCandidate>, ScanError> {
1061    scan_int8_rows(query, codes, factors, row_mask, k, 0, None)
1062}
1063
1064fn scan_int8_rows(
1065    query: &Int8Query,
1066    codes: &[i8],
1067    factors: &[Int8Factors],
1068    row_mask: Option<&roaring::RoaringBitmap>,
1069    k: usize,
1070    first_row: usize,
1071    cancellation: Option<&QueryCancellation<'_>>,
1072) -> Result<Vec<ScanCandidate>, ScanError> {
1073    let dimension = query.len();
1074    if dimension == 0 {
1075        return Err(ScanError::ZeroDimension);
1076    }
1077    if !codes.len().is_multiple_of(dimension) {
1078        return Err(ScanError::RowDataLength {
1079            dimension,
1080            actual: codes.len(),
1081        });
1082    }
1083    let row_count = codes.len() / dimension;
1084    if factors.len() != row_count {
1085        return Err(ScanError::FactorCount {
1086            expected: row_count,
1087            actual: factors.len(),
1088        });
1089    }
1090    let mut selected = BoundedTopK::new(k.min(row_count));
1091    for (local_row, (row, factor)) in codes.chunks_exact(dimension).zip(factors).enumerate() {
1092        check_cancellation_at_row(cancellation, local_row)?;
1093        let row_id = first_row
1094            .checked_add(local_row)
1095            .ok_or(ScanError::ArithmeticOverflow)?;
1096        if row_mask.is_some_and(|mask| u32::try_from(row_id).map_or(true, |id| !mask.contains(id)))
1097        {
1098            continue;
1099        }
1100        let score = dot_int8_query(
1101            query,
1102            Int8Vec {
1103                codes: row,
1104                scale: factor.scale,
1105                offset: factor.offset,
1106            },
1107        )?;
1108        if !score.is_finite() {
1109            return Err(ScanError::NonFiniteScore { row_id });
1110        }
1111        selected.push(ScanCandidate { row_id, score });
1112    }
1113    Ok(selected.into_sorted_with_ties())
1114}
1115
1116fn scan_bit4(
1117    query: &Bit4Query,
1118    codes: &[u8],
1119    factors: &[Bit4Factors],
1120    row_mask: Option<&roaring::RoaringBitmap>,
1121    k: usize,
1122) -> Result<Vec<ScanCandidate>, ScanError> {
1123    scan_bit4_rows(query, codes, factors, row_mask, k, 0, None)
1124}
1125
1126fn scan_bit4_rows(
1127    query: &Bit4Query,
1128    codes: &[u8],
1129    factors: &[Bit4Factors],
1130    row_mask: Option<&roaring::RoaringBitmap>,
1131    k: usize,
1132    first_row: usize,
1133    cancellation: Option<&QueryCancellation<'_>>,
1134) -> Result<Vec<ScanCandidate>, ScanError> {
1135    let dimension = query.len();
1136    if dimension == 0 {
1137        return Err(ScanError::ZeroDimension);
1138    }
1139    let row_width = dimension.div_ceil(2);
1140    if !codes.len().is_multiple_of(row_width) {
1141        return Err(ScanError::RowDataLength {
1142            dimension: row_width,
1143            actual: codes.len(),
1144        });
1145    }
1146    let row_count = codes.len() / row_width;
1147    if factors.len() != row_count {
1148        return Err(ScanError::FactorCount {
1149            expected: row_count,
1150            actual: factors.len(),
1151        });
1152    }
1153    let mut selected = BoundedTopK::new(k.min(row_count));
1154    let mut scores = [0.0_f32; 4];
1155    let mut batch_start = 0_usize;
1156    while batch_start < row_count {
1157        check_cancellation_at_row(cancellation, batch_start)?;
1158        let batch_rows = (row_count - batch_start).min(scores.len());
1159        let batch_end = batch_start
1160            .checked_add(batch_rows)
1161            .ok_or(ScanError::ArithmeticOverflow)?;
1162        let code_start = batch_start
1163            .checked_mul(row_width)
1164            .ok_or(ScanError::ArithmeticOverflow)?;
1165        let code_end = batch_end
1166            .checked_mul(row_width)
1167            .ok_or(ScanError::ArithmeticOverflow)?;
1168        let batch_codes = codes
1169            .get(code_start..code_end)
1170            .ok_or(ScanError::ArithmeticOverflow)?;
1171        let batch_factors = factors
1172            .get(batch_start..batch_end)
1173            .ok_or(ScanError::ArithmeticOverflow)?;
1174        let batch_scores = scores
1175            .get_mut(..batch_rows)
1176            .ok_or(ScanError::ArithmeticOverflow)?;
1177        est_dot_bit4_batch(query, batch_codes, batch_factors, batch_scores)?;
1178        for (batch_row, &score) in batch_scores.iter().enumerate() {
1179            let local_row = batch_start
1180                .checked_add(batch_row)
1181                .ok_or(ScanError::ArithmeticOverflow)?;
1182            let row_id = first_row
1183                .checked_add(local_row)
1184                .ok_or(ScanError::ArithmeticOverflow)?;
1185            if !row_is_allowed(row_mask, row_id) {
1186                continue;
1187            }
1188            if !score.is_finite() {
1189                return Err(ScanError::NonFiniteScore { row_id });
1190            }
1191            selected.push(ScanCandidate { row_id, score });
1192        }
1193        batch_start = batch_end;
1194    }
1195    Ok(selected.into_sorted_with_ties())
1196}
1197
1198#[cfg(test)]
1199#[allow(
1200    clippy::expect_used,
1201    clippy::indexing_slicing,
1202    clippy::panic,
1203    clippy::unwrap_used
1204)]
1205mod tests {
1206    use rand::Rng;
1207    use roaring::RoaringBitmap;
1208    use tempfile::{TempDir, tempdir};
1209
1210    use super::{F32Rows, Int8Factors, ScanOptions, ScanQuery, ScanRequest, ScanRows, top_k};
1211    use crate::lifecycle::{CancelToken, OpenOptions, QueryControl, QueryError, Store};
1212    use crate::quant::{prepare_bit4_query, prepare_int8_query, quantize_bit4};
1213
1214    fn query_store() -> (TempDir, Store) {
1215        let directory = tempdir().expect("query store directory");
1216        let store = Store::open(directory.path(), OpenOptions::default()).expect("query store");
1217        (directory, store)
1218    }
1219
1220    fn pooled_top_k(
1221        store: &Store,
1222        request: ScanRequest<'_>,
1223        k: usize,
1224        options: ScanOptions,
1225    ) -> Result<super::ScanOutcome, QueryError> {
1226        store.top_k_with_options(
1227            request,
1228            k,
1229            options,
1230            QueryControl::Cancel(CancelToken::new()),
1231        )
1232    }
1233
1234    #[test]
1235    fn all_identical_vectors_tie_by_ascending_row_id() {
1236        let (_directory, store) = query_store();
1237        let query = [1.0_f32, -2.0];
1238        let rows = F32Rows::new(vec![1.0_f32, -2.0, 1.0, -2.0, 1.0, -2.0, 1.0, -2.0]);
1239        let request = ScanRequest {
1240            query: ScanQuery::F32(&query),
1241            rows: ScanRows::F32RowMajor(&rows),
1242            row_mask: None,
1243        };
1244
1245        let hits = top_k(request, 3).expect("valid scan");
1246
1247        assert_eq!(
1248            hits.iter().map(|hit| hit.row_id).collect::<Vec<_>>(),
1249            [0, 1, 2]
1250        );
1251        let pooled = pooled_top_k(&store, request, 3, ScanOptions { thread_budget: 2 })
1252            .expect("valid pooled scan");
1253        assert_eq!(
1254            pooled
1255                .candidates
1256                .iter()
1257                .map(|hit| hit.row_id)
1258                .collect::<Vec<_>>(),
1259            [0, 1, 2]
1260        );
1261    }
1262
1263    #[test]
1264    fn truncate_helper_keeps_only_the_kth_total_cmp_score_tie() {
1265        let mut short = vec![2.0_f32];
1266        super::truncate_to_k_with_score_ties(&mut short, 2, |score| *score);
1267        assert_eq!(short, [2.0]);
1268
1269        let mut no_tie = vec![3.0_f32, 2.0, 1.0];
1270        super::truncate_to_k_with_score_ties(&mut no_tie, 2, |score| *score);
1271        assert_eq!(no_tie, [3.0, 2.0]);
1272
1273        let mut boundary_tie = vec![3.0_f32, 2.0, 2.0, 2.0, 1.0];
1274        super::truncate_to_k_with_score_ties(&mut boundary_tie, 2, |score| *score);
1275        assert_eq!(boundary_tie, [3.0, 2.0, 2.0, 2.0]);
1276
1277        let mut all_tied = vec![1.0_f32; 4];
1278        super::truncate_to_k_with_score_ties(&mut all_tied, 1, |score| *score);
1279        assert_eq!(all_tied.len(), 4);
1280
1281        let mut signed_zero = vec![0.0_f32, -0.0];
1282        super::truncate_to_k_with_score_ties(&mut signed_zero, 1, |score| *score);
1283        assert_eq!(signed_zero.len(), 1);
1284        assert_eq!(signed_zero[0].to_bits(), 0.0_f32.to_bits());
1285
1286        super::truncate_to_k_with_score_ties(&mut signed_zero, 0, |score| *score);
1287        assert!(signed_zero.is_empty());
1288    }
1289
1290    #[test]
1291    fn f32_rows_cache_first_non_finite_input() {
1292        let rows = super::F32Rows::new(vec![1.0, f32::NAN, f32::INFINITY]);
1293
1294        assert_eq!(rows.first_non_finite(), Some(1));
1295        assert_eq!(
1296            top_k(
1297                ScanRequest {
1298                    query: ScanQuery::F32(&[1.0, 1.0, 1.0]),
1299                    rows: ScanRows::F32RowMajor(&rows),
1300                    row_mask: None,
1301                },
1302                1,
1303            ),
1304            Err(super::ScanError::NonFiniteInput { index: 4 })
1305        );
1306    }
1307
1308    #[test]
1309    fn bit4_row_major_scans_partial_final_batch() {
1310        let query = prepare_bit4_query(&[1.0, -1.0], 0x05).expect("valid Bit4 query");
1311        let mut codes = Vec::new();
1312        let mut factors = Vec::new();
1313        for magnitude in 1..=5 {
1314            let mut row = [0_u8; 1];
1315            factors.push(
1316                quantize_bit4(&[magnitude as f32, -(magnitude as f32)], &mut row)
1317                    .expect("valid Bit4 row"),
1318            );
1319            codes.extend(row);
1320        }
1321        let request = ScanRequest {
1322            query: ScanQuery::Bit4(&query),
1323            rows: ScanRows::Bit4RowMajor {
1324                codes: &codes,
1325                factors: &factors,
1326            },
1327            row_mask: None,
1328        };
1329
1330        let actual = top_k(request, 5)
1331            .expect("valid Bit4 scan")
1332            .into_iter()
1333            .map(|candidate| candidate.row_id)
1334            .collect::<Vec<_>>();
1335
1336        assert_eq!(actual, [4, 3, 2, 1, 0]);
1337    }
1338
1339    #[test]
1340    fn empty_row_set_returns_no_candidates() {
1341        let query = [1.0_f32, -2.0];
1342        let rows = F32Rows::new(Vec::new());
1343        let request = ScanRequest {
1344            query: ScanQuery::F32(&query),
1345            rows: ScanRows::F32RowMajor(&rows),
1346            row_mask: None,
1347        };
1348
1349        assert!(top_k(request, 4).expect("valid empty scan").is_empty());
1350    }
1351
1352    #[test]
1353    fn prop_row_major_scan_equals_naive_topk() {
1354        let (_directory, store) = query_store();
1355        let mut random =
1356            crate::test_support::seeded_rng("scan::prop_row_major_scan_equals_naive_topk");
1357        let cases = std::env::var("PROPTEST_CASES")
1358            .ok()
1359            .and_then(|value| value.parse::<usize>().ok())
1360            .unwrap_or(256);
1361
1362        for case in 0..cases {
1363            let dimension = if case.is_multiple_of(17) {
1364                1_024
1365            } else {
1366                random.random_range(8..=1_024)
1367            };
1368            let maximum_rows = (32_768 / dimension).clamp(1, 10_000);
1369            let row_count = if case.is_multiple_of(31) {
1370                10_000.min(maximum_rows)
1371            } else {
1372                random.random_range(0..=maximum_rows)
1373            };
1374            let k = random.random_range(0..=row_count);
1375
1376            match case % 4 {
1377                0 => {
1378                    let query = (0..dimension)
1379                        .map(|_| random.random_range(-2.0_f32..=2.0_f32))
1380                        .collect::<Vec<_>>();
1381                    let values = (0..row_count * dimension)
1382                        .map(|_| random.random_range(-2.0_f32..=2.0_f32))
1383                        .collect::<Vec<_>>();
1384                    let rows = F32Rows::new(values);
1385                    let request = ScanRequest {
1386                        query: ScanQuery::F32(&query),
1387                        rows: ScanRows::F32RowMajor(&rows),
1388                        row_mask: None,
1389                    };
1390                    assert_row_major_scan(
1391                        &store,
1392                        request,
1393                        k,
1394                        case,
1395                        ScanComparison::F32 {
1396                            query: &query,
1397                            rows: rows.values(),
1398                        },
1399                    );
1400                }
1401                1 => {
1402                    let query = (0..dimension)
1403                        .map(|_| random_finite_f16(&mut random))
1404                        .collect::<Vec<_>>();
1405                    let rows = (0..row_count * dimension)
1406                        .map(|_| random_finite_f16(&mut random))
1407                        .collect::<Vec<_>>();
1408                    let request = ScanRequest {
1409                        query: ScanQuery::F16(&query),
1410                        rows: ScanRows::F16RowMajor(&rows),
1411                        row_mask: None,
1412                    };
1413                    assert_row_major_scan(
1414                        &store,
1415                        request,
1416                        k,
1417                        case,
1418                        ScanComparison::F16 {
1419                            query: &query,
1420                            rows: &rows,
1421                        },
1422                    );
1423                }
1424                2 => {
1425                    let query_values = (0..dimension)
1426                        .map(|_| random.random_range(-2.0_f32..=2.0_f32))
1427                        .collect::<Vec<_>>();
1428                    let query = prepare_int8_query(&query_values).expect("finite Int8 query");
1429                    let codes = (0..row_count * dimension)
1430                        .map(|_| random.random::<i8>())
1431                        .collect::<Vec<_>>();
1432                    let factors = (0..row_count)
1433                        .map(|_| {
1434                            Int8Factors::new(
1435                                random.random_range(0.001_f32..=0.25_f32),
1436                                random.random_range(-1.0_f32..=1.0_f32),
1437                            )
1438                            .expect("finite factors")
1439                        })
1440                        .collect::<Vec<_>>();
1441                    let request = ScanRequest {
1442                        query: ScanQuery::Int8(&query),
1443                        rows: ScanRows::Int8RowMajor {
1444                            codes: &codes,
1445                            factors: &factors,
1446                        },
1447                        row_mask: None,
1448                    };
1449                    assert_row_major_scan(&store, request, k, case, ScanComparison::Exact);
1450                }
1451                _ => {
1452                    let query_values = (0..dimension)
1453                        .map(|_| random.random_range(-2.0_f32..=2.0_f32))
1454                        .collect::<Vec<_>>();
1455                    let query = prepare_bit4_query(&query_values, random.random())
1456                        .expect("finite Bit4 query");
1457                    let row_bytes = dimension.div_ceil(2);
1458                    let mut codes = (0..row_count * row_bytes)
1459                        .map(|_| random.random::<u8>())
1460                        .collect::<Vec<_>>();
1461                    if !dimension.is_multiple_of(2) {
1462                        for row in codes.chunks_exact_mut(row_bytes) {
1463                            if let Some(last) = row.last_mut() {
1464                                *last &= 0xf0;
1465                            }
1466                        }
1467                    }
1468                    let factor_source = (0..dimension)
1469                        .map(|_| random.random_range(-2.0_f32..=2.0_f32))
1470                        .collect::<Vec<_>>();
1471                    let mut factor_codes = vec![0_u8; row_bytes];
1472                    let factor = quantize_bit4(&factor_source, &mut factor_codes)
1473                        .expect("finite factor source");
1474                    let factors = vec![factor; row_count];
1475                    let request = ScanRequest {
1476                        query: ScanQuery::Bit4(&query),
1477                        rows: ScanRows::Bit4RowMajor {
1478                            codes: &codes,
1479                            factors: &factors,
1480                        },
1481                        row_mask: None,
1482                    };
1483                    assert_row_major_scan(&store, request, k, case, ScanComparison::Exact);
1484                }
1485            }
1486        }
1487    }
1488
1489    fn assert_row_major_scan(
1490        store: &Store,
1491        request: ScanRequest<'_>,
1492        k: usize,
1493        case: usize,
1494        comparison: ScanComparison<'_>,
1495    ) {
1496        let expected = naive_top_k(request, k);
1497        let single = pooled_top_k(store, request, k, ScanOptions { thread_budget: 1 })
1498            .expect("valid single-thread row-major scan")
1499            .candidates;
1500        let parallel = pooled_top_k(
1501            store,
1502            request,
1503            k,
1504            ScanOptions {
1505                thread_budget: case % 12 + 1,
1506            },
1507        )
1508        .expect("valid parallel row-major scan")
1509        .candidates;
1510        assert_eq!(
1511            parallel, single,
1512            "case {case} thread count changed candidates"
1513        );
1514        assert_scan_matches(
1515            request,
1516            &expected,
1517            &parallel,
1518            comparison,
1519            &format!("row-major case {case}"),
1520        );
1521    }
1522
1523    fn naive_top_k(request: ScanRequest<'_>, k: usize) -> Vec<super::ScanCandidate> {
1524        let scalar = crate::kernels::KernelVariant::scalar();
1525        let mut candidates = Vec::new();
1526        match (request.query, request.rows) {
1527            (ScanQuery::F32(query), ScanRows::F32RowMajor(rows)) => {
1528                for (row_id, row) in rows.values().chunks_exact(query.len()).enumerate() {
1529                    candidates.push(super::ScanCandidate {
1530                        row_id,
1531                        score: scalar.dot_f32(query, row),
1532                    });
1533                }
1534            }
1535            (ScanQuery::F16(query), ScanRows::F16RowMajor(rows)) => {
1536                for (row_id, row) in rows.chunks_exact(query.len()).enumerate() {
1537                    candidates.push(super::ScanCandidate {
1538                        row_id,
1539                        score: scalar.dot_f16(query, row),
1540                    });
1541                }
1542            }
1543            (ScanQuery::Int8(query), ScanRows::Int8RowMajor { codes, factors }) => {
1544                for (row_id, (row, factor)) in
1545                    codes.chunks_exact(query.len()).zip(factors).enumerate()
1546                {
1547                    let integer_dot = scalar.dot_i8(query.codes(), row);
1548                    candidates.push(super::ScanCandidate {
1549                        row_id,
1550                        score: query.score_integer_dot(integer_dot, factor.scale, factor.offset),
1551                    });
1552                }
1553            }
1554            (ScanQuery::Bit4(query), ScanRows::Bit4RowMajor { codes, factors }) => {
1555                let row_width = query.len().div_ceil(2);
1556                for (row_id, (row, factor)) in
1557                    codes.chunks_exact(row_width).zip(factors).enumerate()
1558                {
1559                    let mut score = [0.0_f32; 1];
1560                    scalar.score_bit4_prepared_batch(
1561                        query.kernel_parts(),
1562                        row,
1563                        query.len(),
1564                        std::slice::from_ref(factor),
1565                        &mut score,
1566                    );
1567                    candidates.push(super::ScanCandidate {
1568                        row_id,
1569                        score: score[0],
1570                    });
1571                }
1572            }
1573            _ => unreachable!("test requests are scheme-matched"),
1574        }
1575        candidates.sort_by(|left, right| {
1576            right
1577                .score
1578                .total_cmp(&left.score)
1579                .then_with(|| left.row_id.cmp(&right.row_id))
1580        });
1581        candidates.truncate(k.min(candidates.len()));
1582        candidates
1583    }
1584
1585    #[derive(Clone, Copy)]
1586    enum ScanComparison<'a> {
1587        Exact,
1588        F32 { query: &'a [f32], rows: &'a [f32] },
1589        F16 { query: &'a [u16], rows: &'a [u16] },
1590    }
1591
1592    fn assert_scan_matches(
1593        reference_request: ScanRequest<'_>,
1594        expected: &[super::ScanCandidate],
1595        actual: &[super::ScanCandidate],
1596        comparison: ScanComparison<'_>,
1597        context: &str,
1598    ) {
1599        assert_eq!(actual.len(), expected.len(), "{context} length");
1600        assert_own_order(actual, context);
1601        if matches!(comparison, ScanComparison::Exact) {
1602            assert_eq!(actual, expected, "{context}");
1603            return;
1604        }
1605        let row_count = match comparison {
1606            ScanComparison::Exact => 0,
1607            ScanComparison::F32 { query, rows } => rows.len() / query.len(),
1608            ScanComparison::F16 { query, rows } => rows.len() / query.len(),
1609        };
1610        let reference = naive_top_k(reference_request, row_count);
1611        let mut rank_by_id = vec![usize::MAX; row_count];
1612        for (rank, candidate) in reference.iter().enumerate() {
1613            rank_by_id[candidate.row_id] = rank;
1614        }
1615        for (position, (expected_candidate, actual_candidate)) in
1616            expected.iter().zip(actual).enumerate()
1617        {
1618            let tolerance = score_bound(comparison, expected_candidate.row_id)
1619                .max(score_bound(comparison, actual_candidate.row_id));
1620            let error =
1621                (f64::from(actual_candidate.score) - f64::from(expected_candidate.score)).abs();
1622            assert!(
1623                error <= tolerance,
1624                "{context} score position {position}: expected={expected_candidate:?} actual={actual_candidate:?} error={error:?} tolerance={tolerance:?}"
1625            );
1626            let actual_rank = rank_by_id[actual_candidate.row_id];
1627            let (cluster_start, cluster_end) = reference_cluster(&reference, position, comparison);
1628            assert!(
1629                actual_rank >= cluster_start && actual_rank <= cluster_end,
1630                "{context} id escaped near-tie cluster at position {position}: expected={expected_candidate:?} actual={actual_candidate:?} cluster={cluster_start}..={cluster_end} actual_rank={actual_rank}"
1631            );
1632            if actual_rank >= expected.len() && !expected.is_empty() {
1633                let boundary = &reference[expected.len() - 1];
1634                let exchanged = &reference[actual_rank];
1635                let boundary_tolerance = score_bound(comparison, boundary.row_id);
1636                assert!(
1637                    (f64::from(exchanged.score) - f64::from(boundary.score)).abs()
1638                        <= boundary_tolerance,
1639                    "{context} k-boundary exchange exceeded its bound: kth={boundary:?} exchanged={exchanged:?} tolerance={boundary_tolerance:?}"
1640                );
1641            }
1642        }
1643    }
1644
1645    fn reference_cluster(
1646        reference: &[super::ScanCandidate],
1647        position: usize,
1648        comparison: ScanComparison<'_>,
1649    ) -> (usize, usize) {
1650        let mut start = position;
1651        while start > 0
1652            && reference_pair_is_near(&reference[start - 1], &reference[start], comparison)
1653        {
1654            start -= 1;
1655        }
1656        let mut end = position;
1657        while end + 1 < reference.len()
1658            && reference_pair_is_near(&reference[end], &reference[end + 1], comparison)
1659        {
1660            end += 1;
1661        }
1662        (start, end)
1663    }
1664
1665    fn reference_pair_is_near(
1666        left: &super::ScanCandidate,
1667        right: &super::ScanCandidate,
1668        comparison: ScanComparison<'_>,
1669    ) -> bool {
1670        let bound = score_bound(comparison, left.row_id).max(score_bound(comparison, right.row_id));
1671        (f64::from(left.score) - f64::from(right.score)).abs() <= 2.0 * bound
1672    }
1673
1674    fn score_bound(comparison: ScanComparison<'_>, row_id: usize) -> f64 {
1675        match comparison {
1676            ScanComparison::Exact => 0.0,
1677            ScanComparison::F32 { query, rows } => {
1678                rows.get(row_id * query.len()..(row_id + 1) * query.len())
1679                    .unwrap_or_default()
1680                    .iter()
1681                    .zip(query)
1682                    .map(|(&row, &query)| f64::from(row).abs() * f64::from(query).abs())
1683                    .sum::<f64>()
1684                    * 1.0e-5
1685            }
1686            ScanComparison::F16 { query, rows } => {
1687                rows.get(row_id * query.len()..(row_id + 1) * query.len())
1688                    .unwrap_or_default()
1689                    .iter()
1690                    .zip(query)
1691                    .map(|(&row, &query)| f16_to_f64(row).abs() * f16_to_f64(query).abs())
1692                    .sum::<f64>()
1693                    * 1.0e-5
1694            }
1695        }
1696    }
1697
1698    fn assert_own_order(candidates: &[super::ScanCandidate], context: &str) {
1699        for pair in candidates.windows(2) {
1700            let left = pair[0];
1701            let right = pair[1];
1702            assert!(
1703                left.score.total_cmp(&right.score).is_ge(),
1704                "{context} output is not descending: left={left:?} right={right:?}"
1705            );
1706            if left.score.to_bits() == right.score.to_bits() {
1707                assert!(
1708                    left.row_id < right.row_id,
1709                    "{context} bitwise tie is not ascending by row id: left={left:?} right={right:?}"
1710                );
1711            }
1712        }
1713    }
1714
1715    fn random_finite_f16(random: &mut impl Rng) -> u16 {
1716        let sign = if random.random::<bool>() { 0x8000 } else { 0 };
1717        let exponent = random.random_range(0_u16..=30) << 10;
1718        let fraction = random.random_range(0_u16..=0x03ff);
1719        sign | exponent | fraction
1720    }
1721
1722    fn f16_to_f64(bits: u16) -> f64 {
1723        let sign = if bits & 0x8000 == 0 { 1.0 } else { -1.0 };
1724        let exponent = (bits >> 10) & 0x1f;
1725        let fraction = bits & 0x03ff;
1726        match exponent {
1727            0 => sign * f64::from(fraction) * 2.0_f64.powi(-24),
1728            _ => {
1729                sign * (1.0 + f64::from(fraction) / 1_024.0)
1730                    * 2.0_f64.powi(i32::from(exponent) - 15)
1731            }
1732        }
1733    }
1734
1735    #[test]
1736    fn prop_parallel_equals_single_thread() {
1737        let (_directory, store) = query_store();
1738        let mut random =
1739            crate::test_support::seeded_rng("scan::prop_parallel_equals_single_thread");
1740        let cases = std::env::var("PROPTEST_CASES")
1741            .ok()
1742            .and_then(|value| value.parse::<usize>().ok())
1743            .unwrap_or(256);
1744        for case in 0..cases {
1745            let dimension = random.random_range(1..=65);
1746            let row_count = random.random_range(1..=576);
1747            let k = random.random_range(1..=row_count);
1748            let query = (0..dimension)
1749                .map(|_| random.random_range(-2.0_f32..=2.0_f32))
1750                .collect::<Vec<_>>();
1751            let rows = F32Rows::new(if case.is_multiple_of(9) {
1752                vec![1.0_f32; row_count * dimension]
1753            } else {
1754                (0..row_count * dimension)
1755                    .map(|_| random.random_range(-2.0_f32..=2.0_f32))
1756                    .collect::<Vec<_>>()
1757            });
1758            let request = ScanRequest {
1759                query: ScanQuery::F32(&query),
1760                rows: ScanRows::F32RowMajor(&rows),
1761                row_mask: None,
1762            };
1763            let single = pooled_top_k(&store, request, k, ScanOptions { thread_budget: 1 })
1764                .expect("single-thread scan");
1765            let requested = case % 12 + 1;
1766            let parallel = pooled_top_k(
1767                &store,
1768                request,
1769                k,
1770                ScanOptions {
1771                    thread_budget: requested,
1772                },
1773            )
1774            .expect("parallel scan");
1775            assert_eq!(
1776                parallel.candidates.first(),
1777                single.candidates.first(),
1778                "case {case} thread {requested} first candidate"
1779            );
1780            assert_eq!(
1781                parallel.candidates, single.candidates,
1782                "case {case} thread {requested}"
1783            );
1784            assert_eq!(
1785                (parallel.stats.dims_touched, parallel.stats.bytes_read),
1786                (single.stats.dims_touched, single.stats.bytes_read),
1787                "case {case} thread {requested} counters"
1788            );
1789        }
1790    }
1791
1792    #[test]
1793    fn row_major_scan_reports_one_exhaustive_payload_pass() {
1794        let (_directory, store) = query_store();
1795        let dimension = 37;
1796        let row_count = 137;
1797        let query = vec![1.0_f32; dimension];
1798        let rows = F32Rows::new(vec![0.5_f32; row_count * dimension]);
1799        let outcome = pooled_top_k(
1800            &store,
1801            ScanRequest {
1802                query: ScanQuery::F32(&query),
1803                rows: ScanRows::F32RowMajor(&rows),
1804                row_mask: None,
1805            },
1806            10,
1807            ScanOptions { thread_budget: 1 },
1808        )
1809        .expect("exhaustive scan");
1810        assert_eq!(outcome.stats.dims_touched, (row_count * dimension) as u64);
1811        assert_eq!(
1812            outcome.stats.bytes_read,
1813            (row_count * dimension * std::mem::size_of::<f32>()) as u64,
1814            "row-major exhaustive scan must report every loaded payload byte"
1815        );
1816    }
1817
1818    #[test]
1819    fn explicit_thread_budget_is_honoured_up_to_physical_cap() {
1820        let (_directory, store) = query_store();
1821        let capacity = super::physical_thread_capacity().expect("physical thread capacity");
1822        let row_count = capacity.max(8) * 64;
1823        let query = [1.0_f32];
1824        let rows = F32Rows::new(vec![1.0_f32; row_count]);
1825        for requested in 1..=capacity.saturating_add(2) {
1826            let outcome = pooled_top_k(
1827                &store,
1828                ScanRequest {
1829                    query: ScanQuery::F32(&query),
1830                    rows: ScanRows::F32RowMajor(&rows),
1831                    row_mask: None,
1832                },
1833                1,
1834                ScanOptions {
1835                    thread_budget: requested,
1836                },
1837            )
1838            .expect("parallel scan");
1839            assert_eq!(outcome.stats.threads_used, requested.min(capacity));
1840        }
1841    }
1842
1843    #[test]
1844    fn zero_k_returns_no_candidates() {
1845        let query = [0x3c00_u16, 0xbc00];
1846        let rows = [0x3c00_u16, 0xbc00, 0x4000, 0x3800];
1847        let request = ScanRequest {
1848            query: ScanQuery::F16(&query),
1849            rows: ScanRows::F16RowMajor(&rows),
1850            row_mask: None,
1851        };
1852
1853        assert!(top_k(request, 0).expect("valid zero-k scan").is_empty());
1854    }
1855
1856    #[test]
1857    fn dimension_one_scans_exactly() {
1858        let query = [0x3c00_u16];
1859        let rows = [0x4000_u16, 0x3c00, 0xbc00];
1860        let request = ScanRequest {
1861            query: ScanQuery::F16(&query),
1862            rows: ScanRows::F16RowMajor(&rows),
1863            row_mask: None,
1864        };
1865
1866        let hits = top_k(request, 3).expect("valid one-dimensional scan");
1867
1868        assert_eq!(
1869            hits.iter().map(|hit| hit.row_id).collect::<Vec<_>>(),
1870            [0, 1, 2]
1871        );
1872        assert_eq!(
1873            hits.iter().map(|hit| hit.score).collect::<Vec<_>>(),
1874            [2.0, 1.0, -1.0]
1875        );
1876    }
1877
1878    #[test]
1879    fn k_larger_than_row_count_returns_all_rows() {
1880        let query = [1.0_f32, 0.5];
1881        let rows = F32Rows::new(vec![1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0]);
1882        let request = ScanRequest {
1883            query: ScanQuery::F32(&query),
1884            rows: ScanRows::F32RowMajor(&rows),
1885            row_mask: None,
1886        };
1887
1888        let hits = top_k(request, usize::MAX).expect("valid oversized-k scan");
1889
1890        assert_eq!(hits.len(), 3);
1891        assert_eq!(
1892            hits.iter().map(|hit| hit.row_id).collect::<Vec<_>>(),
1893            [0, 1, 2]
1894        );
1895    }
1896
1897    #[test]
1898    fn single_row_returns_that_row() {
1899        let query = [0.25_f32, -0.5, 1.0];
1900        let rows = F32Rows::new(vec![2.0_f32, 1.0, -1.0]);
1901        let request = ScanRequest {
1902            query: ScanQuery::F32(&query),
1903            rows: ScanRows::F32RowMajor(&rows),
1904            row_mask: None,
1905        };
1906
1907        assert_eq!(
1908            top_k(request, 1).expect("valid single-row scan"),
1909            [super::ScanCandidate {
1910                row_id: 0,
1911                score: -1.0,
1912            }]
1913        );
1914    }
1915
1916    fn mask_fixture<'a>(
1917        query: &'a [f32],
1918        rows: &'a F32Rows,
1919        row_mask: Option<&'a RoaringBitmap>,
1920    ) -> ScanRequest<'a> {
1921        ScanRequest {
1922            query: ScanQuery::F32(query),
1923            rows: ScanRows::F32RowMajor(rows),
1924            row_mask,
1925        }
1926    }
1927
1928    #[test]
1929    fn empty_row_mask_returns_nothing() {
1930        let query = [1.0_f32, 0.0];
1931        let rows = F32Rows::new(vec![1.0_f32, 0.0, 2.0, 0.0, 3.0, 0.0]);
1932        let mask = RoaringBitmap::new();
1933
1934        assert!(
1935            top_k(mask_fixture(&query, &rows, Some(&mask)), 3)
1936                .expect("valid empty mask")
1937                .is_empty()
1938        );
1939    }
1940
1941    #[test]
1942    fn full_row_mask_equals_unmasked() {
1943        let query = [1.0_f32, 0.0];
1944        let rows = F32Rows::new(vec![1.0_f32, 0.0, 2.0, 0.0, 3.0, 0.0]);
1945        let mask = RoaringBitmap::from_iter(0_u32..3);
1946
1947        assert_eq!(
1948            top_k(mask_fixture(&query, &rows, Some(&mask)), 3).expect("valid full mask"),
1949            top_k(mask_fixture(&query, &rows, None), 3).expect("valid unmasked scan")
1950        );
1951    }
1952
1953    #[test]
1954    fn alternating_row_mask_returns_exactly_allowed_ids() {
1955        let query = [1.0_f32, 0.0];
1956        let rows = F32Rows::new(vec![
1957            0.0_f32, 0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0, 4.0, 0.0, 5.0, 0.0,
1958        ]);
1959        let mask = RoaringBitmap::from_iter([0_u32, 2, 4]);
1960
1961        let hits =
1962            top_k(mask_fixture(&query, &rows, Some(&mask)), 6).expect("valid alternating mask");
1963        let mut ids = hits.iter().map(|hit| hit.row_id).collect::<Vec<_>>();
1964        ids.sort_unstable();
1965
1966        assert_eq!(ids, [0, 2, 4]);
1967    }
1968
1969    #[test]
1970    fn scan_validation_errors_are_typed_and_actionable() {
1971        assert_eq!(
1972            Int8Factors::new(f32::NAN, 0.0),
1973            Err(super::Int8FactorsError::NonFiniteScale)
1974        );
1975        assert_eq!(
1976            Int8Factors::new(-1.0, 0.0),
1977            Err(super::Int8FactorsError::NegativeScale)
1978        );
1979        assert_eq!(
1980            Int8Factors::new(1.0, f32::INFINITY),
1981            Err(super::Int8FactorsError::NonFiniteOffset)
1982        );
1983        for error in [
1984            super::Int8FactorsError::NonFiniteScale,
1985            super::Int8FactorsError::NegativeScale,
1986            super::Int8FactorsError::NonFiniteOffset,
1987        ] {
1988            assert!(!error.to_string().is_empty());
1989        }
1990
1991        let f32_query = [1.0_f32, 2.0];
1992        let f16_query = [0x3c00_u16, 0x4000];
1993        let int8_query = prepare_int8_query(&f32_query).expect("finite Int8 query");
1994        let bit4_query = prepare_bit4_query(&f32_query, 7).expect("finite Bit4 query");
1995        let int8_factor = Int8Factors::new(1.0, 0.0).expect("valid factor");
1996        let mut bit4_code = [0_u8; 1];
1997        let bit4_factor = quantize_bit4(&f32_query, &mut bit4_code).expect("valid factor");
1998        let partial_f32_rows = F32Rows::new(vec![1.0]);
1999        let finite_f32_rows = F32Rows::new(vec![1.0]);
2000        let overflowing_f32_rows = F32Rows::new(vec![2.0]);
2001
2002        let errors = [
2003            top_k(
2004                ScanRequest {
2005                    query: ScanQuery::F16(&[]),
2006                    rows: ScanRows::F16RowMajor(&[]),
2007                    row_mask: None,
2008                },
2009                1,
2010            )
2011            .expect_err("zero dimension"),
2012            top_k(
2013                ScanRequest {
2014                    query: ScanQuery::F32(&f32_query),
2015                    rows: ScanRows::F32RowMajor(&partial_f32_rows),
2016                    row_mask: None,
2017                },
2018                1,
2019            )
2020            .expect_err("partial f32 row"),
2021            top_k(
2022                ScanRequest {
2023                    query: ScanQuery::F16(&f16_query),
2024                    rows: ScanRows::F16RowMajor(&[0x3c00]),
2025                    row_mask: None,
2026                },
2027                1,
2028            )
2029            .expect_err("partial f16 row"),
2030            top_k(
2031                ScanRequest {
2032                    query: ScanQuery::Int8(&int8_query),
2033                    rows: ScanRows::Int8RowMajor {
2034                        codes: &[1],
2035                        factors: &[],
2036                    },
2037                    row_mask: None,
2038                },
2039                1,
2040            )
2041            .expect_err("partial Int8 row"),
2042            top_k(
2043                ScanRequest {
2044                    query: ScanQuery::Int8(&int8_query),
2045                    rows: ScanRows::Int8RowMajor {
2046                        codes: &[1, 2],
2047                        factors: &[],
2048                    },
2049                    row_mask: None,
2050                },
2051                1,
2052            )
2053            .expect_err("missing Int8 factor"),
2054            top_k(
2055                ScanRequest {
2056                    query: ScanQuery::Bit4(&bit4_query),
2057                    rows: ScanRows::Bit4RowMajor {
2058                        codes: &[0xff, 0xff],
2059                        factors: &[bit4_factor],
2060                    },
2061                    row_mask: None,
2062                },
2063                1,
2064            )
2065            .expect_err("partial Bit4 row"),
2066            top_k(
2067                ScanRequest {
2068                    query: ScanQuery::Bit4(&bit4_query),
2069                    rows: ScanRows::Bit4RowMajor {
2070                        codes: &bit4_code,
2071                        factors: &[],
2072                    },
2073                    row_mask: None,
2074                },
2075                1,
2076            )
2077            .expect_err("missing Bit4 factor"),
2078            top_k(
2079                ScanRequest {
2080                    query: ScanQuery::F32(&f32_query),
2081                    rows: ScanRows::Int8RowMajor {
2082                        codes: &[1, 2],
2083                        factors: &[int8_factor],
2084                    },
2085                    row_mask: None,
2086                },
2087                1,
2088            )
2089            .expect_err("scheme mismatch"),
2090            top_k(
2091                ScanRequest {
2092                    query: ScanQuery::F32(&[f32::NAN]),
2093                    rows: ScanRows::F32RowMajor(&finite_f32_rows),
2094                    row_mask: None,
2095                },
2096                1,
2097            )
2098            .expect_err("non-finite input"),
2099            top_k(
2100                ScanRequest {
2101                    query: ScanQuery::F32(&[f32::MAX]),
2102                    rows: ScanRows::F32RowMajor(&overflowing_f32_rows),
2103                    row_mask: None,
2104                },
2105                1,
2106            )
2107            .expect_err("non-finite score"),
2108        ];
2109        for error in errors {
2110            assert!(!error.to_string().is_empty());
2111        }
2112        assert!(
2113            !super::ScanError::from(crate::quant::QuantError::EmptyVector)
2114                .to_string()
2115                .is_empty()
2116        );
2117    }
2118
2119    #[test]
2120    fn row_masks_filter_every_encoded_scheme() {
2121        let mask = RoaringBitmap::from_iter([1_u32]);
2122
2123        let f16_query = [0x3c00_u16];
2124        let f16_rows = [0x3c00_u16, 0x4000];
2125        assert_eq!(
2126            top_k(
2127                ScanRequest {
2128                    query: ScanQuery::F16(&f16_query),
2129                    rows: ScanRows::F16RowMajor(&f16_rows),
2130                    row_mask: Some(&mask),
2131                },
2132                2,
2133            )
2134            .expect("masked f16 scan")[0]
2135                .row_id,
2136            1
2137        );
2138
2139        let query_values = [1.0_f32];
2140        let int8_query = prepare_int8_query(&query_values).expect("finite Int8 query");
2141        let int8_factors = [
2142            Int8Factors::new(1.0, 0.0).expect("valid factor"),
2143            Int8Factors::new(1.0, 0.0).expect("valid factor"),
2144        ];
2145        assert_eq!(
2146            top_k(
2147                ScanRequest {
2148                    query: ScanQuery::Int8(&int8_query),
2149                    rows: ScanRows::Int8RowMajor {
2150                        codes: &[1, 2],
2151                        factors: &int8_factors,
2152                    },
2153                    row_mask: Some(&mask),
2154                },
2155                2,
2156            )
2157            .expect("masked Int8 scan")[0]
2158                .row_id,
2159            1
2160        );
2161
2162        let bit4_query = prepare_bit4_query(&query_values, 9).expect("finite Bit4 query");
2163        let mut bit4_code = [0_u8; 1];
2164        let factor = quantize_bit4(&query_values, &mut bit4_code).expect("valid factor");
2165        let bit4_codes = [bit4_code[0], bit4_code[0]];
2166        assert_eq!(
2167            top_k(
2168                ScanRequest {
2169                    query: ScanQuery::Bit4(&bit4_query),
2170                    rows: ScanRows::Bit4RowMajor {
2171                        codes: &bit4_codes,
2172                        factors: &[factor, factor],
2173                    },
2174                    row_mask: Some(&mask),
2175                },
2176                2,
2177            )
2178            .expect("masked Bit4 scan")[0]
2179                .row_id,
2180            1
2181        );
2182    }
2183}