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