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