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