Skip to main content

hermes_core/segment/
ann_disk.rs

1//! Merge-native ANN segment format.
2//!
3//! ANN payloads are split into immutable cluster runs. A normal segment merge
4//! copies the three run columns (doc IDs, ordinals, codes) byte-for-byte and
5//! rewrites only the compact run directory with an adjusted document base.
6//! No centroid assignment, payload deserialization, or reserialization occurs
7//! on the merge path.
8
9#[cfg(feature = "native")]
10use std::cmp::Reverse;
11#[cfg(feature = "native")]
12use std::collections::BinaryHeap;
13use std::io;
14#[cfg(feature = "native")]
15use std::io::Write;
16use std::ops::Range;
17
18#[cfg(feature = "native")]
19use byteorder::WriteBytesExt;
20use byteorder::{LittleEndian, ReadBytesExt};
21
22use crate::directories::OwnedBytes;
23use crate::dsl::IvfRoutingMode;
24
25use crate::observe::DenseAnnScanStats;
26#[cfg(feature = "native")]
27use crate::structures::BinaryIvfIndex;
28use crate::structures::vector::index::{BoundedAnnCollector, BoundedUniqueAnnCollector};
29
30/// Combined binary candidates: the retained documents, plus the exact
31/// per-ordinal `(doc_id, ordinal, score)` scores behind them.
32type CombinedBinaryCandidates = (Vec<AnnDocumentCandidate>, Vec<(u32, u16, f32)>);
33/// Distinct-key scan output: `(doc_id, ordinal, score)` hits plus scan diagnostics.
34type ScanResults = (Vec<(u32, u16, f32)>, DenseAnnScanStats);
35
36/// Top-k sink shared by the TQ lane loops: the deduplicating collector for
37/// multi-valued or SOAR-spilled payloads, the heap-only collector when every
38/// `(doc_id, ordinal)` key is unique by construction (single-valued field and
39/// exactly one posting per stored vector).
40trait IvfTqTopK: Sized + Send {
41    fn with_k(k: usize) -> Self;
42    fn insert(&mut self, doc_id: u32, ordinal: u16, score: f32);
43    /// Conservative k-th best score once `k` entries are held.
44    fn pruning_threshold(&self) -> Option<f32>;
45    fn non_finite_dropped(&self) -> usize;
46    #[cfg(feature = "native")]
47    fn merge_from(&mut self, other: Self);
48    fn into_sorted_results(self) -> Vec<(u32, u16, f32)>;
49}
50
51impl IvfTqTopK for BoundedAnnCollector<true, true> {
52    #[inline]
53    fn with_k(k: usize) -> Self {
54        Self::new(k)
55    }
56    #[inline]
57    fn insert(&mut self, doc_id: u32, ordinal: u16, score: f32) {
58        BoundedAnnCollector::insert(self, doc_id, ordinal, score)
59    }
60    #[inline]
61    fn pruning_threshold(&self) -> Option<f32> {
62        BoundedAnnCollector::pruning_threshold(self)
63    }
64    #[inline]
65    fn non_finite_dropped(&self) -> usize {
66        BoundedAnnCollector::non_finite_dropped(self)
67    }
68    #[cfg(feature = "native")]
69    #[inline]
70    fn merge_from(&mut self, other: Self) {
71        BoundedAnnCollector::merge_from(self, other)
72    }
73    #[inline]
74    fn into_sorted_results(self) -> Vec<(u32, u16, f32)> {
75        BoundedAnnCollector::into_sorted_results(self)
76    }
77}
78
79impl IvfTqTopK for BoundedUniqueAnnCollector<true> {
80    #[inline]
81    fn with_k(k: usize) -> Self {
82        Self::new(k)
83    }
84    #[inline]
85    fn insert(&mut self, doc_id: u32, ordinal: u16, score: f32) {
86        BoundedUniqueAnnCollector::insert(self, doc_id, ordinal, score)
87    }
88    #[inline]
89    fn pruning_threshold(&self) -> Option<f32> {
90        BoundedUniqueAnnCollector::pruning_threshold(self)
91    }
92    #[inline]
93    fn non_finite_dropped(&self) -> usize {
94        BoundedUniqueAnnCollector::non_finite_dropped(self)
95    }
96    #[cfg(feature = "native")]
97    #[inline]
98    fn merge_from(&mut self, other: Self) {
99        BoundedUniqueAnnCollector::merge_from(self, other)
100    }
101    #[inline]
102    fn into_sorted_results(self) -> Vec<(u32, u16, f32)> {
103        BoundedUniqueAnnCollector::into_sorted_results(self)
104    }
105}
106
107const ANN_HEADER_MAGIC: u32 = 0x3152_4e41; // "ANR1"
108const ANN_FOOTER_MAGIC: u32 = 0x3146_4e41; // "ANF1"
109const ANN_DISK_VERSION: u16 = 1;
110const ANN_HEADER_SIZE: usize = 56;
111const ANN_RUN_SIZE: usize = 48;
112const ANN_FOOTER_SIZE: usize = 24;
113#[cfg(feature = "native")]
114const COPY_CHUNK: usize = 8 * 1024 * 1024;
115#[cfg(feature = "native")]
116const PREFETCH_COALESCE_GAP: usize = 4 * 1024;
117const BINARY_SCORE_BATCH: usize = 8_192;
118/// Upper bound on the TQ leaf estimate `est⟨q̂,r̂⟩` for a unit residual
119/// direction: `base ≤ ‖recon‖ ≈ 1` plus the QJL term `≤ √(π/2)·γ ≤ 1.26·γ`
120/// with `γ < 1`, kept with slack so pruning can never drop a candidate the
121/// unpruned scan would have kept (pinned by a test).
122const TQ_PRUNE_ESTIMATE_BOUND: f32 = 1.3;
123/// Flat TQ scans fan out across Rayon above this vector count. Rayon folds
124/// chunks into one collector per worker before reducing those collectors, so
125/// temporary top-k memory is bounded by the active worker count rather than
126/// the number of chunks. Small segments stay sequential because fan-out
127/// overhead would dominate.
128#[cfg(feature = "native")]
129const TQ_PARALLEL_SCAN_MIN_VECTORS: usize = 65_536;
130#[cfg(feature = "native")]
131const TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512;
132/// IVF-TQ scans are already parallel across segments. Fan out inside one
133/// segment only when the selected leaves contain enough postings to amortize
134/// Rayon scheduling and worker-local top-k state.
135const IVF_PARALLEL_SCAN_MIN_POSTINGS: usize = 65_536;
136const IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512;
137
138#[derive(Debug, Clone, Copy, PartialEq, Eq)]
139pub(crate) enum AnnKind {
140    // Discriminant 1 was IVF-PQ, removed after IVF-TQ superseded it
141    // (docs/turboquant-quantization.md). Never reuse it.
142    BinaryIvf = 2,
143    /// TurboQuant flat scan: single logical cluster, block-packed codes.
144    TqFlat = 3,
145    /// Trained IVF router with TurboQuant-coded centroid residuals.
146    IvfTq = 4,
147    /// ScaNN float leaves encoded with 4-bit asymmetric hashes.
148    ScannAh = 5,
149    /// ScaNN binary leaves retaining exact packed vectors for Hamming scan.
150    ScannBinary = 6,
151}
152
153impl AnnKind {
154    fn from_u8(value: u8) -> io::Result<Self> {
155        match value {
156            1 => Err(invalid_data(
157                "ANN kind 1 (IVF-PQ) is no longer supported; recreate the index \
158                 with `ivf_tq` and reindex",
159            )),
160            2 => Ok(Self::BinaryIvf),
161            3 => Ok(Self::TqFlat),
162            4 => Ok(Self::IvfTq),
163            5 => Ok(Self::ScannAh),
164            6 => Ok(Self::ScannBinary),
165            _ => Err(invalid_data(format!("unknown ANN kind {value}"))),
166        }
167    }
168}
169
170/// Codes-column byte length for one run. TQ packs vectors into 16-lane
171/// blocks (gammas + dimension-major nibbles), so its column is block-padded
172/// rather than `count * code_size`.
173fn expected_codes_column_len(
174    kind: AnnKind,
175    count: usize,
176    dim: usize,
177    code_size: usize,
178) -> io::Result<usize> {
179    match kind {
180        AnnKind::BinaryIvf | AnnKind::ScannBinary => count
181            .checked_mul(code_size)
182            .ok_or_else(|| invalid_data("ANN code column size overflows usize")),
183        // Single source of truth for the block-packed layouts lives in tq.rs.
184        AnnKind::TqFlat => {
185            crate::structures::vector::quantization::tq_codes_column_len_checked(count, code_size)
186                .ok_or_else(|| invalid_data("TQ code column size overflows usize"))
187        }
188        AnnKind::IvfTq => crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
189            count, code_size,
190        )
191        .ok_or_else(|| invalid_data("IVF-TQ code column size overflows usize")),
192        // For ScaNN AH, `code_size` stores dimensions-per-block. This is the
193        // one encoding parameter not derivable from the fixed ANN header;
194        // the actual byte length remains count-dependent because complete
195        // 32-row FastScan blocks interleave their nibbles.
196        AnnKind::ScannAh => crate::structures::vector::scann::ScannEncoding::AsymmetricHash {
197            dimensions_per_block: u16::try_from(code_size)
198                .map_err(|_| invalid_data("ScaNN AH block dimension exceeds u16"))?,
199            bits_per_code: 4,
200        }
201        .leaf_code_bytes(
202            u32::try_from(dim).map_err(|_| invalid_data("ScaNN dimension exceeds u32"))?,
203            count,
204        )
205        .map_err(|error| invalid_data(error.to_string())),
206    }
207}
208
209#[derive(Debug, Clone, PartialEq, Eq)]
210pub(crate) struct AnnDiskHeader {
211    pub kind: AnnKind,
212    pub routing: IvfRoutingMode,
213    pub dim: usize,
214    pub code_size: usize,
215    pub num_clusters: u32,
216    pub quantizer_version: u64,
217    pub codebook_version: u64,
218    pub vector_count: usize,
219}
220
221#[derive(Debug)]
222struct AnnRun {
223    cluster_id: u32,
224    doc_base: u32,
225    max_doc_id: u32,
226    count: usize,
227    doc_ids: Range<usize>,
228    ordinals: Range<usize>,
229    codes: Range<usize>,
230}
231
232#[derive(Clone, Copy)]
233struct IvfTqScanTask<'a> {
234    run: &'a AnnRun,
235    cluster_dot: f32,
236    first_block: usize,
237    last_block: usize,
238}
239
240#[derive(Clone, Copy)]
241struct BinaryScanTask<'a> {
242    run: &'a AnnRun,
243    first_index: usize,
244    count: usize,
245}
246
247/// Mmap-backed searchable ANN payload. Only the fixed-size run directory is
248/// heap-resident; all corpus-sized columns remain zero-copy file slices.
249pub(crate) struct AnnDiskIndex {
250    // Drop locks before the directory allocation they reference.
251    #[cfg(feature = "native")]
252    heap_pins: crate::segment::pin::HeapPinSet,
253    raw: OwnedBytes,
254    header: AnnDiskHeader,
255    runs: Vec<AnnRun>,
256}
257
258/// Cheap structural health of one ANN payload, computed from the in-memory
259/// run directory in O(runs) — payload bytes are never touched.
260///
261/// Two production failure modes motivated every field here (see
262/// `docs/diagnostics.md`): a 31%-of-vectors leaf built from degenerate
263/// embeddings, and probe read-amplification from byte-copy merges that leave
264/// one logical cluster scattered across many physical extents.
265#[derive(Debug, Clone, Copy, PartialEq)]
266pub struct AnnHealth {
267    /// Vectors across all runs.
268    pub vectors: u64,
269    /// Distinct cluster IDs holding at least one posting.
270    pub clusters_nonempty: u32,
271    /// Codebook size (`num_clusters` from the header).
272    pub clusters_total: u32,
273    /// Run-directory entries; each is one physical extent on disk.
274    pub runs: u32,
275    /// Vectors in the most populated cluster, with its ID.
276    pub largest_cluster: u32,
277    pub largest_cluster_vectors: u64,
278    /// Faiss `imbalance_factor` over non-empty clusters:
279    /// `K · Σ nᵢ² / N²`. 1.0 is perfectly balanced; a value of γ means a
280    /// fixed-`nprobe` probe computes γ× the distances of the balanced
281    /// baseline in expectation.
282    pub imbalance: f64,
283    /// Bytes of the codes columns (what a probe of every leaf would read).
284    pub payload_bytes: u64,
285}
286
287impl AnnHealth {
288    /// Physical extents per non-empty cluster. 1.0 after a rebuild; each
289    /// byte-copy merge multiplies it, and every extent is a potential seek
290    /// when the index is cold.
291    pub fn fragmentation(&self) -> f64 {
292        if self.clusters_nonempty == 0 {
293            return 0.0;
294        }
295        f64::from(self.runs) / f64::from(self.clusters_nonempty)
296    }
297
298    /// Share of all vectors held by the single largest cluster.
299    pub fn largest_cluster_share(&self) -> f64 {
300        if self.vectors == 0 {
301            return 0.0;
302        }
303        self.largest_cluster_vectors as f64 / self.vectors as f64
304    }
305}
306
307/// `largest_cluster_share` above this, on at least [`ANN_SKEW_WARN_MIN_VECTORS`]
308/// vectors, is a scan cliff worth a warning: the production incident value was
309/// 0.31, and a healthy 100k+-cluster codebook sits orders of magnitude lower.
310const ANN_SKEW_WARN_SHARE: f64 = 0.05;
311const ANN_SKEW_WARN_MIN_VECTORS: u64 = 100_000;
312/// A rebuilt segment has fragmentation 1.0; a single 32-way byte-copy merge
313/// can reach 32. Warn once probes pay ~an order of magnitude extra seeks.
314const ANN_FRAGMENTATION_WARN: f64 = 8.0;
315
316/// A document selected by a combiner-aware compressed scan.
317///
318/// This intentionally has no ordinal: `score` combines every value belonging
319/// to the document and must never be reused as an individual vector score by
320/// an exact reranker.
321#[derive(Debug, Clone, Copy, PartialEq)]
322pub(crate) struct AnnDocumentCandidate {
323    pub(crate) doc_id: u32,
324    pub(crate) score: f32,
325}
326
327impl AnnDiskIndex {
328    pub(crate) fn open(
329        raw: OwnedBytes,
330        expected_kind: AnnKind,
331        total_docs: u32,
332    ) -> io::Result<Self> {
333        if raw.len() < ANN_HEADER_SIZE + ANN_FOOTER_SIZE {
334            return Err(invalid_data("ANN payload is shorter than header + footer"));
335        }
336        let bytes = raw.as_slice();
337        let mut header_cursor = std::io::Cursor::new(&bytes[..ANN_HEADER_SIZE]);
338        if header_cursor.read_u32::<LittleEndian>()? != ANN_HEADER_MAGIC {
339            return Err(invalid_data("ANN payload has unsupported header magic"));
340        }
341        let kind = AnnKind::from_u8(header_cursor.read_u8()?)?;
342        if kind != expected_kind {
343            return Err(invalid_data(format!(
344                "ANN payload kind {kind:?} does not match expected {expected_kind:?}"
345            )));
346        }
347        let routing = routing_from_u8(header_cursor.read_u8()?)?;
348        if header_cursor.read_u16::<LittleEndian>()? != ANN_DISK_VERSION {
349            return Err(invalid_data("ANN payload has unsupported format version"));
350        }
351        let dim = header_cursor.read_u32::<LittleEndian>()? as usize;
352        let code_size = header_cursor.read_u32::<LittleEndian>()? as usize;
353        let num_clusters = header_cursor.read_u32::<LittleEndian>()?;
354        check_layout_revision(kind, header_cursor.read_u32::<LittleEndian>()?)?;
355        let quantizer_version = header_cursor.read_u64::<LittleEndian>()?;
356        let codebook_version = header_cursor.read_u64::<LittleEndian>()?;
357        let vector_count = usize::try_from(header_cursor.read_u64::<LittleEndian>()?)
358            .map_err(|_| invalid_data("ANN vector count exceeds usize"))?;
359        if header_cursor.read_u64::<LittleEndian>()? != 0 {
360            return Err(invalid_data("ANN header tail is non-zero"));
361        }
362        let header = AnnDiskHeader {
363            kind,
364            routing,
365            dim,
366            code_size,
367            num_clusters,
368            quantizer_version,
369            codebook_version,
370            vector_count,
371        };
372        validate_header(&header)?;
373
374        let footer_start = bytes.len() - ANN_FOOTER_SIZE;
375        let mut footer_cursor = std::io::Cursor::new(&bytes[footer_start..]);
376        let directory_offset = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
377            .map_err(|_| invalid_data("ANN directory offset exceeds usize"))?;
378        let num_runs = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
379            .map_err(|_| invalid_data("ANN run count exceeds usize"))?;
380        if footer_cursor.read_u32::<LittleEndian>()? != ANN_FOOTER_MAGIC
381            || footer_cursor.read_u32::<LittleEndian>()? != u32::from(ANN_DISK_VERSION)
382        {
383            return Err(invalid_data("ANN payload has unsupported footer"));
384        }
385        if num_runs == 0 {
386            return Err(invalid_data("ANN payload has no cluster runs"));
387        }
388        let directory_len = num_runs
389            .checked_mul(ANN_RUN_SIZE)
390            .ok_or_else(|| invalid_data("ANN directory size overflows usize"))?;
391        if directory_offset < ANN_HEADER_SIZE
392            || directory_offset.checked_add(directory_len) != Some(footer_start)
393        {
394            return Err(invalid_data("ANN directory does not end at the footer"));
395        }
396
397        let mut runs = Vec::with_capacity(num_runs);
398        let mut directory_cursor = std::io::Cursor::new(&bytes[directory_offset..footer_start]);
399        let mut previous_cluster = None;
400        let mut counted_vectors = 0usize;
401        for _ in 0..num_runs {
402            let cluster_id = directory_cursor.read_u32::<LittleEndian>()?;
403            let doc_base = directory_cursor.read_u32::<LittleEndian>()?;
404            let count = directory_cursor.read_u32::<LittleEndian>()? as usize;
405            let max_doc_id = directory_cursor.read_u32::<LittleEndian>()?;
406            let doc_ids_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
407                .map_err(|_| invalid_data("ANN doc-ID offset exceeds usize"))?;
408            let ordinals_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
409                .map_err(|_| invalid_data("ANN ordinal offset exceeds usize"))?;
410            let codes_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
411                .map_err(|_| invalid_data("ANN code offset exceeds usize"))?;
412            let codes_len = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
413                .map_err(|_| invalid_data("ANN code length exceeds usize"))?;
414            if count == 0
415                || cluster_id >= num_clusters
416                || previous_cluster.is_some_and(|previous| previous > cluster_id)
417                || doc_base
418                    .checked_add(max_doc_id)
419                    .is_none_or(|doc_id| doc_id >= total_docs)
420            {
421                return Err(invalid_data("ANN run metadata is invalid"));
422            }
423            previous_cluster = Some(cluster_id);
424            let doc_ids_len = count
425                .checked_mul(std::mem::size_of::<u32>())
426                .ok_or_else(|| invalid_data("ANN doc-ID column size overflows usize"))?;
427            let ordinals_len = count
428                .checked_mul(std::mem::size_of::<u16>())
429                .ok_or_else(|| invalid_data("ANN ordinal column size overflows usize"))?;
430            let expected_codes_len = expected_codes_column_len(kind, count, dim, code_size)?;
431            let doc_ids_end = doc_ids_offset
432                .checked_add(doc_ids_len)
433                .ok_or_else(|| invalid_data("ANN doc-ID range overflows usize"))?;
434            let ordinals_end = ordinals_offset
435                .checked_add(ordinals_len)
436                .ok_or_else(|| invalid_data("ANN ordinal range overflows usize"))?;
437            let codes_end = codes_offset
438                .checked_add(codes_len)
439                .ok_or_else(|| invalid_data("ANN code range overflows usize"))?;
440            if doc_ids_offset < ANN_HEADER_SIZE
441                || ordinals_offset != doc_ids_end
442                || codes_offset != ordinals_end
443                || codes_len != expected_codes_len
444                || codes_end > directory_offset
445            {
446                return Err(invalid_data("ANN run columns are not contiguous/in bounds"));
447            }
448            runs.push(AnnRun {
449                cluster_id,
450                doc_base,
451                max_doc_id,
452                count,
453                doc_ids: doc_ids_offset..doc_ids_end,
454                ordinals: ordinals_offset..ordinals_end,
455                codes: codes_offset..codes_end,
456            });
457            counted_vectors = counted_vectors
458                .checked_add(count)
459                .ok_or_else(|| invalid_data("ANN run vector count overflows usize"))?;
460        }
461        // Validate every doc-ID column once so the scan loops can resolve
462        // IDs with plain reads (`run_doc_id_validated`) instead of a per-lane
463        // bounds/overflow check. A corrupt column fails loud here.
464        for run in &runs {
465            let column = &bytes[run.doc_ids.clone()];
466            if let Some(local_doc_id) = column
467                .chunks_exact(std::mem::size_of::<u32>())
468                .map(|entry| u32::from_le_bytes([entry[0], entry[1], entry[2], entry[3]]))
469                .find(|&local_doc_id| local_doc_id > run.max_doc_id)
470            {
471                return Err(invalid_data(format!(
472                    "ANN run for cluster {} contains local document {local_doc_id} above its declared maximum {}",
473                    run.cluster_id, run.max_doc_id
474                )));
475            }
476        }
477        let mut payload_order: Vec<usize> = (0..runs.len()).collect();
478        payload_order.sort_unstable_by_key(|&index| runs[index].doc_ids.start);
479        let mut expected_payload_offset = ANN_HEADER_SIZE;
480        for index in payload_order {
481            let run = &runs[index];
482            if run.doc_ids.start != expected_payload_offset {
483                return Err(invalid_data("ANN payload runs overlap or contain gaps"));
484            }
485            expected_payload_offset = run.codes.end;
486        }
487        if expected_payload_offset != directory_offset || counted_vectors != vector_count {
488            return Err(invalid_data(
489                "ANN payload coverage/vector count is inconsistent",
490            ));
491        }
492
493        // Clustered queries visit a small set of runs at unrelated offsets.
494        // Disable default mmap readahead for those corpus-sized payloads:
495        // without this, each small run can pull in ~128 KiB and amplify
496        // cold-query IO by an order of magnitude. Their search methods issue
497        // exact WILLNEED ranges before scoring. Flat TQ deliberately scans its
498        // sole cluster and therefore retains a sequential access policy.
499        #[cfg(feature = "native")]
500        raw.madvise_range(
501            ANN_HEADER_SIZE..directory_offset,
502            if kind == AnnKind::TqFlat {
503                libc::MADV_SEQUENTIAL
504            } else {
505                libc::MADV_RANDOM
506            },
507        );
508
509        Ok(Self {
510            #[cfg(feature = "native")]
511            heap_pins: Default::default(),
512            raw,
513            header,
514            runs,
515        })
516    }
517
518    /// Structural health from the run directory alone. O(runs), no payload
519    /// reads — safe to call at every open.
520    pub(crate) fn health(&self) -> AnnHealth {
521        let mut vectors = 0u64;
522        let mut clusters_nonempty = 0u32;
523        let mut payload_bytes = 0u64;
524        let mut largest = (0u32, 0u64);
525        let mut sum_squares = 0f64;
526        // Runs are sorted by cluster ID, so one pass groups them.
527        let mut index = 0usize;
528        while index < self.runs.len() {
529            let cluster_id = self.runs[index].cluster_id;
530            let mut cluster_vectors = 0u64;
531            while index < self.runs.len() && self.runs[index].cluster_id == cluster_id {
532                let run = &self.runs[index];
533                cluster_vectors += run.count as u64;
534                payload_bytes += (run.codes.end - run.codes.start) as u64;
535                index += 1;
536            }
537            vectors += cluster_vectors;
538            clusters_nonempty += 1;
539            sum_squares += (cluster_vectors as f64) * (cluster_vectors as f64);
540            if cluster_vectors > largest.1 {
541                largest = (cluster_id, cluster_vectors);
542            }
543        }
544        let imbalance = if vectors == 0 || clusters_nonempty == 0 {
545            0.0
546        } else {
547            f64::from(clusters_nonempty) * sum_squares / ((vectors as f64) * (vectors as f64))
548        };
549        AnnHealth {
550            vectors,
551            clusters_nonempty,
552            clusters_total: self.header.num_clusters,
553            runs: self.runs.len() as u32,
554            largest_cluster: largest.0,
555            largest_cluster_vectors: largest.1,
556            imbalance,
557            payload_bytes,
558        }
559    }
560
561    /// Log this payload's health, warning on the two known cliff shapes.
562    ///
563    /// Called once per segment open; the caller supplies identity because the
564    /// payload itself does not know its index or field.
565    pub(crate) fn report_health(&self, index_label: &str, field_id: u32, segment_id: u128) {
566        let health = self.health();
567        let share = health.largest_cluster_share();
568        let fragmentation = health.fragmentation();
569        log::info!(
570            "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
571             vectors={} clusters={}/{} runs={} fragmentation={fragmentation:.2} \
572             imbalance={:.2} largest_leaf={:.2}% payload={}",
573            health.vectors,
574            health.clusters_nonempty,
575            health.clusters_total,
576            health.runs,
577            health.imbalance,
578            100.0 * share,
579            crate::format_bytes(health.payload_bytes),
580        );
581        crate::observe::ann_health(
582            index_label,
583            field_id,
584            health.imbalance,
585            fragmentation,
586            share,
587        );
588        if share >= ANN_SKEW_WARN_SHARE && health.vectors >= ANN_SKEW_WARN_MIN_VECTORS {
589            log::warn!(
590                "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
591                 leaf {} holds {:.1}% of {} vectors — every query probing it scans that leaf \
592                 in full; degenerate embeddings collapse into one leaf exactly like this",
593                health.largest_cluster,
594                100.0 * share,
595                health.vectors,
596            );
597        }
598        if fragmentation >= ANN_FRAGMENTATION_WARN {
599            log::warn!(
600                "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
601                 {fragmentation:.1} extents per probed cluster ({} runs / {} clusters) — \
602                 cold probes pay that many seeks; the next merge or vector-generation rewrite \
603                 compacts to 1.0",
604                health.runs,
605                health.clusters_nonempty,
606            );
607        }
608    }
609
610    pub(crate) fn header(&self) -> &AnnDiskHeader {
611        &self.header
612    }
613
614    /// Refuse to pair a segment payload with any global ScaNN model other
615    /// than the exact generation that encoded it.
616    pub(crate) fn validate_scann_generation(
617        &self,
618        config: &crate::structures::vector::scann::ScannConfig,
619        generation: u64,
620        artifact_id: u64,
621    ) -> io::Result<()> {
622        use crate::structures::vector::scann::ScannEncoding;
623
624        let expected_kind = match config.encoding {
625            ScannEncoding::AsymmetricHash { .. } => AnnKind::ScannAh,
626            ScannEncoding::BinaryHamming => AnnKind::ScannBinary,
627        };
628        if self.header.kind != expected_kind
629            || self.header.routing != IvfRoutingMode::Flat
630            || self.header.dim != config.dimension as usize
631            || self.header.num_clusters != config.num_leaves
632            || self.header.quantizer_version != generation
633            || self.header.codebook_version != artifact_id
634        {
635            return Err(invalid_data(
636                "ScaNN ANN payload does not match the global trained generation",
637            ));
638        }
639        match config.encoding {
640            ScannEncoding::AsymmetricHash {
641                dimensions_per_block,
642                bits_per_code: 4,
643            } if self.header.code_size == usize::from(dimensions_per_block) => Ok(()),
644            ScannEncoding::BinaryHamming
645                if self.header.code_size == config.dimension as usize / 8 =>
646            {
647                Ok(())
648            }
649            _ => Err(invalid_data(
650                "ScaNN ANN payload encoding does not match the global trained artifact",
651            )),
652        }
653    }
654
655    /// Validate physical leaf postings against the logical flat-vector count.
656    /// Float ScaNN and primary-only binary ScaNN are exact. Binary spilling
657    /// may add at most one posting per logical vector, with target-fraction
658    /// policies retaining a stricter segment-local cap.
659    #[cfg(feature = "native")]
660    pub(crate) fn validate_scann_posting_count(
661        &self,
662        logical_vectors: usize,
663        soar: Option<&crate::structures::SoarConfig>,
664    ) -> io::Result<()> {
665        let spill_budget = match self.header.kind {
666            AnnKind::ScannAh => 0,
667            AnnKind::ScannBinary
668                if self.header.num_clusters > 1
669                    && soar.is_some_and(|config| config.num_secondary > 0) =>
670            {
671                match soar.and_then(crate::structures::SoarConfig::calibration_target) {
672                    Some(target_fraction) => {
673                        (logical_vectors as f64 * f64::from(target_fraction)).floor() as usize
674                    }
675                    None => logical_vectors,
676                }
677            }
678            AnnKind::ScannBinary => 0,
679            _ => {
680                return Err(invalid_data(
681                    "posting-count validation requires a ScaNN ANN payload",
682                ));
683            }
684        };
685        let maximum = logical_vectors
686            .checked_add(spill_budget)
687            .ok_or_else(|| invalid_data("ScaNN posting-count bound overflows usize"))?;
688        if self.header.vector_count < logical_vectors || self.header.vector_count > maximum {
689            return Err(invalid_data(format!(
690                "ScaNN ANN payload has {} physical postings for {logical_vectors} logical vectors; expected {logical_vectors}..={maximum}",
691                self.header.vector_count,
692            )));
693        }
694        Ok(())
695    }
696
697    pub(crate) fn estimated_heap_bytes(&self) -> usize {
698        std::mem::size_of::<Self>() + self.runs.capacity() * std::mem::size_of::<AnnRun>()
699    }
700
701    #[cfg(feature = "native")]
702    pub(crate) fn pin_lookup_directory(
703        &mut self,
704        mode: crate::segment::pin::PinMode,
705        remaining: &mut u64,
706        report: &mut crate::segment::pin::PinReport,
707    ) {
708        let before = self.heap_pins.report();
709        self.heap_pins
710            .pin_slice(&self.runs, "ANN cluster-run directory", mode, remaining);
711        let after = self.heap_pins.report();
712        report.intended_bytes += after.intended_bytes - before.intended_bytes;
713        report.pinned_bytes += after.pinned_bytes - before.pinned_bytes;
714        report.skipped_budget_bytes += after.skipped_budget_bytes - before.skipped_budget_bytes;
715        report.failed_bytes += after.failed_bytes - before.failed_bytes;
716        report.heap_copy_bytes += after.heap_copy_bytes - before.heap_copy_bytes;
717    }
718
719    fn cluster_runs(&self, cluster_id: u32) -> &[AnnRun] {
720        let start = self.runs.partition_point(|run| run.cluster_id < cluster_id);
721        let end = self
722            .runs
723            .partition_point(|run| run.cluster_id <= cluster_id);
724        &self.runs[start..end]
725    }
726
727    #[cfg(feature = "native")]
728    fn ivf_tq_scan_tasks<'a>(
729        &'a self,
730        plan: &'a crate::structures::TqIvfQueryPlan,
731        block_bytes: usize,
732        chunk_blocks: usize,
733    ) -> Vec<IvfTqScanTask<'a>> {
734        debug_assert!(chunk_blocks > 0);
735        let mut tasks = Vec::new();
736        for (cluster_id, cluster_dot) in plan.cluster_dots() {
737            for run in self.cluster_runs(cluster_id) {
738                let block_count = run.codes.len() / block_bytes;
739                for first_block in (0..block_count).step_by(chunk_blocks) {
740                    tasks.push(IvfTqScanTask {
741                        run,
742                        cluster_dot,
743                        first_block,
744                        last_block: (first_block + chunk_blocks).min(block_count),
745                    });
746                }
747            }
748        }
749        tasks
750    }
751
752    #[cfg(feature = "native")]
753    fn binary_scan_tasks<'a>(&'a self, cluster_ids: &[u32]) -> Vec<BinaryScanTask<'a>> {
754        let mut tasks = Vec::new();
755        for &cluster_id in cluster_ids {
756            for run in self.cluster_runs(cluster_id) {
757                for first_index in (0..run.count).step_by(BINARY_SCORE_BATCH) {
758                    tasks.push(BinaryScanTask {
759                        run,
760                        first_index,
761                        count: BINARY_SCORE_BATCH.min(run.count - first_index),
762                    });
763                }
764            }
765        }
766        tasks
767    }
768
769    /// Prefetch exactly the mmap ranges needed by the selected IVF leaves.
770    ///
771    /// A pure-copy merge preserves each source payload as one physical extent,
772    /// so runs for one logical cluster can be far apart. Sorting by file offset
773    /// lets us coalesce overlaps and page-near ranges without reading through
774    /// unrelated clusters.
775    #[cfg(feature = "native")]
776    fn prefetch_cluster_runs(&self, cluster_ids: &[u32]) {
777        if cluster_ids.is_empty() || !self.raw.is_mmap() {
778            return;
779        }
780        let mut ranges = Vec::with_capacity(cluster_ids.len());
781        for &cluster_id in cluster_ids {
782            ranges.extend(
783                self.cluster_runs(cluster_id)
784                    .iter()
785                    .map(|run| run.doc_ids.start..run.codes.end),
786            );
787        }
788        coalesce_prefetch_ranges(&mut ranges);
789        for range in ranges {
790            self.raw.madvise_range(range, libc::MADV_WILLNEED);
791        }
792    }
793
794    /// Score a flat TQ payload while every value of a document is still in
795    /// hand, combine those approximate scores, and retain document-level
796    /// top-k. TQ build runs preserve `(doc_id, ordinal)` input order, so the
797    /// scratch space is bounded by one document plus the retained heap.
798    pub(crate) fn search_tq_combined_documents(
799        &self,
800        k: usize,
801        plan: &crate::structures::TqQueryPlan,
802        combiner: crate::query::MultiValueCombiner,
803    ) -> io::Result<Vec<AnnDocumentCandidate>> {
804        use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_block_bytes};
805
806        combiner
807            .validate()
808            .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))?;
809        if plan.padded_dim() != self.header.code_size * 2 {
810            return Err(io::Error::new(
811                io::ErrorKind::InvalidInput,
812                "TQ query plan does not match the payload dimension",
813            ));
814        }
815        if k == 0 {
816            return Ok(Vec::new());
817        }
818
819        let block_bytes = tq_block_bytes(self.header.code_size);
820        let bytes = self.raw.as_slice();
821        let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
822        let mut ordinal_scores = Vec::new();
823        let mut scores = [0.0f32; TQ_BLOCK_LANES];
824
825        for run in &self.runs {
826            let mut current_doc = None;
827            let codes = &bytes[run.codes.clone()];
828            for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
829                crate::structures::vector::quantization::tq_score_block(plan, block, &mut scores);
830                let lane_base = block_index * TQ_BLOCK_LANES;
831                let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
832                for (lane, &score) in scores.iter().enumerate().take(lanes) {
833                    let index = lane_base + lane;
834                    let doc_id = run_doc_id_validated(bytes, run, index);
835                    if current_doc.is_some_and(|previous| doc_id < previous) {
836                        return Err(invalid_data("flat TQ run is not grouped by document ID"));
837                    }
838                    if current_doc.is_some_and(|previous| doc_id != previous) {
839                        retain_combined_document(
840                            &mut top_documents,
841                            current_doc.expect("current document is present"),
842                            &ordinal_scores,
843                            combiner,
844                        );
845                        ordinal_scores.clear();
846                    }
847                    current_doc = Some(doc_id);
848                    if score.is_finite() {
849                        ordinal_scores.push((
850                            u32::from(read_u16(bytes, run.ordinals.start + index * 2)),
851                            score,
852                        ));
853                    }
854                }
855            }
856            if let Some(doc_id) = current_doc {
857                retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
858                ordinal_scores.clear();
859            }
860        }
861
862        Ok(top_documents
863            .into_sorted_results()
864            .into_iter()
865            .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
866            .collect())
867    }
868
869    /// Score every posting in the probed IVF-TQ leaves, deduplicate SOAR
870    /// assignments by `(doc_id, ordinal)`, combine every approximate ordinal
871    /// present in the probed leaves, and retain at most `k` documents.
872    ///
873    /// Unlike the Max path, this deliberately performs no individual-score
874    /// pruning: an ordinal that cannot enter vector top-k may still change a
875    /// Sum, Avg, LogSumExp, or WeightedTopK document score.
876    #[cfg(test)]
877    pub(crate) fn search_ivf_tq_combined_documents(
878        &self,
879        k: usize,
880        plan: &crate::structures::TqIvfQueryPlan,
881        combiner: crate::query::MultiValueCombiner,
882    ) -> io::Result<Vec<AnnDocumentCandidate>> {
883        self.search_ivf_tq_combined_documents_with_stats(k, plan, combiner)
884            .map(|(candidates, _)| candidates)
885    }
886
887    /// [`Self::search_ivf_tq_combined_documents`] plus the scan diagnostics:
888    /// every probed posting is materialized here, so `posting_count` is the
889    /// buffer this path allocates.
890    pub(crate) fn search_ivf_tq_combined_documents_with_stats(
891        &self,
892        k: usize,
893        plan: &crate::structures::TqIvfQueryPlan,
894        combiner: crate::query::MultiValueCombiner,
895    ) -> io::Result<(Vec<AnnDocumentCandidate>, DenseAnnScanStats)> {
896        self.search_ivf_tq_combined_documents_impl(
897            k,
898            plan,
899            combiner,
900            IVF_PARALLEL_SCAN_MIN_POSTINGS,
901            IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
902        )
903    }
904
905    #[cfg(test)]
906    fn search_ivf_tq_combined_documents_with_tuning(
907        &self,
908        k: usize,
909        plan: &crate::structures::TqIvfQueryPlan,
910        combiner: crate::query::MultiValueCombiner,
911        parallel_min_postings: usize,
912        parallel_chunk_blocks: usize,
913    ) -> io::Result<Vec<AnnDocumentCandidate>> {
914        self.search_ivf_tq_combined_documents_impl(
915            k,
916            plan,
917            combiner,
918            parallel_min_postings,
919            parallel_chunk_blocks,
920        )
921        .map(|(candidates, _)| candidates)
922    }
923
924    fn search_ivf_tq_combined_documents_impl(
925        &self,
926        k: usize,
927        plan: &crate::structures::TqIvfQueryPlan,
928        combiner: crate::query::MultiValueCombiner,
929        parallel_min_postings: usize,
930        parallel_chunk_blocks: usize,
931    ) -> io::Result<(Vec<AnnDocumentCandidate>, DenseAnnScanStats)> {
932        use crate::structures::vector::quantization::{
933            TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
934        };
935
936        validate_combined_search(combiner)?;
937        let tq_plan = plan.tq_plan();
938        self.validate_ivf_tq_query_plan(plan)?;
939        #[cfg(not(feature = "native"))]
940        let _ = (parallel_min_postings, parallel_chunk_blocks);
941        if k == 0 {
942            return Ok((Vec::new(), DenseAnnScanStats::default()));
943        }
944        #[cfg(feature = "native")]
945        self.prefetch_cluster_runs(&plan.cluster_ids);
946
947        let block_bytes = tq_ivf_block_bytes(self.header.code_size);
948        let bytes = self.raw.as_slice();
949        let posting_count = probed_posting_count(self, &plan.cluster_ids)?;
950        let mut stats = DenseAnnScanStats {
951            posting_count,
952            ..DenseAnnScanStats::default()
953        };
954
955        #[cfg(feature = "native")]
956        if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
957            use rayon::prelude::*;
958            let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks);
959            let (ordinal_scores, scored_blocks, non_finite) = tasks
960                .par_iter()
961                .fold(
962                    || {
963                        (
964                            Vec::with_capacity(parallel_chunk_blocks * TQ_BLOCK_LANES),
965                            0usize,
966                            0usize,
967                        )
968                    },
969                    |(mut ordinal_scores, scored, non_finite), task| {
970                        let (task_scored, task_non_finite) = score_ivf_tq_combined_blocks(
971                            bytes,
972                            tq_plan,
973                            block_bytes,
974                            *task,
975                            &mut ordinal_scores,
976                        );
977                        (
978                            ordinal_scores,
979                            scored + task_scored,
980                            non_finite + task_non_finite,
981                        )
982                    },
983                )
984                .reduce(
985                    || (Vec::new(), 0usize, 0usize),
986                    |(mut left, left_scored, left_non_finite),
987                     (mut right, right_scored, right_non_finite)| {
988                        left.append(&mut right);
989                        (
990                            left,
991                            left_scored + right_scored,
992                            left_non_finite + right_non_finite,
993                        )
994                    },
995                );
996            stats.scored_blocks = scored_blocks;
997            stats.non_finite_dropped = non_finite;
998            stats.parallel = true;
999            return Ok((combine_scored_ordinals(ordinal_scores, k, combiner), stats));
1000        }
1001
1002        let mut ordinal_scores = Vec::new();
1003        ordinal_scores
1004            .try_reserve_exact(posting_count)
1005            .map_err(|_| invalid_data("IVF-TQ combined score buffer allocation failed"))?;
1006        let mut scores = [0.0f32; TQ_BLOCK_LANES];
1007        for (cluster_id, cluster_dot) in plan.cluster_dots() {
1008            for run in self.cluster_runs(cluster_id) {
1009                let codes = &bytes[run.codes.clone()];
1010                for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
1011                    tq_score_ivf_block(tq_plan, block, cluster_dot, &mut scores);
1012                    stats.scored_blocks += 1;
1013                    let lane_base = block_index * TQ_BLOCK_LANES;
1014                    let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1015                    for (lane, &score) in scores.iter().enumerate().take(lanes) {
1016                        if !score.is_finite() {
1017                            stats.non_finite_dropped += 1;
1018                            continue;
1019                        }
1020                        let index = lane_base + lane;
1021                        ordinal_scores.push((
1022                            run_doc_id_validated(bytes, run, index),
1023                            read_u16(bytes, run.ordinals.start + index * 2),
1024                            score,
1025                        ));
1026                    }
1027                }
1028            }
1029        }
1030        Ok((combine_scored_ordinals(ordinal_scores, k, combiner), stats))
1031    }
1032
1033    /// Score packed codes in the selected binary-IVF leaves, deduplicate
1034    /// repeated `(doc_id, ordinal)` assignments, combine per document, and
1035    /// retain at most `k` approximate document candidates.
1036    ///
1037    /// The second return value carries the per-ordinal scores behind the
1038    /// retained documents. Binary leaves store the original packed codes, so
1039    /// those scores are *exact*, and reranking can reuse them instead of
1040    /// re-reading the same codes from flat storage.
1041    pub(crate) fn search_binary_combined_documents(
1042        &self,
1043        k: usize,
1044        query: &[u8],
1045        cluster_ids: &[u32],
1046        combiner: crate::query::MultiValueCombiner,
1047    ) -> io::Result<CombinedBinaryCandidates> {
1048        self.search_binary_combined_documents_with_tuning(
1049            k,
1050            query,
1051            cluster_ids,
1052            combiner,
1053            IVF_PARALLEL_SCAN_MIN_POSTINGS,
1054        )
1055    }
1056
1057    /// Score float ScaNN AH codes in the routed leaves and retain approximate
1058    /// document candidates for the shared exact-flat reranker.
1059    pub(crate) fn search_scann_ah_combined_documents(
1060        &self,
1061        k: usize,
1062        query: &crate::structures::vector::scann::FloatScannQuery,
1063        combiner: crate::query::MultiValueCombiner,
1064    ) -> io::Result<Vec<AnnDocumentCandidate>> {
1065        self.search_scann_ah_combined_documents_with_tuning(
1066            k,
1067            query,
1068            combiner,
1069            IVF_PARALLEL_SCAN_MIN_POSTINGS,
1070        )
1071    }
1072
1073    /// Bounded float ScaNN leaf scan.
1074    ///
1075    /// The FastScan lookup table is built once per query (it lives in the
1076    /// shared `FloatScannQuery`), the kernel is resolved once, and scoring
1077    /// never materializes more than the probed posting count. `Max` streams
1078    /// every lane into a document-level top-k collector and uses its k-th
1079    /// best score as an integer lane cutoff (FAISS-style), so rows that cannot
1080    /// enter the top-k are rejected before any float conversion or doc-ID
1081    /// decode. Other combiners need every ordinal of a document and keep the
1082    /// compact `(doc, ordinal, score)` path, now reserved up front and fanned
1083    /// out across Rayon above the same posting threshold as the binary scan.
1084    fn search_scann_ah_combined_documents_with_tuning(
1085        &self,
1086        k: usize,
1087        query: &crate::structures::vector::scann::FloatScannQuery,
1088        combiner: crate::query::MultiValueCombiner,
1089        parallel_min_postings: usize,
1090    ) -> io::Result<Vec<AnnDocumentCandidate>> {
1091        validate_combined_search(combiner)?;
1092        if self.header.kind != AnnKind::ScannAh {
1093            return Err(invalid_data(
1094                "ScaNN AH query used with a different ANN payload",
1095            ));
1096        }
1097        #[cfg(not(feature = "native"))]
1098        let _ = parallel_min_postings;
1099        if k == 0 {
1100            return Ok(Vec::new());
1101        }
1102        #[cfg(feature = "native")]
1103        self.prefetch_cluster_runs(query.routed_leaves());
1104        let blocks = self.header.dim.div_ceil(self.header.code_size);
1105        if query.fast_scan().blocks() != blocks {
1106            return Err(invalid_data(format!(
1107                "ScaNN query has {} AH blocks but the payload stores {blocks}",
1108                query.fast_scan().blocks()
1109            )));
1110        }
1111        let geometry = ScannLeafGeometry {
1112            packed_block_bytes: query.fast_scan().packed_block_bytes(),
1113            packed_row_bytes: blocks.div_ceil(2),
1114        };
1115        let bytes = self.raw.as_slice();
1116        let posting_count = probed_posting_count(self, query.routed_leaves())?;
1117
1118        if matches!(combiner, crate::query::MultiValueCombiner::Max) {
1119            #[cfg(feature = "native")]
1120            if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1121                use rayon::prelude::*;
1122                let tasks = self.scann_scan_tasks(query, SCANN_PARALLEL_SCAN_CHUNK_ROWS);
1123                let collector = tasks
1124                    .par_iter()
1125                    .try_fold(
1126                        || BoundedAnnCollector::<true, true>::new(k),
1127                        |mut collector, task| {
1128                            scan_scann_ah_task(bytes, query, geometry, *task, &mut collector)?;
1129                            Ok::<_, io::Error>(collector)
1130                        },
1131                    )
1132                    .try_reduce(
1133                        || BoundedAnnCollector::<true, true>::new(k),
1134                        |mut left, right| {
1135                            left.merge_from(right);
1136                            Ok(left)
1137                        },
1138                    )?;
1139                return Ok(scann_document_candidates(collector));
1140            }
1141            let mut collector = BoundedAnnCollector::<true, true>::new(k);
1142            for (leaf, centroid_dot) in query.routed() {
1143                for run in self.cluster_runs(leaf) {
1144                    scan_scann_ah_task(
1145                        bytes,
1146                        query,
1147                        geometry,
1148                        ScannScanTask {
1149                            run,
1150                            centroid_dot,
1151                            first_row: 0,
1152                            count: run.count,
1153                        },
1154                        &mut collector,
1155                    )?;
1156                }
1157            }
1158            return Ok(scann_document_candidates(collector));
1159        }
1160
1161        #[cfg(feature = "native")]
1162        if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1163            use rayon::prelude::*;
1164            let tasks = self.scann_scan_tasks(query, SCANN_PARALLEL_SCAN_CHUNK_ROWS);
1165            let ordinal_scores = tasks
1166                .par_iter()
1167                .try_fold(
1168                    || Vec::with_capacity(SCANN_PARALLEL_SCAN_CHUNK_ROWS),
1169                    |mut ordinal_scores, task| {
1170                        scan_scann_ah_task(bytes, query, geometry, *task, &mut ordinal_scores)?;
1171                        Ok::<_, io::Error>(ordinal_scores)
1172                    },
1173                )
1174                .try_reduce(Vec::new, |mut left, mut right| {
1175                    left.append(&mut right);
1176                    Ok(left)
1177                })?;
1178            return Ok(combine_scored_ordinals(ordinal_scores, k, combiner));
1179        }
1180        let mut ordinal_scores = Vec::new();
1181        ordinal_scores
1182            .try_reserve_exact(posting_count)
1183            .map_err(|_| invalid_data("ScaNN combined score buffer allocation failed"))?;
1184        for (leaf, centroid_dot) in query.routed() {
1185            for run in self.cluster_runs(leaf) {
1186                scan_scann_ah_task(
1187                    bytes,
1188                    query,
1189                    geometry,
1190                    ScannScanTask {
1191                        run,
1192                        centroid_dot,
1193                        first_row: 0,
1194                        count: run.count,
1195                    },
1196                    &mut ordinal_scores,
1197                )?;
1198            }
1199        }
1200        Ok(combine_scored_ordinals(ordinal_scores, k, combiner))
1201    }
1202
1203    /// Split the routed ScaNN leaves into FastScan-group-aligned row ranges.
1204    /// The task that reaches a run's end also owns its row-major tail.
1205    #[cfg(feature = "native")]
1206    fn scann_scan_tasks<'a>(
1207        &'a self,
1208        query: &crate::structures::vector::scann::FloatScannQuery,
1209        chunk_rows: usize,
1210    ) -> Vec<ScannScanTask<'a>> {
1211        debug_assert!(chunk_rows.is_multiple_of(crate::structures::vector::scann::FAST_SCAN_LANES));
1212        let mut tasks = Vec::new();
1213        for (leaf, centroid_dot) in query.routed() {
1214            for run in self.cluster_runs(leaf) {
1215                for first_row in (0..run.count).step_by(chunk_rows) {
1216                    tasks.push(ScannScanTask {
1217                        run,
1218                        centroid_dot,
1219                        first_row,
1220                        count: chunk_rows.min(run.count - first_row),
1221                    });
1222                }
1223            }
1224        }
1225        tasks
1226    }
1227
1228    fn search_binary_combined_documents_with_tuning(
1229        &self,
1230        k: usize,
1231        query: &[u8],
1232        cluster_ids: &[u32],
1233        combiner: crate::query::MultiValueCombiner,
1234        parallel_min_postings: usize,
1235    ) -> io::Result<CombinedBinaryCandidates> {
1236        validate_combined_search(combiner)?;
1237        #[cfg(not(feature = "native"))]
1238        let _ = parallel_min_postings;
1239        if query.len() != self.header.code_size {
1240            return Err(io::Error::new(
1241                io::ErrorKind::InvalidInput,
1242                "binary ANN query has the wrong byte length",
1243            ));
1244        }
1245        if k == 0 {
1246            return Ok((Vec::new(), Vec::new()));
1247        }
1248        #[cfg(feature = "native")]
1249        self.prefetch_cluster_runs(cluster_ids);
1250
1251        let bytes = self.raw.as_slice();
1252        let posting_count = probed_posting_count(self, cluster_ids)?;
1253        #[cfg(feature = "native")]
1254        if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1255            use rayon::prelude::*;
1256            let tasks = self.binary_scan_tasks(cluster_ids);
1257            let (ordinal_scores, _) = tasks
1258                .par_iter()
1259                .try_fold(
1260                    || {
1261                        (
1262                            Vec::with_capacity(BINARY_SCORE_BATCH),
1263                            vec![0.0f32; BINARY_SCORE_BATCH],
1264                        )
1265                    },
1266                    |(mut ordinal_scores, mut scores), task| {
1267                        score_binary_task(
1268                            bytes,
1269                            query,
1270                            self.header.dim,
1271                            self.header.code_size,
1272                            *task,
1273                            &mut scores,
1274                            &mut ordinal_scores,
1275                        )?;
1276                        Ok::<_, io::Error>((ordinal_scores, scores))
1277                    },
1278                )
1279                .try_reduce(
1280                    || (Vec::new(), Vec::new()),
1281                    |(mut left, scores), (mut right, _)| {
1282                        left.append(&mut right);
1283                        Ok((left, scores))
1284                    },
1285                )?;
1286            return Ok(combine_scored_ordinals_retaining(
1287                ordinal_scores,
1288                k,
1289                combiner,
1290            ));
1291        }
1292
1293        let mut score_batch = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
1294        let mut ordinal_scores = Vec::new();
1295        ordinal_scores
1296            .try_reserve_exact(posting_count)
1297            .map_err(|_| invalid_data("binary combined score buffer allocation failed"))?;
1298        score_binary_cluster_runs(
1299            self,
1300            bytes,
1301            query,
1302            cluster_ids,
1303            &mut score_batch,
1304            &mut ordinal_scores,
1305        )?;
1306        Ok(combine_scored_ordinals_retaining(
1307            ordinal_scores,
1308            k,
1309            combiner,
1310        ))
1311    }
1312
1313    /// Score every TQ block against the query plan and keep the top `k`
1314    /// distinct documents by estimated similarity.
1315    #[cfg(test)]
1316    pub(crate) fn search_tq_distinct(
1317        &self,
1318        k: usize,
1319        plan: &crate::structures::TqQueryPlan,
1320    ) -> io::Result<Vec<(u32, u16, f32)>> {
1321        self.search_tq_distinct_generic::<BoundedAnnCollector<true, true>>(k, plan)
1322            .map(|(results, _)| results)
1323    }
1324
1325    /// [`Self::search_tq_distinct`] plus scan diagnostics. `unique_keys`
1326    /// selects the heap-only collector, which is only sound when every
1327    /// `(doc_id, ordinal)` occurs once (single-valued field, one posting per
1328    /// stored vector) — the caller checks both against flat metadata.
1329    pub(crate) fn search_tq_distinct_with_stats(
1330        &self,
1331        k: usize,
1332        plan: &crate::structures::TqQueryPlan,
1333        unique_keys: bool,
1334    ) -> io::Result<ScanResults> {
1335        if unique_keys {
1336            self.search_tq_distinct_generic::<BoundedUniqueAnnCollector<true>>(k, plan)
1337        } else {
1338            self.search_tq_distinct_generic::<BoundedAnnCollector<true, true>>(k, plan)
1339        }
1340    }
1341
1342    fn search_tq_distinct_generic<C: IvfTqTopK>(
1343        &self,
1344        k: usize,
1345        plan: &crate::structures::TqQueryPlan,
1346    ) -> io::Result<ScanResults> {
1347        use crate::structures::vector::quantization::{
1348            TQ_BLOCK_LANES, tq_block_bytes, tq_score_block,
1349        };
1350
1351        if plan.padded_dim() != self.header.code_size * 2 {
1352            return Err(io::Error::new(
1353                io::ErrorKind::InvalidInput,
1354                "TQ query plan does not match the payload dimension",
1355            ));
1356        }
1357        let block_bytes = tq_block_bytes(self.header.code_size);
1358        let bytes = self.raw.as_slice();
1359        let mut stats = DenseAnnScanStats {
1360            posting_count: self.header.vector_count,
1361            ..DenseAnnScanStats::default()
1362        };
1363
1364        // Score a block, then offer only lanes that can still enter top-k.
1365        // `insert` re-checks, so a stale threshold is merely conservative;
1366        // NaN fails the comparison and reaches `insert`, which counts it.
1367        let score_lanes = |collector: &mut C,
1368                           run: &AnnRun,
1369                           block_index: usize,
1370                           scores: &[f32; TQ_BLOCK_LANES]| {
1371            let lane_base = block_index * TQ_BLOCK_LANES;
1372            let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1373            let mut threshold = collector.pruning_threshold();
1374            for (lane, &score) in scores.iter().enumerate().take(lanes) {
1375                if threshold.is_some_and(|threshold| score < threshold) {
1376                    continue;
1377                }
1378                let index = lane_base + lane;
1379                collector.insert(
1380                    run_doc_id_validated(bytes, run, index),
1381                    read_u16(bytes, run.ordinals.start + index * 2),
1382                    score,
1383                );
1384                threshold = collector.pruning_threshold();
1385            }
1386        };
1387
1388        // Large flat scans are CPU-bound on the LUT16 kernel. Fold chunks into
1389        // worker-local collectors and reduce them directly: materializing one
1390        // top-k Vec per chunk would make temporary memory O(chunks * k) on
1391        // large segments instead of O(workers * k).
1392        #[cfg(feature = "native")]
1393        if self.header.vector_count >= TQ_PARALLEL_SCAN_MIN_VECTORS {
1394            use rayon::prelude::*;
1395            let (collector, scored_blocks) = self
1396                .runs
1397                .par_iter()
1398                .flat_map(|run| {
1399                    let codes = &bytes[run.codes.clone()];
1400                    let blocks = codes.len() / block_bytes;
1401                    (0..blocks.div_ceil(TQ_PARALLEL_SCAN_CHUNK_BLOCKS))
1402                        .into_par_iter()
1403                        .map(move |chunk| (run, chunk * TQ_PARALLEL_SCAN_CHUNK_BLOCKS, blocks))
1404                })
1405                .fold(
1406                    || (C::with_k(k), 0usize),
1407                    |(mut collector, scored), (run, first_block, total_blocks)| {
1408                        let codes = &bytes[run.codes.clone()];
1409                        let last_block =
1410                            (first_block + TQ_PARALLEL_SCAN_CHUNK_BLOCKS).min(total_blocks);
1411                        let mut scores = [0.0f32; TQ_BLOCK_LANES];
1412                        for block_index in first_block..last_block {
1413                            let block = &codes[block_index * block_bytes..][..block_bytes];
1414                            tq_score_block(plan, block, &mut scores);
1415                            score_lanes(&mut collector, run, block_index, &scores);
1416                        }
1417                        (collector, scored + (last_block - first_block))
1418                    },
1419                )
1420                .reduce(
1421                    || (C::with_k(k), 0usize),
1422                    |(mut collector, left_scored), (partial, right_scored)| {
1423                        collector.merge_from(partial);
1424                        (collector, left_scored + right_scored)
1425                    },
1426                );
1427            stats.scored_blocks = scored_blocks;
1428            stats.non_finite_dropped = collector.non_finite_dropped();
1429            stats.parallel = true;
1430            return Ok((collector.into_sorted_results(), stats));
1431        }
1432
1433        let mut collector = C::with_k(k);
1434        let mut scores = [0.0f32; TQ_BLOCK_LANES];
1435        for run in &self.runs {
1436            let codes = &bytes[run.codes.clone()];
1437            for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
1438                tq_score_block(plan, block, &mut scores);
1439                stats.scored_blocks += 1;
1440                score_lanes(&mut collector, run, block_index, &scores);
1441            }
1442        }
1443        stats.non_finite_dropped = collector.non_finite_dropped();
1444        Ok((collector.into_sorted_results(), stats))
1445    }
1446
1447    /// Score the probed IVF-TQ leaves and keep the top `k` distinct
1448    /// documents by estimated cosine similarity
1449    /// (`⟨q̂,c⟩ + scale·⟨q̂,r̂⟩`).
1450    ///
1451    /// Blocks whose best possible estimate cannot beat the running k-th score
1452    /// terminate their run. The supported cosine generation guarantees that
1453    /// residual scales are stored in descending order.
1454    #[cfg(test)]
1455    pub(crate) fn search_ivf_tq_distinct(
1456        &self,
1457        k: usize,
1458        plan: &crate::structures::TqIvfQueryPlan,
1459    ) -> io::Result<Vec<(u32, u16, f32)>> {
1460        self.search_ivf_tq_distinct_generic::<BoundedAnnCollector<true, true>>(
1461            k,
1462            plan,
1463            IVF_PARALLEL_SCAN_MIN_POSTINGS,
1464            IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
1465        )
1466        .map(|(results, _)| results)
1467    }
1468
1469    /// [`Self::search_ivf_tq_distinct`] plus scan diagnostics. `unique_keys`
1470    /// selects the heap-only collector; it is only sound when the field is
1471    /// single-valued *and* the payload holds exactly one posting per stored
1472    /// vector (no SOAR spill) — the caller checks both against flat metadata.
1473    pub(crate) fn search_ivf_tq_distinct_with_stats(
1474        &self,
1475        k: usize,
1476        plan: &crate::structures::TqIvfQueryPlan,
1477        unique_keys: bool,
1478    ) -> io::Result<ScanResults> {
1479        if unique_keys {
1480            self.search_ivf_tq_distinct_generic::<BoundedUniqueAnnCollector<true>>(
1481                k,
1482                plan,
1483                IVF_PARALLEL_SCAN_MIN_POSTINGS,
1484                IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
1485            )
1486        } else {
1487            self.search_ivf_tq_distinct_generic::<BoundedAnnCollector<true, true>>(
1488                k,
1489                plan,
1490                IVF_PARALLEL_SCAN_MIN_POSTINGS,
1491                IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
1492            )
1493        }
1494    }
1495
1496    #[cfg(test)]
1497    fn search_ivf_tq_distinct_with_tuning(
1498        &self,
1499        k: usize,
1500        plan: &crate::structures::TqIvfQueryPlan,
1501        parallel_min_postings: usize,
1502        parallel_chunk_blocks: usize,
1503    ) -> io::Result<Vec<(u32, u16, f32)>> {
1504        self.search_ivf_tq_distinct_generic::<BoundedAnnCollector<true, true>>(
1505            k,
1506            plan,
1507            parallel_min_postings,
1508            parallel_chunk_blocks,
1509        )
1510        .map(|(results, _)| results)
1511    }
1512
1513    fn search_ivf_tq_distinct_generic<C: IvfTqTopK>(
1514        &self,
1515        k: usize,
1516        plan: &crate::structures::TqIvfQueryPlan,
1517        parallel_min_postings: usize,
1518        parallel_chunk_blocks: usize,
1519    ) -> io::Result<ScanResults> {
1520        use crate::structures::vector::quantization::tq_ivf_block_bytes;
1521
1522        let tq_plan = plan.tq_plan();
1523        self.validate_ivf_tq_query_plan(plan)?;
1524        #[cfg(not(feature = "native"))]
1525        let _ = (parallel_min_postings, parallel_chunk_blocks);
1526        if k == 0 {
1527            return Ok((Vec::new(), DenseAnnScanStats::default()));
1528        }
1529        #[cfg(feature = "native")]
1530        self.prefetch_cluster_runs(&plan.cluster_ids);
1531        let block_bytes = tq_ivf_block_bytes(self.header.code_size);
1532        let bytes = self.raw.as_slice();
1533        let posting_count = probed_posting_count(self, &plan.cluster_ids)?;
1534        let mut stats = DenseAnnScanStats {
1535            posting_count,
1536            ..DenseAnnScanStats::default()
1537        };
1538
1539        // Score one chunk first so every worker starts with a top-k backed by
1540        // real documents. Its k-th score is therefore a safe pruning floor.
1541        // Chunking also exposes parallelism when a skewed leaf is one large
1542        // physical run. Nested Rayon stays on the caller's bounded search pool.
1543        #[cfg(feature = "native")]
1544        if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1545            use rayon::prelude::*;
1546            let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks);
1547
1548            if let Some((pilot, remaining)) = tasks.split_first() {
1549                let mut pilot_collector = C::with_k(k);
1550                let (pilot_pruned, pilot_scored) =
1551                    score_ivf_tq_blocks(bytes, tq_plan, block_bytes, *pilot, &mut pilot_collector);
1552                let pilot_non_finite = pilot_collector.non_finite_dropped();
1553                let seed = pilot_collector.into_sorted_results();
1554                let seeded_collector = || {
1555                    let mut collector = C::with_k(k);
1556                    for &(doc_id, ordinal, score) in &seed {
1557                        collector.insert(doc_id, ordinal, score);
1558                    }
1559                    collector
1560                };
1561                let (collector, pruned_blocks, scored_blocks) = remaining
1562                    .par_iter()
1563                    .fold(
1564                        || (seeded_collector(), 0usize, 0usize),
1565                        |(mut collector, pruned, scored), task| {
1566                            let (task_pruned, task_scored) = score_ivf_tq_blocks(
1567                                bytes,
1568                                tq_plan,
1569                                block_bytes,
1570                                *task,
1571                                &mut collector,
1572                            );
1573                            (collector, pruned + task_pruned, scored + task_scored)
1574                        },
1575                    )
1576                    .reduce(
1577                        || (seeded_collector(), 0usize, 0usize),
1578                        |(mut collector, left_pruned, left_scored),
1579                         (partial, right_pruned, right_scored)| {
1580                            collector.merge_from(partial);
1581                            (
1582                                collector,
1583                                left_pruned + right_pruned,
1584                                left_scored + right_scored,
1585                            )
1586                        },
1587                    );
1588                stats.pruned_blocks = pilot_pruned + pruned_blocks;
1589                stats.scored_blocks = pilot_scored + scored_blocks;
1590                stats.non_finite_dropped = pilot_non_finite + collector.non_finite_dropped();
1591                stats.parallel = true;
1592                log_ivf_tq_pruning(stats.pruned_blocks, stats.scored_blocks, true);
1593                return Ok((collector.into_sorted_results(), stats));
1594            }
1595        }
1596
1597        let mut collector = C::with_k(k);
1598        for (cluster_id, cluster_dot) in plan.cluster_dots() {
1599            for run in self.cluster_runs(cluster_id) {
1600                let block_count = run.codes.len() / block_bytes;
1601                let (run_pruned, run_scored) = score_ivf_tq_blocks(
1602                    bytes,
1603                    tq_plan,
1604                    block_bytes,
1605                    IvfTqScanTask {
1606                        run,
1607                        cluster_dot,
1608                        first_block: 0,
1609                        last_block: block_count,
1610                    },
1611                    &mut collector,
1612                );
1613                stats.pruned_blocks += run_pruned;
1614                stats.scored_blocks += run_scored;
1615            }
1616        }
1617        stats.non_finite_dropped = collector.non_finite_dropped();
1618        log_ivf_tq_pruning(stats.pruned_blocks, stats.scored_blocks, false);
1619        Ok((collector.into_sorted_results(), stats))
1620    }
1621
1622    fn validate_ivf_tq_query_plan(
1623        &self,
1624        plan: &crate::structures::TqIvfQueryPlan,
1625    ) -> io::Result<()> {
1626        if self.header.kind != AnnKind::IvfTq
1627            || !crate::structures::is_ivf_tq_cosine_generation(self.header.quantizer_version)
1628        {
1629            return Err(invalid_data(
1630                "legacy raw IVF-TQ payloads cannot be searched; rebuild the index",
1631            ));
1632        }
1633        if plan.tq_plan().padded_dim() != self.header.code_size * 2
1634            || plan.quantizer_version != self.header.quantizer_version
1635            || plan.fingerprint != self.header.codebook_version
1636        {
1637            return Err(io::Error::new(
1638                io::ErrorKind::InvalidInput,
1639                "IVF-TQ query plan does not match the payload generation",
1640            ));
1641        }
1642        Ok(())
1643    }
1644
1645    pub(crate) fn search_binary_clusters<const BY_DOCUMENT: bool>(
1646        &self,
1647        query: &[u8],
1648        k: usize,
1649        cluster_ids: &[u32],
1650    ) -> io::Result<Vec<(u32, u16, f32)>> {
1651        self.search_binary_clusters_with_tuning::<BY_DOCUMENT>(
1652            query,
1653            k,
1654            cluster_ids,
1655            IVF_PARALLEL_SCAN_MIN_POSTINGS,
1656        )
1657    }
1658
1659    fn search_binary_clusters_with_tuning<const BY_DOCUMENT: bool>(
1660        &self,
1661        query: &[u8],
1662        k: usize,
1663        cluster_ids: &[u32],
1664        parallel_min_postings: usize,
1665    ) -> io::Result<Vec<(u32, u16, f32)>> {
1666        #[cfg(not(feature = "native"))]
1667        let _ = parallel_min_postings;
1668        if query.len() != self.header.code_size {
1669            return Err(io::Error::new(
1670                io::ErrorKind::InvalidInput,
1671                "binary ANN query has the wrong byte length",
1672            ));
1673        }
1674        if k == 0 {
1675            return Ok(Vec::new());
1676        }
1677        #[cfg(feature = "native")]
1678        self.prefetch_cluster_runs(cluster_ids);
1679        let bytes = self.raw.as_slice();
1680
1681        // A probe plan returns each leaf once. IVF has one posting per vector;
1682        // binary ScaNN may intentionally spill a vector into a second leaf.
1683        debug_assert!(
1684            BY_DOCUMENT || self.header.kind == AnnKind::ScannBinary || {
1685                let mut seen = rustc_hash::FxHashSet::default();
1686                cluster_ids
1687                    .iter()
1688                    .all(|cluster_id| seen.insert(*cluster_id))
1689            },
1690            "an IVF probe plan must not repeat a cluster",
1691        );
1692
1693        #[cfg(feature = "native")]
1694        if rayon::current_num_threads() > 1
1695            && probed_posting_count(self, cluster_ids)? >= parallel_min_postings
1696        {
1697            use rayon::prelude::*;
1698            let tasks = self.binary_scan_tasks(cluster_ids);
1699            let (collector, _) = tasks
1700                .par_iter()
1701                .try_fold(
1702                    || {
1703                        (
1704                            BoundedAnnCollector::<BY_DOCUMENT, true>::new(k),
1705                            vec![0.0f32; BINARY_SCORE_BATCH],
1706                        )
1707                    },
1708                    |(mut collector, mut scores), task| {
1709                        score_binary_task(
1710                            bytes,
1711                            query,
1712                            self.header.dim,
1713                            self.header.code_size,
1714                            *task,
1715                            &mut scores,
1716                            &mut collector,
1717                        )?;
1718                        Ok::<_, io::Error>((collector, scores))
1719                    },
1720                )
1721                .try_reduce(
1722                    || (BoundedAnnCollector::<BY_DOCUMENT, true>::new(k), Vec::new()),
1723                    |(mut collector, scores), (partial, _)| {
1724                        collector.merge_from(partial);
1725                        Ok((collector, scores))
1726                    },
1727                )?;
1728            return Ok(collector.into_sorted_results());
1729        }
1730
1731        let mut scores = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
1732        if BY_DOCUMENT || self.header.kind == AnnKind::ScannBinary {
1733            let mut collector = BoundedAnnCollector::<BY_DOCUMENT, true>::new(k);
1734            score_binary_cluster_runs(
1735                self,
1736                bytes,
1737                query,
1738                cluster_ids,
1739                &mut scores,
1740                &mut collector,
1741            )?;
1742            return Ok(collector.into_sorted_results());
1743        }
1744
1745        // Avoid the deduplication hash map in the serial single-value path.
1746        let mut collector = BoundedUniqueAnnCollector::<true>::new(k);
1747        score_binary_cluster_runs(self, bytes, query, cluster_ids, &mut scores, &mut collector)?;
1748        Ok(collector.into_sorted_results())
1749    }
1750}
1751
1752/// Score one contiguous block range from an IVF-TQ run. A task may start in
1753/// the middle of a run because residual scales descend across the complete
1754/// run: the first block is still an upper bound for the remainder of the task.
1755///
1756/// Returns `(pruned_blocks, scored_blocks)`.
1757fn score_ivf_tq_blocks<C: IvfTqTopK>(
1758    bytes: &[u8],
1759    plan: &crate::structures::TqQueryPlan,
1760    block_bytes: usize,
1761    task: IvfTqScanTask<'_>,
1762    collector: &mut C,
1763) -> (usize, usize) {
1764    use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block};
1765
1766    let run = task.run;
1767    let codes = &bytes[run.codes.clone()];
1768    let mut scores = [0.0f32; TQ_BLOCK_LANES];
1769    let mut scored_blocks = 0usize;
1770    let mut threshold = collector.pruning_threshold();
1771    for block_index in task.first_block..task.last_block {
1772        let block = &codes[block_index * block_bytes..][..block_bytes];
1773        if let Some(threshold) = threshold
1774            && task.cluster_dot + tq_ivf_block_max_scale(block) * TQ_PRUNE_ESTIMATE_BOUND
1775                <= threshold
1776        {
1777            return (task.last_block - block_index, scored_blocks);
1778        }
1779        scored_blocks += 1;
1780        tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores);
1781        let lane_base = block_index * TQ_BLOCK_LANES;
1782        let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1783        for (lane, &score) in scores.iter().enumerate().take(lanes) {
1784            // Reject non-competitive lanes before touching the doc-ID and
1785            // ordinal columns. Strict `<` leaves exact ties to `insert`, which
1786            // breaks them by key exactly as before; NaN fails the comparison
1787            // and reaches `insert`, which counts it.
1788            if threshold.is_some_and(|threshold| score < threshold) {
1789                continue;
1790            }
1791            let index = lane_base + lane;
1792            collector.insert(
1793                run_doc_id_validated(bytes, run, index),
1794                read_u16(bytes, run.ordinals.start + index * 2),
1795                score,
1796            );
1797            threshold = collector.pruning_threshold();
1798        }
1799    }
1800    (0, scored_blocks)
1801}
1802
1803/// Returns `(scored_blocks, non_finite_dropped)`.
1804#[cfg(feature = "native")]
1805fn score_ivf_tq_combined_blocks(
1806    bytes: &[u8],
1807    plan: &crate::structures::TqQueryPlan,
1808    block_bytes: usize,
1809    task: IvfTqScanTask<'_>,
1810    ordinal_scores: &mut Vec<(u32, u16, f32)>,
1811) -> (usize, usize) {
1812    use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block};
1813
1814    let run = task.run;
1815    let codes = &bytes[run.codes.clone()];
1816    let mut scores = [0.0f32; TQ_BLOCK_LANES];
1817    let mut non_finite = 0usize;
1818    for block_index in task.first_block..task.last_block {
1819        let block = &codes[block_index * block_bytes..][..block_bytes];
1820        tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores);
1821        let lane_base = block_index * TQ_BLOCK_LANES;
1822        let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1823        for (lane, &score) in scores.iter().enumerate().take(lanes) {
1824            if !score.is_finite() {
1825                non_finite += 1;
1826                continue;
1827            }
1828            let index = lane_base + lane;
1829            ordinal_scores.push((
1830                run_doc_id_validated(bytes, run, index),
1831                read_u16(bytes, run.ordinals.start + index * 2),
1832                score,
1833            ));
1834        }
1835    }
1836    (task.last_block - task.first_block, non_finite)
1837}
1838
1839fn log_ivf_tq_pruning(pruned_blocks: usize, scored_blocks: usize, parallel: bool) {
1840    if pruned_blocks > 0 {
1841        log::debug!(
1842            "[search_ivf_tq] pruned {pruned_blocks} of {} blocks via scale bounds ({})",
1843            pruned_blocks + scored_blocks,
1844            if parallel { "parallel" } else { "serial" },
1845        );
1846    }
1847}
1848
1849#[cfg(feature = "native")]
1850fn coalesce_prefetch_ranges(ranges: &mut Vec<Range<usize>>) {
1851    if ranges.len() < 2 {
1852        return;
1853    }
1854    ranges.sort_unstable_by_key(|range| range.start);
1855    let mut output_len = 1usize;
1856    for input_index in 1..ranges.len() {
1857        let next_start = ranges[input_index].start;
1858        let next_end = ranges[input_index].end;
1859        let previous = &mut ranges[output_len - 1];
1860        if next_start <= previous.end.saturating_add(PREFETCH_COALESCE_GAP) {
1861            previous.end = previous.end.max(next_end);
1862        } else {
1863            ranges[output_len] = next_start..next_end;
1864            output_len += 1;
1865        }
1866    }
1867    ranges.truncate(output_len);
1868}
1869
1870trait AnnScoreSink {
1871    fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32);
1872}
1873
1874impl<const BY_DOCUMENT: bool> AnnScoreSink for BoundedAnnCollector<BY_DOCUMENT, true> {
1875    #[inline]
1876    fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1877        self.insert(doc_id, ordinal, score);
1878    }
1879}
1880
1881impl AnnScoreSink for BoundedUniqueAnnCollector<true> {
1882    #[inline]
1883    fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1884        self.insert(doc_id, ordinal, score);
1885    }
1886}
1887
1888impl AnnScoreSink for Vec<(u32, u16, f32)> {
1889    #[inline]
1890    fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1891        if score.is_finite() {
1892            self.push((doc_id, ordinal, score));
1893        }
1894    }
1895}
1896
1897fn probed_posting_count(index: &AnnDiskIndex, cluster_ids: &[u32]) -> io::Result<usize> {
1898    cluster_ids.iter().try_fold(0usize, |count, &cluster_id| {
1899        index
1900            .cluster_runs(cluster_id)
1901            .iter()
1902            .try_fold(count, |count, run| {
1903                count
1904                    .checked_add(run.count)
1905                    .ok_or_else(|| invalid_data("ANN probed posting count overflows usize"))
1906            })
1907    })
1908}
1909
1910fn score_binary_cluster_runs(
1911    index: &AnnDiskIndex,
1912    bytes: &[u8],
1913    query: &[u8],
1914    cluster_ids: &[u32],
1915    scores: &mut [f32],
1916    collector: &mut impl AnnScoreSink,
1917) -> io::Result<()> {
1918    for &cluster_id in cluster_ids {
1919        for run in index.cluster_runs(cluster_id) {
1920            score_binary_run(
1921                bytes,
1922                run,
1923                query,
1924                index.header.dim,
1925                index.header.code_size,
1926                scores,
1927                collector,
1928            )?;
1929        }
1930    }
1931    Ok(())
1932}
1933
1934fn score_binary_run(
1935    bytes: &[u8],
1936    run: &AnnRun,
1937    query: &[u8],
1938    dim_bits: usize,
1939    code_size: usize,
1940    scores: &mut [f32],
1941    collector: &mut impl AnnScoreSink,
1942) -> io::Result<()> {
1943    for batch_start in (0..run.count).step_by(BINARY_SCORE_BATCH) {
1944        score_binary_task(
1945            bytes,
1946            query,
1947            dim_bits,
1948            code_size,
1949            BinaryScanTask {
1950                run,
1951                first_index: batch_start,
1952                count: BINARY_SCORE_BATCH.min(run.count - batch_start),
1953            },
1954            scores,
1955            collector,
1956        )?;
1957    }
1958    Ok(())
1959}
1960
1961fn score_binary_task(
1962    bytes: &[u8],
1963    query: &[u8],
1964    dim_bits: usize,
1965    code_size: usize,
1966    task: BinaryScanTask<'_>,
1967    scores: &mut [f32],
1968    collector: &mut impl AnnScoreSink,
1969) -> io::Result<()> {
1970    let code_start = task.run.codes.start + task.first_index * code_size;
1971    let code_end = code_start + task.count * code_size;
1972    crate::structures::simd::batch_hamming_scores(
1973        query,
1974        &bytes[code_start..code_end],
1975        code_size,
1976        dim_bits,
1977        &mut scores[..task.count],
1978    );
1979    for (batch_index, &score) in scores.iter().enumerate().take(task.count) {
1980        let index = task.first_index + batch_index;
1981        collector.insert_score(
1982            run_doc_id_validated(bytes, task.run, index),
1983            read_u16(bytes, task.run.ordinals.start + index * 2),
1984            score,
1985        );
1986    }
1987    Ok(())
1988}
1989
1990#[inline]
1991fn retain_combined_document(
1992    collector: &mut BoundedAnnCollector<true, true>,
1993    doc_id: u32,
1994    ordinal_scores: &[(u32, f32)],
1995    combiner: crate::query::MultiValueCombiner,
1996) {
1997    if !ordinal_scores.is_empty() {
1998        collector.insert(doc_id, 0, combiner.combine(ordinal_scores));
1999    }
2000}
2001
2002fn validate_combined_search(combiner: crate::query::MultiValueCombiner) -> io::Result<()> {
2003    combiner
2004        .validate()
2005        .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))
2006}
2007
2008/// Total non-finite leaf scores dropped by `combine_scored_ordinals_retaining`
2009/// (also exported as `hermes_ann_non_finite_scores_dropped_total`).
2010pub(crate) static NON_FINITE_ANN_SCORES_DROPPED: std::sync::atomic::AtomicU64 =
2011    std::sync::atomic::AtomicU64::new(0);
2012static NON_FINITE_ANN_SCORE_DROP_EVENTS: std::sync::atomic::AtomicU64 =
2013    std::sync::atomic::AtomicU64::new(0);
2014
2015/// Rows per parallel ScaNN scan task: 512 FastScan groups, matching the
2016/// IVF-TQ chunking so worker-local top-k state stays small.
2017#[cfg(feature = "native")]
2018const SCANN_PARALLEL_SCAN_CHUNK_ROWS: usize =
2019    512 * crate::structures::vector::scann::FAST_SCAN_LANES;
2020
2021/// Byte geometry of one ScaNN AH run under a fixed block count.
2022#[derive(Clone, Copy)]
2023struct ScannLeafGeometry {
2024    /// Bytes per complete 32-row FastScan group.
2025    packed_block_bytes: usize,
2026    /// Bytes per row-major tail row.
2027    packed_row_bytes: usize,
2028}
2029
2030/// One FastScan-aligned row range of a routed ScaNN leaf run.
2031#[derive(Clone, Copy)]
2032struct ScannScanTask<'a> {
2033    run: &'a AnnRun,
2034    centroid_dot: f32,
2035    first_row: usize,
2036    count: usize,
2037}
2038
2039/// Score sinks that can expose their current k-th best score so FastScan can
2040/// reject integer lane sums before converting them.
2041trait ScannPruneThreshold {
2042    fn prune_threshold(&self) -> Option<f32>;
2043}
2044
2045impl ScannPruneThreshold for BoundedAnnCollector<true, true> {
2046    #[inline]
2047    fn prune_threshold(&self) -> Option<f32> {
2048        self.pruning_threshold()
2049    }
2050}
2051
2052impl ScannPruneThreshold for Vec<(u32, u16, f32)> {
2053    #[inline]
2054    fn prune_threshold(&self) -> Option<f32> {
2055        // Every ordinal may change a Sum/Avg/LogSumExp/WeightedTopK document
2056        // score, so nothing can be pruned on this path.
2057        None
2058    }
2059}
2060
2061/// Score one row range of a ScaNN AH run. Complete 32-row groups go through
2062/// the integer FastScan kernel; the run's row-major tail (only reached by the
2063/// task that ends at `run.count`) uses the float packed-row scorer.
2064fn scan_scann_ah_task(
2065    bytes: &[u8],
2066    query: &crate::structures::vector::scann::FloatScannQuery,
2067    geometry: ScannLeafGeometry,
2068    task: ScannScanTask<'_>,
2069    sink: &mut (impl AnnScoreSink + ScannPruneThreshold),
2070) -> io::Result<()> {
2071    use crate::structures::vector::scann::FAST_SCAN_LANES;
2072
2073    let run = task.run;
2074    let codes = &bytes[run.codes.clone()];
2075    let fast = query.fast_scan();
2076    let full_rows = run.count / FAST_SCAN_LANES * FAST_SCAN_LANES;
2077    let end = task.first_row + task.count;
2078    debug_assert!(task.first_row.is_multiple_of(FAST_SCAN_LANES) || task.first_row >= full_rows);
2079    let mut lanes = [0i32; FAST_SCAN_LANES];
2080    let mut row = task.first_row;
2081    while row < full_rows.min(end) {
2082        let start = (row / FAST_SCAN_LANES) * geometry.packed_block_bytes;
2083        fast.accumulate_block(
2084            &codes[start..start + geometry.packed_block_bytes],
2085            &mut lanes,
2086        )
2087        .map_err(|error| invalid_data(error.to_string()))?;
2088        let cutoff = sink
2089            .prune_threshold()
2090            .map(|threshold| fast.lane_cutoff(threshold, task.centroid_dot));
2091        for (lane, &sum) in lanes.iter().enumerate() {
2092            if cutoff.is_some_and(|cutoff| sum <= cutoff) {
2093                continue;
2094            }
2095            let index = row + lane;
2096            sink.insert_score(
2097                run_doc_id(bytes, run, index)?,
2098                read_u16(bytes, run.ordinals.start + index * 2),
2099                fast.lane_score(sum, task.centroid_dot),
2100            );
2101        }
2102        row += FAST_SCAN_LANES;
2103    }
2104    let tail_bytes = (full_rows / FAST_SCAN_LANES) * geometry.packed_block_bytes;
2105    while row < end {
2106        let start = tail_bytes + (row - full_rows) * geometry.packed_row_bytes;
2107        let score = query
2108            .ah_query()
2109            .score_packed(
2110                &codes[start..start + geometry.packed_row_bytes],
2111                task.centroid_dot,
2112            )
2113            .map_err(|error| invalid_data(error.to_string()))?;
2114        sink.insert_score(
2115            run_doc_id(bytes, run, row)?,
2116            read_u16(bytes, run.ordinals.start + row * 2),
2117            score,
2118        );
2119        row += 1;
2120    }
2121    Ok(())
2122}
2123
2124fn scann_document_candidates(
2125    collector: BoundedAnnCollector<true, true>,
2126) -> Vec<AnnDocumentCandidate> {
2127    collector
2128        .into_sorted_results()
2129        .into_iter()
2130        .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
2131        .collect()
2132}
2133
2134/// Sort arbitrary probed-run output into complete documents, retain the best
2135/// estimate for every duplicated `(doc_id, ordinal)` SOAR assignment, then
2136/// apply the requested combiner. The corpus-dependent scratch is one compact
2137/// tuple per probed posting and never expands to unprobed leaves or raw
2138/// vectors; final retained state is O(k).
2139fn combine_scored_ordinals(
2140    scores: Vec<(u32, u16, f32)>,
2141    k: usize,
2142    combiner: crate::query::MultiValueCombiner,
2143) -> Vec<AnnDocumentCandidate> {
2144    combine_scored_ordinals_retaining(scores, k, combiner).0
2145}
2146
2147/// As [`combine_scored_ordinals`], additionally returning the deduplicated
2148/// per-ordinal scores of the retained documents, sorted by `(doc_id, ordinal)`.
2149///
2150/// Only callers whose leaf scores are exact may reuse those numbers; TQ block
2151/// scores are estimates and must still be reranked against raw vectors.
2152fn combine_scored_ordinals_retaining(
2153    mut scores: Vec<(u32, u16, f32)>,
2154    k: usize,
2155    combiner: crate::query::MultiValueCombiner,
2156) -> CombinedBinaryCandidates {
2157    if k == 0 || scores.is_empty() {
2158        return (Vec::new(), Vec::new());
2159    }
2160    let scored = scores.len();
2161    scores.retain(|entry| entry.2.is_finite());
2162    let dropped = scored - scores.len();
2163    if dropped > 0 {
2164        // A non-finite leaf score means a degenerate stored vector or
2165        // codebook; dropping it is the right ranking decision but never a
2166        // silent one.
2167        NON_FINITE_ANN_SCORES_DROPPED
2168            .fetch_add(dropped as u64, std::sync::atomic::Ordering::Relaxed);
2169        crate::observe::ann_non_finite_scores_dropped(dropped as u64);
2170        crate::structures::vector::scann::warn_rate_limited(
2171            &NON_FINITE_ANN_SCORE_DROP_EVENTS,
2172            |events| {
2173                format!(
2174                    "[ann] dropped {dropped} of {scored} non-finite leaf scores before combining \
2175                     ({events} scans affected so far; total dropped {})",
2176                    NON_FINITE_ANN_SCORES_DROPPED.load(std::sync::atomic::Ordering::Relaxed)
2177                )
2178            },
2179        );
2180    }
2181    let by_document_ordinal = |left: &(u32, u16, f32), right: &(u32, u16, f32)| {
2182        left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
2183    };
2184    #[cfg(feature = "native")]
2185    if rayon::current_num_threads() > 1 && scores.len() >= IVF_PARALLEL_SCAN_MIN_POSTINGS {
2186        use rayon::prelude::*;
2187        scores.par_sort_unstable_by(by_document_ordinal);
2188    } else {
2189        scores.sort_unstable_by(by_document_ordinal);
2190    }
2191    #[cfg(not(feature = "native"))]
2192    scores.sort_unstable_by(by_document_ordinal);
2193
2194    // Compact duplicates in place. IVF-TQ SOAR copies can have slightly
2195    // different residual estimates; preserving the highest one matches the
2196    // legacy Max collector's duplicate semantics without double-counting it
2197    // for additive combiners.
2198    let mut unique_len = 0usize;
2199    for read_index in 0..scores.len() {
2200        let candidate = scores[read_index];
2201        if unique_len > 0
2202            && scores[unique_len - 1].0 == candidate.0
2203            && scores[unique_len - 1].1 == candidate.1
2204        {
2205            if candidate.2.total_cmp(&scores[unique_len - 1].2).is_gt() {
2206                scores[unique_len - 1].2 = candidate.2;
2207            }
2208        } else {
2209            scores[unique_len] = candidate;
2210            unique_len += 1;
2211        }
2212    }
2213    scores.truncate(unique_len);
2214
2215    let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
2216    let mut current_doc = None;
2217    let mut ordinal_scores = Vec::new();
2218    for &(doc_id, ordinal, score) in &scores {
2219        if current_doc.is_some_and(|current| current != doc_id) {
2220            retain_combined_document(
2221                &mut top_documents,
2222                current_doc.expect("current document is present"),
2223                &ordinal_scores,
2224                combiner,
2225            );
2226            ordinal_scores.clear();
2227        }
2228        current_doc = Some(doc_id);
2229        ordinal_scores.push((u32::from(ordinal), score));
2230    }
2231    if let Some(doc_id) = current_doc {
2232        retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
2233    }
2234
2235    let candidates: Vec<AnnDocumentCandidate> = top_documents
2236        .into_sorted_results()
2237        .into_iter()
2238        .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
2239        .collect();
2240
2241    // Narrow the per-ordinal scores to the retained documents. `scores` is
2242    // already sorted by document, so this is one pass over a sorted ID set
2243    // rather than a hash lookup per posting.
2244    let mut retained_ids: Vec<u32> = candidates
2245        .iter()
2246        .map(|candidate| candidate.doc_id)
2247        .collect();
2248    retained_ids.sort_unstable();
2249    let mut retained = Vec::with_capacity(scores.len().min(retained_ids.len().saturating_mul(4)));
2250    let mut cursor = 0usize;
2251    for entry in scores {
2252        while cursor < retained_ids.len() && retained_ids[cursor] < entry.0 {
2253            cursor += 1;
2254        }
2255        if retained_ids.get(cursor) == Some(&entry.0) {
2256            retained.push(entry);
2257        }
2258    }
2259    (candidates, retained)
2260}
2261
2262#[cfg(feature = "native")]
2263pub(crate) fn write_built_binary_ivf(
2264    index: &BinaryIvfIndex,
2265    routing: IvfRoutingMode,
2266    writer: &mut (impl Write + ?Sized),
2267) -> io::Result<u64> {
2268    let runs: Vec<_> = index
2269        .clusters
2270        .iter()
2271        .map(|(cluster_id, cluster)| BuildRun {
2272            cluster_id: *cluster_id,
2273            doc_ids: &cluster.doc_ids,
2274            ordinals: &cluster.ordinals,
2275            codes: &cluster.codes,
2276        })
2277        .collect();
2278    write_built_runs(
2279        AnnDiskHeader {
2280            kind: AnnKind::BinaryIvf,
2281            routing,
2282            dim: index.dim_bits,
2283            code_size: index.dim_bits.div_ceil(8),
2284            num_clusters: index.num_clusters,
2285            quantizer_version: index.quantizer_version,
2286            codebook_version: 0,
2287            vector_count: index.len(),
2288        },
2289        &runs,
2290        writer,
2291    )
2292}
2293
2294/// Serialize ScaNN segment-local leaf runs into Hermes's merge-native ANN
2295/// format. The global model is not embedded: its generation and content
2296/// fingerprint occupy the header's compatibility slots, so an ordinary merge
2297/// can reject mixed models before copying any corpus bytes.
2298#[cfg(feature = "native")]
2299#[cfg_attr(not(test), allow(dead_code))]
2300pub(crate) fn write_built_scann(
2301    payload: &crate::structures::vector::scann::ScannSegmentPayload,
2302    writer: &mut (impl Write + ?Sized),
2303) -> io::Result<u64> {
2304    use crate::structures::vector::scann::ScannEncoding;
2305
2306    let (kind, code_size) = match payload.encoding {
2307        ScannEncoding::AsymmetricHash {
2308            dimensions_per_block,
2309            bits_per_code: 4,
2310        } => (AnnKind::ScannAh, usize::from(dimensions_per_block)),
2311        ScannEncoding::AsymmetricHash { .. } => {
2312            return Err(invalid_data(
2313                "ScaNN ANN disk format supports only 4-bit AH codes",
2314            ));
2315        }
2316        ScannEncoding::BinaryHamming => (
2317            AnnKind::ScannBinary,
2318            usize::try_from(payload.dimension)
2319                .map_err(|_| invalid_data("ScaNN dimension exceeds usize"))?
2320                .div_ceil(8),
2321        ),
2322    };
2323    let vector_count = payload.runs().iter().try_fold(0usize, |total, run| {
2324        total
2325            .checked_add(run.row_count as usize)
2326            .ok_or_else(|| invalid_data("ScaNN vector count overflows usize"))
2327    })?;
2328    let header = AnnDiskHeader {
2329        kind,
2330        // ScaNN owns its multi-level routing geometry in the global artifact;
2331        // Flat is the reserved sentinel in the shared ANN header.
2332        routing: IvfRoutingMode::Flat,
2333        dim: usize::try_from(payload.dimension)
2334            .map_err(|_| invalid_data("ScaNN dimension exceeds usize"))?,
2335        code_size,
2336        num_clusters: payload.num_leaves,
2337        quantizer_version: payload.generation,
2338        codebook_version: payload.artifact_id,
2339        vector_count,
2340    };
2341    validate_header(&header)?;
2342    write_header(writer, &header)?;
2343
2344    let mut offset = ANN_HEADER_SIZE as u64;
2345    let mut records = Vec::with_capacity(payload.runs().len());
2346    let mut previous_leaf = None;
2347    for run in payload.runs() {
2348        let count = run.row_count as usize;
2349        let expected_docs = count
2350            .checked_mul(4)
2351            .ok_or_else(|| invalid_data("ScaNN doc-ID column size overflows usize"))?;
2352        let expected_ordinals = count
2353            .checked_mul(2)
2354            .ok_or_else(|| invalid_data("ScaNN ordinal column size overflows usize"))?;
2355        if count == 0
2356            || run.leaf_id >= header.num_clusters
2357            || previous_leaf.is_some_and(|leaf| leaf > run.leaf_id)
2358            || run.doc_ids_le.len() != expected_docs
2359            || run.ordinals_le.len() != expected_ordinals
2360            || run.codes.len()
2361                != expected_codes_column_len(header.kind, count, header.dim, header.code_size)?
2362        {
2363            return Err(invalid_data("ScaNN ANN leaf run columns are inconsistent"));
2364        }
2365        previous_leaf = Some(run.leaf_id);
2366        let max_doc_id = run
2367            .doc_ids_le
2368            .chunks_exact(4)
2369            .map(|bytes| u32::from_le_bytes(bytes.try_into().unwrap()))
2370            .max()
2371            .unwrap_or(0);
2372        if run
2373            .doc_base
2374            .checked_add(max_doc_id)
2375            .is_none_or(|doc_id| doc_id >= payload.doc_count)
2376        {
2377            return Err(invalid_data(
2378                "ScaNN ANN leaf run document ID is out of range",
2379            ));
2380        }
2381
2382        let doc_ids_offset = offset;
2383        writer.write_all(&run.doc_ids_le)?;
2384        offset = checked_advance(offset, run.doc_ids_le.len())?;
2385        let ordinals_offset = offset;
2386        writer.write_all(&run.ordinals_le)?;
2387        offset = checked_advance(offset, run.ordinals_le.len())?;
2388        let codes_offset = offset;
2389        writer.write_all(&run.codes)?;
2390        offset = checked_advance(offset, run.codes.len())?;
2391        records.push(RunRecord {
2392            cluster_id: run.leaf_id,
2393            doc_base: run.doc_base,
2394            count: run.row_count,
2395            max_doc_id,
2396            doc_ids_offset,
2397            ordinals_offset,
2398            codes_offset,
2399            codes_len: u64::try_from(run.codes.len())
2400                .map_err(|_| invalid_data("ScaNN code output size exceeds u64"))?,
2401        });
2402    }
2403    if records.is_empty() {
2404        return Err(invalid_data("cannot write an empty ScaNN ANN payload"));
2405    }
2406    finish_layout(writer, offset, &records)
2407}
2408
2409/// Serialize a populated IVF-TQ build: one run per non-empty leaf, codes
2410/// block-packed per run (scales + gammas + dimension-major nibbles).
2411#[cfg(feature = "native")]
2412pub(crate) fn write_built_ivf_tq(
2413    index: &crate::structures::IvfTqIndex,
2414    num_clusters: u32,
2415    writer: &mut (impl Write + ?Sized),
2416) -> io::Result<u64> {
2417    use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_pack_ivf_block};
2418
2419    if !crate::structures::is_ivf_tq_cosine_generation(index.centroids_version) {
2420        return Err(invalid_data(
2421            "legacy raw IVF-TQ generations cannot be serialized; rebuild the index",
2422        ));
2423    }
2424    let codec = index.codec();
2425    let padded_dim = codec.padded_dim();
2426    let mut clusters: Vec<_> = index.clusters.iter().collect();
2427    clusters.sort_unstable_by_key(|(cluster_id, _)| **cluster_id);
2428    // Emit every cluster in descending residual-scale order: block maxima
2429    // then decrease monotonically, so the scan's per-block score bound can
2430    // stop a run at the first block that cannot beat the running k-th score.
2431    struct PackedCluster {
2432        cluster_id: u32,
2433        doc_ids: Vec<u32>,
2434        ordinals: Vec<u16>,
2435        codes: Vec<u8>,
2436    }
2437    let packed: Vec<PackedCluster> = clusters
2438        .iter()
2439        .map(|&(&cluster_id, cluster)| {
2440            let count = cluster.doc_ids.len();
2441            let mut order: Vec<usize> = (0..count).collect();
2442            order.sort_by(|&a, &b| {
2443                cluster.scales[b]
2444                    .total_cmp(&cluster.scales[a])
2445                    .then_with(|| a.cmp(&b))
2446            });
2447            let doc_ids: Vec<u32> = order.iter().map(|&i| cluster.doc_ids[i]).collect();
2448            let ordinals: Vec<u16> = order.iter().map(|&i| cluster.ordinals[i]).collect();
2449            let scales: Vec<f32> = order.iter().map(|&i| cluster.scales[i]).collect();
2450            let gammas: Vec<f32> = order.iter().map(|&i| cluster.gammas[i]).collect();
2451            let mut codes = Vec::with_capacity(
2452                crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
2453                    count,
2454                    codec.code_size(),
2455                )
2456                .unwrap_or_default(),
2457            );
2458            for block_start in (0..count).step_by(TQ_BLOCK_LANES) {
2459                let lanes = TQ_BLOCK_LANES.min(count - block_start);
2460                let rows: Vec<&[u8]> = order[block_start..block_start + lanes]
2461                    .iter()
2462                    .map(|&row| &cluster.rows[row * padded_dim..(row + 1) * padded_dim])
2463                    .collect();
2464                tq_pack_ivf_block(
2465                    &rows,
2466                    &scales[block_start..block_start + lanes],
2467                    &gammas[block_start..block_start + lanes],
2468                    padded_dim,
2469                    &mut codes,
2470                );
2471            }
2472            PackedCluster {
2473                cluster_id,
2474                doc_ids,
2475                ordinals,
2476                codes,
2477            }
2478        })
2479        .collect();
2480    let runs: Vec<_> = packed
2481        .iter()
2482        .map(|cluster| BuildRun {
2483            cluster_id: cluster.cluster_id,
2484            doc_ids: &cluster.doc_ids,
2485            ordinals: &cluster.ordinals,
2486            codes: &cluster.codes,
2487        })
2488        .collect();
2489    write_built_runs(
2490        AnnDiskHeader {
2491            kind: AnnKind::IvfTq,
2492            routing: index.routing,
2493            dim: index.dim,
2494            code_size: codec.code_size(),
2495            num_clusters,
2496            quantizer_version: index.centroids_version,
2497            codebook_version: codec.fingerprint(),
2498            vector_count: index.len(),
2499        },
2500        &runs,
2501        writer,
2502    )
2503}
2504
2505/// View a populated TQ builder as a single extra merge run (cluster 0).
2506#[cfg(feature = "native")]
2507pub(crate) fn tq_builder_extra_run(builder: &crate::structures::TqFlatBuilder) -> BuildRun<'_> {
2508    BuildRun {
2509        cluster_id: 0,
2510        doc_ids: &builder.doc_ids,
2511        ordinals: &builder.ordinals,
2512        codes: &builder.codes,
2513    }
2514}
2515
2516/// Serialize a populated TQ builder as a single-run payload.
2517#[cfg(feature = "native")]
2518pub(crate) fn write_built_tq_flat(
2519    builder: &crate::structures::TqFlatBuilder,
2520    writer: &mut (impl Write + ?Sized),
2521) -> io::Result<u64> {
2522    let codec = builder.codec();
2523    let runs = [BuildRun {
2524        cluster_id: 0,
2525        doc_ids: &builder.doc_ids,
2526        ordinals: &builder.ordinals,
2527        codes: &builder.codes,
2528    }];
2529    write_built_runs(
2530        AnnDiskHeader {
2531            kind: AnnKind::TqFlat,
2532            routing: IvfRoutingMode::Flat,
2533            dim: codec.dim(),
2534            code_size: codec.code_size(),
2535            num_clusters: 1,
2536            quantizer_version: codec.fingerprint(),
2537            codebook_version: 0,
2538            vector_count: builder.len(),
2539        },
2540        &runs,
2541        writer,
2542    )
2543}
2544
2545#[cfg(feature = "native")]
2546pub(crate) struct BuildRun<'a> {
2547    cluster_id: u32,
2548    doc_ids: &'a [u32],
2549    ordinals: &'a [u16],
2550    codes: &'a [u8],
2551}
2552
2553#[cfg(feature = "native")]
2554struct RunRecord {
2555    cluster_id: u32,
2556    doc_base: u32,
2557    count: u32,
2558    max_doc_id: u32,
2559    doc_ids_offset: u64,
2560    ordinals_offset: u64,
2561    codes_offset: u64,
2562    codes_len: u64,
2563}
2564
2565#[cfg(feature = "native")]
2566fn write_built_runs(
2567    header: AnnDiskHeader,
2568    runs: &[BuildRun<'_>],
2569    writer: &mut (impl Write + ?Sized),
2570) -> io::Result<u64> {
2571    if runs.is_empty() || header.vector_count == 0 {
2572        return Err(invalid_data("cannot write an empty ANN payload"));
2573    }
2574    validate_header(&header)?;
2575    write_header(writer, &header)?;
2576    let mut offset = ANN_HEADER_SIZE as u64;
2577    let mut records = Vec::with_capacity(runs.len());
2578    let mut counted = 0usize;
2579    let mut scratch = Vec::new();
2580    let mut previous_cluster = None;
2581    for run in runs {
2582        let count = run.doc_ids.len();
2583        if count == 0
2584            || run.cluster_id >= header.num_clusters
2585            || previous_cluster.is_some_and(|cluster| cluster >= run.cluster_id)
2586            || run.ordinals.len() != count
2587            || run.codes.len()
2588                != expected_codes_column_len(header.kind, count, header.dim, header.code_size)?
2589        {
2590            return Err(invalid_data("ANN build run columns are inconsistent"));
2591        }
2592        previous_cluster = Some(run.cluster_id);
2593        let count_u32 = u32::try_from(count)
2594            .map_err(|_| invalid_data("ANN cluster run exceeds u32 vectors"))?;
2595        let max_doc_id = run.doc_ids.iter().copied().max().unwrap_or(0);
2596        let doc_ids_offset = offset;
2597        write_u32_column(writer, run.doc_ids, &mut scratch)?;
2598        offset = offset
2599            .checked_add(
2600                u64::try_from(count)
2601                    .ok()
2602                    .and_then(|count| count.checked_mul(4))
2603                    .ok_or_else(|| invalid_data("ANN doc-ID output size overflows u64"))?,
2604            )
2605            .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2606        let ordinals_offset = offset;
2607        write_u16_column(writer, run.ordinals, &mut scratch)?;
2608        offset = offset
2609            .checked_add(
2610                u64::try_from(count)
2611                    .ok()
2612                    .and_then(|count| count.checked_mul(2))
2613                    .ok_or_else(|| invalid_data("ANN ordinal output size overflows u64"))?,
2614            )
2615            .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2616        let codes_offset = offset;
2617        writer.write_all(run.codes)?;
2618        offset = offset
2619            .checked_add(
2620                u64::try_from(run.codes.len())
2621                    .map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
2622            )
2623            .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2624        records.push(RunRecord {
2625            cluster_id: run.cluster_id,
2626            doc_base: 0,
2627            count: count_u32,
2628            max_doc_id,
2629            doc_ids_offset,
2630            ordinals_offset,
2631            codes_offset,
2632            codes_len: u64::try_from(run.codes.len())
2633                .map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
2634        });
2635        counted = counted
2636            .checked_add(count)
2637            .ok_or_else(|| invalid_data("ANN vector count overflow"))?;
2638    }
2639    if counted != header.vector_count {
2640        return Err(invalid_data("ANN header/build vector counts disagree"));
2641    }
2642    finish_layout(writer, offset, &records)
2643}
2644
2645/// Pure-copy normal merge. Corpus-sized source columns are never decoded or
2646/// rewritten; only this compact directory is regenerated with adjusted bases.
2647#[cfg(all(feature = "native", test))]
2648pub(crate) fn write_merged_ann(
2649    sources: &[(&AnnDiskIndex, u32)],
2650    writer: &mut (impl Write + ?Sized),
2651) -> io::Result<u64> {
2652    write_merged_ann_impl(sources, &[], writer, None)
2653}
2654
2655#[cfg(feature = "native")]
2656pub(crate) fn write_merged_ann_cancellable(
2657    sources: &[(&AnnDiskIndex, u32)],
2658    writer: &mut (impl Write + ?Sized),
2659    cancellation: Option<&std::sync::atomic::AtomicBool>,
2660) -> io::Result<u64> {
2661    write_merged_ann_impl(sources, &[], writer, cancellation)
2662}
2663
2664/// Physical extents per non-empty cluster the byte-copy merge of `sources`
2665/// would produce.
2666///
2667/// Byte-copy preserves every source run, so the merged fragmentation is
2668/// `total runs / distinct non-empty clusters` — computable exactly from the
2669/// in-memory directories before writing a byte. The merge policy compacts
2670/// when this crosses its threshold instead of letting probe read
2671/// amplification grow another generation.
2672#[cfg(feature = "native")]
2673pub(crate) fn predicted_merge_fragmentation(sources: &[(&AnnDiskIndex, u32)]) -> f64 {
2674    let mut total_runs = 0usize;
2675    // Distinct clusters via a k-way sorted walk over the (already
2676    // cluster-sorted) directories — no allocation proportional to clusters.
2677    let mut cursors: Vec<std::iter::Peekable<std::slice::Iter<'_, AnnRun>>> = sources
2678        .iter()
2679        .map(|(source, _)| {
2680            total_runs += source.runs.len();
2681            source.runs.iter().peekable()
2682        })
2683        .collect();
2684    let mut distinct = 0usize;
2685    while let Some(cluster) = cursors
2686        .iter_mut()
2687        .filter_map(|cursor| cursor.peek().map(|run| run.cluster_id))
2688        .min()
2689    {
2690        distinct += 1;
2691        for cursor in &mut cursors {
2692            while cursor.peek().is_some_and(|run| run.cluster_id == cluster) {
2693                cursor.next();
2694            }
2695        }
2696    }
2697    if distinct == 0 {
2698        0.0
2699    } else {
2700        total_runs as f64 / distinct as f64
2701    }
2702}
2703
2704/// Doc IDs rewritten per scratch flush during compaction (256 KiB of u32s).
2705#[cfg(feature = "native")]
2706const DOC_ID_REWRITE_CHUNK: usize = 64 * 1024;
2707
2708/// Cluster-major compacting merge for exact-binary and ScaNN-AH payloads.
2709///
2710/// The byte-copy merge keeps each source payload as one physical extent, so a
2711/// logical cluster's postings scatter across up to `sources.len()` extents —
2712/// and another factor per earlier merge generation. Every extent is a
2713/// potential seek when the index is cold; production measured the array
2714/// IOPS-bound at 32 KB/read from exactly this. This writer instead gathers
2715/// each cluster's runs from all sources and emits **one contiguous run per
2716/// cluster**, restoring the freshly-built layout (fragmentation 1.0).
2717///
2718/// Cost: the same total payload bytes the byte-copy merge already streams,
2719/// plus one `u32` add per posting — document IDs are rewritten absolute
2720/// (`doc_base = 0`) because runs from different sources cannot share a single
2721/// directory entry otherwise. Ordinals and binary codes are copied verbatim.
2722/// ScaNN AH codes are decoded one row at a time into a single 32-row scratch
2723/// block and repacked, because FastScan block/tail boundaries are run-local.
2724///
2725/// TQ payloads also pack codes into fixed-lane blocks, but their quantized
2726/// representation is intentionally outside this ScaNN compactor; TQ merges
2727/// stay byte-copy.
2728///
2729/// Every supported merge whose prediction is fragmented takes this path.
2730/// Measured (interleaved best-of-3, pre-faulted buffers, aarch64): byte-copy
2731/// 38.0 GiB/s vs compaction 32.4 GiB/s — ~17% more CPU on a stage that is a
2732/// rounding error of merge wall-clock (a production dense stage is ~0.7s of
2733/// a 20s+ merge), so there is no threshold below which byte-copy is worth
2734/// the fragmentation it leaves behind.
2735#[cfg(feature = "native")]
2736pub(crate) fn write_compacted_ann_cancellable(
2737    sources: &[(&AnnDiskIndex, u32)],
2738    writer: &mut (impl Write + ?Sized),
2739    cancellation: Option<&std::sync::atomic::AtomicBool>,
2740) -> io::Result<u64> {
2741    let Some((first, _)) = sources.first() else {
2742        return Err(invalid_data("cannot compact an empty ANN source list"));
2743    };
2744    if !matches!(
2745        first.header.kind,
2746        AnnKind::BinaryIvf | AnnKind::ScannBinary | AnnKind::ScannAh
2747    ) {
2748        return Err(invalid_data(
2749            "ANN run compaction is only defined for exact binary and ScaNN AH payloads",
2750        ));
2751    }
2752    let mut header = first.header.clone();
2753    header.vector_count = 0;
2754    for &(source, _) in sources {
2755        if !headers_compatible(&first.header, &source.header) {
2756            return Err(invalid_data(
2757                "ANN compaction sources use incompatible generations",
2758            ));
2759        }
2760        header.vector_count = header
2761            .vector_count
2762            .checked_add(source.header.vector_count)
2763            .ok_or_else(|| invalid_data("compacted ANN vector count overflows usize"))?;
2764    }
2765    validate_header(&header)?;
2766    write_header(writer, &header)?;
2767
2768    let code_size = header.code_size;
2769    let mut offset = ANN_HEADER_SIZE as u64;
2770    let mut records: Vec<RunRecord> = Vec::new();
2771    let mut scratch = Vec::new();
2772    let mut cursors: Vec<usize> = vec![0; sources.len()];
2773
2774    loop {
2775        if cancellation.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed)) {
2776            return Err(io::Error::new(
2777                io::ErrorKind::Interrupted,
2778                "ANN compaction cancelled",
2779            ));
2780        }
2781        // Next cluster = minimum un-consumed cluster ID across sources.
2782        let Some(cluster_id) = sources
2783            .iter()
2784            .zip(&cursors)
2785            .filter_map(|((source, _), &cursor)| source.runs.get(cursor).map(|run| run.cluster_id))
2786            .min()
2787        else {
2788            break;
2789        };
2790
2791        // Pass 1: doc IDs, rewritten absolute. Sources are visited in
2792        // segment order and each source's same-cluster runs in directory
2793        // order, which is ascending document ranges — so the output column
2794        // stays sorted like a built segment's.
2795        let mut count = 0usize;
2796        let mut max_doc_id = 0u32;
2797        let doc_ids_offset = offset;
2798        for (source_index, &(source, segment_base)) in sources.iter().enumerate() {
2799            let mut cursor = cursors[source_index];
2800            while let Some(run) = source
2801                .runs
2802                .get(cursor)
2803                .filter(|run| run.cluster_id == cluster_id)
2804            {
2805                let base = run
2806                    .doc_base
2807                    .checked_add(segment_base)
2808                    .ok_or_else(|| invalid_data("compacted ANN document base overflows u32"))?;
2809                let bytes = source.raw.as_slice();
2810                // Chunked rewrite: peak scratch stays at 256 KiB no matter how
2811                // large the run — the production incident had a single run of
2812                // 20M postings, and buffering it whole would be an 80 MB spike
2813                // in the middle of a merge.
2814                for chunk_start in (0..run.count).step_by(DOC_ID_REWRITE_CHUNK) {
2815                    let chunk_end = (chunk_start + DOC_ID_REWRITE_CHUNK).min(run.count);
2816                    scratch.clear();
2817                    scratch.reserve((chunk_end - chunk_start) * 4);
2818                    for index in chunk_start..chunk_end {
2819                        let doc_id = run_doc_id_with_base(bytes, run, index, base)?;
2820                        max_doc_id = max_doc_id.max(doc_id);
2821                        scratch.extend_from_slice(&doc_id.to_le_bytes());
2822                    }
2823                    writer.write_all(&scratch)?;
2824                    if cancellation
2825                        .is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed))
2826                    {
2827                        return Err(io::Error::new(
2828                            io::ErrorKind::Interrupted,
2829                            "ANN compaction cancelled",
2830                        ));
2831                    }
2832                }
2833                offset = checked_advance(offset, run.count * 4)?;
2834                count = count
2835                    .checked_add(run.count)
2836                    .ok_or_else(|| invalid_data("compacted ANN run count overflows usize"))?;
2837                cursor += 1;
2838            }
2839        }
2840
2841        // Pass 2: ordinals, verbatim.
2842        let ordinals_offset = offset;
2843        for (source_index, &(source, _)) in sources.iter().enumerate() {
2844            let mut cursor = cursors[source_index];
2845            while let Some(run) = source
2846                .runs
2847                .get(cursor)
2848                .filter(|run| run.cluster_id == cluster_id)
2849            {
2850                copy_range(writer, &source.raw, run.ordinals.clone(), cancellation)?;
2851                offset = checked_advance(offset, run.ordinals.len())?;
2852                cursor += 1;
2853            }
2854        }
2855
2856        // Pass 3: codes. Exact binary rows concatenate directly. ScaNN AH
2857        // blocks are run-relative, so repack a continuous output stream.
2858        let codes_offset = offset;
2859        if header.kind == AnnKind::ScannAh {
2860            let blocks = header.dim.div_ceil(code_size);
2861            let lanes = crate::structures::vector::scann::FAST_SCAN_LANES;
2862            let mut unpacked = Vec::with_capacity(lanes * blocks);
2863            let mut packed = Vec::with_capacity(blocks * lanes / 2);
2864            for (source_index, &(source, _)) in sources.iter().enumerate() {
2865                let mut cursor = cursors[source_index];
2866                while let Some(run) = source
2867                    .runs
2868                    .get(cursor)
2869                    .filter(|run| run.cluster_id == cluster_id)
2870                {
2871                    let bytes = &source.raw.as_slice()[run.codes.clone()];
2872                    for row in 0..run.count {
2873                        unpack_scann_ah_row(bytes, run.count, blocks, row, &mut unpacked)?;
2874                        if unpacked.len() == lanes * blocks {
2875                            packed.clear();
2876                            crate::structures::vector::scann::pack_fast_scan_block(
2877                                &unpacked,
2878                                blocks,
2879                                &mut packed,
2880                            )
2881                            .map_err(|error| invalid_data(error.to_string()))?;
2882                            writer.write_all(&packed)?;
2883                            offset = checked_advance(offset, packed.len())?;
2884                            unpacked.clear();
2885                            if cancellation
2886                                .is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed))
2887                            {
2888                                return Err(io::Error::new(
2889                                    io::ErrorKind::Interrupted,
2890                                    "ANN compaction cancelled",
2891                                ));
2892                            }
2893                        }
2894                    }
2895                    cursor += 1;
2896                }
2897            }
2898            // One compacted run has one row-major tail, not one tail per
2899            // source run. Pack each remaining row's adjacent block nibbles.
2900            if !unpacked.is_empty() {
2901                packed.clear();
2902                for row in unpacked.chunks_exact(blocks) {
2903                    for pair in row.chunks(2) {
2904                        packed.push(pair[0] | (pair.get(1).copied().unwrap_or(0) << 4));
2905                    }
2906                }
2907                writer.write_all(&packed)?;
2908                offset = checked_advance(offset, packed.len())?;
2909            }
2910        } else {
2911            for (source_index, &(source, _)) in sources.iter().enumerate() {
2912                let mut cursor = cursors[source_index];
2913                while let Some(run) = source
2914                    .runs
2915                    .get(cursor)
2916                    .filter(|run| run.cluster_id == cluster_id)
2917                {
2918                    copy_range(writer, &source.raw, run.codes.clone(), cancellation)?;
2919                    offset = checked_advance(offset, run.codes.len())?;
2920                    cursor += 1;
2921                }
2922            }
2923        }
2924        let expected_codes_len =
2925            expected_codes_column_len(header.kind, count, header.dim, header.code_size)?;
2926        if offset - codes_offset != expected_codes_len as u64 {
2927            return Err(invalid_data("compacted ANN code column length mismatch"));
2928        }
2929
2930        // Consume this cluster's runs from every cursor.
2931        for (source_index, &(source, _)) in sources.iter().enumerate() {
2932            while source
2933                .runs
2934                .get(cursors[source_index])
2935                .is_some_and(|run| run.cluster_id == cluster_id)
2936            {
2937                cursors[source_index] += 1;
2938            }
2939        }
2940
2941        records.push(RunRecord {
2942            cluster_id,
2943            doc_base: 0,
2944            count: u32::try_from(count)
2945                .map_err(|_| invalid_data("compacted ANN run exceeds u32 vectors"))?,
2946            max_doc_id,
2947            doc_ids_offset,
2948            ordinals_offset,
2949            codes_offset,
2950            codes_len: u64::try_from(expected_codes_column_len(
2951                header.kind,
2952                count,
2953                header.dim,
2954                code_size,
2955            )?)
2956            .map_err(|_| invalid_data("compacted ANN code length exceeds u64"))?,
2957        });
2958    }
2959
2960    if records.is_empty() {
2961        return Err(invalid_data("cannot compact an ANN payload with no runs"));
2962    }
2963    finish_layout(writer, offset, &records)
2964}
2965
2966/// Append one ScaNN-AH row as unpacked 4-bit block codes. Complete 32-row
2967/// groups use the FastScan v2 block-pair layout (`packed_code_position` is
2968/// the single definition); each run's final partial group is row-major.
2969#[cfg(feature = "native")]
2970fn unpack_scann_ah_row(
2971    codes: &[u8],
2972    count: usize,
2973    blocks: usize,
2974    row: usize,
2975    output: &mut Vec<u8>,
2976) -> io::Result<()> {
2977    use crate::structures::vector::scann::{
2978        FAST_SCAN_LANES, packed_block_bytes, packed_code_position,
2979    };
2980
2981    let lanes = FAST_SCAN_LANES;
2982    let full_rows = count / lanes * lanes;
2983    let full_block_bytes = packed_block_bytes(blocks)
2984        .ok_or_else(|| invalid_data("ScaNN AH block size overflows usize"))?;
2985    let tail_row_bytes = blocks.div_ceil(2);
2986    for block in 0..blocks {
2987        let (byte_offset, high) = if row < full_rows {
2988            let (offset, high) = packed_code_position(row % lanes, block);
2989            ((row / lanes) * full_block_bytes + offset, high)
2990        } else {
2991            (
2992                (full_rows / lanes) * full_block_bytes
2993                    + (row - full_rows) * tail_row_bytes
2994                    + block / 2,
2995                !block.is_multiple_of(2),
2996            )
2997        };
2998        let byte = *codes
2999            .get(byte_offset)
3000            .ok_or_else(|| invalid_data("ScaNN AH row exceeds its code column"))?;
3001        output.push(if high { byte >> 4 } else { byte & 0x0f });
3002    }
3003    Ok(())
3004}
3005
3006/// [`run_doc_id`] against an explicit base, for rewriting IDs absolute.
3007#[cfg(feature = "native")]
3008fn run_doc_id_with_base(bytes: &[u8], run: &AnnRun, index: usize, base: u32) -> io::Result<u32> {
3009    let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
3010    if local_doc_id > run.max_doc_id {
3011        return Err(invalid_data(
3012            "ANN run contains a document above its declared maximum",
3013        ));
3014    }
3015    base.checked_add(local_doc_id)
3016        .ok_or_else(|| invalid_data("compacted ANN document ID overflows u32"))
3017}
3018
3019/// [`write_merged_ann`] plus freshly built runs appended to the payload —
3020/// used when some merge sources predate the field's current format and were
3021/// re-encoded while every compatible source is still byte-copied.
3022#[cfg(feature = "native")]
3023pub(crate) fn write_merged_ann_with_extra(
3024    sources: &[(&AnnDiskIndex, u32)],
3025    extra: &[BuildRun<'_>],
3026    writer: &mut (impl Write + ?Sized),
3027    cancellation: Option<&std::sync::atomic::AtomicBool>,
3028) -> io::Result<u64> {
3029    write_merged_ann_impl(sources, extra, writer, cancellation)
3030}
3031
3032#[cfg(feature = "native")]
3033fn write_merged_ann_impl(
3034    sources: &[(&AnnDiskIndex, u32)],
3035    extra: &[BuildRun<'_>],
3036    writer: &mut (impl Write + ?Sized),
3037    cancellation: Option<&std::sync::atomic::AtomicBool>,
3038) -> io::Result<u64> {
3039    let Some((first, _)) = sources.first() else {
3040        return Err(invalid_data("cannot merge an empty ANN source list"));
3041    };
3042    if first.header.kind == AnnKind::IvfTq
3043        && !crate::structures::is_ivf_tq_cosine_generation(first.header.quantizer_version)
3044    {
3045        return Err(invalid_data(
3046            "legacy raw IVF-TQ generations cannot be merged; rebuild the index",
3047        ));
3048    }
3049    let mut header = first.header.clone();
3050    header.vector_count = 0;
3051    for &(source, _) in sources {
3052        if !headers_compatible(&first.header, &source.header) {
3053            return Err(invalid_data(
3054                "ANN merge sources use incompatible generations",
3055            ));
3056        }
3057        header.vector_count = header
3058            .vector_count
3059            .checked_add(source.header.vector_count)
3060            .ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
3061    }
3062    for run in extra {
3063        if run.doc_ids.is_empty()
3064            || run.cluster_id >= header.num_clusters
3065            || run.ordinals.len() != run.doc_ids.len()
3066            || run.codes.len()
3067                != expected_codes_column_len(
3068                    header.kind,
3069                    run.doc_ids.len(),
3070                    header.dim,
3071                    header.code_size,
3072                )?
3073        {
3074            return Err(invalid_data("extra ANN merge run columns are inconsistent"));
3075        }
3076        header.vector_count = header
3077            .vector_count
3078            .checked_add(run.doc_ids.len())
3079            .ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
3080    }
3081    validate_header(&header)?;
3082    write_header(writer, &header)?;
3083    let mut offset = ANN_HEADER_SIZE as u64;
3084    let run_capacity = sources.iter().try_fold(extra.len(), |count, (source, _)| {
3085        count
3086            .checked_add(source.runs.len())
3087            .ok_or_else(|| invalid_data("merged ANN run count overflows usize"))
3088    })?;
3089    let mut output_payload_starts = Vec::with_capacity(sources.len());
3090    for &(source, _) in sources {
3091        let payload_end = source
3092            .runs
3093            .iter()
3094            .map(|run| run.codes.end)
3095            .max()
3096            .ok_or_else(|| invalid_data("ANN source has no payload runs"))?;
3097        let output_payload_start = offset;
3098        output_payload_starts.push(output_payload_start);
3099        copy_range(
3100            writer,
3101            &source.raw,
3102            ANN_HEADER_SIZE..payload_end,
3103            cancellation,
3104        )?;
3105        offset = checked_advance(offset, payload_end - ANN_HEADER_SIZE)?;
3106    }
3107
3108    // Extra runs' columns are appended after the copied extents so the
3109    // payload region stays contiguous for open()'s coverage validation.
3110    let mut extra_records = Vec::with_capacity(extra.len());
3111    let mut scratch = Vec::new();
3112    for run in extra {
3113        let count = run.doc_ids.len();
3114        let doc_ids_offset = offset;
3115        write_u32_column(writer, run.doc_ids, &mut scratch)?;
3116        offset = checked_advance(offset, count * 4)?;
3117        let ordinals_offset = offset;
3118        write_u16_column(writer, run.ordinals, &mut scratch)?;
3119        offset = checked_advance(offset, count * 2)?;
3120        let codes_offset = offset;
3121        writer.write_all(run.codes)?;
3122        offset = checked_advance(offset, run.codes.len())?;
3123        extra_records.push(RunRecord {
3124            cluster_id: run.cluster_id,
3125            doc_base: 0,
3126            count: u32::try_from(count)
3127                .map_err(|_| invalid_data("extra ANN run exceeds u32 vectors"))?,
3128            max_doc_id: run.doc_ids.iter().copied().max().unwrap_or(0),
3129            doc_ids_offset,
3130            ordinals_offset,
3131            codes_offset,
3132            codes_len: u64::try_from(run.codes.len())
3133                .map_err(|_| invalid_data("extra ANN code length exceeds u64"))?,
3134        });
3135    }
3136
3137    // Every source directory is already cluster-sorted. Merge those compact
3138    // directories (plus the extra records) directly into the output with
3139    // O(source count) heap memory; the corpus payload extents above remain
3140    // untouched and source-contiguous. Extra records use pseudo source index
3141    // `sources.len()` so ties stay deterministic.
3142    let directory_offset = offset;
3143    let mut pending = BinaryHeap::with_capacity(sources.len() + 1);
3144    for (source_index, (source, _)) in sources.iter().enumerate() {
3145        pending.push(Reverse((source.runs[0].cluster_id, source_index, 0usize)));
3146    }
3147    if let Some(first_extra) = extra_records.first() {
3148        pending.push(Reverse((first_extra.cluster_id, sources.len(), 0usize)));
3149    }
3150    let mut written_runs = 0usize;
3151    while let Some(Reverse((_, source_index, run_index))) = pending.pop() {
3152        if source_index == sources.len() {
3153            write_run_record(writer, &extra_records[run_index])?;
3154            written_runs += 1;
3155            if let Some(next) = extra_records.get(run_index + 1) {
3156                pending.push(Reverse((next.cluster_id, source_index, run_index + 1)));
3157            }
3158            continue;
3159        }
3160        let (source, segment_base) = sources[source_index];
3161        let run = &source.runs[run_index];
3162        write_run_record(
3163            writer,
3164            &RunRecord {
3165                cluster_id: run.cluster_id,
3166                doc_base: run
3167                    .doc_base
3168                    .checked_add(segment_base)
3169                    .ok_or_else(|| invalid_data("merged ANN document base overflows u32"))?,
3170                count: u32::try_from(run.count)
3171                    .map_err(|_| invalid_data("ANN source run exceeds u32 vectors"))?,
3172                max_doc_id: run.max_doc_id,
3173                doc_ids_offset: relocate_payload_offset(
3174                    output_payload_starts[source_index],
3175                    run.doc_ids.start,
3176                )?,
3177                ordinals_offset: relocate_payload_offset(
3178                    output_payload_starts[source_index],
3179                    run.ordinals.start,
3180                )?,
3181                codes_offset: relocate_payload_offset(
3182                    output_payload_starts[source_index],
3183                    run.codes.start,
3184                )?,
3185                codes_len: u64::try_from(run.codes.len())
3186                    .map_err(|_| invalid_data("ANN source code length exceeds u64"))?,
3187            },
3188        )?;
3189        written_runs = written_runs
3190            .checked_add(1)
3191            .ok_or_else(|| invalid_data("merged ANN run count overflows usize"))?;
3192        let next_run_index = run_index + 1;
3193        if let Some(next_run) = source.runs.get(next_run_index) {
3194            pending.push(Reverse((next_run.cluster_id, source_index, next_run_index)));
3195        }
3196    }
3197    if written_runs != run_capacity {
3198        return Err(invalid_data("merged ANN directory lost source runs"));
3199    }
3200    finish_footer(writer, directory_offset, written_runs)
3201}
3202
3203#[cfg(feature = "native")]
3204fn relocate_payload_offset(output_payload_start: u64, source_offset: usize) -> io::Result<u64> {
3205    let relative = source_offset
3206        .checked_sub(ANN_HEADER_SIZE)
3207        .ok_or_else(|| invalid_data("ANN source offset precedes its payload"))?;
3208    output_payload_start
3209        .checked_add(
3210            u64::try_from(relative)
3211                .map_err(|_| invalid_data("ANN source relative offset exceeds u64"))?,
3212        )
3213        .ok_or_else(|| invalid_data("merged ANN payload offset overflows u64"))
3214}
3215
3216#[cfg(feature = "native")]
3217fn headers_compatible(left: &AnnDiskHeader, right: &AnnDiskHeader) -> bool {
3218    left.kind == right.kind
3219        && left.routing == right.routing
3220        && left.dim == right.dim
3221        && left.code_size == right.code_size
3222        && left.num_clusters == right.num_clusters
3223        && left.quantizer_version == right.quantizer_version
3224        && left.codebook_version == right.codebook_version
3225}
3226
3227#[cfg(feature = "native")]
3228fn finish_layout(
3229    writer: &mut (impl Write + ?Sized),
3230    directory_offset: u64,
3231    records: &[RunRecord],
3232) -> io::Result<u64> {
3233    for record in records {
3234        write_run_record(writer, record)?;
3235    }
3236    finish_footer(writer, directory_offset, records.len())
3237}
3238
3239#[cfg(feature = "native")]
3240fn write_run_record(writer: &mut (impl Write + ?Sized), record: &RunRecord) -> io::Result<()> {
3241    writer.write_u32::<LittleEndian>(record.cluster_id)?;
3242    writer.write_u32::<LittleEndian>(record.doc_base)?;
3243    writer.write_u32::<LittleEndian>(record.count)?;
3244    writer.write_u32::<LittleEndian>(record.max_doc_id)?;
3245    writer.write_u64::<LittleEndian>(record.doc_ids_offset)?;
3246    writer.write_u64::<LittleEndian>(record.ordinals_offset)?;
3247    writer.write_u64::<LittleEndian>(record.codes_offset)?;
3248    writer.write_u64::<LittleEndian>(record.codes_len)?;
3249    Ok(())
3250}
3251
3252#[cfg(feature = "native")]
3253fn finish_footer(
3254    writer: &mut (impl Write + ?Sized),
3255    directory_offset: u64,
3256    num_records: usize,
3257) -> io::Result<u64> {
3258    writer.write_u64::<LittleEndian>(directory_offset)?;
3259    writer.write_u64::<LittleEndian>(
3260        u64::try_from(num_records).map_err(|_| invalid_data("ANN run count exceeds u64"))?,
3261    )?;
3262    writer.write_u32::<LittleEndian>(ANN_FOOTER_MAGIC)?;
3263    writer.write_u32::<LittleEndian>(u32::from(ANN_DISK_VERSION))?;
3264    let tail_size = num_records
3265        .checked_mul(ANN_RUN_SIZE)
3266        .and_then(|size| size.checked_add(ANN_FOOTER_SIZE))
3267        .and_then(|size| u64::try_from(size).ok())
3268        .ok_or_else(|| invalid_data("ANN final tail size overflows u64"))?;
3269    directory_offset
3270        .checked_add(tail_size)
3271        .ok_or_else(|| invalid_data("ANN final size overflows u64"))
3272}
3273
3274/// Code-layout revision stamped into the header's former reserved `u32`.
3275///
3276/// Kinds whose leaf encoding changed shape without changing byte length need
3277/// an explicit gate: a FastScan v1 (`ScannAh`) column has exactly the v2
3278/// length for even block counts, so nothing else in the container would
3279/// notice a stale generation (docs/fast-scan-layout-v2.md). Every other kind
3280/// keeps the field at zero, preserving the old "reserved must be zero" check.
3281fn expected_layout_revision(kind: AnnKind) -> u32 {
3282    match kind {
3283        AnnKind::ScannAh => u32::from(crate::structures::vector::scann::FAST_SCAN_LAYOUT_VERSION),
3284        AnnKind::BinaryIvf | AnnKind::TqFlat | AnnKind::IvfTq | AnnKind::ScannBinary => 0,
3285    }
3286}
3287
3288fn check_layout_revision(kind: AnnKind, found: u32) -> io::Result<()> {
3289    let expected = expected_layout_revision(kind);
3290    if found == expected {
3291        return Ok(());
3292    }
3293    Err(invalid_data(match kind {
3294        AnnKind::ScannAh => format!(
3295            "ScaNN AH payload uses FastScan layout revision {found}; this build reads revision \
3296             {expected} (docs/fast-scan-layout-v2.md). Rebuild the float ScaNN generation for \
3297             this field before searching it"
3298        ),
3299        _ => format!("ANN header reserved field is non-zero ({found}) for kind {kind:?}"),
3300    }))
3301}
3302
3303#[cfg(feature = "native")]
3304fn write_header(writer: &mut (impl Write + ?Sized), header: &AnnDiskHeader) -> io::Result<()> {
3305    writer.write_u32::<LittleEndian>(ANN_HEADER_MAGIC)?;
3306    writer.write_u8(header.kind as u8)?;
3307    writer.write_u8(routing_to_u8(header.routing))?;
3308    writer.write_u16::<LittleEndian>(ANN_DISK_VERSION)?;
3309    writer.write_u32::<LittleEndian>(
3310        u32::try_from(header.dim).map_err(|_| invalid_data("ANN dimension exceeds u32"))?,
3311    )?;
3312    writer.write_u32::<LittleEndian>(
3313        u32::try_from(header.code_size).map_err(|_| invalid_data("ANN code size exceeds u32"))?,
3314    )?;
3315    writer.write_u32::<LittleEndian>(header.num_clusters)?;
3316    writer.write_u32::<LittleEndian>(expected_layout_revision(header.kind))?;
3317    writer.write_u64::<LittleEndian>(header.quantizer_version)?;
3318    writer.write_u64::<LittleEndian>(header.codebook_version)?;
3319    writer.write_u64::<LittleEndian>(
3320        u64::try_from(header.vector_count)
3321            .map_err(|_| invalid_data("ANN vector count exceeds u64"))?,
3322    )?;
3323    writer.write_u64::<LittleEndian>(0)?;
3324    Ok(())
3325}
3326
3327#[cfg(feature = "native")]
3328fn write_u32_column(
3329    writer: &mut (impl Write + ?Sized),
3330    values: &[u32],
3331    scratch: &mut Vec<u8>,
3332) -> io::Result<()> {
3333    for chunk in values.chunks(64 * 1024) {
3334        scratch.clear();
3335        scratch.reserve(chunk.len() * 4);
3336        for value in chunk {
3337            scratch.extend_from_slice(&value.to_le_bytes());
3338        }
3339        writer.write_all(scratch)?;
3340    }
3341    Ok(())
3342}
3343
3344#[cfg(feature = "native")]
3345fn write_u16_column(
3346    writer: &mut (impl Write + ?Sized),
3347    values: &[u16],
3348    scratch: &mut Vec<u8>,
3349) -> io::Result<()> {
3350    for chunk in values.chunks(64 * 1024) {
3351        scratch.clear();
3352        scratch.reserve(chunk.len() * 2);
3353        for value in chunk {
3354            scratch.extend_from_slice(&value.to_le_bytes());
3355        }
3356        writer.write_all(scratch)?;
3357    }
3358    Ok(())
3359}
3360
3361#[cfg(feature = "native")]
3362fn copy_range(
3363    writer: &mut (impl Write + ?Sized),
3364    bytes: &OwnedBytes,
3365    range: Range<usize>,
3366    cancellation: Option<&std::sync::atomic::AtomicBool>,
3367) -> io::Result<()> {
3368    if range.is_empty() {
3369        return Ok(());
3370    }
3371    let range_end = range.end;
3372    let mut chunk_start = range.start;
3373    let first_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
3374    bytes.madvise_range(chunk_start..first_end, libc::MADV_WILLNEED);
3375    while chunk_start < range_end {
3376        if cancellation
3377            .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Relaxed))
3378        {
3379            return Err(io::Error::new(
3380                io::ErrorKind::Interrupted,
3381                "ANN merge copy cancelled",
3382            ));
3383        }
3384        let chunk_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
3385        let next_end = chunk_end.saturating_add(COPY_CHUNK).min(range_end);
3386        if chunk_end < next_end {
3387            // Keep one bounded window of IO in flight while the current
3388            // window is copied. The query mapping remains MADV_RANDOM.
3389            bytes.madvise_range(chunk_end..next_end, libc::MADV_WILLNEED);
3390        }
3391        writer.write_all(&bytes.as_slice()[chunk_start..chunk_end])?;
3392        chunk_start = chunk_end;
3393    }
3394    Ok(())
3395}
3396
3397#[cfg(feature = "native")]
3398fn checked_advance(offset: u64, length: usize) -> io::Result<u64> {
3399    offset
3400        .checked_add(
3401            u64::try_from(length).map_err(|_| invalid_data("ANN copy length exceeds u64"))?,
3402        )
3403        .ok_or_else(|| invalid_data("ANN output offset overflows u64"))
3404}
3405
3406fn validate_header(header: &AnnDiskHeader) -> io::Result<()> {
3407    if header.kind == AnnKind::IvfTq
3408        && !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
3409    {
3410        return Err(invalid_data(
3411            "IVF-TQ payload uses an unsupported legacy generation; rebuild the index",
3412        ));
3413    }
3414    if header.dim == 0
3415        || header.code_size == 0
3416        || header.num_clusters == 0
3417        || header.quantizer_version == 0
3418        || header.vector_count == 0
3419        || (header.kind == AnnKind::BinaryIvf
3420            && (header.codebook_version != 0
3421                || !header.dim.is_multiple_of(8)
3422                || header.code_size != header.dim.div_ceil(8)))
3423        || (header.kind == AnnKind::ScannBinary
3424            && (header.codebook_version == 0
3425                || header.routing != IvfRoutingMode::Flat
3426                || !header.dim.is_multiple_of(8)
3427                || header.code_size != header.dim / 8))
3428        || (header.kind == AnnKind::TqFlat
3429            && (header.codebook_version != 0
3430                || header.num_clusters != 1
3431                || header.routing != IvfRoutingMode::Flat
3432                || header.code_size * 2
3433                    != crate::structures::vector::quantization::tq_padded_dim(header.dim)))
3434        // IVF-TQ: quantizer_version is the trained centroid generation and
3435        // codebook_version carries the (nonzero) TQ codec fingerprint.
3436        || (header.kind == AnnKind::IvfTq
3437            && (header.codebook_version == 0
3438                || header.code_size * 2
3439                    != crate::structures::vector::quantization::tq_padded_dim(header.dim)))
3440        || (header.kind == AnnKind::ScannAh
3441            && (header.codebook_version == 0
3442                || header.routing != IvfRoutingMode::Flat
3443                || u16::try_from(header.code_size).is_err()
3444                || u32::try_from(header.dim).is_err()
3445                || crate::structures::vector::scann::ScannEncoding::AsymmetricHash {
3446                    dimensions_per_block: header.code_size as u16,
3447                    bits_per_code: 4,
3448                }
3449                .row_code_bytes(u32::try_from(header.dim).unwrap_or(0))
3450                .is_err()))
3451    {
3452        return Err(invalid_data("ANN header contains invalid metadata"));
3453    }
3454    Ok(())
3455}
3456
3457fn read_u32(bytes: &[u8], offset: usize) -> u32 {
3458    u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
3459}
3460
3461fn read_u16(bytes: &[u8], offset: usize) -> u16 {
3462    u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap())
3463}
3464
3465#[inline]
3466fn tq_ivf_block_max_scale(block: &[u8]) -> f32 {
3467    // Supported cosine-generation writers sort the complete run by descending
3468    // residual scale, so lane zero is both this block's maximum and an upper
3469    // bound for every following block.
3470    f32::from_le_bytes(
3471        block[..size_of::<f32>()]
3472            .try_into()
3473            .expect("scale is one f32"),
3474    )
3475}
3476
3477/// Doc ID for `index` in a run whose doc-ID column was validated at
3478/// [`AnnDiskIndex::open`]: every local ID is `<= max_doc_id` and
3479/// `doc_base + max_doc_id < total_docs`, so the sum cannot overflow. The scan
3480/// loops use this; [`run_doc_id`] remains the checked form for merge/rewrite
3481/// paths and tests.
3482#[inline]
3483fn run_doc_id_validated(bytes: &[u8], run: &AnnRun, index: usize) -> u32 {
3484    debug_assert!(index < run.count);
3485    run.doc_base + read_u32(bytes, run.doc_ids.start + index * 4)
3486}
3487
3488fn run_doc_id(bytes: &[u8], run: &AnnRun, index: usize) -> io::Result<u32> {
3489    let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
3490    if local_doc_id > run.max_doc_id {
3491        return Err(invalid_data(
3492            "ANN run contains a document above its declared maximum",
3493        ));
3494    }
3495    run.doc_base
3496        .checked_add(local_doc_id)
3497        .ok_or_else(|| invalid_data("ANN run document ID overflows u32"))
3498}
3499
3500#[cfg(feature = "native")]
3501fn routing_to_u8(routing: IvfRoutingMode) -> u8 {
3502    match routing {
3503        IvfRoutingMode::Auto => 0,
3504        IvfRoutingMode::Flat => 1,
3505        IvfRoutingMode::TwoLevel => 2,
3506        IvfRoutingMode::Hnsw => 3,
3507    }
3508}
3509
3510fn routing_from_u8(value: u8) -> io::Result<IvfRoutingMode> {
3511    match value {
3512        0 => Ok(IvfRoutingMode::Auto),
3513        1 => Ok(IvfRoutingMode::Flat),
3514        2 => Ok(IvfRoutingMode::TwoLevel),
3515        3 => Ok(IvfRoutingMode::Hnsw),
3516        _ => Err(invalid_data(format!("unknown ANN routing mode {value}"))),
3517    }
3518}
3519
3520fn invalid_data(message: impl Into<String>) -> io::Error {
3521    io::Error::new(io::ErrorKind::InvalidData, message.into())
3522}
3523
3524#[cfg(all(test, feature = "native"))]
3525mod tests {
3526    use super::*;
3527
3528    /// Compaction must produce a payload indistinguishable from a fresh
3529    /// build: one run per cluster, fragmentation 1.0, absolute doc IDs — and
3530    /// return exactly the results the byte-copy merge of the same sources
3531    /// returns, for every cluster.
3532    #[test]
3533    fn compacted_merge_matches_byte_copy_and_resets_fragmentation() {
3534        // Segment A: clusters {0: 3 docs, 5: 2 docs}. Segment B: {0: 2, 2: 1}.
3535        // Merging A+B and then merging that result with A again produces
3536        // multi-generation fragmentation (cluster 0 in three extents).
3537        let a0_docs = [0u32, 1, 2];
3538        let a0_ords = [0u16, 0, 1];
3539        let a0_codes = [0x11u8, 0x22, 0x33];
3540        let a5_docs = [3u32, 4];
3541        let a5_ords = [0u16; 2];
3542        let a5_codes = [0x44u8, 0x55];
3543        let a_runs = [
3544            BuildRun {
3545                cluster_id: 0,
3546                doc_ids: &a0_docs,
3547                ordinals: &a0_ords,
3548                codes: &a0_codes,
3549            },
3550            BuildRun {
3551                cluster_id: 5,
3552                doc_ids: &a5_docs,
3553                ordinals: &a5_ords,
3554                codes: &a5_codes,
3555            },
3556        ];
3557        let mut header = binary_header(5);
3558        header.num_clusters = 8;
3559        let mut a_bytes = Vec::new();
3560        write_built_runs(header.clone(), &a_runs, &mut a_bytes).unwrap();
3561        let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 5).unwrap();
3562
3563        let b0_docs = [0u32, 2];
3564        let b0_ords = [1u16, 0];
3565        let b0_codes = [0x66u8, 0x77];
3566        let b2_docs = [1u32];
3567        let b2_ords = [0u16];
3568        let b2_codes = [0x88u8];
3569        let b_runs = [
3570            BuildRun {
3571                cluster_id: 0,
3572                doc_ids: &b0_docs,
3573                ordinals: &b0_ords,
3574                codes: &b0_codes,
3575            },
3576            BuildRun {
3577                cluster_id: 2,
3578                doc_ids: &b2_docs,
3579                ordinals: &b2_ords,
3580                codes: &b2_codes,
3581            },
3582        ];
3583        let mut header_b = binary_header(3);
3584        header_b.num_clusters = 8;
3585        let mut b_bytes = Vec::new();
3586        write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
3587        let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 3).unwrap();
3588
3589        // Generation 1: byte-copy A (docs 0..5) + B (docs 5..8).
3590        let mut gen1_bytes = Vec::new();
3591        write_merged_ann(&[(&a, 0), (&b, 5)], &mut gen1_bytes).unwrap();
3592        let gen1 = AnnDiskIndex::open(OwnedBytes::new(gen1_bytes), AnnKind::BinaryIvf, 8).unwrap();
3593
3594        // Generation 2 sources: gen1 (docs 0..8) + A again (docs 8..13).
3595        let sources: [(&AnnDiskIndex, u32); 2] = [(&gen1, 0), (&a, 8)];
3596        let predicted = predicted_merge_fragmentation(&sources);
3597        // 4 runs (gen1) + 2 runs (a) over 3 distinct clusters {0, 2, 5}.
3598        assert!((predicted - 2.0).abs() < 1e-9, "{predicted}");
3599
3600        let mut copied_bytes = Vec::new();
3601        write_merged_ann(&sources, &mut copied_bytes).unwrap();
3602        let copied =
3603            AnnDiskIndex::open(OwnedBytes::new(copied_bytes), AnnKind::BinaryIvf, 13).unwrap();
3604        let mut compacted_bytes = Vec::new();
3605        write_compacted_ann_cancellable(&sources, &mut compacted_bytes, None).unwrap();
3606        let compacted =
3607            AnnDiskIndex::open(OwnedBytes::new(compacted_bytes), AnnKind::BinaryIvf, 13).unwrap();
3608
3609        // Byte-copy carries the fragmentation forward; compaction resets it.
3610        let copied_health = copied.health();
3611        let compacted_health = compacted.health();
3612        assert!((copied_health.fragmentation() - 2.0).abs() < 1e-9);
3613        assert!((compacted_health.fragmentation() - 1.0).abs() < 1e-9);
3614        assert_eq!(compacted_health.runs, 3, "one run per non-empty cluster");
3615        assert_eq!(copied_health.vectors, compacted_health.vectors);
3616        assert_eq!(copied_health.payload_bytes, compacted_health.payload_bytes);
3617        assert_eq!(
3618            copied_health.largest_cluster_vectors,
3619            compacted_health.largest_cluster_vectors
3620        );
3621
3622        // Every cluster returns identical (doc, ordinal, score) results.
3623        for cluster in 0..8u32 {
3624            let query = [0x5Au8];
3625            let from_copy = copied
3626                .search_binary_clusters::<false>(&query, 16, &[cluster])
3627                .unwrap();
3628            let from_compact = compacted
3629                .search_binary_clusters::<false>(&query, 16, &[cluster])
3630                .unwrap();
3631            assert_eq!(from_copy, from_compact, "cluster {cluster} diverged");
3632        }
3633
3634        // Doc IDs are absolute now: every directory entry has doc_base 0.
3635        assert!(compacted.runs.iter().all(|run| run.doc_base == 0));
3636
3637        // A compacted payload is indistinguishable from a built one, so it
3638        // must remain a valid source for future ordinary byte-copy merges.
3639        let mut generation3 = Vec::new();
3640        write_merged_ann(&[(&compacted, 0), (&b, 13)], &mut generation3).unwrap();
3641        let generation3 =
3642            AnnDiskIndex::open(OwnedBytes::new(generation3), AnnKind::BinaryIvf, 16).unwrap();
3643        assert_eq!(generation3.health().vectors, 16);
3644        // And the third A copy's docs landed at offset 8.
3645        let all: Vec<(u32, u16, f32)> = compacted
3646            .search_binary_clusters::<false>(&[0x5A], 32, &[0, 2, 5])
3647            .unwrap();
3648        let mut docs: Vec<u32> = all.iter().map(|&(doc, _, _)| doc).collect();
3649        docs.sort_unstable();
3650        assert_eq!(docs, (0..=12).collect::<Vec<u32>>());
3651    }
3652
3653    #[test]
3654    fn compacted_scann_ah_repacks_blocks_across_run_boundaries() {
3655        use crate::structures::vector::scann::{FAST_SCAN_LANES, pack_fast_scan_block};
3656
3657        const DIM: usize = 8;
3658        const DIMS_PER_BLOCK: usize = 2;
3659        const BLOCKS: usize = DIM / DIMS_PER_BLOCK;
3660
3661        let make_rows = |start: usize, count: usize| -> Vec<u8> {
3662            (0..count)
3663                .flat_map(|row| {
3664                    (0..BLOCKS).map(move |block| ((start + row * 3 + block * 5) & 0x0f) as u8)
3665                })
3666                .collect()
3667        };
3668        let encode = |rows: &[u8]| -> Vec<u8> {
3669            let count = rows.len() / BLOCKS;
3670            let full_rows = count / FAST_SCAN_LANES * FAST_SCAN_LANES;
3671            let mut codes = Vec::new();
3672            for group in rows[..full_rows * BLOCKS].chunks_exact(FAST_SCAN_LANES * BLOCKS) {
3673                pack_fast_scan_block(group, BLOCKS, &mut codes).unwrap();
3674            }
3675            for row in rows[full_rows * BLOCKS..].chunks_exact(BLOCKS) {
3676                for pair in row.chunks(2) {
3677                    codes.push(pair[0] | (pair.get(1).copied().unwrap_or(0) << 4));
3678                }
3679            }
3680            codes
3681        };
3682        let make_source = |rows: &[u8]| {
3683            let count = rows.len() / BLOCKS;
3684            let docs: Vec<u32> = (0..count as u32).collect();
3685            let ordinals = vec![0u16; count];
3686            let codes = encode(rows);
3687            let run = BuildRun {
3688                cluster_id: 0,
3689                doc_ids: &docs,
3690                ordinals: &ordinals,
3691                codes: &codes,
3692            };
3693            let header = AnnDiskHeader {
3694                kind: AnnKind::ScannAh,
3695                routing: IvfRoutingMode::Flat,
3696                dim: DIM,
3697                code_size: DIMS_PER_BLOCK,
3698                num_clusters: 1,
3699                quantizer_version: 41,
3700                codebook_version: 73,
3701                vector_count: count,
3702            };
3703            let mut bytes = Vec::new();
3704            write_built_runs(header, &[run], &mut bytes).unwrap();
3705            AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, count as u32).unwrap()
3706        };
3707
3708        // The first source has a complete FastScan group plus a tail; the
3709        // second tail completes a new cross-source group plus an output tail.
3710        let left_rows = make_rows(1, 35);
3711        let right_rows = make_rows(9, 30);
3712        let left = make_source(&left_rows);
3713        let right = make_source(&right_rows);
3714        let mut bytes = Vec::new();
3715        write_compacted_ann_cancellable(&[(&left, 0), (&right, 35)], &mut bytes, None).unwrap();
3716        let compacted = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, 65).unwrap();
3717
3718        assert_eq!(compacted.runs.len(), 1);
3719        assert!((compacted.health().fragmentation() - 1.0).abs() < 1e-9);
3720        let run = &compacted.runs[0];
3721        let codes = &compacted.raw.as_slice()[run.codes.clone()];
3722        let mut decoded = Vec::new();
3723        for row in 0..run.count {
3724            unpack_scann_ah_row(codes, run.count, BLOCKS, row, &mut decoded).unwrap();
3725        }
3726        let mut expected = left_rows;
3727        expected.extend_from_slice(&right_rows);
3728        assert_eq!(decoded, expected);
3729        for row in 0..65 {
3730            assert_eq!(
3731                run_doc_id(compacted.raw.as_slice(), run, row).unwrap(),
3732                row as u32
3733            );
3734        }
3735    }
3736
3737    /// Throughput comparison, prod-shaped: 320-byte codes, 4 sources.
3738    /// Ignored: run with `cargo test --release -- --ignored ann_merge_throughput --nocapture`.
3739    #[test]
3740    #[ignore]
3741    fn ann_merge_throughput_byte_copy_vs_compaction() {
3742        let code_size = 320usize;
3743        let clusters = 4_096u32;
3744        let vectors_per_source = 262_144usize;
3745        let sources_count = 4usize;
3746
3747        let mut sources_bytes = Vec::new();
3748        for source_index in 0..sources_count {
3749            let mut per_cluster: Vec<(Vec<u32>, Vec<u16>, Vec<u8>)> = Vec::new();
3750            let vectors_per_cluster = vectors_per_source / clusters as usize;
3751            let mut doc = 0u32;
3752            for cluster in 0..clusters {
3753                let mut docs = Vec::with_capacity(vectors_per_cluster);
3754                let mut ords = Vec::with_capacity(vectors_per_cluster);
3755                let mut codes = Vec::with_capacity(vectors_per_cluster * code_size);
3756                for _ in 0..vectors_per_cluster {
3757                    docs.push(doc);
3758                    ords.push(0u16);
3759                    codes.extend(std::iter::repeat_n(
3760                        (doc ^ cluster ^ source_index as u32) as u8,
3761                        code_size,
3762                    ));
3763                    doc += 1;
3764                }
3765                per_cluster.push((docs, ords, codes));
3766            }
3767            let runs: Vec<BuildRun<'_>> = per_cluster
3768                .iter()
3769                .enumerate()
3770                .map(|(cluster, (docs, ords, codes))| BuildRun {
3771                    cluster_id: cluster as u32,
3772                    doc_ids: docs,
3773                    ordinals: ords,
3774                    codes,
3775                })
3776                .collect();
3777            let header = AnnDiskHeader {
3778                kind: AnnKind::BinaryIvf,
3779                routing: IvfRoutingMode::Hnsw,
3780                dim: code_size * 8,
3781                code_size,
3782                num_clusters: clusters,
3783                quantizer_version: 42,
3784                codebook_version: 0,
3785                vector_count: vectors_per_source,
3786            };
3787            let mut bytes = Vec::new();
3788            write_built_runs(header, &runs, &mut bytes).unwrap();
3789            sources_bytes.push(bytes);
3790        }
3791        let sources_open: Vec<AnnDiskIndex> = sources_bytes
3792            .iter()
3793            .map(|bytes| {
3794                AnnDiskIndex::open(
3795                    OwnedBytes::new(bytes.clone()),
3796                    AnnKind::BinaryIvf,
3797                    (vectors_per_source * sources_count) as u32,
3798                )
3799                .unwrap()
3800            })
3801            .collect();
3802        let sources: Vec<(&AnnDiskIndex, u32)> = sources_open
3803            .iter()
3804            .enumerate()
3805            .map(|(index, source)| (source, (index * vectors_per_source) as u32))
3806            .collect();
3807        let payload_bytes = sources_bytes.iter().map(Vec::len).sum::<usize>();
3808
3809        // Fault the output buffer in before timing anything: the first pass
3810        // over a fresh Vec pays demand paging + kernel page zeroing for the
3811        // whole capacity, which the earlier version of this bench silently
3812        // charged to whichever writer ran first (making compaction look 2×
3813        // faster than the byte-copy purely by running second).
3814        let mut out = vec![0u8; payload_bytes + (1 << 20)];
3815        out.clear();
3816        // Interleave rounds and keep the best of each so neither path is
3817        // systematically first.
3818        let mut copy_secs = f64::INFINITY;
3819        let mut compact_secs = f64::INFINITY;
3820        let mut compacted_bytes = Vec::new();
3821        for _ in 0..3 {
3822            out.clear();
3823            let start = std::time::Instant::now();
3824            write_merged_ann(&sources, &mut out).unwrap();
3825            copy_secs = copy_secs.min(start.elapsed().as_secs_f64());
3826
3827            out.clear();
3828            let start = std::time::Instant::now();
3829            write_compacted_ann_cancellable(&sources, &mut out, None).unwrap();
3830            compact_secs = compact_secs.min(start.elapsed().as_secs_f64());
3831            compacted_bytes = out.clone();
3832        }
3833        let compacted = AnnDiskIndex::open(
3834            OwnedBytes::new(compacted_bytes),
3835            AnnKind::BinaryIvf,
3836            (vectors_per_source * sources_count) as u32,
3837        )
3838        .unwrap();
3839        assert!((compacted.health().fragmentation() - 1.0).abs() < 1e-9);
3840
3841        let gib = payload_bytes as f64 / (1u64 << 30) as f64;
3842        println!(
3843            "ann merge {:.2} GiB: byte-copy {:.3}s ({:.2} GiB/s), compaction {:.3}s \
3844             ({:.2} GiB/s), overhead {:.1}%",
3845            gib,
3846            copy_secs,
3847            gib / copy_secs,
3848            compact_secs,
3849            gib / compact_secs,
3850            100.0 * (compact_secs - copy_secs) / copy_secs,
3851        );
3852    }
3853
3854    /// Compaction is undefined for block-packed TQ codes and must refuse.
3855    #[test]
3856    fn compaction_refuses_non_binary_payloads() {
3857        // A binary payload whose header is rewritten to the TQ kind would not
3858        // validate, so exercise the guard through the real gate: any source
3859        // list whose first header is not BinaryIvf is refused before a byte
3860        // is written. Reuse a TQ payload from the flat-TQ writer used by the
3861        // pruning tests.
3862        let codec = std::sync::Arc::new(crate::structures::TqCodec::new(8));
3863        let mut builder = crate::structures::TqFlatBuilder::new(codec);
3864        builder
3865            .add_batch(
3866                &[(0, 0), (1, 0)],
3867                &[
3868                    1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, //
3869                    0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
3870                ],
3871            )
3872            .unwrap();
3873        builder.finish();
3874        let mut bytes = Vec::new();
3875        write_built_tq_flat(&builder, &mut bytes).unwrap();
3876        let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
3877        let error = write_compacted_ann_cancellable(&[(&disk, 0)], &mut Vec::new(), None)
3878            .expect_err("TQ payloads must not compact");
3879        assert!(error.to_string().contains("binary"), "{error}");
3880    }
3881
3882    /// Health math on a hand-built payload: two runs of one cluster (a
3883    /// byte-copy merge shape) plus one dominant leaf, checked against the
3884    /// Faiss imbalance definition computed by hand.
3885    #[test]
3886    fn ann_health_measures_skew_and_fragmentation() {
3887        // Segment A: 6 vectors in cluster 0 and 2 in cluster 3; segment B: 2
3888        // more in cluster 0. A byte-copy merge preserves each source extent,
3889        // producing the fragmented shape (two physical runs for cluster 0)
3890        // that build alone can never emit.
3891        let a0_docs = [0u32, 1, 2, 3, 4, 5];
3892        let a0_ords = [0u16; 6];
3893        let a0_codes = [0xAAu8; 6];
3894        let a3_docs = [6u32, 7];
3895        let a3_ords = [0u16; 2];
3896        let a3_codes = [0x0Fu8; 2];
3897        let a_runs = [
3898            BuildRun {
3899                cluster_id: 0,
3900                doc_ids: &a0_docs,
3901                ordinals: &a0_ords,
3902                codes: &a0_codes,
3903            },
3904            BuildRun {
3905                cluster_id: 3,
3906                doc_ids: &a3_docs,
3907                ordinals: &a3_ords,
3908                codes: &a3_codes,
3909            },
3910        ];
3911        let mut header = binary_header(8);
3912        header.num_clusters = 8;
3913        let mut a_bytes = Vec::new();
3914        write_built_runs(header, &a_runs, &mut a_bytes).unwrap();
3915        let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 8).unwrap();
3916
3917        let b0_docs = [0u32, 1];
3918        let b0_ords = [0u16; 2];
3919        let b0_codes = [0xBBu8; 2];
3920        let b_runs = [BuildRun {
3921            cluster_id: 0,
3922            doc_ids: &b0_docs,
3923            ordinals: &b0_ords,
3924            codes: &b0_codes,
3925        }];
3926        let mut header_b = binary_header(2);
3927        header_b.num_clusters = 8;
3928        let mut b_bytes = Vec::new();
3929        write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
3930        let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 2).unwrap();
3931
3932        let mut merged_bytes = Vec::new();
3933        write_merged_ann(&[(&a, 0), (&b, 8)], &mut merged_bytes).unwrap();
3934        let disk =
3935            AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::BinaryIvf, 10).unwrap();
3936
3937        let health = disk.health();
3938        assert_eq!(health.vectors, 10);
3939        assert_eq!(health.clusters_nonempty, 2);
3940        assert_eq!(health.clusters_total, 8);
3941        assert_eq!(health.runs, 3);
3942        assert_eq!(health.largest_cluster, 0);
3943        assert_eq!(health.largest_cluster_vectors, 8);
3944        assert!((health.largest_cluster_share() - 0.8).abs() < 1e-9);
3945        // 3 runs over 2 non-empty clusters.
3946        assert!((health.fragmentation() - 1.5).abs() < 1e-9);
3947        // Faiss: K * sum(n_i^2) / N^2 = 2 * (64 + 4) / 100 = 1.36
3948        assert!(
3949            (health.imbalance - 1.36).abs() < 1e-9,
3950            "{}",
3951            health.imbalance
3952        );
3953        // codes columns: 6 + 2 + 2 bytes at code_size 1.
3954        assert_eq!(health.payload_bytes, 10);
3955    }
3956
3957    #[test]
3958    fn ann_health_is_balanced_at_one() {
3959        let docs: Vec<Vec<u32>> = (0..4).map(|c| vec![c * 2, c * 2 + 1]).collect();
3960        let ords = [0u16; 2];
3961        let codes = [0x55u8; 2];
3962        let runs: Vec<BuildRun<'_>> = docs
3963            .iter()
3964            .enumerate()
3965            .map(|(cluster, doc_ids)| BuildRun {
3966                cluster_id: cluster as u32,
3967                doc_ids,
3968                ordinals: &ords,
3969                codes: &codes,
3970            })
3971            .collect();
3972        let mut header = binary_header(8);
3973        header.num_clusters = 4;
3974        let mut bytes = Vec::new();
3975        write_built_runs(header, &runs, &mut bytes).unwrap();
3976        let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 8).unwrap();
3977        let health = disk.health();
3978        assert!((health.imbalance - 1.0).abs() < 1e-9);
3979        assert!((health.fragmentation() - 1.0).abs() < 1e-9);
3980        assert!((health.largest_cluster_share() - 0.25).abs() < 1e-9);
3981    }
3982
3983    fn binary_header(vector_count: usize) -> AnnDiskHeader {
3984        AnnDiskHeader {
3985            kind: AnnKind::BinaryIvf,
3986            routing: IvfRoutingMode::Hnsw,
3987            dim: 8,
3988            code_size: 1,
3989            num_clusters: 2,
3990            quantizer_version: 42,
3991            codebook_version: 0,
3992            vector_count,
3993        }
3994    }
3995
3996    fn payload_end(index: &AnnDiskIndex) -> usize {
3997        index.runs.iter().map(|run| run.codes.end).max().unwrap()
3998    }
3999
4000    #[test]
4001    fn ann_prefetch_ranges_are_sorted_and_only_merge_page_near_extents() {
4002        let mut ranges = vec![
4003            15_000..16_000,
4004            0..1_000,
4005            9_000..10_000,
4006            1_000..2_000,
4007            7_000..8_000,
4008        ];
4009        coalesce_prefetch_ranges(&mut ranges);
4010        assert_eq!(ranges, [0..2_000, 7_000..10_000, 15_000..16_000]);
4011    }
4012
4013    #[test]
4014    fn normal_merge_copies_ann_payload_columns_byte_for_byte() {
4015        let first_doc_0 = [0u32];
4016        let first_doc_1 = [1u32];
4017        let first_ord_0 = [0u16];
4018        let first_ord_1 = [2u16];
4019        let first_code_0 = [0x00u8];
4020        let first_code_1 = [0xffu8];
4021        let first_runs = [
4022            BuildRun {
4023                cluster_id: 0,
4024                doc_ids: &first_doc_0,
4025                ordinals: &first_ord_0,
4026                codes: &first_code_0,
4027            },
4028            BuildRun {
4029                cluster_id: 1,
4030                doc_ids: &first_doc_1,
4031                ordinals: &first_ord_1,
4032                codes: &first_code_1,
4033            },
4034        ];
4035        let mut first_bytes = Vec::new();
4036        write_built_runs(binary_header(2), &first_runs, &mut first_bytes).unwrap();
4037        let first = AnnDiskIndex::open(OwnedBytes::new(first_bytes.clone()), AnnKind::BinaryIvf, 2)
4038            .unwrap();
4039
4040        let second_docs = [0u32, 1u32];
4041        let second_ords = [1u16, 0u16];
4042        let second_codes = [0x0fu8, 0xf0u8];
4043        let second_runs = [BuildRun {
4044            cluster_id: 0,
4045            doc_ids: &second_docs,
4046            ordinals: &second_ords,
4047            codes: &second_codes,
4048        }];
4049        let mut second_bytes = Vec::new();
4050        write_built_runs(binary_header(2), &second_runs, &mut second_bytes).unwrap();
4051        let second =
4052            AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::BinaryIvf, 2)
4053                .unwrap();
4054
4055        let mut merged_bytes = Vec::new();
4056        write_merged_ann(&[(&first, 0), (&second, 2)], &mut merged_bytes).unwrap();
4057        let merged =
4058            AnnDiskIndex::open(OwnedBytes::new(merged_bytes.clone()), AnnKind::BinaryIvf, 4)
4059                .unwrap();
4060
4061        let mut expected_payload = first_bytes[ANN_HEADER_SIZE..payload_end(&first)].to_vec();
4062        expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
4063        assert_eq!(
4064            &merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
4065            expected_payload.as_slice(),
4066            "normal merge must not decode or rewrite any corpus-sized ANN column",
4067        );
4068
4069        let mut docs: Vec<u32> = merged
4070            .search_binary_clusters::<false>(&[0], 4, &[0, 1])
4071            .unwrap()
4072            .into_iter()
4073            .map(|result| result.0)
4074            .collect();
4075        docs.sort_unstable();
4076        assert_eq!(docs, [0, 1, 2, 3]);
4077        let serial = merged
4078            .search_binary_clusters_with_tuning::<false>(&[0], 4, &[0, 1], usize::MAX)
4079            .unwrap();
4080        let parallel = rayon::ThreadPoolBuilder::new()
4081            .num_threads(4)
4082            .build()
4083            .unwrap()
4084            .install(|| merged.search_binary_clusters_with_tuning::<false>(&[0], 4, &[0, 1], 1))
4085            .unwrap();
4086        assert_eq!(parallel, serial, "parallel binary top-k changed results");
4087
4088        // A merged source's directory is cluster-sorted while its payload is
4089        // source-order. A later merge must follow physical offsets and still
4090        // preserve every source column byte-for-byte.
4091        let mut second_merge_bytes = Vec::new();
4092        write_merged_ann(&[(&merged, 0), (&first, 4)], &mut second_merge_bytes).unwrap();
4093        let second_merge = AnnDiskIndex::open(
4094            OwnedBytes::new(second_merge_bytes.clone()),
4095            AnnKind::BinaryIvf,
4096            6,
4097        )
4098        .unwrap();
4099        let mut expected_second_payload =
4100            merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)].to_vec();
4101        expected_second_payload
4102            .extend_from_slice(&first_bytes[ANN_HEADER_SIZE..payload_end(&first)]);
4103        assert_eq!(
4104            &second_merge_bytes[ANN_HEADER_SIZE..payload_end(&second_merge)],
4105            expected_second_payload.as_slice(),
4106        );
4107        let mut docs: Vec<u32> = second_merge
4108            .search_binary_clusters::<false>(&[0], 6, &[0, 1])
4109            .unwrap()
4110            .into_iter()
4111            .map(|result| result.0)
4112            .collect();
4113        docs.sort_unstable();
4114        assert_eq!(docs, [0, 1, 2, 3, 4, 5]);
4115    }
4116
4117    #[test]
4118    fn legacy_ivf_tq_payload_is_rejected_while_opening() {
4119        let dim = 8;
4120        let marked_version = crate::structures::mark_ivf_tq_cosine_generation(7);
4121        let centroids = crate::structures::CoarseCentroids {
4122            num_clusters: 1,
4123            dim,
4124            centroids: vec![0.0; dim],
4125            version: marked_version,
4126            soar_config: None,
4127            routing_index: None,
4128        };
4129        let mut bytes = crate::segment::ann_build::build_ivf_tq(
4130            dim,
4131            IvfRoutingMode::Flat,
4132            &centroids,
4133            &[(0, 0)],
4134            &[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
4135        )
4136        .unwrap();
4137
4138        // The fixed header stores quantizer_version at bytes 24..32.
4139        bytes[24..32].copy_from_slice(&7u64.to_le_bytes());
4140        let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
4141            .err()
4142            .expect("legacy IVF-TQ payload must fail while opening")
4143            .to_string();
4144        assert!(error.contains("unsupported legacy generation"), "{error}");
4145    }
4146
4147    #[test]
4148    fn binary_combined_search_deduplicates_soar_and_bounds_document_results() {
4149        // doc 0 / ordinal 0 occurs in both leaves, as it can under SOAR.
4150        // Without `(doc, ordinal)` dedup it would beat doc 1 for both Sum and
4151        // default LogSumExp; with correct dedup doc 1's two distinct values win.
4152        let cluster_0_docs = [0u32, 0, 1];
4153        let cluster_0_ordinals = [0u16, 1, 0];
4154        let cluster_0_codes = [0x00u8, 0xff, 0x03];
4155        let cluster_1_docs = [0u32, 1, 2];
4156        let cluster_1_ordinals = [0u16, 1, 0];
4157        let cluster_1_codes = [0x00u8, 0x0c, 0xf0];
4158        let runs = [
4159            BuildRun {
4160                cluster_id: 0,
4161                doc_ids: &cluster_0_docs,
4162                ordinals: &cluster_0_ordinals,
4163                codes: &cluster_0_codes,
4164            },
4165            BuildRun {
4166                cluster_id: 1,
4167                doc_ids: &cluster_1_docs,
4168                ordinals: &cluster_1_ordinals,
4169                codes: &cluster_1_codes,
4170            },
4171        ];
4172        let mut bytes = Vec::new();
4173        write_built_runs(binary_header(6), &runs, &mut bytes).unwrap();
4174        let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 3).unwrap();
4175        let parallel_pool = rayon::ThreadPoolBuilder::new()
4176            .num_threads(4)
4177            .build()
4178            .unwrap();
4179        let serial_documents = disk
4180            .search_binary_clusters_with_tuning::<true>(&[0], 3, &[0, 1], usize::MAX)
4181            .unwrap();
4182        let parallel_documents = parallel_pool
4183            .install(|| disk.search_binary_clusters_with_tuning::<true>(&[0], 3, &[0, 1], 1))
4184            .unwrap();
4185        assert_eq!(parallel_documents, serial_documents);
4186
4187        // Under Sum the SOAR duplicate decides the winner outright: doc 0
4188        // counted twice (2.0) would beat doc 1's two distinct values (1.5);
4189        // deduplicated, doc 1 wins.
4190        let sum = crate::query::MultiValueCombiner::Sum;
4191        let (result, probed) = disk
4192            .search_binary_combined_documents(1, &[0], &[0, 1], sum)
4193            .unwrap();
4194        let parallel_sum = parallel_pool
4195            .install(|| disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], sum, 1))
4196            .unwrap();
4197        assert_eq!(parallel_sum, (result.clone(), probed.clone()));
4198        assert_eq!(result.len(), 1, "combined search must honor k");
4199        assert_eq!(
4200            result[0].doc_id, 1,
4201            "SOAR duplicate changed {sum:?} ranking: {result:?}",
4202        );
4203        // Exact leaf scores are handed back only for the retained document,
4204        // deduplicated and sorted, so reranking can skip re-reading them.
4205        assert_eq!(
4206            probed
4207                .iter()
4208                .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
4209                .collect::<Vec<_>>(),
4210            vec![(1, 0), (1, 1)],
4211            "{sum:?}",
4212        );
4213
4214        // Under the default smooth-max combiner doc 0's perfect ordinal wins
4215        // regardless of the duplicate, so dedup is pinned through the exact
4216        // score instead: it must equal the combiner over the two *distinct*
4217        // ordinals (1.0 and 0.0) — a double-counted (doc 0, ordinal 0) would
4218        // inflate it.
4219        let smooth_max = crate::query::MultiValueCombiner::default();
4220        let (result, probed) = disk
4221            .search_binary_combined_documents(1, &[0], &[0, 1], smooth_max)
4222            .unwrap();
4223        let parallel_smooth_max = parallel_pool
4224            .install(|| {
4225                disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], smooth_max, 1)
4226            })
4227            .unwrap();
4228        assert_eq!(parallel_smooth_max, (result.clone(), probed.clone()));
4229        assert_eq!(result.len(), 1, "combined search must honor k");
4230        assert_eq!(result[0].doc_id, 0, "{result:?}");
4231        let expected = smooth_max.combine(&[(0, 1.0), (1, 0.0)]);
4232        assert!(
4233            (result[0].score - expected).abs() < 1e-6,
4234            "SOAR duplicate leaked into {smooth_max:?}: got {}, expected {expected}",
4235            result[0].score,
4236        );
4237        assert_eq!(
4238            probed
4239                .iter()
4240                .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
4241                .collect::<Vec<_>>(),
4242            vec![(0, 0), (0, 1)],
4243            "{smooth_max:?}",
4244        );
4245
4246        let (top_two, probed) = disk
4247            .search_binary_combined_documents(
4248                2,
4249                &[0],
4250                &[0, 1],
4251                crate::query::MultiValueCombiner::Sum,
4252            )
4253            .unwrap();
4254        assert_eq!(
4255            top_two
4256                .iter()
4257                .map(|candidate| candidate.doc_id)
4258                .collect::<Vec<_>>(),
4259            vec![1, 0],
4260        );
4261        assert_eq!(top_two.len(), 2, "full probing must still return at most k");
4262        // doc 0 ordinal 0 was probed twice (a SOAR duplicate) and must appear
4263        // once, with the two retained documents in ascending order.
4264        assert_eq!(
4265            probed
4266                .iter()
4267                .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
4268                .collect::<Vec<_>>(),
4269            vec![(0, 0), (0, 1), (1, 0), (1, 1)],
4270        );
4271    }
4272
4273    #[test]
4274    fn combined_ordinal_reduction_handles_out_of_order_runs_for_every_combiner() {
4275        let out_of_order_with_duplicate = vec![
4276            (7, 1, 0.4),
4277            (3, 0, 0.8),
4278            (7, 0, 0.6),
4279            (3, 1, 0.2),
4280            (7, 1, 0.5), // higher SOAR estimate replaces 0.4, never adds to it
4281        ];
4282        for combiner in [
4283            crate::query::MultiValueCombiner::Max,
4284            crate::query::MultiValueCombiner::Sum,
4285            crate::query::MultiValueCombiner::Avg,
4286            crate::query::MultiValueCombiner::default(),
4287            crate::query::MultiValueCombiner::WeightedTopK { k: 2, decay: 0.7 },
4288        ] {
4289            let actual = combine_scored_ordinals(out_of_order_with_duplicate.clone(), 2, combiner);
4290            let mut expected = vec![
4291                AnnDocumentCandidate {
4292                    doc_id: 3,
4293                    score: combiner.combine(&[(0, 0.8), (1, 0.2)]),
4294                },
4295                AnnDocumentCandidate {
4296                    doc_id: 7,
4297                    score: combiner.combine(&[(0, 0.6), (1, 0.5)]),
4298                },
4299            ];
4300            expected.sort_unstable_by(|left, right| {
4301                right
4302                    .score
4303                    .total_cmp(&left.score)
4304                    .then_with(|| left.doc_id.cmp(&right.doc_id))
4305            });
4306            assert_eq!(actual, expected, "combiner {combiner:?}");
4307        }
4308    }
4309
4310    fn build_tq_payload(dim: usize, count: usize, seed: u64) -> (Vec<u8>, Vec<Vec<f32>>) {
4311        let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
4312        let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
4313        let mut state = seed;
4314        let mut vectors = Vec::new();
4315        let mut flat = Vec::new();
4316        for _ in 0..count {
4317            let vector: Vec<f32> = (0..dim)
4318                .map(|_| {
4319                    state = state
4320                        .wrapping_mul(6364136223846793005)
4321                        .wrapping_add(1442695040888963407);
4322                    ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
4323                })
4324                .collect();
4325            flat.extend_from_slice(&vector);
4326            vectors.push(vector);
4327        }
4328        let labels: Vec<(u32, u16)> = (0..count).map(|index| (index as u32, 0)).collect();
4329        builder.add_batch(&labels, &flat).unwrap();
4330        builder.finish();
4331        let mut bytes = Vec::new();
4332        write_built_tq_flat(&builder, &mut bytes).unwrap();
4333        (bytes, vectors)
4334    }
4335
4336    #[test]
4337    fn tq_combined_scan_ranks_complete_documents_instead_of_individual_values() {
4338        let dim = 8;
4339        let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
4340        let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
4341        let labels = [(0u32, 0u16), (1, 0), (1, 1)];
4342        let vectors = [
4343            1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // doc 0: best single value
4344            0.8, 0.6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // doc 1: two good values
4345            0.8, -0.6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
4346        ];
4347        builder.add_batch(&labels, &vectors).unwrap();
4348        builder.finish();
4349        let mut bytes = Vec::new();
4350        write_built_tq_flat(&builder, &mut bytes).unwrap();
4351        let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
4352        let query = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
4353        let plan = crate::structures::TqQueryPlan::build(&codec, &query);
4354
4355        let max = disk.search_tq_distinct(1, &plan).unwrap();
4356        let sum = disk
4357            .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
4358            .unwrap();
4359        let default_combiner = disk
4360            .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::default())
4361            .unwrap();
4362        assert_eq!(max[0].0, 0, "fixture must favor doc 0 by Max: {max:?}");
4363        assert_eq!(
4364            sum[0].doc_id, 1,
4365            "two complete values must make doc 1 win by Sum: {sum:?}"
4366        );
4367        // The default combiner is a smooth maximum: doc 0's single best value
4368        // (1.0) outranks doc 1's two 0.8 values — value count alone no longer
4369        // wins. Sum above is what pins that both of doc 1's values were seen.
4370        assert_eq!(
4371            default_combiner[0].doc_id, 0,
4372            "default smooth-max must follow the best value: {default_combiner:?}"
4373        );
4374        assert!(sum[0].score > max[0].2);
4375
4376        // Pure-copy upgrade merges may append a re-encoded older segment
4377        // after copied runs, so global run order need not follow doc IDs.
4378        // Documents remain complete within each run and must still combine.
4379        let (single_bytes, _) = build_tq_payload(dim, 1, 91);
4380        let single = AnnDiskIndex::open(OwnedBytes::new(single_bytes), AnnKind::TqFlat, 1).unwrap();
4381        let mut merged_bytes = Vec::new();
4382        write_merged_ann(&[(&disk, 1), (&single, 0)], &mut merged_bytes).unwrap();
4383        let merged = AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::TqFlat, 3).unwrap();
4384        let merged_sum = merged
4385            .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
4386            .unwrap();
4387        assert_eq!(merged_sum[0].doc_id, 2);
4388    }
4389
4390    #[test]
4391    fn tq_payload_roundtrip_search_and_pure_copy_merge() {
4392        let dim = 20; // pads to 32; exercises padding + partial final block
4393        let count = 21;
4394        let (bytes, vectors) = build_tq_payload(dim, count, 42);
4395        let index = AnnDiskIndex::open(
4396            OwnedBytes::new(bytes.clone()),
4397            AnnKind::TqFlat,
4398            count as u32,
4399        )
4400        .unwrap();
4401        assert_eq!(index.header().vector_count, count);
4402
4403        // The stored estimate must rank an exact-duplicate query's own row first.
4404        let codec = crate::structures::TqCodec::new(dim);
4405        for target in [0usize, 7, 20] {
4406            let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[target]);
4407            let results = index.search_tq_distinct(3, &plan).unwrap();
4408            assert_eq!(
4409                results[0].0, target as u32,
4410                "query duplicating vector {target} must rank it first: {results:?}"
4411            );
4412        }
4413
4414        // Ordinary merge must not decode or rewrite the corpus columns.
4415        let (second_bytes, _) = build_tq_payload(dim, 5, 77);
4416        let second =
4417            AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::TqFlat, 5).unwrap();
4418        let mut merged_bytes = Vec::new();
4419        write_merged_ann(&[(&index, 0), (&second, count as u32)], &mut merged_bytes).unwrap();
4420        let merged = AnnDiskIndex::open(
4421            OwnedBytes::new(merged_bytes.clone()),
4422            AnnKind::TqFlat,
4423            count as u32 + 5,
4424        )
4425        .unwrap();
4426        let mut expected_payload = bytes[ANN_HEADER_SIZE..payload_end(&index)].to_vec();
4427        expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
4428        assert_eq!(
4429            &merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
4430            expected_payload.as_slice(),
4431            "TQ merge must be a pure byte copy of the source columns",
4432        );
4433        let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[7]);
4434        let results = merged.search_tq_distinct(1, &plan).unwrap();
4435        assert_eq!(results[0].0, 7, "merged payload must keep doc bases");
4436    }
4437
4438    /// The scan loops read doc IDs without a per-lane maximum check, so the
4439    /// open-time column validation is what keeps a corrupt payload from
4440    /// producing out-of-range document IDs.
4441    /// A FastScan v1 `ScannAh` column has the same byte length as v2 for even
4442    /// block counts, so the container header must carry the layout revision
4443    /// and refuse stale generations loudly instead of misreading nibbles.
4444    #[test]
4445    fn open_refuses_scann_ah_payload_with_pre_v2_fast_scan_layout() {
4446        use crate::structures::vector::scann::{
4447            FAST_SCAN_LANES, FAST_SCAN_LAYOUT_VERSION, pack_fast_scan_block,
4448        };
4449        const BLOCKS: usize = 4;
4450        const DIMS_PER_BLOCK: usize = 2;
4451        const DIM: usize = BLOCKS * DIMS_PER_BLOCK;
4452        // Header layout: magic u32, kind u8, routing u8, version u16, dim u32,
4453        // code_size u32, num_clusters u32, then the layout-revision u32.
4454        const REVISION_OFFSET: usize = 20;
4455
4456        let count = FAST_SCAN_LANES;
4457        let rows: Vec<u8> = (0..count * BLOCKS).map(|i| (i % 16) as u8).collect();
4458        let mut codes = Vec::new();
4459        pack_fast_scan_block(&rows, BLOCKS, &mut codes).unwrap();
4460        let docs: Vec<u32> = (0..count as u32).collect();
4461        let ordinals = vec![0u16; count];
4462        let run = BuildRun {
4463            cluster_id: 0,
4464            doc_ids: &docs,
4465            ordinals: &ordinals,
4466            codes: &codes,
4467        };
4468        let header = AnnDiskHeader {
4469            kind: AnnKind::ScannAh,
4470            routing: IvfRoutingMode::Flat,
4471            dim: DIM,
4472            code_size: DIMS_PER_BLOCK,
4473            num_clusters: 1,
4474            quantizer_version: 1,
4475            codebook_version: 1,
4476            vector_count: count,
4477        };
4478        let mut bytes = Vec::new();
4479        write_built_runs(header, &[run], &mut bytes).unwrap();
4480
4481        let stamped = u32::from_le_bytes(
4482            bytes[REVISION_OFFSET..REVISION_OFFSET + 4]
4483                .try_into()
4484                .unwrap(),
4485        );
4486        assert_eq!(stamped, u32::from(FAST_SCAN_LAYOUT_VERSION));
4487        AnnDiskIndex::open(
4488            OwnedBytes::new(bytes.clone()),
4489            AnnKind::ScannAh,
4490            count as u32,
4491        )
4492        .expect("current layout revision opens");
4493
4494        // A v1 build left the field at zero.
4495        let mut legacy = bytes;
4496        legacy[REVISION_OFFSET..REVISION_OFFSET + 4].copy_from_slice(&0u32.to_le_bytes());
4497        let error =
4498            match AnnDiskIndex::open(OwnedBytes::new(legacy), AnnKind::ScannAh, count as u32) {
4499                Ok(_) => panic!("a pre-v2 FastScan payload must be refused at open"),
4500                Err(error) => error,
4501            };
4502        let message = error.to_string();
4503        assert!(
4504            message.contains("FastScan layout revision 0") && message.contains("Rebuild"),
4505            "unexpected error: {message}"
4506        );
4507
4508        // Every other kind still requires the field to be zero.
4509        let (tq, _) = build_tq_payload(8, 40, 5);
4510        let mut tq_bad = tq;
4511        tq_bad[REVISION_OFFSET..REVISION_OFFSET + 4].copy_from_slice(&7u32.to_le_bytes());
4512        let error = AnnDiskIndex::open(OwnedBytes::new(tq_bad), AnnKind::TqFlat, 40)
4513            .err()
4514            .expect("non-zero reserved field must be refused");
4515        assert!(
4516            error.to_string().contains("reserved field is non-zero"),
4517            "unexpected error: {error}"
4518        );
4519    }
4520
4521    #[test]
4522    fn open_rejects_doc_id_column_above_run_maximum() {
4523        let dim = 8;
4524        let count = 40;
4525        let (bytes, _) = build_tq_payload(dim, count, 5);
4526        let disk = AnnDiskIndex::open(
4527            OwnedBytes::new(bytes.clone()),
4528            AnnKind::TqFlat,
4529            count as u32,
4530        )
4531        .unwrap();
4532        let run = &disk.runs[0];
4533        // Corrupt one local doc ID to exceed the run's declared maximum
4534        // without touching the directory (so every other check still passes).
4535        let mut corrupt = bytes;
4536        let offset = run.doc_ids.start + 3 * 4;
4537        corrupt[offset..offset + 4].copy_from_slice(&(run.max_doc_id + 1).to_le_bytes());
4538        let error =
4539            match AnnDiskIndex::open(OwnedBytes::new(corrupt), AnnKind::TqFlat, count as u32) {
4540                Ok(_) => panic!("corrupt doc-ID column must fail at open"),
4541                Err(error) => error,
4542            };
4543        assert!(
4544            error.to_string().contains("above its declared maximum"),
4545            "unexpected error: {error}"
4546        );
4547    }
4548
4549    #[test]
4550    fn tq_parallel_fold_matches_a_sequential_scan() {
4551        use crate::structures::vector::quantization::{
4552            TQ_BLOCK_LANES, tq_block_bytes, tq_score_block,
4553        };
4554
4555        // Exactly the production fan-out threshold exercises the parallel
4556        // fold/reduce path without relying on a test-only configuration.
4557        let dim = 8;
4558        let count = TQ_PARALLEL_SCAN_MIN_VECTORS;
4559        let (bytes, vectors) = build_tq_payload(dim, count, 87);
4560        let disk =
4561            AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, count as u32).unwrap();
4562        let codec = crate::structures::TqCodec::new(dim);
4563        let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[count / 3]);
4564        let k = 31;
4565
4566        let block_bytes = tq_block_bytes(disk.header().code_size);
4567        assert_eq!(
4568            disk.runs.len(),
4569            1,
4570            "the regression must parallelize chunks inside one run"
4571        );
4572        assert!(
4573            disk.runs[0].codes.len() / block_bytes > TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
4574            "the single run must span multiple parallel chunks"
4575        );
4576        let raw = disk.raw.as_slice();
4577        let mut reference = BoundedAnnCollector::<true, true>::new(k);
4578        let mut scores = [0.0f32; TQ_BLOCK_LANES];
4579        for run in &disk.runs {
4580            let codes = &raw[run.codes.clone()];
4581            for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
4582                tq_score_block(&plan, block, &mut scores);
4583                let lane_base = block_index * TQ_BLOCK_LANES;
4584                let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
4585                for (lane, &score) in scores.iter().enumerate().take(lanes) {
4586                    let index = lane_base + lane;
4587                    reference.insert(
4588                        run_doc_id(raw, run, index).unwrap(),
4589                        read_u16(raw, run.ordinals.start + index * 2),
4590                        score,
4591                    );
4592                }
4593            }
4594        }
4595
4596        assert_eq!(
4597            disk.search_tq_distinct(k, &plan).unwrap(),
4598            reference.into_sorted_results(),
4599        );
4600    }
4601
4602    #[test]
4603    fn ivf_tq_scale_pruning_matches_the_unpruned_scan() {
4604        use crate::structures::vector::ivf::{CoarseCentroids, CoarseConfig};
4605        use crate::structures::vector::quantization::{
4606            TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
4607        };
4608        use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
4609
4610        let dim = 32;
4611        let count = 400usize;
4612        let codec = std::sync::Arc::new(TqCodec::new(dim));
4613        let mut state = 5u64;
4614        let mut next = move || {
4615            state = state
4616                .wrapping_mul(6364136223846793005)
4617                .wrapping_add(1442695040888963407);
4618            ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
4619        };
4620        let vectors: Vec<Vec<f32>> = (0..count)
4621            .map(|_| {
4622                let mut v: Vec<f32> = (0..dim).map(|_| next()).collect();
4623                let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
4624                v.iter_mut().for_each(|x| *x /= norm);
4625                v
4626            })
4627            .collect();
4628        let mut centroids = CoarseCentroids::train(&CoarseConfig::new(dim, 8), &vectors, "test");
4629        centroids.version = crate::structures::mark_ivf_tq_cosine_generation(centroids.version);
4630        let mut index = IvfTqIndex::new(
4631            dim,
4632            crate::dsl::IvfRoutingMode::Flat,
4633            centroids.version,
4634            std::sync::Arc::clone(&codec),
4635        );
4636        let mut scratch = TqIvfEncodeScratch::default();
4637        for (i, vector) in vectors.iter().enumerate() {
4638            index.add_vector(
4639                &centroids,
4640                (i / 2) as u32,
4641                (i % 2) as u16,
4642                vector,
4643                &mut scratch,
4644            );
4645        }
4646        let mut bytes = Vec::new();
4647        write_built_ivf_tq(&index, centroids.num_clusters, &mut bytes).unwrap();
4648        let disk =
4649            AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, count as u32).unwrap();
4650
4651        // The supported cosine generation promises descending residual scales,
4652        // enabling the reader to terminate the run at the first losing block.
4653        let block_bytes = tq_ivf_block_bytes(disk.header().code_size);
4654        let raw = disk.raw.as_slice();
4655        for run in &disk.runs {
4656            let codes = &raw[run.codes.clone()];
4657            let mut previous_scale = f32::INFINITY;
4658            for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
4659                let lane_base = block_index * TQ_BLOCK_LANES;
4660                let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
4661                let mut block_scales = block[..TQ_BLOCK_LANES * size_of::<f32>()]
4662                    .chunks_exact(size_of::<f32>())
4663                    .take(lanes)
4664                    .map(|lane| f32::from_le_bytes(lane.try_into().unwrap()));
4665                let first_scale = block_scales.next().unwrap();
4666                assert_eq!(tq_ivf_block_max_scale(block), first_scale);
4667                assert!(first_scale <= previous_scale);
4668                previous_scale = first_scale;
4669                for scale in block_scales {
4670                    assert!(scale <= previous_scale);
4671                    previous_scale = scale;
4672                }
4673            }
4674        }
4675        let k = 10;
4676        let parallel_pool = rayon::ThreadPoolBuilder::new()
4677            .num_threads(4)
4678            .build()
4679            .unwrap();
4680        for query_seed in [1u64, 9, 42] {
4681            let mut qstate = query_seed;
4682            let mut qnext = move || {
4683                qstate = qstate
4684                    .wrapping_mul(6364136223846793005)
4685                    .wrapping_add(1442695040888963407);
4686                ((qstate >> 33) as f32 / (1u64 << 31) as f32) - 0.5
4687            };
4688            let query: Vec<f32> = (0..dim).map(|_| qnext()).collect();
4689            let plan = TqIvfQueryPlan::build(
4690                &centroids,
4691                &codec,
4692                &query,
4693                8,
4694                crate::dsl::IvfRoutingMode::Flat,
4695            );
4696            // Unpruned reference: score every block of every probed run with
4697            // the identical kernel and collector.
4698            let mut reference = BoundedAnnCollector::<true, true>::new(k);
4699            let mut unpruned_scores = Vec::new();
4700            let mut scores = [0.0f32; TQ_BLOCK_LANES];
4701            for (cluster_id, cluster_dot) in plan.cluster_dots() {
4702                for run in disk.cluster_runs(cluster_id) {
4703                    let codes = &raw[run.codes.clone()];
4704                    for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
4705                        tq_score_ivf_block(plan.tq_plan(), block, cluster_dot, &mut scores);
4706                        let lane_base = block_index * TQ_BLOCK_LANES;
4707                        let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
4708                        for (lane, &score) in scores.iter().enumerate().take(lanes) {
4709                            let idx = lane_base + lane;
4710                            reference.insert(
4711                                run_doc_id(raw, run, idx).unwrap(),
4712                                read_u16(raw, run.ordinals.start + idx * 2),
4713                                score,
4714                            );
4715                            unpruned_scores.push((
4716                                run_doc_id(raw, run, idx).unwrap(),
4717                                read_u16(raw, run.ordinals.start + idx * 2),
4718                                score,
4719                            ));
4720                        }
4721                    }
4722                }
4723            }
4724
4725            let pruned = disk.search_ivf_tq_distinct(k, &plan).unwrap();
4726            let forced_parallel = parallel_pool.install(|| {
4727                // Small chunks force several tasks from this compact test
4728                // payload instead of allocating a production-sized one.
4729                disk.search_ivf_tq_distinct_with_tuning(k, &plan, 1, 4)
4730                    .unwrap()
4731            });
4732            let reference = reference.into_sorted_results();
4733            assert_eq!(
4734                pruned, reference,
4735                "scale-bound pruning must not change the estimated top-k (seed {query_seed})"
4736            );
4737            assert_eq!(
4738                forced_parallel, reference,
4739                "parallel IVF-TQ pruning must match the unpruned top-k (seed {query_seed})"
4740            );
4741            for combiner in [
4742                crate::query::MultiValueCombiner::Max,
4743                crate::query::MultiValueCombiner::Sum,
4744                crate::query::MultiValueCombiner::Avg,
4745                crate::query::MultiValueCombiner::default(),
4746                crate::query::MultiValueCombiner::WeightedTopK { k: 3, decay: 0.7 },
4747            ] {
4748                let expected = combine_scored_ordinals(unpruned_scores.clone(), k, combiner);
4749                let combined = disk
4750                    .search_ivf_tq_combined_documents(k, &plan, combiner)
4751                    .unwrap();
4752                let parallel_combined = parallel_pool
4753                    .install(|| {
4754                        disk.search_ivf_tq_combined_documents_with_tuning(k, &plan, combiner, 1, 4)
4755                    })
4756                    .unwrap();
4757                assert_eq!(
4758                    combined, expected,
4759                    "combined IVF-TQ scan diverged from the unpruned reference \
4760                     for {combiner:?} (seed {query_seed})",
4761                );
4762                assert_eq!(
4763                    parallel_combined, expected,
4764                    "parallel combined IVF-TQ scan diverged from the unpruned reference \
4765                     for {combiner:?} (seed {query_seed})",
4766                );
4767                assert!(combined.len() <= k);
4768            }
4769        }
4770    }
4771
4772    /// Focused single-segment scan benchmark for the adaptive IVF-TQ fan-out.
4773    /// Run with:
4774    ///
4775    /// `cargo test --release -p hermes-core ivf_tq_parallel_scan_benchmark -- --ignored --nocapture`
4776    ///
4777    /// `IVF_SCAN_BENCH_DOCS`, `IVF_SCAN_BENCH_DIM`, `IVF_SCAN_BENCH_ITERS`,
4778    /// and `IVF_SCAN_BENCH_THREADS` override the defaults.
4779    #[test]
4780    #[ignore]
4781    fn ivf_tq_parallel_scan_benchmark() {
4782        use crate::structures::vector::ivf::CoarseCentroids;
4783        use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
4784
4785        fn env_usize(name: &str, default: usize) -> usize {
4786            std::env::var(name)
4787                .ok()
4788                .map(|value| {
4789                    value
4790                        .parse::<usize>()
4791                        .unwrap_or_else(|_| panic!("{name} must be a positive integer"))
4792                })
4793                .unwrap_or(default)
4794                .max(1)
4795        }
4796
4797        fn median_ms(samples: &mut [f64]) -> f64 {
4798            samples.sort_unstable_by(f64::total_cmp);
4799            samples[samples.len() / 2]
4800        }
4801
4802        let dim = env_usize("IVF_SCAN_BENCH_DIM", 128);
4803        let count = env_usize("IVF_SCAN_BENCH_DOCS", 262_144);
4804        let iterations = env_usize("IVF_SCAN_BENCH_ITERS", 30);
4805        let threads = env_usize("IVF_SCAN_BENCH_THREADS", num_cpus::get().min(8));
4806        let codec = std::sync::Arc::new(TqCodec::new(dim));
4807        let version = crate::structures::mark_ivf_tq_cosine_generation(0x51ca_0001);
4808        let mut centroid = vec![0.0f32; dim];
4809        centroid[0] = 1.0;
4810        let centroids = CoarseCentroids {
4811            num_clusters: 1,
4812            dim,
4813            centroids: centroid,
4814            version,
4815            soar_config: None,
4816            routing_index: None,
4817        };
4818        let mut index = IvfTqIndex::new(
4819            dim,
4820            crate::dsl::IvfRoutingMode::Flat,
4821            version,
4822            std::sync::Arc::clone(&codec),
4823        );
4824        let mut scratch = TqIvfEncodeScratch::default();
4825        let mut state = 0x5eed_u64;
4826        let mut vector = vec![0.0f32; dim];
4827        for posting_index in 0..count {
4828            for value in &mut vector {
4829                state = state
4830                    .wrapping_mul(6364136223846793005)
4831                    .wrapping_add(1442695040888963407);
4832                *value = ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5;
4833            }
4834            let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
4835            vector.iter_mut().for_each(|value| *value /= norm);
4836            index.add_vector(
4837                &centroids,
4838                u32::try_from(posting_index / 2).unwrap(),
4839                u16::try_from(posting_index % 2).unwrap(),
4840                &vector,
4841                &mut scratch,
4842            );
4843        }
4844        let mut bytes = Vec::new();
4845        write_built_ivf_tq(&index, 1, &mut bytes).unwrap();
4846        let disk = AnnDiskIndex::open(
4847            OwnedBytes::new(bytes),
4848            AnnKind::IvfTq,
4849            u32::try_from(count.div_ceil(2)).unwrap(),
4850        )
4851        .unwrap();
4852        let query = vec![1.0f32; dim];
4853        let plan = TqIvfQueryPlan::build(
4854            &centroids,
4855            &codec,
4856            &query,
4857            1,
4858            crate::dsl::IvfRoutingMode::Flat,
4859        );
4860        let pool = rayon::ThreadPoolBuilder::new()
4861            .num_threads(threads)
4862            .build()
4863            .unwrap();
4864        let serial_pool = rayon::ThreadPoolBuilder::new()
4865            .num_threads(1)
4866            .build()
4867            .unwrap();
4868
4869        let serial = pool
4870            .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512))
4871            .unwrap();
4872        let parallel = pool
4873            .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512))
4874            .unwrap();
4875        assert_eq!(parallel, serial);
4876
4877        let mut serial_ms = Vec::with_capacity(iterations);
4878        let mut parallel_ms = Vec::with_capacity(iterations);
4879        for _ in 0..iterations {
4880            let start = std::time::Instant::now();
4881            std::hint::black_box(
4882                pool.install(|| {
4883                    disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512)
4884                })
4885                .unwrap(),
4886            );
4887            serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4888
4889            let start = std::time::Instant::now();
4890            std::hint::black_box(
4891                pool.install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512))
4892                    .unwrap(),
4893            );
4894            parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4895        }
4896        let serial_p50 = median_ms(&mut serial_ms);
4897        let parallel_p50 = median_ms(&mut parallel_ms);
4898        println!(
4899            "IVF-TQ distinct scan: postings={count} dim={dim} threads={threads} \
4900             serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x",
4901            serial_p50 / parallel_p50,
4902        );
4903
4904        let combiner = crate::query::MultiValueCombiner::Sum;
4905        let serial_combined = serial_pool
4906            .install(|| {
4907                disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4908            })
4909            .unwrap();
4910        let parallel_combined = pool
4911            .install(|| {
4912                disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4913            })
4914            .unwrap();
4915        assert_eq!(parallel_combined, serial_combined);
4916        let mut serial_combined_ms = Vec::with_capacity(iterations);
4917        let mut parallel_combined_ms = Vec::with_capacity(iterations);
4918        for _ in 0..iterations {
4919            let start = std::time::Instant::now();
4920            std::hint::black_box(
4921                serial_pool
4922                    .install(|| {
4923                        disk.search_ivf_tq_combined_documents_with_tuning(
4924                            20, &plan, combiner, 1, 512,
4925                        )
4926                    })
4927                    .unwrap(),
4928            );
4929            serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4930
4931            let start = std::time::Instant::now();
4932            std::hint::black_box(
4933                pool.install(|| {
4934                    disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4935                })
4936                .unwrap(),
4937            );
4938            parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4939        }
4940        let serial_combined_p50 = median_ms(&mut serial_combined_ms);
4941        let parallel_combined_p50 = median_ms(&mut parallel_combined_ms);
4942        println!(
4943            "IVF-TQ combined scan: postings={count} dim={dim} threads={threads} \
4944             serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \
4945             speedup={:.2}x",
4946            serial_combined_p50 / parallel_combined_p50,
4947        );
4948    }
4949
4950    /// Focused binary-IVF benchmark for single-value top-k and multi-value
4951    /// combined scoring. Environment overrides use the `BINARY_SCAN_BENCH_*`
4952    /// prefix with `POSTINGS`, `DIM_BITS`, `ITERS`, and `THREADS` suffixes.
4953    #[test]
4954    #[ignore]
4955    fn binary_ivf_parallel_scan_benchmark() {
4956        fn env_usize(name: &str, default: usize) -> usize {
4957            std::env::var(name)
4958                .ok()
4959                .map(|value| value.parse::<usize>().unwrap())
4960                .unwrap_or(default)
4961                .max(1)
4962        }
4963
4964        fn median_ms(samples: &mut [f64]) -> f64 {
4965            samples.sort_unstable_by(f64::total_cmp);
4966            samples[samples.len() / 2]
4967        }
4968
4969        let count = env_usize("BINARY_SCAN_BENCH_POSTINGS", 131_072);
4970        let dim_bits = env_usize("BINARY_SCAN_BENCH_DIM_BITS", 2_048);
4971        assert!(dim_bits.is_multiple_of(8));
4972        let code_size = dim_bits / 8;
4973        let iterations = env_usize("BINARY_SCAN_BENCH_ITERS", 30);
4974        let threads = env_usize("BINARY_SCAN_BENCH_THREADS", num_cpus::get().min(8));
4975        let mut state = 0xb1a4_5eed_u64;
4976        let mut codes = vec![0u8; count * code_size];
4977        for code in &mut codes {
4978            state = state
4979                .wrapping_mul(6364136223846793005)
4980                .wrapping_add(1442695040888963407);
4981            *code = (state >> 56) as u8;
4982        }
4983        let query = vec![0xa5u8; code_size];
4984        let unique_docs: Vec<u32> = (0..count)
4985            .map(|index| u32::try_from(index).unwrap())
4986            .collect();
4987        let unique_ordinals = vec![0u16; count];
4988        let multi_docs: Vec<u32> = (0..count)
4989            .map(|index| u32::try_from(index / 2).unwrap())
4990            .collect();
4991        let multi_ordinals: Vec<u16> = (0..count)
4992            .map(|index| u16::try_from(index % 2).unwrap())
4993            .collect();
4994        let header = AnnDiskHeader {
4995            kind: AnnKind::BinaryIvf,
4996            routing: IvfRoutingMode::Flat,
4997            dim: dim_bits,
4998            code_size,
4999            num_clusters: 1,
5000            quantizer_version: 0xb1a4,
5001            codebook_version: 0,
5002            vector_count: count,
5003        };
5004        let mut unique_bytes = Vec::new();
5005        write_built_runs(
5006            header.clone(),
5007            &[BuildRun {
5008                cluster_id: 0,
5009                doc_ids: &unique_docs,
5010                ordinals: &unique_ordinals,
5011                codes: &codes,
5012            }],
5013            &mut unique_bytes,
5014        )
5015        .unwrap();
5016        let unique_disk = AnnDiskIndex::open(
5017            OwnedBytes::new(unique_bytes),
5018            AnnKind::BinaryIvf,
5019            u32::try_from(count).unwrap(),
5020        )
5021        .unwrap();
5022        let mut multi_bytes = Vec::new();
5023        write_built_runs(
5024            header,
5025            &[BuildRun {
5026                cluster_id: 0,
5027                doc_ids: &multi_docs,
5028                ordinals: &multi_ordinals,
5029                codes: &codes,
5030            }],
5031            &mut multi_bytes,
5032        )
5033        .unwrap();
5034        let multi_disk = AnnDiskIndex::open(
5035            OwnedBytes::new(multi_bytes),
5036            AnnKind::BinaryIvf,
5037            u32::try_from(count.div_ceil(2)).unwrap(),
5038        )
5039        .unwrap();
5040        let pool = rayon::ThreadPoolBuilder::new()
5041            .num_threads(threads)
5042            .build()
5043            .unwrap();
5044        let serial_pool = rayon::ThreadPoolBuilder::new()
5045            .num_threads(1)
5046            .build()
5047            .unwrap();
5048
5049        let serial = pool
5050            .install(|| {
5051                unique_disk.search_binary_clusters_with_tuning::<false>(
5052                    &query,
5053                    20,
5054                    &[0],
5055                    usize::MAX,
5056                )
5057            })
5058            .unwrap();
5059        let parallel = pool
5060            .install(|| {
5061                unique_disk.search_binary_clusters_with_tuning::<false>(&query, 20, &[0], 1)
5062            })
5063            .unwrap();
5064        assert_eq!(parallel, serial);
5065        let mut serial_ms = Vec::with_capacity(iterations);
5066        let mut parallel_ms = Vec::with_capacity(iterations);
5067        for _ in 0..iterations {
5068            let start = std::time::Instant::now();
5069            std::hint::black_box(
5070                pool.install(|| {
5071                    unique_disk.search_binary_clusters_with_tuning::<false>(
5072                        &query,
5073                        20,
5074                        &[0],
5075                        usize::MAX,
5076                    )
5077                })
5078                .unwrap(),
5079            );
5080            serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
5081
5082            let start = std::time::Instant::now();
5083            std::hint::black_box(
5084                pool.install(|| {
5085                    unique_disk.search_binary_clusters_with_tuning::<false>(&query, 20, &[0], 1)
5086                })
5087                .unwrap(),
5088            );
5089            parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
5090        }
5091        let serial_p50 = median_ms(&mut serial_ms);
5092        let parallel_p50 = median_ms(&mut parallel_ms);
5093        println!(
5094            "binary-IVF distinct scan: postings={count} dim_bits={dim_bits} threads={threads} \
5095             serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x",
5096            serial_p50 / parallel_p50,
5097        );
5098
5099        let combiner = crate::query::MultiValueCombiner::Sum;
5100        let serial_combined = serial_pool
5101            .install(|| {
5102                multi_disk.search_binary_combined_documents_with_tuning(
5103                    20,
5104                    &query,
5105                    &[0],
5106                    combiner,
5107                    1,
5108                )
5109            })
5110            .unwrap();
5111        let parallel_combined = pool
5112            .install(|| {
5113                multi_disk.search_binary_combined_documents_with_tuning(
5114                    20,
5115                    &query,
5116                    &[0],
5117                    combiner,
5118                    1,
5119                )
5120            })
5121            .unwrap();
5122        assert_eq!(parallel_combined, serial_combined);
5123        let mut serial_combined_ms = Vec::with_capacity(iterations);
5124        let mut parallel_combined_ms = Vec::with_capacity(iterations);
5125        for _ in 0..iterations {
5126            let start = std::time::Instant::now();
5127            std::hint::black_box(
5128                serial_pool
5129                    .install(|| {
5130                        multi_disk.search_binary_combined_documents_with_tuning(
5131                            20,
5132                            &query,
5133                            &[0],
5134                            combiner,
5135                            1,
5136                        )
5137                    })
5138                    .unwrap(),
5139            );
5140            serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
5141
5142            let start = std::time::Instant::now();
5143            std::hint::black_box(
5144                pool.install(|| {
5145                    multi_disk.search_binary_combined_documents_with_tuning(
5146                        20,
5147                        &query,
5148                        &[0],
5149                        combiner,
5150                        1,
5151                    )
5152                })
5153                .unwrap(),
5154            );
5155            parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
5156        }
5157        let serial_combined_p50 = median_ms(&mut serial_combined_ms);
5158        let parallel_combined_p50 = median_ms(&mut parallel_combined_ms);
5159        println!(
5160            "binary-IVF combined scan: postings={count} dim_bits={dim_bits} threads={threads} \
5161             serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \
5162             speedup={:.2}x",
5163            serial_combined_p50 / parallel_combined_p50,
5164        );
5165    }
5166
5167    #[test]
5168    fn open_rejects_tq_payload_with_inconsistent_geometry() {
5169        let (bytes, _) = build_tq_payload(20, 4, 9);
5170        // code_size (header bytes 12..16) is P/2 = 16 for dim 20; corrupt to 15.
5171        let mut corrupted = bytes.clone();
5172        corrupted[12..16].copy_from_slice(&15u32.to_le_bytes());
5173        assert!(
5174            AnnDiskIndex::open(OwnedBytes::new(corrupted), AnnKind::TqFlat, 4).is_err(),
5175            "TQ header with code_size != padded_dim/2 must be refused"
5176        );
5177
5178        // A block-padded TQ column must not validate under another kind.
5179        let (short_bytes, _) = build_tq_payload(20, 4, 9);
5180        let mut wrong_kind = short_bytes.clone();
5181        wrong_kind[4] = AnnKind::BinaryIvf as u8;
5182        assert!(
5183            AnnDiskIndex::open(OwnedBytes::new(wrong_kind), AnnKind::BinaryIvf, 4).is_err(),
5184            "TQ block-padded columns must not validate under another kind"
5185        );
5186
5187        // The retired IVF-PQ discriminant must be refused loudly.
5188        let (legacy_bytes, _) = build_tq_payload(20, 4, 9);
5189        let mut legacy_kind = legacy_bytes.clone();
5190        legacy_kind[4] = 1;
5191        let Err(error) = AnnDiskIndex::open(OwnedBytes::new(legacy_kind), AnnKind::TqFlat, 4)
5192        else {
5193            panic!("retired IVF-PQ payloads must not open");
5194        };
5195        assert!(
5196            error.to_string().contains("IVF-PQ"),
5197            "error must name the retired format: {error}"
5198        );
5199        assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 4).is_ok());
5200    }
5201
5202    #[test]
5203    fn open_rejects_old_or_out_of_range_ann_payloads() {
5204        let mut legacy = vec![0u8; ANN_HEADER_SIZE + ANN_FOOTER_SIZE];
5205        legacy[..4].copy_from_slice(b"old!");
5206        assert!(AnnDiskIndex::open(OwnedBytes::new(legacy), AnnKind::BinaryIvf, 1).is_err());
5207
5208        let docs = [0u32];
5209        let ordinals = [0u16];
5210        let codes = [0u8];
5211        let runs = [BuildRun {
5212            cluster_id: 0,
5213            doc_ids: &docs,
5214            ordinals: &ordinals,
5215            codes: &codes,
5216        }];
5217        let mut bytes = Vec::new();
5218        write_built_runs(binary_header(1), &runs, &mut bytes).unwrap();
5219        let footer = bytes.len() - ANN_FOOTER_SIZE;
5220        let directory = usize::try_from(u64::from_le_bytes(
5221            bytes[footer..footer + 8].try_into().unwrap(),
5222        ))
5223        .unwrap();
5224        bytes[directory + 12..directory + 16].copy_from_slice(&10u32.to_le_bytes());
5225        assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 1).is_err());
5226    }
5227
5228    fn scann_binary_artifact(
5229        generation: u64,
5230    ) -> crate::structures::vector::scann::ScannTrainedArtifact {
5231        use crate::structures::vector::scann::{
5232            ScannConfig, ScannEncoding, ScannRoutingLevel, ScannTrainedArtifact,
5233        };
5234        ScannTrainedArtifact::new(
5235            generation,
5236            100_000,
5237            ScannConfig {
5238                dimension: 16,
5239                tree_levels: 1,
5240                num_leaves: 2,
5241                encoding: ScannEncoding::BinaryHamming,
5242            },
5243            vec![ScannRoutingLevel {
5244                centroid_count: 2,
5245                centroid_codes: vec![0, 0, 0xff, 0xff],
5246                minimums: Vec::new(),
5247                steps: Vec::new(),
5248                child_offsets: Vec::new(),
5249            }],
5250            None,
5251        )
5252        .unwrap()
5253    }
5254
5255    #[test]
5256    fn scann_binary_ann_round_trip_validates_exact_global_generation() {
5257        use crate::structures::vector::scann::{
5258            ScannEncoding, ScannLeafRun, ScannSegmentPayload, ScannTrainedArtifactView,
5259        };
5260
5261        let artifact = scann_binary_artifact(41);
5262        let primary = ScannLeafRun::from_rows(
5263            0,
5264            0,
5265            &[0, 1],
5266            &[0, 2],
5267            vec![0x12, 0x34, 0xab, 0xcd],
5268            ScannEncoding::BinaryHamming,
5269            16,
5270        )
5271        .unwrap();
5272        let secondary = ScannLeafRun::from_rows(
5273            1,
5274            0,
5275            &[0],
5276            &[0],
5277            vec![0x12, 0x34],
5278            ScannEncoding::BinaryHamming,
5279            16,
5280        )
5281        .unwrap();
5282        let payload = ScannSegmentPayload::new(&artifact, 2, vec![primary, secondary]).unwrap();
5283        let mut bytes = Vec::new();
5284        write_built_scann(&payload, &mut bytes).unwrap();
5285
5286        let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannBinary, 2).unwrap();
5287        assert_eq!(disk.header.quantizer_version, artifact.generation);
5288        assert_eq!(disk.header.codebook_version, artifact.artifact_id);
5289        assert_eq!(disk.header.vector_count, 3);
5290        let artifact_bytes = artifact.to_bytes().unwrap();
5291        let view = ScannTrainedArtifactView::parse(&artifact_bytes).unwrap();
5292        disk.validate_scann_generation(&view.config, view.generation, view.artifact_id)
5293            .unwrap();
5294        let selective = crate::structures::SoarConfig::new().target_spill_fraction(0.5);
5295        disk.validate_scann_posting_count(2, Some(&selective))
5296            .unwrap();
5297        assert!(
5298            disk.validate_scann_posting_count(2, None)
5299                .unwrap_err()
5300                .to_string()
5301                .contains("physical postings")
5302        );
5303        let too_small_budget = crate::structures::SoarConfig::new().target_spill_fraction(0.49);
5304        assert!(
5305            disk.validate_scann_posting_count(2, Some(&too_small_budget))
5306                .is_err()
5307        );
5308
5309        let serial = disk
5310            .search_binary_clusters_with_tuning::<false>(&[0x12, 0x34], 2, &[0, 1], usize::MAX)
5311            .unwrap();
5312        let parallel = rayon::ThreadPoolBuilder::new()
5313            .num_threads(2)
5314            .build()
5315            .unwrap()
5316            .install(|| {
5317                disk.search_binary_clusters_with_tuning::<false>(&[0x12, 0x34], 2, &[0, 1], 1)
5318            })
5319            .unwrap();
5320        assert_eq!(parallel, serial);
5321        assert_eq!(serial.len(), 2, "secondary posting duplicated a result");
5322        assert_eq!(serial[0], (0, 0, 1.0));
5323
5324        let other = scann_binary_artifact(42);
5325        let other_bytes = other.to_bytes().unwrap();
5326        let other_view = ScannTrainedArtifactView::parse(&other_bytes).unwrap();
5327        let error = disk
5328            .validate_scann_generation(
5329                &other_view.config,
5330                other_view.generation,
5331                other_view.artifact_id,
5332            )
5333            .unwrap_err();
5334        assert!(error.to_string().contains("global trained generation"));
5335    }
5336
5337    #[test]
5338    fn scann_ah_ann_round_trip_preserves_fastscan_tail_geometry() {
5339        use crate::structures::vector::scann::{
5340            ScannAhCodebook, ScannConfig, ScannEncoding, ScannLeafRun, ScannRoutingLevel,
5341            ScannSegmentPayload, ScannTrainedArtifact, ScannTrainedArtifactView,
5342        };
5343
5344        let encoding = ScannEncoding::AsymmetricHash {
5345            dimensions_per_block: 2,
5346            bits_per_code: 4,
5347        };
5348        let artifact = ScannTrainedArtifact::new(
5349            51,
5350            100_000,
5351            ScannConfig {
5352                dimension: 8,
5353                tree_levels: 1,
5354                num_leaves: 2,
5355                encoding,
5356            },
5357            vec![ScannRoutingLevel {
5358                centroid_count: 2,
5359                centroid_codes: vec![0; 16],
5360                minimums: vec![0.0; 8],
5361                steps: vec![1.0; 8],
5362                child_offsets: Vec::new(),
5363            }],
5364            Some(ScannAhCodebook {
5365                dimensions_per_block: 2,
5366                centers_per_block: 16,
5367                centers: vec![0.0; 4 * 16 * 2],
5368            }),
5369        )
5370        .unwrap();
5371        let docs: Vec<u32> = (0..33).collect();
5372        let ordinals = vec![0u16; docs.len()];
5373        let codes = vec![0x5a; encoding.leaf_code_bytes(8, docs.len()).unwrap()];
5374        let run = ScannLeafRun::from_rows(0, 0, &docs, &ordinals, codes, encoding, 8).unwrap();
5375        let payload = ScannSegmentPayload::new(&artifact, 33, vec![run]).unwrap();
5376        let mut bytes = Vec::new();
5377        write_built_scann(&payload, &mut bytes).unwrap();
5378
5379        let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, 33).unwrap();
5380        assert_eq!(
5381            disk.runs[0].codes.len(),
5382            encoding.leaf_code_bytes(8, 33).unwrap()
5383        );
5384        let artifact_bytes = artifact.to_bytes().unwrap();
5385        let view = ScannTrainedArtifactView::parse(&artifact_bytes).unwrap();
5386        disk.validate_scann_generation(&view.config, view.generation, view.artifact_id)
5387            .unwrap();
5388    }
5389
5390    #[test]
5391    fn scann_binary_merge_copies_codes_and_rebases_without_retraining() {
5392        use crate::structures::vector::scann::{ScannEncoding, ScannLeafRun, ScannSegmentPayload};
5393
5394        let artifact = scann_binary_artifact(61);
5395        let make = |codes: Vec<u8>| {
5396            let secondary_code = codes[..2].to_vec();
5397            let primary = ScannLeafRun::from_rows(
5398                0,
5399                0,
5400                &[0, 1],
5401                &[0, 0],
5402                codes,
5403                ScannEncoding::BinaryHamming,
5404                16,
5405            )
5406            .unwrap();
5407            let secondary = ScannLeafRun::from_rows(
5408                1,
5409                0,
5410                &[0],
5411                &[0],
5412                secondary_code,
5413                ScannEncoding::BinaryHamming,
5414                16,
5415            )
5416            .unwrap();
5417            let payload = ScannSegmentPayload::new(&artifact, 2, vec![primary, secondary]).unwrap();
5418            let mut bytes = Vec::new();
5419            write_built_scann(&payload, &mut bytes).unwrap();
5420            AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannBinary, 2).unwrap()
5421        };
5422        let left = make(vec![1, 2, 3, 4]);
5423        let right = make(vec![5, 6, 7, 8]);
5424        let left_codes = left.raw.as_slice()[left.runs[0].codes.clone()].to_vec();
5425        let right_codes = right.raw.as_slice()[right.runs[0].codes.clone()].to_vec();
5426
5427        let mut merged_bytes = Vec::new();
5428        write_merged_ann(&[(&left, 0), (&right, 2)], &mut merged_bytes).unwrap();
5429        let merged =
5430            AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::ScannBinary, 4).unwrap();
5431        assert_eq!(merged.header.vector_count, 6);
5432        merged
5433            .validate_scann_posting_count(
5434                4,
5435                Some(&crate::structures::SoarConfig::new().target_spill_fraction(0.5)),
5436            )
5437            .unwrap();
5438        assert_eq!(merged.runs[0].doc_base, 0);
5439        assert_eq!(merged.runs[1].doc_base, 2);
5440        assert_eq!(
5441            &merged.raw.as_slice()[merged.runs[0].codes.clone()],
5442            left_codes
5443        );
5444        assert_eq!(
5445            &merged.raw.as_slice()[merged.runs[1].codes.clone()],
5446            right_codes
5447        );
5448        assert_eq!(
5449            merged
5450                .search_binary_clusters::<false>(&[0, 0], 4, &[0, 1])
5451                .unwrap()
5452                .len(),
5453            4,
5454            "merge exposed primary and secondary postings as separate results",
5455        );
5456
5457        let other_artifact = scann_binary_artifact(62);
5458        let other_run = ScannLeafRun::from_rows(
5459            0,
5460            0,
5461            &[0],
5462            &[0],
5463            vec![9, 10],
5464            ScannEncoding::BinaryHamming,
5465            16,
5466        )
5467        .unwrap();
5468        let other_payload = ScannSegmentPayload::new(&other_artifact, 1, vec![other_run]).unwrap();
5469        let mut other_bytes = Vec::new();
5470        write_built_scann(&other_payload, &mut other_bytes).unwrap();
5471        let other =
5472            AnnDiskIndex::open(OwnedBytes::new(other_bytes), AnnKind::ScannBinary, 1).unwrap();
5473        let error = write_merged_ann(&[(&left, 0), (&other, 2)], &mut Vec::new()).unwrap_err();
5474        assert!(error.to_string().contains("incompatible generations"));
5475    }
5476}